113 lines
3.8 KiB
Python
113 lines
3.8 KiB
Python
"""同步 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")
|