如何利用python来提取SQL语句中的表名称

news2024/9/21 4:20:52

1.介绍

在某些场景下，我们可能需要从一个复杂的SQL语句中提取对应的表名称，在这样的场景下，我们如果在python中处理的话，就需要用到SQLparse这个库。
SQLparse 是一个用于解析 SQL 查询语句的 Python 库。它可以将复杂的 SQL 查询转换为易于分析的结构，并提供了许多功能来检索、修改和分析 SQL 查询中的元素。其中一个常见的用例是从查询中提取表名称。

2.安装 SQLparse

首先，确保您已经安装了 SQLparse。您可以使用 pip 进行安装：

pip install sqlparse

3.解析 SQL 查询语句

使用 SQLparse 解析 SQL 查询语句非常简单。只需导入库并将查询语句作为字符串传递给 sqlparse.parse() 函数。

4.提取表名称

提取表名称通常涉及到遍历解析后的 SQL 结构以查找关键元素，这些元素通常表示表名称。SQLparse 使用了一种抽象语法树（AST）的结构，可以轻松访问查询中的各个部分。

示例代码
下面是一个示例代码，演示如何使用 SQLparse 从简单的SQL 查询语句中得到表名称：

import sqlparse
# 定义一个 SQL 查询语句
sql_query = "SELECT * FROM customers WHERE age > 30;"
# 解析 SQL 查询语句
parsed_query = sqlparse.parse(sql_query)
# 遍历解析出表名
table_names=[]
for stmt in parsed_query:
    for token in stmt.tokens:
        if isinstance(token, sqlparse.sql.Identifier) :
            table_name= token.get_real_name()
            table_names.append(table_name)
print(table_names)

代码运行后，结果如下：

在上述代码中，我们首先定义了一个 SQL 查询语句。然后，我们使用 SQLparse 解析查询语句，将其转换为 AST 结构。接下来，我们遍历 AST 结构以查找关键元素，例如表名称。但是敏感和聪明的你，会发现用上述代码在实际场景中时会出现各种各样的问题。

5.实际场景中的问题

5.1 无法处理嵌套类型的语句

如在处理下面这个SQL语句时，无法提取嵌套在里面的表名称

    SELECT o.id, o.name FROM orders o
    LEFT JOIN products p ON o.product_code = p.product_code
    WHERE o.customer_id IN (
        SELECT id FROM customers 
        WHERE country = 'USA'
    )
    OR p.category = 'Electronics'

简单的嵌套，无法识别：

5.2 无法处理重命名表的语句

如下面的SQL语句，会出现提取错误的情况：

with co as 
(SELECT o.id, o.name FROM orders o
    LEFT JOIN products p ON o.product_code = p.product_code
    WHERE o.customer_id IN (
        SELECT id FROM customers 
        WHERE country = 'USA'
    )
    OR p.category = 'Electronics')
select * from co

提取结果如下：

5.3 多个表名重复时的去重问题

在处理基础的SQL语句时，会出现多个相同的表

    SELECT o.id, o.name FROM orders o
    LEFT JOIN orders o2 on o.id = o2.pre_id

处理后，截图如下：

5.4 其他问题

有一些场景下，还会提取出字段信息夹杂其中的情况，或者有一些存在表空间的情况会进行重命名，等等，听到这里，你是不是觉得处理起来令人头大？别慌，接着往下看。

6.终极解决方案

6.1 解决思路

经过本人的一番思考后，采用sqlparse模块和re模块相结合后，惊喜的发现，能够灵活的从绝大部分的sql语句中正确地提取出表名称来。
解决思路如下：

1.先将sql语句进行规范
2.获取left和join关键字后面的表名
3.对表名进行各种处理
4.去重后输出结果

6.2 实际代码

import sqlparse
import re


def format_sql(sql_content):
    '''将sql语句进行规范化，并去除sql中的注释，输入和输出均为字符串'''
    parse_str=sqlparse.format(sql_content,reindent=True,strip_comments=True)
    return parse_str

def extract_table_names(sql_query):
    '''从sql中提取对应的表名称，输出为列表'''
    table_names = set()
    # 解析SQL语句
    parsed = sqlparse.parse(sql_query)
    # 正则表达式模式，用于匹配表名
    table_name_pattern = r'\bFROM\s+([^\s\(\)\,]+)|\bJOIN\s+([^\s\(\)\,]+)'
    # with 子句判断
    with_pattern = r'with\s+(\w+)\s+as'
    remove_with_name=[]
    
    # 遍历解析后的语句块
    for statement in parsed:
        # 转换为字符串
        statement_str =  str(statement).lower()
        
        #将字符串中的特殊语法置空
        statement_str = re.sub('(substring|extract)\s*\(((.|\s)*?)\)','',statement_str)
        
        # 查找匹配的表名
        matches = re.findall(table_name_pattern, statement_str, re.IGNORECASE)
        
        for match in matches:
            # 提取非空的表名部分
            for name in match:
                #if name and name not in not_contain_list:
                if name :
                    # 对于可能包含命名空间的情况，只保留最后一部分作为表名
                    table_name = name.split('.')[-1]
                    #去除表名中的特殊符号
                    table_name = re.sub('("|`|\'|;)','',table_name)
                    table_names.add(table_name)
        #处理特殊的with语句
        if 'with' in statement_str:
            match = re.search(with_pattern, statement_str)
            if match:
                result = match.group(1)
                remove_with_name.append(result)
    table_list = list(table_names)
    #移除多余的表名
    if remove_with_name:
        table_list =list(set(table_list)-set(remove_with_name))
    return table_list
sql_query='''with co as 
(SELECT o.id, o.name FROM orders o
    LEFT JOIN products p ON o.product_code = p.product_code
    WHERE o.customer_id IN (
        SELECT id FROM customers 
        WHERE country = 'USA'
    )
    OR p.category = 'Electronics')
select * from co'''
parse_str=format_sql(sql_query)
table_names = extract_table_names(parse_str)
# 打印提取的表名
print(table_names)