"""构造并 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