"""从 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