Mutual_Fund/nl2sql/schema.py

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