114 lines
4.1 KiB
Python
114 lines
4.1 KiB
Python
"""将 MySQL information_schema 结果标准化为 NL2SQL 元数据 chunk。"""
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
from typing import Any
|
|
|
|
|
|
def _value(row: dict[str, Any], name: str, default: Any = None) -> Any:
|
|
if name in row:
|
|
return row[name]
|
|
upper = name.upper()
|
|
lower = name.lower()
|
|
return row.get(upper, row.get(lower, default))
|
|
|
|
|
|
def normalize_table_row(row: dict[str, Any]) -> dict[str, Any] | None:
|
|
table_name = str(_value(row, "table_name", "") or "").strip()
|
|
table_type = str(_value(row, "table_type", "BASE TABLE") or "").upper()
|
|
if not table_name or table_type != "BASE TABLE":
|
|
return None
|
|
return {
|
|
"table_name": table_name,
|
|
"table_comment": str(_value(row, "table_comment", "") or "").strip(),
|
|
"is_valid": True,
|
|
}
|
|
|
|
|
|
def normalize_column_row(row: dict[str, Any]) -> dict[str, Any]:
|
|
nullable = str(_value(row, "is_nullable", "NO") or "NO").upper() == "YES"
|
|
return {
|
|
"table_name": str(_value(row, "table_name", "") or "").strip(),
|
|
"field_name": str(
|
|
_value(row, "column_name", _value(row, "field_name", "")) or ""
|
|
).strip(),
|
|
"column_comment": str(_value(row, "column_comment", "") or "").strip(),
|
|
"data_type": str(_value(row, "data_type", "") or "").strip(),
|
|
"is_nullable": nullable,
|
|
"ordinal_position": int(_value(row, "ordinal_position", 0) or 0),
|
|
}
|
|
|
|
|
|
def _hash_text(text: str) -> str:
|
|
return hashlib.sha256(text.encode("utf-8")).hexdigest()
|
|
|
|
|
|
def _chunk_id(chunk_type: str, table_name: str, field_name: str = "") -> str:
|
|
raw = f"{chunk_type}:{table_name}:{field_name}"
|
|
return hashlib.sha256(raw.encode("utf-8")).hexdigest()[:64]
|
|
|
|
|
|
def build_metadata_chunks(
|
|
tables: list[dict[str, Any]], columns: list[dict[str, Any]]
|
|
) -> list[dict[str, Any]]:
|
|
valid_tables = sorted(
|
|
(table for table in (normalize_table_row(row) for row in tables) if table),
|
|
key=lambda item: item["table_name"],
|
|
)
|
|
valid_table_names = {table["table_name"] for table in valid_tables}
|
|
normalized_columns = sorted(
|
|
(
|
|
column
|
|
for column in (normalize_column_row(row) for row in columns)
|
|
if column["table_name"] in valid_table_names and column["field_name"]
|
|
),
|
|
key=lambda item: (item["table_name"], item["ordinal_position"], item["field_name"]),
|
|
)
|
|
|
|
chunks: list[dict[str, Any]] = []
|
|
for table in valid_tables:
|
|
text = f"表名：{table['table_name']}\n表说明：{table['table_comment']}"
|
|
chunks.append(
|
|
{
|
|
"id": _chunk_id("table_meta", table["table_name"]),
|
|
"chunk_type": "table_meta",
|
|
"table_name": table["table_name"],
|
|
"field_name": "",
|
|
"text": text,
|
|
"content_hash": _hash_text(text),
|
|
"is_valid": table["is_valid"],
|
|
"is_deprecated": False,
|
|
}
|
|
)
|
|
|
|
for column in normalized_columns:
|
|
text = (
|
|
f"表名：{column['table_name']}\n"
|
|
f"字段名：{column['field_name']}\n"
|
|
f"字段说明：{column['column_comment']}\n"
|
|
f"字段类型：{column['data_type']}\n"
|
|
f"允许为空：{'是' if column['is_nullable'] else '否'}"
|
|
)
|
|
chunks.append(
|
|
{
|
|
"id": _chunk_id(
|
|
"field_meta", column["table_name"], column["field_name"]
|
|
),
|
|
"chunk_type": "field_meta",
|
|
"table_name": column["table_name"],
|
|
"field_name": column["field_name"],
|
|
"text": text,
|
|
"content_hash": _hash_text(text),
|
|
"is_valid": True,
|
|
"is_deprecated": False,
|
|
}
|
|
)
|
|
return chunks
|
|
|
|
|
|
def metadata_signature(chunk: dict[str, Any]) -> str:
|
|
"""返回用于比较 chunk 内容的稳定签名字符串。"""
|
|
payload = {key: chunk[key] for key in ("chunk_type", "table_name", "field_name", "text")}
|
|
return json.dumps(payload, ensure_ascii=False, sort_keys=True, separators=(",", ":"))
|