63 lines
1.7 KiB
Python
63 lines
1.7 KiB
Python
"""读取当前 MySQL 元数据并同步到 Milvus。"""
|
|
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.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 nl2sql.metadata_sync import 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
|
|
"""
|
|
)
|
|
|
|
|
|
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
|
|
|
|
|
|
async def synchronize() -> int:
|
|
try:
|
|
tables, columns = await load_information_schema()
|
|
return await sync_metadata(milvus_client(), tables, columns)
|
|
finally:
|
|
await mysql.dispose()
|
|
await milvus_db.dispose()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
print(f"upserted {asyncio.run(synchronize())} NL2SQL metadata chunks")
|