Files
Mutual_Fund/scripts/check_nl2sql_consistency.py
T
2026-09-14 10:57:48 +08:00

63 lines
2.6 KiB
Python

"""检查 MySQL 数据字典和 Milvus NL2SQL 向量数量的一致性。"""
from __future__ import annotations
import asyncio
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from sqlalchemy import text
from config import database
from config.database.milvus import client
from config.database.mysql import get_session_factory
from config.settings import settings
from nl2sql.metadata import build_metadata_chunks
from nl2sql.milvus_collections import NL2SQL_COLLECTION
from model.base import Base
from scripts.sync_nl2sql_metadata import collect_orm_tables, merge_orm_metadata_rows
async def collect_consistency() -> dict:
table_sql = text(
"SELECT TABLE_NAME, TABLE_COMMENT, TABLE_TYPE FROM information_schema.tables WHERE TABLE_SCHEMA=:db"
)
column_sql = text(
"SELECT TABLE_NAME, COLUMN_NAME, COLUMN_COMMENT, DATA_TYPE, IS_NULLABLE, ORDINAL_POSITION "
"FROM information_schema.columns WHERE TABLE_SCHEMA=:db ORDER BY TABLE_NAME, ORDINAL_POSITION"
)
try:
async with get_session_factory()() as session:
tables = [dict(row) for row in (await session.execute(table_sql, {"db": settings.mysql.database})).mappings()]
columns = [dict(row) for row in (await session.execute(column_sql, {"db": settings.mysql.database})).mappings()]
# 与同步脚本保持同一口径:ORM 表是唯一权威名单,数据库中存在但
# 未被 model/ 定义的表不应影响 NL2SQL 元数据一致性判断。
orm_tables = collect_orm_tables()
db_tables = {
str(row.get("TABLE_NAME") or row.get("table_name") or "").strip()
for row in tables
}
missing_tables = sorted(orm_tables - db_tables)
expected_tables, expected_columns = merge_orm_metadata_rows(
tables,
columns,
allowed_tables=orm_tables,
missing_table_objects=(Base.metadata.tables[name] for name in missing_tables),
)
expected = len(build_metadata_chunks(expected_tables, expected_columns))
rows = await client().query(
collection_name=NL2SQL_COLLECTION,
filter="is_valid == true and is_deprecated == false",
output_fields=["id"],
)
actual = len(rows or [])
return {"expected_chunks": expected, "actual_valid_vectors": actual, "consistent": expected == actual}
finally:
await database.mysql.dispose()
await database.milvus.dispose()
if __name__ == "__main__":
print(asyncio.run(collect_consistency()))