SQL查询字段提取脚本适配多Case/多Select...From语句问题求助

问题背景 #

现有Python脚本提取SQL查询字段时，普通场景能正常工作，但碰到多CASE语句或嵌套SELECT...FROM子查询时就会失效。当前脚本代码如下：

import re

def remove_comments(sql):
    # 移除单行注释
    sql = re.sub(r'--.*', '', sql)
    # 移除多行注释
    sql = re.sub(r'/\*.*?\*/', '', sql, flags=re.DOTALL)
    return sql

def extract_fields(sql):
    # 提取SELECT子句中的字段（支持表别名和AS别名）
    pattern = r'\s*([\w\.]+(?:\s+AS\s+[\w]+)?)\s*(?:,|$)'
    
    fields = re.findall(pattern, sql)
    
    # 清理字段名，移除别名
    cleaned_fields = [field.split()[0] for field in fields]
    
    # 过滤数值、CTE相关关键字，去重
    valid_fields = [field for field in cleaned_fields if not field.isdigit() and not field.lower().startswith('cte')]
    valid_fields = list(set(valid_fields))
    
    return valid_fields

需求：要准确提取SQL主SELECT子句里的目标字段（比如类似user.id、user.status这类真实引用的字段），忽略CASE表达式的语法关键字、子查询内的SELECT字段等无关内容。

现有脚本的核心问题 #

1. 正则表达式分不清主SELECT和子查询/CASE里的嵌套SELECT，会错误把子查询内的字段也提取出来
2. 无法识别CASE表达式里包裹的真实字段，要么漏掉这些字段，要么错误匹配表达式里的语法片段
3. 单纯靠字符串匹配处理不了SQL的语法结构，比如括号嵌套、复杂表达式场景

改进后的代码 #

import re

def remove_comments(sql):
    sql = re.sub(r'--.*', '', sql)
    sql = re.sub(r'/\*.*?\*/', '', sql, flags=re.DOTALL)
    # 移除多余空白，简化后续处理
    sql = re.sub(r'\s+', ' ', sql).strip()
    return sql

def extract_main_select_fields(sql):
    sql_clean = remove_comments(sql)
    
    # 替换嵌套括号内的内容为占位符，避免干扰主SELECT定位
    def replace_nested_parens(s):
        stack = []
        result = []
        for char in s:
            if char == '(':
                stack.append(len(result))
                result.append('')
            elif char == ')':
                if stack:
                    start_idx = stack.pop()
                    result = result[:start_idx] + ['[NESTED]']
            else:
                if stack:
                    result[-1] += char
                else:
                    result.append(char)
        return ''.join(result)
    
    sql_no_nested = replace_nested_parens(sql_clean)
    
    # 匹配主SELECT子句：从SELECT开始到第一个FROM/UNION/INTERSECT/EXCEPT
    main_select_match = re.search(r'(?i)^SELECT\s+(.*?)\s+(?:FROM|UNION|INTERSECT|EXCEPT)', sql_no_nested)
    if not main_select_match:
        return []
    
    main_select_content = main_select_match.group(1)
    
    # 从CASE表达式中提取真实字段
    def extract_fields_from_case(case_expr):
        case_fields = re.findall(r'([\w\.]+)', case_expr)
        # 过滤CASE语法关键字
        return [f for f in case_fields if f.lower() not in ('case', 'when', 'then', 'else', 'end')]
    
    # 分割主SELECT的各个项，处理逗号分隔的内容
    select_items = re.split(r',\s*(?![^\[]*\])', main_select_content)
    valid_fields = set()
    
    for item in select_items:
        # 处理CASE表达式项
        if re.search(r'(?i)CASE', item):
            case_match = re.search(r'(?i)CASE.*?END', sql_clean)
            if case_match:
                case_fields = extract_fields_from_case(case_match.group())
                valid_fields.update(case_fields)
        else:
            # 处理普通字段（支持AS别名）
            field_match = re.search(r'([\w\.]+)(?:\s+AS\s+[\w]+)?', item)
            if field_match:
                field = field_match.group(1)
                if not field.isdigit() and not field.lower().startswith('cte'):
                    valid_fields.add(field)
    
    return list(valid_fields)

关键优化点 #

• 嵌套内容隔离：把子查询、函数等嵌套括号里的内容换成占位符，避免干扰主SELECT子句的定位
• 主SELECT精准提取：只处理从SELECT开头到第一个FROM/UNION等关键字的内容，排除子查询内的SELECT
• CASE字段解析：单独处理CASE表达式，提取里面真实引用的字段，而非把整个CASE当作一个字段
• 自动去重过滤：保留原有的数值、CTE关键字过滤逻辑，用集合自动去重

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