Mutual_Fund/nl2sql/metadata_sync.py

"""构造并 upsert NL2SQL 元数据行。"""
from __future__ import annotations

import time
from collections.abc import Awaitable, Callable, Iterable
from typing import Any

from nl2sql.embedding import embed_texts
from nl2sql.metadata import build_metadata_chunks
from nl2sql.milvus_collections import NL2SQL_COLLECTION

# ORM 已定义但数据库尚未建表的表，在表说明后追加此标注，
# 让 NL2SQL 生成侧知道该表暂不可查询（权威 Schema 校验也会兜底拒绝）。
MISSING_TABLE_MARKER = "（注意：当前数据库中尚未建表）"


def _escape_filter_value(value: str) -> str:
    return value.replace("\\", "\\\\").replace('"', '\\"')


def _row_table_name(row: dict[str, Any]) -> str:
    """兼容 information_schema 大小写键名，取行内表名。"""
    value = (
        row.get("table_name")
        or row.get("TABLE_NAME")
        or row.get("Table_name")
        or ""
    )
    return str(value).strip()


def _filter_rows_by_tables(
    rows: list[dict[str, Any]], allowed_tables: set[str]
) -> list[dict[str, Any]]:
    return [row for row in rows if _row_table_name(row) in allowed_tables]


async def _delete_tables_outside_allowlist(
    milvus_client, allowed_tables: set[str]
) -> list[str]:
    """删除集合中不在允许名单内的表元数据 chunk，返回被清理的表名。"""
    existing_rows = await milvus_client.query(
        collection_name=NL2SQL_COLLECTION,
        filter="is_valid == true",
        output_fields=["table_name"],
        limit=16384,
    )
    existing_tables = {
        str(row.get("table_name") or "").strip()
        for row in existing_rows or []
        if str(row.get("table_name") or "").strip()
    }
    removed: list[str] = []
    for table_name in sorted(existing_tables - allowed_tables):
        await milvus_client.delete(
            collection_name=NL2SQL_COLLECTION,
            filter=f'table_name == "{_escape_filter_value(table_name)}"',
        )
        removed.append(table_name)
    return removed


def build_orm_metadata_rows(
    tables: Iterable[Any],
) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
    """把 SQLAlchemy Table 定义合成 information_schema 风格的元数据行。

    用于 model/ 已定义、但数据库尚未建表的表：
    - 表注释取 ORM ``__table_args__`` 的 comment（ORM 无则空串）；
    - 字段注释取列 comment（ORM 通常未标注，则为空串）；
    - 字段类型/可空性从 ORM 列定义推导。
    """
    table_rows: list[dict[str, Any]] = []
    column_rows: list[dict[str, Any]] = []
    for table in tables:
        table_rows.append(
            {
                "TABLE_NAME": str(table.name).strip(),
                "TABLE_COMMENT": str(table.comment or "").strip(),
                "TABLE_TYPE": "BASE TABLE",
            }
        )
        for position, column in enumerate(table.columns, start=1):
            column_rows.append(
                {
                    "TABLE_NAME": str(table.name).strip(),
                    "COLUMN_NAME": str(column.name).strip(),
                    "COLUMN_COMMENT": str(column.comment or "").strip(),
                    "DATA_TYPE": str(column.type).strip().lower(),
                    "IS_NULLABLE": "YES" if column.nullable else "NO",
                    "ORDINAL_POSITION": position,
                }
            )
    return table_rows, column_rows


