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

Question

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


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() 中可以做什么吗？生成的函数应该能够识别别名并删除 table aliases/identifiers（例如“sales.”或“p.”）。

提前致谢！

Answer 1

我使用库 sqlparse 做过类似的事情。基本上，这个库接受您的 SQL 查询并将其标记化。完成后，您可以搜索 select 查询令牌并解析基础令牌。在代码中，读起来像

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 :)")

此代码也适用于使用常见 table 表达式的查询，即它仅 return 最后的 select 列。根据 SQL 方言和您使用的引号字符，您可能必须调整 col.value.split(" ")[-1].strip("`").rpartition('. ')[-1]

Answer 2

试用SQLGlot

它比 sqlparse 更容易，更不容易出错。

import sqlglot
import sqlglot.expressions as exp

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 = []

for expression in sqlglot.parse_one(query).find(exp.Select).args["expressions"]:
    if isinstance(expression, exp.Alias):
        column_names.append(expression.text("alias"))
    elif isinstance(expression, exp.Column):
        column_names.append(expression.text("this"))

print(column_names)

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

How to extract column names from SQL query using Python

python

sql

parsing

extract