Skip to content

Commit da84d2c

Browse files
committed
fix: extract_tables_from_sql incorrectly includes CTE names as table names
#1278
1 parent deff40b commit da84d2c

1 file changed

Lines changed: 7 additions & 2 deletions

File tree

backend/apps/chat/task/llm.py

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -69,15 +69,20 @@
6969

7070

7171
def extract_tables_from_sql(sql: str, ds_type: str = None) -> set:
72-
"""从 SQL 中提取表名(使用 sqlglot 解析,可信)"""
72+
"""从 SQL 中提取真实表名(排除 CTE 别名)"""
7373
tables = set()
7474
dialect = get_sqlglot_dialect(ds_type)
7575
try:
7676
statements = sqlglot.parse(sql, dialect=dialect)
7777
for stmt in statements:
7878
if stmt:
79+
# 收集 CTE 别名,排除嵌套 CTE
80+
cte_names = set()
81+
for cte in stmt.find_all(exp.CTE):
82+
if cte.alias:
83+
cte_names.add(cte.alias)
7984
for table in stmt.find_all(exp.Table):
80-
if table.name:
85+
if table.name and table.name not in cte_names:
8186
tables.add(table.name)
8287
except Exception:
8388
pass

0 commit comments

Comments
 (0)