def merge_orm_metadata_rows(
    table_rows: list[dict[str, Any]],
    column_rows: list[dict[str, Any]],
    *,
    allowed_tables: set[str],
    missing_table_objects: Iterable[Any] = (),
) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
    """以 model/ ORM 全量为名单主体合并元数据行。

    - 库里真实存在的 ORM 表：沿用 information_schema 行（列注释更全）；
    - 库里没有的 ORM 表（missing_table_objects）：用 ORM 定义合成行，
      并在表说明追加 MISSING_TABLE_MARKER 标注"尚未建表"；
    - 名单之外（DB-only）的行在此丢弃，交由 sync_metadata 从 Milvus 清理。
    """
    allowed = {str(name).strip() for name in allowed_tables if str(name).strip()}
    merged_tables = [row for row in table_rows if _row_table_name(row) in allowed]
    merged_columns = [row for row in column_rows if _row_table_name(row) in allowed]

    missing_table_rows, missing_column_rows = build_orm_metadata_rows(
        missing_table_objects
    )
    for row in missing_table_rows:
        comment = str(row.get("TABLE_COMMENT") or "").strip()
        row["TABLE_COMMENT"] = f"{comment}{MISSING_TABLE_MARKER}".strip()
    merged_tables.extend(missing_table_rows)
    merged_columns.extend(missing_column_rows)
    return merged_tables, merged_columns


async def prepare_metadata_rows(
    chunks: list[dict[str, Any]],
    *,
    embedder: Callable[[list[str]], Awaitable[list[list[float]]]] = embed_texts,
    timestamp: int | None = None,
) -> list[dict[str, Any]]:
    if not chunks:
        return []
    vectors = await embedder([chunk["text"] for chunk in chunks])
    if len(vectors) != len(chunks):
        raise ValueError("Embedding result count must match metadata chunk count")
    created_at = int(time.time()) if timestamp is None else timestamp
    return [
        {**chunk, "vector": vector, "created_at": created_at}
        for chunk, vector in zip(chunks, vectors)
    ]


async def sync_metadata(
    milvus_client,
    table_rows: list[dict[str, Any]],
    column_rows: list[dict[str, Any]],
    *,
    embedder: Callable[[list[str]], Awaitable[list[list[float]]]] = embed_texts,
    timestamp: int | None = None,
    allowed_tables: set[str] | None = None,
) -> int:
    """构造并 upsert 元数据 chunk。

    allowed_tables 提供时：仅同步名单内的表，并删除集合中名单之外的
    表元数据（ Milvus 里只保留主数据源认可的表，如 model/ ORM 覆盖的表）。
    """
    if allowed_tables is not None:
        allowed = {str(name).strip() for name in allowed_tables if str(name).strip()}
        table_rows = _filter_rows_by_tables(table_rows, allowed)
        column_rows = _filter_rows_by_tables(column_rows, allowed)

    chunks = build_metadata_chunks(table_rows, column_rows)
    chunks_by_table: dict[str, list[dict[str, Any]]] = {}
    for chunk in chunks:
        chunks_by_table.setdefault(chunk["table_name"], []).append(chunk)

    updated_count = 0
    for table_name, table_chunks in chunks_by_table.items():
        escaped_name = _escape_filter_value(table_name)
        existing_rows = await milvus_client.query(
            collection_name=NL2SQL_COLLECTION,
            filter=f'table_name == "{escaped_name}"',
            output_fields=["id", "content_hash", "is_valid", "is_deprecated"],
        )
        desired_signature = {
            (chunk["id"], chunk["content_hash"], chunk["is_valid"], chunk["is_deprecated"])
            for chunk in table_chunks
        }
        existing_signature = {
            (
                row.get("id"),
                row.get("content_hash"),
                row.get("is_valid", True),
                row.get("is_deprecated", False),
            )
            for row in existing_rows or []
        }
        if desired_signature == existing_signature:
            continue

        if existing_rows:
            await milvus_client.delete(
                collection_name=NL2SQL_COLLECTION,
                filter=f'table_name == "{escaped_name}"',
            )
        rows = await prepare_metadata_rows(
            table_chunks, embedder=embedder, timestamp=timestamp
        )
        if rows:
            await milvus_client.upsert(collection_name=NL2SQL_COLLECTION, data=rows)
            updated_count += len(rows)

    if allowed_tables is not None:
        await _delete_tables_outside_allowlist(milvus_client, allowed)
    return updated_count