如何使用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]

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)

Answer 3

我希望您有数据的列名称列表。 并且您有一个字符串查询。 对于查询中的每个单词，在列名称列表中查找相同的单词。 因此，通过这种方式，您必须能够获取选择查询中的所有列