316 lines
10 KiB
Python
316 lines
10 KiB
Python
|
"""
|
|||
|
|
Schema管理器 - 管理数据库Schema的加载、索引和检索
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import json
|
|||
|
|
import hashlib
|
|||
|
|
from pathlib import Path
|
|||
|
|
from typing import List, Dict, Optional, Tuple
|
|||
|
|
import logging
|
|||
|
|
|
|||
|
|
from .models import DatabaseSchema, Table
|
|||
|
|
from .loader import SchemaLoader, load_schema_from_json, load_schema_from_g3sb
|
|||
|
|
|
|||
|
|
logger = logging.getLogger(__name__)
|
|||
|
|
|
|||
|
|
|
|||
|
|
class SchemaManager:
|
|||
|
|
"""
|
|||
|
|
Schema管理器
|
|||
|
|
|
|||
|
|
功能：
|
|||
|
|
1. 加载和解析Schema文件
|
|||
|
|
2. 管理表关系
|
|||
|
|
3. 提供表检索接口
|
|||
|
|
4. 生成紧凑的Schema描述
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
def __init__(self, schema: DatabaseSchema):
|
|||
|
|
"""
|
|||
|
|
初始化Schema管理器
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
schema: DatabaseSchema对象
|
|||
|
|
"""
|
|||
|
|
self.schema = schema
|
|||
|
|
self._table_dict = {tbl.name: tbl for tbl in schema.tables}
|
|||
|
|
self._embedding_index = None # 向量索引（延迟初始化）
|
|||
|
|
self._cache = {}
|
|||
|
|
|
|||
|
|
@classmethod
|
|||
|
|
def load_from_json(
|
|||
|
|
cls, json_path: str, g3sb_meta_path: Optional[str] = None
|
|||
|
|
) -> "SchemaManager":
|
|||
|
|
"""
|
|||
|
|
从JSON文件加载Schema
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
json_path: JSON文件路径
|
|||
|
|
g3sb_meta_path: 可选，G3SB table_meta.json（与 table_structure 配对）
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
SchemaManager实例
|
|||
|
|
"""
|
|||
|
|
schema = load_schema_from_json(json_path, g3sb_meta_path=g3sb_meta_path)
|
|||
|
|
return cls(schema)
|
|||
|
|
|
|||
|
|
@classmethod
|
|||
|
|
def load_from_g3sb(cls, meta_path: str, structure_path: str) -> "SchemaManager":
|
|||
|
|
"""
|
|||
|
|
从G3SB格式加载Schema
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
meta_path: table_meta.json路径
|
|||
|
|
structure_path: table_structure.json路径
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
SchemaManager实例
|
|||
|
|
"""
|
|||
|
|
schema = load_schema_from_g3sb(meta_path, structure_path)
|
|||
|
|
return cls(schema)
|
|||
|
|
|
|||
|
|
@classmethod
|
|||
|
|
def from_dict(cls, data: Dict) -> "SchemaManager":
|
|||
|
|
"""
|
|||
|
|
从字典创建SchemaManager
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
data: 包含database和tables的字典
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
SchemaManager实例
|
|||
|
|
"""
|
|||
|
|
loader = SchemaLoader()
|
|||
|
|
schema = loader._parse_table(data) if isinstance(data, dict) else None
|
|||
|
|
if not schema:
|
|||
|
|
raise ValueError("Invalid schema data")
|
|||
|
|
return cls(schema)
|
|||
|
|
|
|||
|
|
# ========== 查询接口 ==========
|
|||
|
|
|
|||
|
|
def get_table(self, table_name: str) -> Optional[Table]:
|
|||
|
|
"""获取指定表"""
|
|||
|
|
return self._table_dict.get(table_name)
|
|||
|
|
|
|||
|
|
def list_tables(self) -> List[str]:
|
|||
|
|
"""获取所有表名"""
|
|||
|
|
return list(self._table_dict.keys())
|
|||
|
|
|
|||
|
|
def get_tables(self) -> List[Table]:
|
|||
|
|
"""获取所有Table对象"""
|
|||
|
|
return list(self._table_dict.values())
|
|||
|
|
|
|||
|
|
def get_related_tables(self, table_name: str, depth: int = 1) -> List[Table]:
|
|||
|
|
"""
|
|||
|
|
获取关联表（通过外键关系）
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
table_name: 起始表名
|
|||
|
|
depth: 递归深度（1=直接关联，2=间接关联...）
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
相关表列表
|
|||
|
|
"""
|
|||
|
|
result = set()
|
|||
|
|
visited = set()
|
|||
|
|
|
|||
|
|
def _traverse(tname: str, current_depth: int):
|
|||
|
|
if tname not in self._table_dict or current_depth > depth:
|
|||
|
|
return
|
|||
|
|
if tname in visited:
|
|||
|
|
return
|
|||
|
|
visited.add(tname)
|
|||
|
|
|
|||
|
|
table = self._table_dict[tname]
|
|||
|
|
for fk in table.foreign_keys:
|
|||
|
|
result.add(fk.ref_table)
|
|||
|
|
_traverse(fk.ref_table, current_depth + 1)
|
|||
|
|
|
|||
|
|
# 反向外键（被引用的表）
|
|||
|
|
for other_table in self._table_dict.values():
|
|||
|
|
for fk in other_table.foreign_keys:
|
|||
|
|
if fk.ref_table == tname:
|
|||
|
|
result.add(other_table.name)
|
|||
|
|
_traverse(other_table.name, current_depth + 1)
|
|||
|
|
|
|||
|
|
_traverse(table_name, 0)
|
|||
|
|
return [self._table_dict[t] for t in result if t in self._table_dict]
|
|||
|
|
|
|||
|
|
# ========== Schema生成 ==========
|
|||
|
|
|
|||
|
|
def to_compact_string(
|
|||
|
|
self,
|
|||
|
|
table_names: List[str] = None,
|
|||
|
|
include_columns: bool = True,
|
|||
|
|
max_columns_per_table: int = 15,
|
|||
|
|
) -> str:
|
|||
|
|
"""
|
|||
|
|
生成紧凑的Schema描述字符串（用于LLM输入）
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
table_names: 指定包含的表（None表示所有表）
|
|||
|
|
include_columns: 是否包含字段详情
|
|||
|
|
max_columns_per_table: 每表最多显示字段数
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
Schema描述字符串
|
|||
|
|
"""
|
|||
|
|
if table_names is None:
|
|||
|
|
tables = self.get_tables()
|
|||
|
|
else:
|
|||
|
|
missing = [n for n in table_names if n not in self._table_dict]
|
|||
|
|
if missing:
|
|||
|
|
logger.warning(
|
|||
|
|
"以下表名不在已加载Schema中，已从本次Schema片段中省略: %s",
|
|||
|
|
missing,
|
|||
|
|
)
|
|||
|
|
tables = [t for t in self.get_tables() if t.name in table_names]
|
|||
|
|
|
|||
|
|
lines = []
|
|||
|
|
lines.append(f"数据库: {self.schema.name}")
|
|||
|
|
lines.append(f"涉及表数: {len(tables)}")
|
|||
|
|
lines.append("=" * 60)
|
|||
|
|
lines.append("")
|
|||
|
|
|
|||
|
|
for table in tables:
|
|||
|
|
lines.append(f"【表】{table.name}")
|
|||
|
|
if table.comment:
|
|||
|
|
lines.append(f" 描述: {table.comment}")
|
|||
|
|
|
|||
|
|
if include_columns and table.columns:
|
|||
|
|
lines.append(f" 字段:")
|
|||
|
|
for col in table.columns[:max_columns_per_table]:
|
|||
|
|
pk_marker = " [PK]" if col.is_primary_key else ""
|
|||
|
|
null_marker = " NOT NULL" if not col.nullable else ""
|
|||
|
|
comment = f" -- {col.comment}" if col.comment else ""
|
|||
|
|
lines.append(f" {col.name}: {col.data_type}{pk_marker}{null_marker}{comment}")
|
|||
|
|
|
|||
|
|
if len(table.columns) > max_columns_per_table:
|
|||
|
|
lines.append(f" ... 还有 {len(table.columns) - max_columns_per_table} 个字段")
|
|||
|
|
|
|||
|
|
# 外键关系
|
|||
|
|
if table.foreign_keys:
|
|||
|
|
lines.append(" 外键:")
|
|||
|
|
for fk in table.foreign_keys[:5]:
|
|||
|
|
lines.append(
|
|||
|
|
f" {', '.join(fk.columns)} → {fk.ref_table}({', '.join(fk.ref_columns)})"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
lines.append("")
|
|||
|
|
|
|||
|
|
return "\n".join(lines)
|
|||
|
|
|
|||
|
|
def to_filtered_schema(self, table_names: List[str]) -> DatabaseSchema:
|
|||
|
|
"""
|
|||
|
|
根据表名过滤，生成子集Schema
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
table_names: 要保留的表名列表
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
新的DatabaseSchema对象
|
|||
|
|
"""
|
|||
|
|
filtered_tables = [t for t in self.get_tables() if t.name in table_names]
|
|||
|
|
return DatabaseSchema(
|
|||
|
|
name=f"{self.schema.name}_filtered",
|
|||
|
|
tables=filtered_tables,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# ========== 向量索引支持 ==========
|
|||
|
|
|
|||
|
|
def set_embeddings(self, table_embeddings: Dict[str, List[float]]):
|
|||
|
|
"""
|
|||
|
|
为表设置预计算的embedding向量
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
table_embeddings: {table_name: embedding_vector}
|
|||
|
|
"""
|
|||
|
|
for table in self.get_tables():
|
|||
|
|
if table.name in table_embeddings:
|
|||
|
|
table.embedding = table_embeddings[table.name]
|
|||
|
|
|
|||
|
|
def get_table_for_embedding(
|
|||
|
|
self,
|
|||
|
|
max_columns: int = 48,
|
|||
|
|
max_col_comment_chars: int = 160,
|
|||
|
|
max_total_chars: int = 8000,
|
|||
|
|
) -> List[Dict]:
|
|||
|
|
"""
|
|||
|
|
获取用于embedding的表信息列表（单条文本同时承载 meta + structure 的可检索信息）。
|
|||
|
|
|
|||
|
|
- 表级：表名、表注释（来自 table_meta 合并后的 Table.comment）
|
|||
|
|
- 列级：字段名、类型、主键标记、列注释（来自 table_structure 解析后的 Column）
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
[{"name": "...", "text": "..."}, ...]
|
|||
|
|
"""
|
|||
|
|
result = []
|
|||
|
|
for table in self.get_tables():
|
|||
|
|
text_parts = [f"表名: {table.name}"]
|
|||
|
|
if table.comment:
|
|||
|
|
text_parts.append(f"描述: {table.comment}")
|
|||
|
|
|
|||
|
|
col_segments: List[str] = []
|
|||
|
|
shown = table.columns[:max_columns]
|
|||
|
|
for col in shown:
|
|||
|
|
seg = f"{col.name} {col.data_type}"
|
|||
|
|
if col.is_primary_key:
|
|||
|
|
seg += " PK"
|
|||
|
|
if col.comment:
|
|||
|
|
c = col.comment.strip().replace("\n", " ")
|
|||
|
|
if len(c) > max_col_comment_chars:
|
|||
|
|
c = c[: max_col_comment_chars - 1] + "…"
|
|||
|
|
seg += f" — {c}"
|
|||
|
|
col_segments.append(seg)
|
|||
|
|
|
|||
|
|
if col_segments:
|
|||
|
|
text_parts.append("字段: " + "；".join(col_segments))
|
|||
|
|
if len(table.columns) > max_columns:
|
|||
|
|
rest = len(table.columns) - max_columns
|
|||
|
|
text_parts.append(f"另有{rest}个字段未列出（共{len(table.columns)}列）")
|
|||
|
|
|
|||
|
|
text = "。".join(text_parts)
|
|||
|
|
if len(text) > max_total_chars:
|
|||
|
|
text = text[: max_total_chars - 1] + "…"
|
|||
|
|
|
|||
|
|
result.append({
|
|||
|
|
"name": table.name,
|
|||
|
|
"text": text,
|
|||
|
|
})
|
|||
|
|
return result
|
|||
|
|
|
|||
|
|
# ========== 缓存支持 ==========
|
|||
|
|
|
|||
|
|
def get_cached_result(self, key: str) -> Optional[str]:
|
|||
|
|
"""获取缓存结果"""
|
|||
|
|
return self._cache.get(key)
|
|||
|
|
|
|||
|
|
def set_cached_result(self, key: str, value: str, ttl: int = 3600):
|
|||
|
|
"""设置缓存结果"""
|
|||
|
|
self._cache[key] = value
|
|||
|
|
# TODO: 可扩展为Redis缓存
|
|||
|
|
|
|||
|
|
# ========== 统计信息 ==========
|
|||
|
|
|
|||
|
|
def get_statistics(self) -> Dict:
|
|||
|
|
"""获取Schema统计信息"""
|
|||
|
|
total_columns = sum(len(t.columns) for t in self.get_tables())
|
|||
|
|
tables_with_pk = sum(1 for t in self.get_tables() if t.primary_keys)
|
|||
|
|
tables_with_fk = sum(1 for t in self.get_tables() if t.foreign_keys)
|
|||
|
|
|
|||
|
|
return {
|
|||
|
|
"database": self.schema.name,
|
|||
|
|
"total_tables": len(self.get_tables()),
|
|||
|
|
"total_columns": total_columns,
|
|||
|
|
"avg_columns_per_table": total_columns / len(self.get_tables()) if self.get_tables() else 0,
|
|||
|
|
"tables_with_primary_key": tables_with_pk,
|
|||
|
|
"tables_with_foreign_key": tables_with_fk,
|
|||
|
|
"table_names": self.list_tables(),
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
def __len__(self) -> int:
|
|||
|
|
return len(self.get_tables())
|
|||
|
|
|
|||
|
|
def __repr__(self) -> str:
|
|||
|
|
return f"SchemaManager(database={self.schema.name}, tables={len(self)})"
|