BrandCultivationDataAnalysis/tets.py

import re

def replace_select_content(sql):
    """
    将SELECT和FROM之间的内容替换为SELECT COUNT(原内容) FROM
    """
    pattern = r'SELECT\s+(.*?)\s+FROM'
    
    def replace_func(match):
        content = match.group(1)  # 获取SELECT和FROM之间的内容
        # 如果是DISTINCT或ALL等特殊关键字，需要特殊处理
        if re.match(r'^(DISTINCT|ALL)\s+', content, re.IGNORECASE):
            # 保留DISTINCT/ALL关键字
            return f'SELECT COUNT({content}) FROM'
        else:
            return f'SELECT COUNT({content}) FROM'
    
    # 使用re.sub进行替换，re.IGNORECASE忽略大小写
    result = re.sub(pattern, replace_func, sql, flags=re.IGNORECASE)
    return result

# 测试用例
test_sqls = [
    "SELECT name FROM users",
    "SELECT name, age, high FROM users WHERE age > 20",
    "SELECT DISTINCT name FROM users",
    "SELECT * FROM employees",
    "select id, name from students where grade = 'A'",
    "SELECT name FROM (SELECT * FROM users) as t",
]

for sql in test_sqls:
    result = replace_select_content(sql)
    print(f"原始SQL: {sql}")
    print(f"替换后: {result}\n")