如何从SQL SELECT语句提取列名？解决sqlparse适配问题

问题背景 #

需要从SQL SELECT语句中提取生成的列名，用于构建INSERT ... SELECT语句。最初使用sqlparse编写的代码无法正确解析部分场景的列别名，现需适配两种典型查询场景：

原代码实现 #

import sqlparse

def get_new_col_list(sql: str) -> str:
    parsed = sqlparse.parse(sql)
    id_lists = filter(lambda x: isinstance(x, sqlparse.sql.IdentifierList), parsed[0].tokens)
    names = []
    for id_list in id_lists:
        sub_names = []
        for i in id_list.get_identifiers():
            if not isinstance(i, sqlparse.sql.Identifier):
                print("no id ---->", i, i.value, i.flatten, sep="\n")
                continue
            sub_names.append(i.get_name())
        names.extend(sub_names)
    return names

无法适配的场景1 #

执行如下SQL时，原代码无法识别yy列：

select
  p.x xx
  , 'SOMETHING' yy
from
  polls p
;

期望返回结果：["xx", "yy"]

无法适配的场景2 #

执行带复杂表达式的SQL时，原代码无法提取colname别名：

select
a aa
, p pp
, cast(case
    when x like 'TT%' then 'TT'
    when x like 'PP%' then 'PP'
    else 'TT'
  end as xtype)
  as colname
from
 table

期望返回结果：["aa", "pp", "colname"]

解决方案 #

修改代码逻辑，手动定位SELECT与FROM之间的字段区域，遍历每个字段项并从后往前识别别名（支持AS 别名和直接空格后跟别名两种格式），同时兼容普通列、字面量、复杂表达式等多种场景：

import sqlparse
from sqlparse.sql import Identifier, Parenthesis, Literal
from sqlparse.tokens import Keyword, DML, Whitespace, Newline, Punctuation

def get_select_columns(sql: str) -> list[str]:
    parsed = sqlparse.parse(sql)[0]
    
    # 定位SELECT关键字的位置
    select_idx = None
    for idx, token in enumerate(parsed.tokens):
        if token.ttype is DML and token.value.upper() == 'SELECT':
            select_idx = idx
            break
    if select_idx is None:
        return []
    
    columns = []
    current_col_tokens = []
    
    # 遍历SELECT之后的token，直到遇到FROM关键字
    for token in parsed.tokens[select_idx+1:]:
        if token.ttype is Keyword and token.value.upper() == 'FROM':
            break
        # 跳过空白、换行、逗号等分隔符，遇到时先处理当前收集的字段项
        if isinstance(token, (Whitespace, Newline)) or token.ttype == Punctuation:
            if current_col_tokens:
                columns.append(process_column_tokens(current_col_tokens))
                current_col_tokens = []
            continue
        current_col_tokens.append(token)
    
    # 处理最后一个未被分隔符触发的字段项
    if current_col_tokens:
        columns.append(process_column_tokens(current_col_tokens))
    
    return columns

def process_column_tokens(tokens: list) -> str:
    # 从后往前查找别名
    alias = None
    for i in range(len(tokens)-1, -1, -1):
        token = tokens[i]
        # 处理AS关键字后的别名
        if token.ttype is Keyword and token.value.upper() == 'AS':
            if i + 1 < len(tokens) and isinstance(tokens[i+1], Identifier):
                alias = tokens[i+1].get_name()
                break
        # 处理直接空格后跟的别名
        elif isinstance(token, Identifier):
            prev_token = tokens[i-1] if i > 0 else None
            if prev_token and isinstance(prev_token, Whitespace):
                alias = token.get_name()
                break
    
    if alias:
        return alias
    
    # 无别名时返回字段本身的名称（兼容普通列场景）
    for token in tokens:
        if isinstance(token, Identifier):
            return token.get_name()
    
    # 极端情况返回表达式字符串（可选，根据需求调整）
    return sqlparse.sql.TokenList(tokens).value.strip()

测试验证 #

• 针对场景1的SQL，调用get_select_columns(sql)返回["xx", "yy"]
• 针对场景2的SQL，调用函数返回["aa", "pp", "colname"]

核心逻辑说明 #

1. 手动定位SELECT与FROM的区间，避免依赖sqlparse对IdentifierList的解析限制
2. 从后往前遍历字段的token，优先识别别名（别名通常位于字段项末尾）
3. 兼容AS 别名和直接空格别名两种语法，同时支持普通列、字面量、复杂表达式等多种字段类型

内容的提问来源于stack exchange，提问作者Arpad Horvath -- Слава Україні