""" 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)})"