89 lines
3.0 KiB
Python
89 lines
3.0 KiB
Python
"""从 MySQL information_schema 加载并校验 NL2SQL 权威 Schema。"""
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
import logging
|
|
from typing import Any
|
|
|
|
from sqlalchemy import bindparam, text
|
|
|
|
from nl2sql.metadata import normalize_column_row, normalize_table_row
|
|
|
|
logger = logging.getLogger("nl2sql.schema")
|
|
SCHEMA_CACHE_TTL = 60
|
|
|
|
|
|
def _schema_cache_key(database: str, candidates: list[str]) -> str:
|
|
payload = json.dumps([database, candidates], ensure_ascii=False, separators=(",", ":"))
|
|
return "nl2sql:schema:" + hashlib.sha256(payload.encode()).hexdigest()
|
|
|
|
|
|
TABLES_SQL = text(
|
|
"""
|
|
SELECT TABLE_NAME, TABLE_COMMENT, TABLE_TYPE
|
|
FROM information_schema.tables
|
|
WHERE TABLE_SCHEMA = :database
|
|
AND TABLE_NAME IN :candidate_tables
|
|
"""
|
|
).bindparams(bindparam("candidate_tables", expanding=True))
|
|
|
|
COLUMNS_SQL = text(
|
|
"""
|
|
SELECT TABLE_NAME, COLUMN_NAME, COLUMN_COMMENT, DATA_TYPE,
|
|
IS_NULLABLE, ORDINAL_POSITION
|
|
FROM information_schema.columns
|
|
WHERE TABLE_SCHEMA = :database
|
|
AND TABLE_NAME IN :candidate_tables
|
|
ORDER BY TABLE_NAME, ORDINAL_POSITION
|
|
"""
|
|
).bindparams(bindparam("candidate_tables", expanding=True))
|
|
|
|
|
|
async def load_authoritative_schema(
|
|
session,
|
|
*,
|
|
database: str,
|
|
candidate_tables: set[str] | list[str],
|
|
redis=None,
|
|
cache_ttl: int = SCHEMA_CACHE_TTL,
|
|
) -> dict[str, list[dict[str, Any]]]:
|
|
"""从权威元数据源加载候选表，并剔除不存在的表和孤立字段。"""
|
|
candidates = {str(name).strip() for name in candidate_tables if str(name).strip()}
|
|
if not candidates:
|
|
return {"tables": [], "columns": []}
|
|
|
|
candidate_list = sorted(candidates)
|
|
cache_key = _schema_cache_key(database, candidate_list)
|
|
if redis is not None:
|
|
try:
|
|
cached = await redis.get(cache_key)
|
|
if cached:
|
|
return json.loads(cached)
|
|
except Exception: # noqa: BLE001 缓存异常回退权威查询
|
|
logger.warning("Schema 缓存读取失败", exc_info=True)
|
|
|
|
params = {"database": database, "candidate_tables": candidate_list}
|
|
table_result = await session.execute(TABLES_SQL, params)
|
|
column_result = await session.execute(COLUMNS_SQL, params)
|
|
|
|
tables = [
|
|
normalized
|
|
for row in table_result.mappings()
|
|
if (normalized := normalize_table_row(dict(row))) is not None
|
|
and normalized["table_name"] in candidates
|
|
]
|
|
valid_table_names = {table["table_name"] for table in tables}
|
|
columns = []
|
|
for row in column_result.mappings():
|
|
normalized = normalize_column_row(dict(row))
|
|
if normalized["table_name"] in valid_table_names and normalized["field_name"]:
|
|
columns.append(normalized)
|
|
schema = {"tables": tables, "columns": columns}
|
|
if redis is not None:
|
|
try:
|
|
await redis.set(cache_key, json.dumps(schema, ensure_ascii=False), ex=cache_ttl)
|
|
except Exception: # noqa: BLE001 缓存异常不阻断查询
|
|
logger.warning("Schema 缓存写入失败", exc_info=True)
|
|
return schema
|