使用sqlparse提取SQL列名报错：'Identifier'对象无get_identifiers属性

我用下面的代码从SQL查询语句中提取列名，但触发了错误：

import sqlparse

def get_query_columns(sql):
    stmt = sqlparse.parse(sql)[0]
    columns = []
    column_identifiers = []

    # get column_identifieres
    in_select = False
    for token in stmt.tokens:
        if isinstance(token, sqlparse.sql.Comment):
            continue
        if str(token).lower() == 'select':
            in_select = True
        elif in_select and token.ttype is None:
            for identifier in token.get_identifiers():
                column_identifiers.append(identifier)
            break

    # get column names
    for column_identifier in column_identifiers:
        columns.append(column_identifier.get_name())

    return columns


dfff['SQL_TEXT'].apply(get_query_columns)

报错信息：

Input In [1069], in get_query_columns(sql)
     14     in_select = True
     15 elif in_select and token.ttype is None:
---> 16     for identifier in token.get_identifiers():
     17         column_identifiers.append(identifier)
     18     break

AttributeError: 'Identifier' object has no attribute 'get_identifiers'

错误原因 #

当in_select为True时，代码默认后续的token是包含多个列的IdentifierList对象，但实际场景中，这个token可能是单个列对应的Identifier对象——这类对象没有get_identifiers()方法，因此触发属性错误。

修正后的代码 #

import sqlparse
from sqlparse.sql import IdentifierList, Identifier

def get_query_columns(sql):
    # 处理空SQL的情况
    if not sql.strip():
        return []
    
    stmt = sqlparse.parse(sql)[0]
    columns = []
    column_identifiers = []
    in_select = False

    for token in stmt.tokens:
        # 跳过注释
        if isinstance(token, sqlparse.sql.Comment):
            continue
        
        # 标记进入SELECT子句
        if str(token).lower() == 'select':
            in_select = True
            continue
        
        # 在SELECT子句中处理列
        if in_select:
            # 处理逗号分隔的列列表
            if isinstance(token, IdentifierList):
                for identifier in token.get_identifiers():
                    column_identifiers.append(identifier)
            # 处理单个列的情况
            elif isinstance(token, Identifier):
                column_identifiers.append(token)
            
            # 遇到FROM就停止（避免处理后续内容）
            if str(token).lower() == 'from':
                break
    
    # 提取列名
    for identifier in column_identifiers:
        # 处理带别名的列，比如`user_id as uid`，只取原始列名
        if hasattr(identifier, 'get_real_name'):
            columns.append(identifier.get_real_name())
        else:
            columns.append(identifier.get_name())
    
    return columns


dfff['SQL_TEXT'].apply(get_query_columns)

关键修正点 #

1. 类型判断兼容：区分IdentifierList（多列）和Identifier（单列）两种情况，避免直接调用不存在的方法
2. 增加终止条件：遇到FROM关键字就停止遍历，防止误处理SELECT子句之外的内容
3. 容错处理：增加空SQL的判断，避免解析空字符串时出错
4. 优化列名提取：支持提取带别名的列的原始名称，提升适配性

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