From da84d2cbc527d68ff0bcb893c7e0dcd15d96ad9b Mon Sep 17 00:00:00 2001 From: ulleo Date: Fri, 17 Jul 2026 16:26:43 +0800 Subject: [PATCH] fix: extract_tables_from_sql incorrectly includes CTE names as table names #1278 --- backend/apps/chat/task/llm.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/backend/apps/chat/task/llm.py b/backend/apps/chat/task/llm.py index 308e2ad6..e9ef8042 100644 --- a/backend/apps/chat/task/llm.py +++ b/backend/apps/chat/task/llm.py @@ -69,15 +69,20 @@ def extract_tables_from_sql(sql: str, ds_type: str = None) -> set: - """从 SQL 中提取表名（使用 sqlglot 解析，可信）""" + """从 SQL 中提取真实表名（排除 CTE 别名）""" tables = set() dialect = get_sqlglot_dialect(ds_type) try: statements = sqlglot.parse(sql, dialect=dialect) for stmt in statements: if stmt: + # 收集 CTE 别名，排除嵌套 CTE + cte_names = set() + for cte in stmt.find_all(exp.CTE): + if cte.alias: + cte_names.add(cte.alias) for table in stmt.find_all(exp.Table): - if table.name: + if table.name and table.name not in cte_names: tables.add(table.name) except Exception: pass