207 lines
7.7 KiB
Python
207 lines
7.7 KiB
Python
"""构造并 upsert NL2SQL 元数据行。"""
|
||
from __future__ import annotations
|
||
|
||
import time
|
||
from collections.abc import Awaitable, Callable, Iterable
|
||
from typing import Any
|
||
|
||
from nl2sql.embedding import embed_texts
|
||
from nl2sql.metadata import build_metadata_chunks
|
||
from nl2sql.milvus_collections import NL2SQL_COLLECTION
|
||
|
||
# ORM 已定义但数据库尚未建表的表，在表说明后追加此标注，
|
||
# 让 NL2SQL 生成侧知道该表暂不可查询（权威 Schema 校验也会兜底拒绝）。
|
||
MISSING_TABLE_MARKER = "（注意：当前数据库中尚未建表）"
|
||
|
||
|
||
def _escape_filter_value(value: str) -> str:
|
||
return value.replace("\\", "\\\\").replace('"', '\\"')
|
||
|
||
|
||
def _row_table_name(row: dict[str, Any]) -> str:
|
||
"""兼容 information_schema 大小写键名，取行内表名。"""
|
||
value = (
|
||
row.get("table_name")
|
||
or row.get("TABLE_NAME")
|
||
or row.get("Table_name")
|
||
or ""
|
||
)
|
||
return str(value).strip()
|
||
|
||
|
||
def _filter_rows_by_tables(
|
||
rows: list[dict[str, Any]], allowed_tables: set[str]
|
||
) -> list[dict[str, Any]]:
|
||
return [row for row in rows if _row_table_name(row) in allowed_tables]
|
||
|
||
|
||
async def _delete_tables_outside_allowlist(
|
||
milvus_client, allowed_tables: set[str]
|
||
) -> list[str]:
|
||
"""删除集合中不在允许名单内的表元数据 chunk，返回被清理的表名。"""
|
||
existing_rows = await milvus_client.query(
|
||
collection_name=NL2SQL_COLLECTION,
|
||
filter="is_valid == true",
|
||
output_fields=["table_name"],
|
||
limit=16384,
|
||
)
|
||
existing_tables = {
|
||
str(row.get("table_name") or "").strip()
|
||
for row in existing_rows or []
|
||
if str(row.get("table_name") or "").strip()
|
||
}
|
||
removed: list[str] = []
|
||
for table_name in sorted(existing_tables - allowed_tables):
|
||
await milvus_client.delete(
|
||
collection_name=NL2SQL_COLLECTION,
|
||
filter=f'table_name == "{_escape_filter_value(table_name)}"',
|
||
)
|
||
removed.append(table_name)
|
||
return removed
|
||
|
||
|
||
def build_orm_metadata_rows(
|
||
tables: Iterable[Any],
|
||
) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
|
||
"""把 SQLAlchemy Table 定义合成 information_schema 风格的元数据行。
|
||
|
||
用于 model/ 已定义、但数据库尚未建表的表：
|
||
- 表注释取 ORM ``__table_args__`` 的 comment（ORM 无则空串）；
|
||
- 字段注释取列 comment（ORM 通常未标注，则为空串）；
|
||
- 字段类型/可空性从 ORM 列定义推导。
|
||
"""
|
||
table_rows: list[dict[str, Any]] = []
|
||
column_rows: list[dict[str, Any]] = []
|
||
for table in tables:
|
||
table_rows.append(
|
||
{
|
||
"TABLE_NAME": str(table.name).strip(),
|
||
"TABLE_COMMENT": str(table.comment or "").strip(),
|
||
"TABLE_TYPE": "BASE TABLE",
|
||
}
|
||
)
|
||
for position, column in enumerate(table.columns, start=1):
|
||
column_rows.append(
|
||
{
|
||
"TABLE_NAME": str(table.name).strip(),
|
||
"COLUMN_NAME": str(column.name).strip(),
|
||
"COLUMN_COMMENT": str(column.comment or "").strip(),
|
||
"DATA_TYPE": str(column.type).strip().lower(),
|
||
"IS_NULLABLE": "YES" if column.nullable else "NO",
|
||
"ORDINAL_POSITION": position,
|
||
}
|
||
)
|
||
return table_rows, column_rows
|
||
|
||
|
||
def merge_orm_metadata_rows(
|
||
table_rows: list[dict[str, Any]],
|
||
column_rows: list[dict[str, Any]],
|
||
*,
|
||
allowed_tables: set[str],
|
||
missing_table_objects: Iterable[Any] = (),
|
||
) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
|
||
"""以 model/ ORM 全量为名单主体合并元数据行。
|
||
|
||
- 库里真实存在的 ORM 表：沿用 information_schema 行（列注释更全）；
|
||
- 库里没有的 ORM 表（missing_table_objects）：用 ORM 定义合成行，
|
||
并在表说明追加 MISSING_TABLE_MARKER 标注"尚未建表"；
|
||
- 名单之外（DB-only）的行在此丢弃，交由 sync_metadata 从 Milvus 清理。
|
||
"""
|
||
allowed = {str(name).strip() for name in allowed_tables if str(name).strip()}
|
||
merged_tables = [row for row in table_rows if _row_table_name(row) in allowed]
|
||
merged_columns = [row for row in column_rows if _row_table_name(row) in allowed]
|
||
|
||
missing_table_rows, missing_column_rows = build_orm_metadata_rows(
|
||
missing_table_objects
|
||
)
|
||
for row in missing_table_rows:
|
||
comment = str(row.get("TABLE_COMMENT") or "").strip()
|
||
row["TABLE_COMMENT"] = f"{comment}{MISSING_TABLE_MARKER}".strip()
|
||
merged_tables.extend(missing_table_rows)
|
||
merged_columns.extend(missing_column_rows)
|
||
return merged_tables, merged_columns
|
||
|
||
|
||
async def prepare_metadata_rows(
|
||
chunks: list[dict[str, Any]],
|
||
*,
|
||
embedder: Callable[[list[str]], Awaitable[list[list[float]]]] = embed_texts,
|
||
timestamp: int | None = None,
|
||
) -> list[dict[str, Any]]:
|
||
if not chunks:
|
||
return []
|
||
vectors = await embedder([chunk["text"] for chunk in chunks])
|
||
if len(vectors) != len(chunks):
|
||
raise ValueError("Embedding result count must match metadata chunk count")
|
||
created_at = int(time.time()) if timestamp is None else timestamp
|
||
return [
|
||
{**chunk, "vector": vector, "created_at": created_at}
|
||
for chunk, vector in zip(chunks, vectors)
|
||
]
|
||
|
||
|
||
async def sync_metadata(
|
||
milvus_client,
|
||
table_rows: list[dict[str, Any]],
|
||
column_rows: list[dict[str, Any]],
|
||
*,
|
||
embedder: Callable[[list[str]], Awaitable[list[list[float]]]] = embed_texts,
|
||
timestamp: int | None = None,
|
||
allowed_tables: set[str] | None = None,
|
||
) -> int:
|
||
"""构造并 upsert 元数据 chunk。
|
||
|
||
allowed_tables 提供时：仅同步名单内的表，并删除集合中名单之外的
|
||
表元数据（ Milvus 里只保留主数据源认可的表，如 model/ ORM 覆盖的表）。
|
||
"""
|
||
if allowed_tables is not None:
|
||
allowed = {str(name).strip() for name in allowed_tables if str(name).strip()}
|
||
table_rows = _filter_rows_by_tables(table_rows, allowed)
|
||
column_rows = _filter_rows_by_tables(column_rows, allowed)
|
||
|
||
chunks = build_metadata_chunks(table_rows, column_rows)
|
||
chunks_by_table: dict[str, list[dict[str, Any]]] = {}
|
||
for chunk in chunks:
|
||
chunks_by_table.setdefault(chunk["table_name"], []).append(chunk)
|
||
|
||
updated_count = 0
|
||
for table_name, table_chunks in chunks_by_table.items():
|
||
escaped_name = _escape_filter_value(table_name)
|
||
existing_rows = await milvus_client.query(
|
||
collection_name=NL2SQL_COLLECTION,
|
||
filter=f'table_name == "{escaped_name}"',
|
||
output_fields=["id", "content_hash", "is_valid", "is_deprecated"],
|
||
)
|
||
desired_signature = {
|
||
(chunk["id"], chunk["content_hash"], chunk["is_valid"], chunk["is_deprecated"])
|
||
for chunk in table_chunks
|
||
}
|
||
existing_signature = {
|
||
(
|
||
row.get("id"),
|
||
row.get("content_hash"),
|
||
row.get("is_valid", True),
|
||
row.get("is_deprecated", False),
|
||
)
|
||
for row in existing_rows or []
|
||
}
|
||
if desired_signature == existing_signature:
|
||
continue
|
||
|
||
if existing_rows:
|
||
await milvus_client.delete(
|
||
collection_name=NL2SQL_COLLECTION,
|
||
filter=f'table_name == "{escaped_name}"',
|
||
)
|
||
rows = await prepare_metadata_rows(
|
||
table_chunks, embedder=embedder, timestamp=timestamp
|
||
)
|
||
if rows:
|
||
await milvus_client.upsert(collection_name=NL2SQL_COLLECTION, data=rows)
|
||
updated_count += len(rows)
|
||
|
||
if allowed_tables is not None:
|
||
await _delete_tables_outside_allowlist(milvus_client, allowed)
|
||
return updated_count
|