Mutual_Fund/scripts/sync_nl2sql_metadata.py

"""同步 NL2SQL 元数据到 Milvus（表名单以 model/ 目录的 ORM 为主）。

主数据源策略：
- 表名单 = model/ 下 SQLAlchemy ORM 定义的全部表（Base.metadata），
  不再与数据库求交集——数据库尚未建表的 ORM 表也纳入元数据，
  其表/字段信息从 ORM 定义合成，并在表说明标注"尚未建表"；
- 数据库里真实存在的 ORM 表，表/列中文注释仍取自 information_schema
  （列注释只存在于库中，ORM 代码里没有字段级注释）；
- 名单之外（库里有、model/ 没定义）的表元数据 chunk 会被从 Milvus 删除。
"""
from __future__ import annotations

import asyncio
import importlib
import pkgutil
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from sqlalchemy import text

from config.database import milvus as milvus_db
from config.database import mysql
from config.database.milvus import client as milvus_client
from config.database.mysql import get_session_factory
from config.settings import settings
from model.base import Base
from nl2sql.metadata_sync import merge_orm_metadata_rows, sync_metadata


TABLES_SQL = text(
    """
    SELECT TABLE_NAME, TABLE_COMMENT, TABLE_TYPE
    FROM information_schema.tables
    WHERE TABLE_SCHEMA = :database
    """
)

COLUMNS_SQL = text(
    """
    SELECT TABLE_NAME, COLUMN_NAME, COLUMN_COMMENT, DATA_TYPE,
           IS_NULLABLE, ORDINAL_POSITION
    FROM information_schema.columns
    WHERE TABLE_SCHEMA = :database
    ORDER BY TABLE_NAME, ORDINAL_POSITION
    """
)


def collect_orm_tables() -> set[str]:
    """导入 model/ 全部模块，从 Base.metadata 收集 ORM 表名。"""
    import model

    for module_info in pkgutil.iter_modules(model.__path__):
        importlib.import_module(f"model.{module_info.name}")
    return set(Base.metadata.tables.keys())


async def load_information_schema() -> tuple[list[dict], list[dict]]:
    async with get_session_factory()() as session:
        tables = [
            dict(row)
            for row in (await session.execute(TABLES_SQL, {"database": settings.mysql.database})).mappings()
        ]
        columns = [
            dict(row)
            for row in (await session.execute(COLUMNS_SQL, {"database": settings.mysql.database})).mappings()
        ]
    return tables, columns


def _table_name(row: dict) -> str:
    return str(row.get("TABLE_NAME") or row.get("table_name") or "").strip()


async def synchronize() -> tuple[int, set[str], set[str], set[str]]:
    """返回 (upserted, allowed_tables, 名单外表, 未建表的 ORM 表)。"""
    orm_tables = collect_orm_tables()
    try:
        tables, columns = await load_information_schema()
        db_tables = {_table_name(row) for row in tables}
        missing_tables = sorted(orm_tables - db_tables)
        table_rows, column_rows = merge_orm_metadata_rows(
            tables,
            columns,
            allowed_tables=orm_tables,
            missing_table_objects=(
                Base.metadata.tables[name] for name in missing_tables
            ),
        )
        upserted = await sync_metadata(
            milvus_client(),
            table_rows,
            column_rows,
            allowed_tables=orm_tables,
        )
        dropped = db_tables - orm_tables
        return upserted, orm_tables, dropped, set(missing_tables)
    finally:
        await mysql.dispose()
        await milvus_db.dispose()


if __name__ == "__main__":
    upserted, allowed, dropped, missing = asyncio.run(synchronize())
    print(f"ORM 表名单（model/ 全量）: {len(allowed)} 张")
    print(f"其中数据库尚未建表（用 ORM 定义合成元数据）: {len(missing)} 张")
    if missing:
        print(f"  {sorted(missing)}")
    print(f"已排除的非 ORM 表: {sorted(dropped) if dropped else '无'}")
    print(f"upserted {upserted} NL2SQL metadata chunks")