from langchain_core.prompts import ChatPromptTemplate from langgraph.graph import START, StateGraph, END from llmops.agents.datadev.llm import get_llm, get_llm_coder from typing import List from typing_extensions import TypedDict from pydantic import BaseModel, Field from llmops.agents.datadev.memory.memory_saver_with_expiry2 import MemorySaverWithExpiry from langchain_core.output_parsers import PydanticOutputParser from config import linage_agent_config from llmops.agents.datadev.tools.timeit import timeit import asyncio from datetime import datetime class Column(BaseModel): """ 字段信息 """ col_name: str = Field(description="字段名称") def __eq__(self, other): if isinstance(other, Column): return self.col_name == other.col_name return False class ColumnDep(BaseModel): """ 目标表字段信息，包含依赖关系 """ col_name: str = Field(description="目标表字段名称") col_comment: str = Field(description="目标表字段说明") from_cols: List[str] = Field(description="字段来源信息，格式是 源库名.源表名.源字段名 或 中间表名.中间表字段名") dep_type: List[str] = Field(description="字段获取方式 1：直取 2：函数 3：表达式") desp: List[str] = Field(description="详细解释选取来源字段的原因") class SourceTable(BaseModel): """ 单个来源信息表信息 """ database: str = Field(description="来源数据库名称") table_name: str = Field(description="来源表名称") table_name_alias: str = Field(default="", description="来源表别名") col_list: List[Column] = Field(description="来源表的字段集合") is_temp: bool = Field(description="是否是临时表") def __eq__(self, other): if isinstance(other, SourceTable): return self.database == other.database and self.table_name == other.table_name return False class SourceTableList(BaseModel): """ 所有来源信息表信息 """ source_tables: List[SourceTable] = Field(description="所有来源表信息") class WhereCondItem(BaseModel): """ 单个表过滤条件信息 """ database: str = Field(description="数据库名") table_name: str = Field(description="表原名") col_name: str = Field(description="条件字段名称") operator: str = Field(default="=", description="条件操作符，如=,>,<,in") value: str = Field(description="条件值") class WhereCondList(BaseModel): """ 表过滤条件集合 """ where_list: List[WhereCondItem] = Field(description="所有的where条件") class JoinColumnPair(BaseModel): """关联字段对""" main_table_alias: str = Field(description="主表别名") main_column: str = Field(description="主表关联字段") join_table_alias: str = Field(description="从表别名") join_column: str = Field(description="从表关联字段") operator: str = Field(default="=", description="关联操作符，如=,>,<") class JoinRelation(BaseModel): """完整的JOIN关系""" main_database: str = Field(description="主表数据库名") main_table: str = Field(description="主表物理表名") main_alias: str = Field(description="主表别名") join_type: str = Field(description="JOIN类型：INNER/LEFT/RIGHT/FULL") join_database: str = Field(description="从表数据库名") join_table: str = Field(description="从表物理表名") join_alias: str = Field(description="从表别名") column_pairs: List[JoinColumnPair] = Field(description="关联字段对列表") class JoinRelationList(BaseModel): """ join关系集合 """ join_list: List[JoinRelation] = Field(description="所有的join关系") class TargetTable(BaseModel): """ 目标表信息 """ database: str = Field(default="", description="目标表的数据库名称") table_name: str = Field(description="目标表名称") col_size: int = Field(description="目标表字段数量") col_list: List[str] = Field(description="目标表字段集合") col_dep: List[ColumnDep] = Field(description="目标表字段依赖信息") src_table_size: int = Field(description="依赖来源表数量") class CreateTable(BaseModel): """ 建表信息 """ database: str = Field(description="数据库名") table: str = Field(description="表名") col_list: list[str] = Field(description="字段名集合") source_table_list: list[str] = Field(description="来源表集合") class AgentState(TypedDict): """ Agent状态 """ question: str session: str dialect: str sql: str # 带行号的SQL语句 sql_type: str # sql操作类型 lineage: dict status: str error: str stop: bool file_name: str # 脚本文件 target_table: dict # 目标表信息 source_tables: list[dict] # 来源表信息 where_list: list[dict] # where条件信息 join_list: list[dict] # 表关联信息 class SqlLineageAgent: """ 解析用户提供的SQL语句/SQL脚本中的数据血缘关系 """ def __init__(self): # 读取配置文件 c = linage_agent_config # 数据库方言 self.db_dialect = c["dialect"] # 并发度 self.concurrency = c["concurrency"] # 大模型实例 self.llm = get_llm() self.llm_coder = get_llm_coder() self.memory = MemorySaverWithExpiry(expire_time=600, clean_interval=60) # 构建图 self.graph = self._build_graph() # SQL 语句中设置的变量 self.var_list: list[str] = [] # 最终的来源表集合 self.final_source_table_list: list[dict] = [] # 最终的中间表集合 self.final_mid_table_list: list[dict] = [] # 最终的where集合 self.final_where_list: list[dict] = [] # 最终的join集合 self.final_join_list: list[dict] = [] # 最终的目标表集合 self.final_target_list: list[dict] = [] # 最终目标表，需合并 self.final_target_table = {} # 目标表 格式 库名.表名 self.target_table = "" # 内部建表信息 self.ct_list = [] def _build_graph(self): # 构造计算流程 graph_builder = StateGraph(AgentState) # 添加节点 graph_builder.add_node("_sql_type", self._sql_type) graph_builder.add_node("_extract_set", self._extract_set) graph_builder.add_node("_extract_create_table", self._extract_create_table) graph_builder.add_node("_invalid", self._invalid) graph_builder.add_node("source", self.__extract_source) graph_builder.add_node("target", self.__extract_target) graph_builder.add_node("where", self.__extract_where) graph_builder.add_node("join", self.__extract_join) graph_builder.add_node("merge", self.__merge_result) # 添加边 graph_builder.add_edge(START, "_sql_type") graph_builder.add_conditional_edges("_sql_type", path=self.__parallel_route2, path_map={ "_invalid": "_invalid", "_extract_set": "_extract_set", "_extract_create_table": "_extract_create_table", "where": "where", "join": "join" }) graph_builder.add_edge("_invalid", END) graph_builder.add_edge("_extract_set", END) graph_builder.add_edge("_extract_set", END) graph_builder.add_edge("_extract_create_table", END) graph_builder.add_edge("where", "source") graph_builder.add_edge("join", "source") graph_builder.add_edge("source", "target") graph_builder.add_edge("target", "merge") graph_builder.add_edge("merge", END) return graph_builder.compile(checkpointer=self.memory) def _sql_type(self, state: AgentState): """ 判断SQL的操作类型 """ template = """ 你是数据库 {dialect} 专家，对 SQL语句: {sql} 中的操作进行分类, 直接返回分类ID(数字) ### 操作分类标准如下： 1:create, 如建表、建视图等 2:insert, 向表中插入数据 3:set、设置参数操作、变量赋值操作 4:其它 不要重复用户问题，不要中间分析过程，直接回答。 """ pt = ChatPromptTemplate.from_template(template) chain = pt | self.llm response = chain.invoke({"dialect": state["dialect"], "sql": state["question"]}) return {"sql_type": response.content} def __parallel_route(self, state: AgentState): """并发节点""" status = state["status"] if status == "invalid" or status == "canceled" or state["sql_type"] == "3": # 无效问题 或 用户中止 return END return "continue" def __parallel_route2(self, state: AgentState): """并发节点""" if state["sql_type"] == "4": # 非create和非insert操作 return "_invalid" if state["sql_type"] == "3": # 设置参数 return "_extract_set" if state["sql_type"] == "1": # 建表 return "_extract_create_table" return ["where", "join"] def _invalid(self, state: AgentState): """ sql_type == 4，不关注的语句，直接跳过 """ return { "lineage": { "source_tables": [], "join_list": [], "where_list": [], "target_table": {} } } async def _extract_set(self, state: AgentState): """ 提取 SQL 中设置的参数变量 """ if state["sql_type"] != "3": return state template = """ 你是数据库 {dialect} SQL分析专家，从SQL语句: {sql} 中提取 参数设置的全部内容，包括变量赋值、注释。 不要重复用户问题，不需要中间思考过程，直接回答。 """ pt = ChatPromptTemplate.from_template(template) chain = pt | self.llm_coder response = await chain.ainvoke({"dialect": state["dialect"], "sql": state["question"]}) self.var_list.append(response.content) return { "lineage": { "source_tables": [], "join_list": [], "where_list": [], "target_table": {} } } async def _extract_create_table(self, state: AgentState): """ 提取建表信息 """ if state["sql_type"] != "1": return state template = """ 你是数据库 {dialect} SQL分析专家，从SQL语句: {sql} 中提取 建表信息。 ### 提取要求 1、提取建表对应的数据库名（无则留空） 2、提取建表对应的表名 3、提取建表对应的字段名 4、提取建表对应的来源表名，来源表名格式是 来源库库名.来源表名 ### 输出规范 {format_instructions} 不要重复用户问题，不需要中间思考过程，直接回答。 """ parser = PydanticOutputParser(pydantic_object=CreateTable) pt = ChatPromptTemplate.from_template(template).partial(format_instructions=parser.get_format_instructions()) chain = pt | self.llm_coder | parser response = await chain.ainvoke({"dialect": state["dialect"], "sql": state["question"]}) # print(f"extract-create-table: {response}") self.ct_list.append(response.model_dump()) return { "lineage": { "source_tables": [], "join_list": [], "where_list": [], "target_table": {} } } async def _parse(self, semaphore, sql: str, session: str, dialect: str, file_name: str = "SQL_FILE.SQL"): """ 解析SQL（片段）的血缘关系 """ async with semaphore: # 进入信号量区域（限制并发） config = {"configurable": {"thread_id": session}} # 调用任务流 response = await self.graph.ainvoke({ "question": sql, "session": session, "dialect": dialect or self.db_dialect, "file_name": file_name, "stop": False, "error": "", "sql_type": "0", "status": "success"}, config) return response async def find_target_table(self, sql: str, dialect: str): """ 找出目标表 """ template = """ 你是SQL数据血缘分析专家，参照数据库方言{dialect}，仔细分析SQL语句，找出目标表并返回。 ### SQL片段如下 {sql} ### 目标表标准 1、目标表一定是出现在 insert into 后的 表 2、最后 insert into 后的表为目标表 3、目标表只有一张 ### 核心要求 1、目标表的返回格式是 数据库名.表名 (如果无数据库名，则留空) 2、SQL语句中可能包括多个insert into语句，但要从全局分析，找出最终插入的目标表 不要重复用户问题、不要中间分析过程，直接回答。 """ pt = ChatPromptTemplate.from_template(template) chain = pt | self.llm_coder response = await chain.ainvoke({"sql": sql, "dialect": dialect}) return response.content.strip().upper() def split_sql_statements(self, sql): """ 将SQL语句按分号分段，同时正确处理字符串和注释中的分号 参数: sql: 包含一个或多个SQL语句的字符串 返回: 分割后的SQL语句列表 """ # 状态变量 in_single_quote = False in_double_quote = False in_line_comment = False in_block_comment = False escape_next = False statements = [] current_start = 0 i = 0 while i < len(sql): char = sql[i] # 处理转义字符 if escape_next: escape_next = False i += 1 continue # 处理字符串和注释中的情况（这些地方的分号不应该作为分隔符） if in_line_comment: if char == '\n': in_line_comment = False elif in_block_comment: if char == '*' and i + 1 < len(sql) and sql[i + 1] == '/': in_block_comment = False i += 1 # 跳过下一个字符 elif in_single_quote: if char == "'": in_single_quote = False elif char == '\\': escape_next = True elif in_double_quote: if char == '"': in_double_quote = False elif char == '\\': escape_next = True else: # 不在字符串或注释中，检查是否进入这些状态 if char == "'": in_single_quote = True elif char == '"': in_double_quote = True elif char == '-' and i + 1 < len(sql) and sql[i + 1] == '-': in_line_comment = True i += 1 # 跳过下一个字符 elif char == '/' and i + 1 < len(sql) and sql[i + 1] == '*': in_block_comment = True i += 1 # 跳过下一个字符 elif char == ';': # 找到真正的分号分隔符 statement = sql[current_start:i].strip() if statement: statements.append(statement) current_start = i + 1 i += 1 # 添加最后一个语句（如果没有以分号结尾） if current_start < len(sql): statement = sql[current_start:].strip() if statement: statements.append(statement) return statements async def find_target_table2(self, results: list[dict]) -> str: """ 根据各段解析结果查找出目标表 规则：一个段的目标表如果没有出现在其它各段的来源表中，即为目标表 """ target_table = "" exists = False for i, item in enumerate(results): tt = item["lineage"]["target_table"] # 取一个目标表 db = tt.get("database", "").strip().upper() table = tt.get("table_name", "").strip().upper() exists = False if len(table) > 0: # 从源表里面查找 for j, item2 in enumerate(results): if i != j: # 不跟自身比 st_list = item2["lineage"]["source_tables"] for st in st_list: sdb = st.get("database", "").strip().upper() stable = st.get("table_name", "").strip().upper() if db == sdb and table == stable: # 目标表在源表中存在 exists = True break if exists: break if not exists: # 说明目标表不在其它各段的源表中存在 target_table = ".".join([db, table]) break return target_table.strip().upper() async def ask_question(self, sql_content: str, session: str, dialect: str, owner: str, sql_file: str = "SQL_FILE.SQL"): """ 功能：根据用户问题，解析SQL中的数据血缘关系。先分段解析，再组织合并。 :param sql_content: SQL内容 :param session: 会话 :param dialect: 数据库方言 :param owner: 脚本平台属主 :param var_map: 变量关系 :sql_file: 脚本名称 """ # 通过session保持记忆 t1 = datetime.now() # 最终解析结果 final_result = {} final_result["task_id"] = session final_result["file_name"] = sql_file final_result["owner"] = owner final_result["status"] = "success" final_result["error"] = "" lineage = {} try: # 根据全局SQL找出目标表 self.target_table = await self.find_target_table(sql_content, dialect) print(f"大模型查找目标表：{self.target_table}, {len(self.target_table)}") # 对SQL分段，并发执行解析 sql_list = self.split_sql_statements(sql=sql_content) print("SQL分段数量：", len(sql_list)) # 信号量，控制并发量 semaphore = asyncio.Semaphore(self.concurrency) task_list = [self._parse(semaphore=semaphore, sql=sql, session=session, dialect=dialect, file_name=sql_file) for sql in sql_list] # 并发分段解析 results = await asyncio.gather(*task_list) if len(self.target_table) == 0: # 从各段解析结果中查找出目标表 self.target_table = await self.find_target_table2(results) or "" print(f"从血缘关系中找出目标表：{self.target_table}") if not self.target_table: # 未出现目标表, 直接返回 final_result["status"] = "error", final_result["error"] = "未找到目标表" return final_result for response in results: # 来源表、where、join、target self._merge_source_tables(response["lineage"]["source_tables"]) self._merge_where(response["lineage"].get("where_list",[])) self._merge_join(response["lineage"]["join_list"]) # 合并目标表 self._merge_target(response["lineage"]["target_table"]) # 增加中间表标识 self._add_mid_table_label(self.final_source_table_list, self.final_mid_table_list) # 补充中间表依赖字段 self._add_mid_table_col(self.final_mid_table_list, self.final_target_list) # 多目标表合并 self.final_target_table["src_table_size"] = len(self.final_source_table_list) lineage["target_table"] = self.final_target_table lineage["source_tables"] = self.final_source_table_list lineage["mid_table_list"] = self.final_mid_table_list lineage["join_list"] = self.final_join_list lineage["where_list"] = self.final_where_list final_result["lineage"] = lineage except Exception as e: final_result["status"] = "error" final_result["error"] = str(e) print(str(e)) t2 = datetime.now() print("总共用时(sec)：", (t2-t1).seconds) parse_time = (t2 - t1).seconds # 解析时间（秒） final_result["parse_time"] = parse_time final_result["parse_end_time"] = t2.strftime("%Y-%m-%d %H:%M:%S") return final_result def _merge_source_tables(self, source_table_list: list[dict]): """ 合并分段来源表 """ if len(source_table_list) > 0: # 目标源表的keys table_keys = [(t.get("database","").strip().upper(), t.get("table_name","").strip().upper()) for t in self.final_source_table_list] for st in source_table_list: # 将字典转换为元组 db = st.get("database","").strip().upper() table = st.get("table_name","").strip().upper() key = (db, table) if key not in table_keys: table_keys.append(key) self.final_source_table_list.append(st) else: # 合并字段 col_list = st.get("col_list", []) # 找出相同的元素, 库名，表名相同 e = next(filter(lambda item: item["database"].upper()==db and item["table_name"].upper()==table, self.final_source_table_list), None) if e: final_col_list = e.get("col_list",[]) col_keys = [(c.get("col_name","").upper()) for c in final_col_list] for c in col_list: ck = (c["col_name"].upper()) if ck not in col_keys and len(c) > 0: col_keys.append(ck) final_col_list.append(c) @timeit def _merge_where(self, where_list: list[dict]): """ 最终合并所有的where条件 """ if where_list and len(where_list) > 0: self.final_where_list += where_list def _merge_join(self, join_list: list[dict]): """ 最终合并所有的join条件 """ if len(join_list) > 0: self.final_join_list += join_list def _merge_target(self, mid_target_table: dict): """ 合并中间目标表，只有 名称是 target_table_name 的表才是真目标表 目标表可能出现多次（如多次插入），需要合并 字段信息(字段名称和字段来源) 对于非目标表的中间目标表，放到中间表集合中 """ if not mid_target_table: # 中间目标表为空 return else: # 将临时目标表与目标表做合并 db = mid_target_table.get("database", "").strip().upper() table = mid_target_table.get("table_name", "").strip().upper() col_size = mid_target_table.get("col_list",[]) if len(table) == 0 or len(col_size) == 0: # 表名为空 或 无字段 直接返回 return if len(db) == 0: # 缺少数据库名, 检查表名中是否带库名 arr = table.split('.') if len(arr) == 2: # 说明表名中当库名，解析有错，修正 mid_target_table["database"] = arr[0].strip().upper() mid_target_table["table_name"] = arr[1].strip().upper() db = arr[0].strip().upper() table = arr[1].strip().upper() elif len(arr) == 1: # 无库名，只有表名，使用真目标表库名替换 arr = self.target_table.split(".") if len(arr) == 2: db = mid_target_table["database"] = arr[0].strip().upper() key = ".".join([db, table]) if key == self.target_table: # 找到最终目标表 print(f"合并目标表：{mid_target_table}") if not self.final_target_table: # 目标表为空（首次） self.final_target_table = mid_target_table return # 合并字段和字段来源信息 # 合并新字段 new_col_list = [] # 新字段集 col_list = mid_target_table.get("col_list", []) for new_col in col_list: # 检查字段是否存在、忽略大小写 if new_col.strip().upper() not in [col.strip().upper() for col in self.final_target_table.get("col_list", [])]: # 不存在 print(f"合并新字段：{new_col}") self.final_target_table.get("col_list", []).append(new_col.strip().upper()) new_col_list.append(new_col.strip().upper()) # 合并字段来源信息 col_dep_list = mid_target_table.get("col_dep", []) for col_dep in col_dep_list: col_name = col_dep["col_name"].strip().upper() if col_name in new_col_list: # 新字段, 直接添加 self.final_target_table.get("col_dep",[]).append(col_dep) else: # 合并来源信息 # 找到同名字段，合并来源字段 for dep in self.final_target_table.get("col_dep", []): cn = dep["col_name"].strip().upper() if col_name == cn: print(f"合并来源字段, {col_name}") m_from_col = col_dep.get('from_cols', []) f_from_col = dep.get("from_cols", []) print(f"中间表来源字段:{m_from_col}") print(f"目标表来源字段:{f_from_col}") # 将临时中间表的来源字段合并至目标表的来源字段 self._merge_col_dep(m_from_col, f_from_col) else: # 加入中间表集合中 print(f"加入中间表:{mid_target_table}") self.final_mid_table_list.append(mid_target_table) print(f"final tt:{self.final_target_table}") def _merge_col_dep(self, mid_from_cols: list[str], final_from_cols: list[str]): """ 合并临时中间表来源字段 到目标表来源字段中 """ for from_col in mid_from_cols: if len(from_col.split(".")) > 0: # 忽略常量,NULL值 if from_col.strip().upper() not in [fc.strip().upper() for fc in final_from_cols]: print(f"合并字段：{from_col}") final_from_cols.append(from_col.strip().upper()) @staticmethod @timeit def _add_mid_table_label(source_table_list: list[dict], mid_table_list: list[dict]): """ 给最终的来源表增加中间表标识 """ if len(source_table_list) > 0 and len(mid_table_list) > 0: mid_table_keys = [(t["database"].strip().upper(), t["table_name"].strip().upper()) for t in mid_table_list] for st in source_table_list: st["is_temp"] = False # 将字典转换为元组 key = (st["database"].strip().upper(), st["table_name"].strip().upper()) if key in mid_table_keys: # 是中间表 st["is_temp"] = True @staticmethod @timeit def _add_mid_table_col(mid_table_list: list[dict], target_table_list: list[dict]): """ 给中间表补充来源字段，使用对应的目标表col_dep去替换 """ if len(mid_table_list) > 0 and len(target_table_list) > 0: for mt in mid_table_list: for tt in target_table_list: tt["is_temp"] = False # 目标表的数据库名与表名相同 tdn = tt["database"].strip().upper() ttn = tt["table_name"].strip().upper() if len(tdn) == 0 and len(ttn) == 0: # 全部为空 tt["is_temp"] = True continue if mt["database"].strip().upper()==tdn and mt["table_name"].strip().upper()==ttn: mt["col_dep"] = tt["col_dep"] tt["is_temp"] = True if len(mt["col_list"]) == 0: mt["col_list"] = [c.upper() for c in tt["col_list"]] mt["col_size"] = len(mt["col_list"]) @timeit def _merge_final_target_table(self, seg_target_table_list: list[dict], target_table_name: str): """ 合并分段目标表内容 """ # 可能存在多个相同的真正目标表（多次插入目标表情况） for tt in seg_target_table_list: tt["is_target"] = False database = tt["database"].strip() table = tt["table_name"].strip() if len(database) > 0: dt = ".".join([database,table]) else: dt = table if dt.upper() == target_table_name.strip().upper(): tt["is_target"] = True # 合并真正的目标表 i = 0 result = {} final_col_list: list[str] = [] final_col_dep_list: list[dict] = [] # 合并目标表信息 for target_table in seg_target_table_list: if target_table["is_target"]: if i == 0: result["database"] = target_table["database"] result["table_name"] = target_table["table_name"] i = 1 # 合并列信息 col_list = target_table["col_list"] for c in col_list: if c not in final_col_list and len(c.strip()) > 0: final_col_list.append(c.upper()) # 合并列来源信息 col_dep_list = target_table["col_dep"] for col_dep in col_dep_list: cn = col_dep["col_name"].strip().upper() from_cols = col_dep["from_cols"] existed = False if len(cn) > 0: for fcd in final_col_dep_list: if cn == fcd["col_name"].upper(): existed = True # 合并来源字段 final_from_cols = fcd["from_cols"] fcd["from_cols"] += [c.upper() for c in from_cols if c.upper() not in final_from_cols] if not existed: final_col_dep_list.append({ "col_name": cn.upper(), "col_comment": col_dep["col_comment"], "from_cols": [c.upper() for c in from_cols] }) result["col_size"] = len(final_col_list) result["col_list"] = final_col_list result["col_dep"] = final_col_dep_list return result async def __check(self, state: AgentState): """ 检查用户问题是否是有效的SQL语句/SQL片段 """ template = """ 你是 数据库 {dialect} SQL分析助手， 仔细分析SQL语句: {sql} ，判断是否语法正确，如果是 直接返回 1，否则说明错误地方，并返回 0。 ### 分析要求 1. 判断子查询的正确性 2. 判断整体语句的正确性 3. 如果存在语法错误，给出说明 不要重复用户问题、不要分析过程、直接回答。 """ pt = ChatPromptTemplate.from_template(template) chain = pt | self.llm stopped = False response = "" config = {"configurable": {"thread_id": state["session"]}} async for chunk in chain.astream({"sql": state["question"], "dialect": state["dialect"]}, config=config): current_state = self.memory.get(config) if current_state and current_state["channel_values"]["stop"]: stopped = True break response += chunk.content if stopped: return {"stop": True, "error": "用户中止", "status": "canceled", "lineage": {}} if response == "0": return {"status": "invalid", "error": "无效问题或SQL存在语法错误，请重新描述！", "lineage": {}} return {"status": "success"} async def __merge_result(self, state: AgentState): """ 合并所有节点结果，输出血缘关系 """ result = {} result["file_name"] = state["file_name"] source_tables = state["source_tables"] target_table = state["target_table"] # 设置目标表依赖的来源表数量和字段数量 target_table["src_table_size"] = len(source_tables) target_table["col_size"] = len(target_table["col_dep"]) result["target_table"] = target_table result["source_tables"] = source_tables result["where_list"] = state.get("where_list",[]) result["join_list"] = state["join_list"] return {"lineage": result} async def __extract_source(self, state: AgentState): """ 根据SQL，提取其中的来源库表信息 """ dt1 = datetime.now() # 默认DB default_db = "TMP" if len(self.target_table) > 0: arr = self.target_table.split(".") if len(arr) == 2: # 使用识别出的目标表db作为默认DB default_db = arr[0] template = """ 你是SQL数据血缘分析专家，仔细分析SQL片段和上下文信息, 从SQL片段中提取出 来源表信息。 ### SQL和上下文信息如下 - SQL片段：{sql} - 数据库方言: {dialect} - 变量设置信息：{var} - 已有的建表信息: {create_table_info} ### 核心要求： 1. 提取 来源表数据库名称（无则使用 {default_db} 填充） 2. 提取 来源表名，包括 from子句、子查询、嵌套子查询、union语句、with语句中出现的物理表名(注意：不要带库名前缀) 3. 提取 来源表别名（无则留空） 4. 提取来源表的所有字段信息，提取来自几种： - 从select中提取，如果带有表别名，则转换成物理表名 - 从where过滤条件中提取 - 从关联条件(inner join,left join,right join,full join)中提取 - 从group中提取 - 从函数参数中提取,比如函数 COALESCE(T13.ECIF_CUST_ID,T2.RELAT_PTY_ID,'') AS ECIF_CUST_ID，提取出T13.ECIF_CUST_ID和T2.RELAT_PTY_ID - 从表达式参数中提取 5. 判断来源表是否是临时表 6. 不要包含目标表（目标表指最终插入的表） ### 输出规范 {format_instructions} 不要重复用户问题，不要分析中间过程，直接给出答案。 """ parser = PydanticOutputParser(pydantic_object=SourceTableList) pt = ChatPromptTemplate.from_template(template=template).partial(format_instructions=parser.get_format_instructions()) chain = pt | self.llm_coder | parser answer = {} config = {"configurable": {"thread_id": state["session"]}} stopped = False async for chunk in chain.astream({ "sql": state["question"], "dialect": state["dialect"], "var": self.var_list or [], "default_db": default_db, "create_table_info": self.ct_list }): current_state = self.memory.get(config) if current_state and current_state["channel_values"]["stop"]: stopped = True break # 设置目标表字段数量,来源表数量 # 合并相同源表 table_list: list[SourceTable] = [] for table in chunk.source_tables: if table not in table_list: table_list.append(table) else: idx = table_list.index(table) tt = table_list[idx] # 合并相同的列 for col in table.col_list: if col not in tt.col_list: tt.col_list.append(col) answer["source_tables"] = [table.model_dump() for table in table_list] dt2 = datetime.now() print("提取源表用时(sec):", (dt2-dt1).seconds) if stopped: return {"status": "canceled", "error": "用户中止", "lineage": {}} return answer async def __extract_target(self, state: AgentState): """ 根据SQL，提取其中的目标表信息 """ if state["sql_type"] == "1": return {"target_table": {"database":"", "table_name": "", "col_list": [], "col_size": 0, "col_dep": []}} # 默认DB default_db = "TMP" if len(self.target_table) > 0: arr = self.target_table.split(".") if len(arr) == 2: # 使用识别出的目标表db作为默认DB default_db = arr[0] template = """ 你是SQL数据血缘分析专家，仔细分析SQL片段和以下上下文： - SQL片段: {sql} - 来源表信息: {source_tables} - 数据库方言: {dialect} - 变量设置信息：{var} - 已有的建表信息: {create_table_info} ### 核心要求： 1. **目标表识别**： - 提取 INSERT INTO 后的表作为目标表 - 格式：`数据库名.表名`（未指定库名则使用 {default_db} 填充) 2. **字段依赖分析**： - 目标表字段必须溯源到来源表的物理字段 - 字段获取方式分类： - `直取`：直接引用源字段（如 `src.col`） - `函数`：通过函数转换（如 `COALESCE(col1, col2)`） - `表达式`：计算表达式（如 `col1 + col2`） - `常量`：固定值（如 `'2023'`） - `子查询`：嵌套SELECT结果 - 字段来源格式：`源库.源表.字段名`（不要使用表别名，如果源库未指定则使用 {default_db}替换） - 字段来源信息只能包含 库名、表名和字段信息，如果经过函数或表达式处理，则从参数中提取出具体的字段信息 3. **关键约束**： - 每个目标字段必须有详细的来源说明 - 函数参数需完整展开（如 `COALESCE(T1.id, T2.id)` 需解析所有参数） - 表名必须是 英文字符 构成（非汉字构成） ### 输出规范 {format_instructions} 请直接输出JSON格式结果，无需解释过程。 """ parser = PydanticOutputParser(pydantic_object=TargetTable) pt = ChatPromptTemplate.from_template(template).partial(format_instructions=parser.get_format_instructions()) chain = pt | self.llm_coder | parser answer = {} config = {"configurable": {"thread_id": state["session"]}} dt1 = datetime.now() print("开始解析目标表...") stopped = False async for chunk in chain.astream({ "sql": state["question"], "dialect": state["dialect"], "source_tables": state["source_tables"], "var": self.var_list or [], "default_db": default_db, "default_db": default_db, "create_table_info": self.ct_list }): current_state = self.memory.get(config) if current_state and current_state["channel_values"]["stop"]: stopped = True break answer["target_table"] = chunk.model_dump() if stopped: return {"status": "canceled", "error": "用户中止", "lineage": {}} dt2 = datetime.now() print("提取目标表用时(sec)：", (dt2 - dt1).seconds) print(f"target:{answer}") return answer async def __extract_where(self, state: AgentState): """ 根据SQL，提取来源表 where 信息 """ # 默认DB default_db = "TMP" if len(self.target_table) > 0: arr = self.target_table.split(".") if len(arr) == 2: # 使用识别出的目标表db作为默认DB default_db = arr[0] dt1 = datetime.now() template = """ 你是SQL数据血缘分析专家，仔细分析 SQL语句和上下文，提取 where 条件信息。 ### SQL语句和上下文如下 - SQL语句: {sql} - 数据库方言: {dialect} - 变量设置信息: {var} ### 提取要求： 1. 提取 所有 来源表 中的 where 条件信息，包含 where 中出现的 表原名（非别名）、 表所在数据库名（如果没有则使用 {default_db} 填充）、字段名、条件操作符和条件值 2. 字段名不能包括函数，比如 length(INT_ORG_NO)，提取的字段名是 INT_ORG_NO 3. 如果SQL语句中存在变量，则根据变量设置信息转换成实际值 ### 输出规范 {format_instructions} 不要重复用户问题，不要分析中间过程，直接给出答案。 """ parser = PydanticOutputParser(pydantic_object=WhereCondList) pt = ChatPromptTemplate.from_template(template).partial(format_instructions=parser.get_format_instructions()) chain = pt | self.llm_coder | parser answer = {} config = {"configurable": {"thread_id": state["session"]}} stopped = False async for chunk in chain.astream({ "sql": state["question"], "dialect": state["dialect"], "var": self.var_list or [], "default_db": default_db }): current_state = self.memory.get(config) if current_state and current_state["channel_values"]["stop"]: stopped = True break # 设置目标表字段数量,来源表数量 dt2 = datetime.now() if stopped: return {"status": "canceled", "error": "用户中止", "lineage": {}} return answer async def __extract_join(self, state: AgentState): """ 根据SQL，提取 表关联 信息 """ # 默认DB default_db = "TMP" if len(self.target_table) > 0: arr = self.target_table.split(".") if len(arr) == 2: # 使用识别出的目标表db作为默认DB default_db = arr[0] dt1 = datetime.now() template = """ 你是SQL数据血缘分析专家，负责提取完整的JOIN关系。仔细分析SQL片段： - SQL片段: {sql} - 数据库方言: {dialect} ### 关键要求： 1. **SQL语句中必须存在关联条件(包括inner、left join、right join、full join)** 2. **识别JOIN结构**： - 提取FROM子句中的主表信息（包括 数据库名和物理表名）,如果无数据库名，则使用 {default_db} 填充 - 提取每个JOIN子句中的从表信息（包括 数据库名和物理表名）,如果无数据库名，则使用 {j_default_db} 填充 - 明确标注JOIN类型（INNER/LEFT/RIGHT/FULL） 3. **关联字段提取**： - 必须提取每对关联字段的完整信息： * 主表端：表别名.字段名 * 从表端：表别名.字段名 - 明确标注关联操作符（=, >, <等） - 多条件JOIN拆分为多个字段对 4. **特殊场景处理**： - 子查询作为表时，表名填写子查询别名 - 隐式JOIN（WHERE条件）需转换为显式JOIN结构 - 多表JOIN时保持原始顺序 - 如果关联字段是 常量，则关联字段取空 - 如果关联字段经函数处，则提取参数中的 字段，如果参数是常量，则关联字段取空 ### 输出规范 {format_instructions} 不要重复用户问题，不要中间思考过程，直接回答。 """ parser = PydanticOutputParser(pydantic_object=JoinRelationList) pt = ChatPromptTemplate.from_template(template).partial(format_instructions=parser.get_format_instructions()) chain = pt | self.llm_coder | parser answer = {} config = {"configurable": {"thread_id": state["session"]}} stopped = False async for chunk in chain.astream({ "sql": state["question"], "dialect": state["dialect"], "default_db": default_db, "j_default_db": default_db }): current_state = self.memory.get(config) if current_state and current_state["channel_values"]["stop"]: stopped = True break # 设置目标表字段数量,来源表数量 answer["join_list"] = [join.model_dump() for join in chunk.join_list] dt2 = datetime.now() if stopped: return {"status": "canceled", "error": "用户中止", "lineage": {}} return answer def trigger_stop(self, session: str): """外部触发中止""" config = {"configurable": {"thread_id": session}} current_state = self.memory.get(config) if current_state: current_state["channel_values"]["stop"] = True async def main(sql_content: str, sql_file: str, dialect: str, owner: str): agent = SqlLineageAgent() result = await agent.ask_question(sql_content=sql_content, session="s-1", dialect=dialect, sql_file=sql_file, owner=owner) return result def read_var_file(var_file: str): """ 读取变量文件 :param var_file: 变量文件(格式excel) """ import pandas as pd # 读取变量excel文件 result = [] df = pd.read_excel(var_file) for row in df.itertuples(): print(f"变量 {row.Index + 1}: {row.变量}, {row.含义}") item = {} arr = row.变量.split("=") item["var_name"] = arr[0] item["var_value"] = arr[1] item["var_comment"] = row.含义 result.append(item) return result def replace_vars(sql_file: str, var_list: list[dict]): """ 替换 sql_file 中的变量 :param sql_file: SQL文件 :param var_list: 变量集合 [{'var_name':xx,'var_value':xxx,'var_comment':xx}] """ new_content = "" encodings = ['utf-8', 'latin-1', 'iso-8859-1', 'cp1252', 'gbk', 'gb2312'] for encoding in encodings: try: with open(sql_file, 'r', encoding=encoding) as f: # 读取文件内容 content = f.read() # 替换变量 for var in var_list: name = var["var_name"] value = var["var_value"] new_content = content.replace(name, value) content = new_content break except UnicodeDecodeError: print(f"文件{os.path.basename(f.name)} 编码 {encoding} 不支持") continue else: print("无法找到合适的编码") raise Exception(f"无法读取文件 {sql_file}, 未找到合适的编码.") return new_content def get_ddl_list(ddl_dir: str, var_list: list[dict]): """ 获取 DDL目录中定义的所有建表DDL，并替换其中的变量 :param ddl_dir: DDL目录,目录结构 ddl/db/xx.ddl :param var_list: 变量集合 [{'var_name':xx,'var_value':xxx,'var_comment':xx}] """ file_list = [] # 遍历目录，获取所有.ddl文件 for root, dirs, files in os.walk(ddl_dir): for file in files: if file.endswith(".ddl"): file_path = os.path.join(root, file) file_list.append(file_path) print(f"DDL文件数量：{len(file_list)}") # 读取ddl文件，替换其中的变量 result = [] for file in file_list: fp = Path(file) # 文件的名称规则：db__table，解析出数据库和表名 name = fp.name.replace(".ddl", "") arr = name.split("__") db = arr[0] table = arr[1] # key key = f"{db.upper()}.{table.upper()}" ddl_content = "" encodings = ['utf-8', 'latin-1', 'iso-8859-1', 'cp1252', 'gbk', 'gb2312'] for encoding in encodings: try: with open(file, 'r', encoding=encoding) as f: # 读取文件内容 content = f.read() # 替换变量 for var in var_list: name = var["var_name"] value = var["var_value"] new_content = content.replace(name, value) content = new_content result.append({key: content}) break except UnicodeDecodeError: print(f"文件{os.path.basename(f.name)} 编码 {encoding} 不支持") continue else: print("无法找到合适的编码") raise Exception(f"无法读取文件 {file}, 未找到合适的编码.") print(f"result:{result[0]}") return result if __name__ == "__main__": import sys import os from pathlib import Path import json # 参数检查 if len(sys.argv) <= 4: print(f"usage python sql_lineage_agent_xmgj.py var_file sql_file dialect owner") exit(1) # 变量文件 var_file = sys.argv[1] # sql脚本文件或目录 sql_file = sys.argv[2] # 数据库方言 dialect = sys.argv[3] # 属主 owner = sys.argv[4] print(f"1、解析SQL文件/目录:{sql_file}") print(f"2、SQL文件属主平台:{owner}") print(f"3、变量文件:{var_file}") print(f"4、SQL数据库方言: {dialect}") # 检查变量文件是否存在 var_fp = Path(var_file) if not var_fp.exists(): raise FileNotFoundError(f"变量文件 {var_file} 不存在，请指定正确路径.") # 检查SQL脚本文件是否存在 sql_fp = Path(sql_file) if not sql_fp.exists(): raise FileNotFoundError(f"SQL脚本文件 {sql_file} 不存在，请指定正确路径.") # 读取变量文件，获取变量值 var_list = read_var_file(var_file) # 读取SQL文件, 替换变量 sql_content = replace_vars(sql_file=sql_file, var_list=var_list) # 解析SQL血缘 result = asyncio.run(main(sql_content=sql_content, sql_file=sql_fp.name, dialect=dialect, owner=owner)) # 将结果写入同级目录下，文件后缀为.json target_file = Path(sql_fp.parent.absolute() / (sql_fp.name + ".json")) print(f"写目标文件:{target_file}") # 写文件 with open(target_file, 'w', encoding="utf-8") as t: t.write(json.dumps(result, ensure_ascii=False, indent=2))