47 lines
1.8 KiB
Python
47 lines
1.8 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
|
|
|
|
|
|
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()]
|
|
expected = len(build_metadata_chunks(tables, 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()))
|