如何使用 Python 从 SQL 查询中提取列名

Question

我想直接从 SQL 语句中提取结果表的列名：

query = """

select 
    sales.order_id as id, 
    p.product_name, 
    sum(p.price) as sales_volume 
from sales
right join products as p 
    on sales.product_id=p.product_id
group by id, p.product_name;

"""

column_names = parse_sql(query)
# column_names:
# ['id', 'product_name', 'sales_volume']

知道在parse_sql() 中做什么吗？生成的函数应该能够识别别名并删除表别名/标识符（例如“sales.”或“p.”）。

提前致谢！

Answer 1

我已经使用库 sqlparse 完成了类似的操作。基本上，这个库接受您的 SQL 查询并将其标记化。完成后，您可以搜索选择查询标记并解析底层标记。 在代码中，读起来像

import sqlparse
def find_selected_columns(query) -> list[str]:
    tokens = sqlparse.parse(query)[0].tokens
    found_select = False
    for token in tokens:
        if found_select:
            if isinstance(token, sqlparse.sql.IdentifierList):
                return [
                    col.value.split(" ")[-1].strip("`").rpartition('.')[-1]
                    for col in token.tokens
                    if isinstance(col, sqlparse.sql.Identifier)
                ]
        else:
            found_select = token.match(sqlparse.tokens.Keyword.DML, ["select", "SELECT"])
    raise Exception("Could not find a select statement. Weired query :)")

此代码也适用于使用公用表表达式的查询，即它只返回最终选择列。 根据您使用的 SQL 方言和引号字符，您可能需要调整 col.value.split(" ")[-1].strip("`").rpartition('.')[- 1]