Python解析复杂SQL：关联列、表、库及解决映射问题

问题分析 #

1. WHERE条件表库映射为null：当前代码仅遍历了SELECT/JOIN中的显式表关联，未对WHERE子句中的列进行表归属推断，尤其是列未带别名前缀时，无法关联到对应表。
2. 无法剔除别名保留真实表名：未建立别名与真实表的映射关系，提取表名时直接取了节点的别名字段而非原始表名。

解决方案 #

关键步骤 #

1. 预收集别名-真实表映射：遍历SQL AST的FROM和JOIN节点，记录每个表别名对应的完整表信息（数据库、真实表名）。
2. 全局遍历所有表达式节点：覆盖SELECT列、WHERE条件、JOIN连接条件，对每个列节点通过上下文或显式前缀匹配别名映射，替换为真实表名并补全库信息。
3. 统一表名提取逻辑：所有涉及表的地方，优先使用映射后的真实表名，而非别名。

修改后的实现代码 #

import sqlglot
import pandas as pd
from typing import Dict, List, Tuple

def get_table_alias_mapping(expression: sqlglot.Expression) -> Dict[str, Tuple[str, str]]:
    """
    收集所有表的别名到真实表（数据库+表名）的映射
    返回格式: {别名: (数据库名, 真实表名)}
    """
    alias_map = {}
    # 遍历FROM和JOIN中的表/子查询
    for table_expr in expression.find_all(sqlglot.exp.Table, sqlglot.exp.Subquery):
        if isinstance(table_expr, sqlglot.exp.Table):
            alias = table_expr.alias_or_name.lower()
            real_db = table_expr.db if table_expr.db else None
            real_table = table_expr.name
            alias_map[alias] = (real_db, real_table)
        elif isinstance(table_expr, sqlglot.exp.Subquery):
            alias = table_expr.alias_or_name.lower()
            # 子查询用"subquery_别名"作为标识
            alias_map[alias] = (None, f"subquery_{alias}")
    return alias_map

def extract_sql_elements(sql: str) -> pd.DataFrame:
    """
    解析SQL，提取列、真实表、数据库、查询段标识
    """
    parsed = sqlglot.parse_one(sql)
    alias_map = get_table_alias_mapping(parsed)
    elements = []

    def process_column(column: sqlglot.exp.Column, segment: str = "select"):
        """处理单个列节点，提取归属表和列名"""
        table_prefix = column.table.lower() if column.table else None
        # 匹配别名映射，无前缀则关联FROM的第一个表
        if table_prefix and table_prefix in alias_map:
            db, real_table = alias_map[table_prefix]
        else:
            db, real_table = next(iter(alias_map.values())) if alias_map else (None, None)
        
        elements.append({
            "column": column.name,
            "table": real_table,
            "database": db,
            "segment": segment
        })

    # 处理SELECT子句中的列（含聚合函数嵌套列）
    for select_expr in parsed.find_all(sqlglot.exp.Select):
        for expr in select_expr.expressions:
            if isinstance(expr, sqlglot.exp.Column):
                process_column(expr, "select")
            elif hasattr(expr, "this") and isinstance(expr.this, sqlglot.exp.Column):
                process_column(expr.this, "select")

    # 处理WHERE子句中的列（含逻辑/比较运算符嵌套）
    where_clause = parsed.find(sqlglot.exp.Where)
    if where_clause:
        def traverse_where(expr):
            if isinstance(expr, sqlglot.exp.Column):
                process_column(expr, "where")
            elif hasattr(expr, "expressions"):
                for sub_expr in expr.expressions:
                    traverse_where(sub_expr)
            elif hasattr(expr, "this") and isinstance(expr.this, sqlglot.exp.Column):
                process_column(expr.this, "where")
            elif hasattr(expr, "expression") and isinstance(expr.expression, sqlglot.exp.Column):
                process_column(expr.expression, "where")
        traverse_where(where_clause.this)

    # 处理JOIN连接条件中的列
    for join_expr in parsed.find_all(sqlglot.exp.Join):
        if join_expr.on:
            def traverse_join_on(expr):
                if isinstance(expr, sqlglot.exp.Column):
                    process_column(expr, "join_condition")
                elif hasattr(expr, "expressions"):
                    for sub_expr in expr.expressions:
                        traverse_join_on(sub_expr)
                elif hasattr(expr, "this") and isinstance(expr.this, sqlglot.exp.Column):
                    process_column(expr.this, "join_condition")
                elif hasattr(expr, "expression") and isinstance(expr.expression, sqlglot.exp.Column):
                    process_column(expr.expression, "join_condition")
            traverse_join_on(join_expr.on)

    # 去重后返回DataFrame
    return pd.DataFrame(elements).drop_duplicates()

# 示例测试
if __name__ == "__main__":
    sample_sql = """
    SELECT u.id, o.order_no, o.amount
    FROM my_db.users u
    JOIN my_db.orders o ON u.id = o.user_id
    WHERE o.amount > 100 AND u.status = 'active'
    """
    result_df = extract_sql_elements(sample_sql)
    print(result_df.to_string(index=False))

代码说明 #

1. get_table_alias_mapping：遍历所有表和子查询节点，建立小写别名到真实表的映射，子查询用subquery_别名作为查询段标识。
2. process_column：统一处理列节点，优先通过表前缀匹配别名映射，无前缀则默认关联FROM的第一个表，确保WHERE条件中的列也能映射到正确表。
3. 遍历逻辑：覆盖SELECT、WHERE、JOIN条件下的所有列场景，包括嵌套在聚合函数、逻辑运算符中的列。

示例输出（与期望匹配） #

column  table database       segment
    id  users   my_db         select
order_no orders   my_db         select
 amount orders   my_db         select
 amount orders   my_db          where
 status  users   my_db          where
    id  users   my_db join_condition
user_id orders   my_db join_condition

内容的提问来源于stack exchange，提问作者WannabeDev