"""从 MySQL information_schema 加载并校验 NL2SQL 权威 Schema。""" from __future__ import annotations from typing import Any from sqlalchemy import bindparam, text from nl2sql.metadata import normalize_column_row, normalize_table_row 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], ) -> 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": []} params = {"database": database, "candidate_tables": sorted(candidates)} 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) return {"tables": tables, "columns": columns}