如何用Python从SQL中提取出涉及到的表名、列名?


=Start=

缘由:

当前需要用 Python 或 Java 对 SQL语句 进行解析,以提取出其中的 表名、列名 等信息,方便后续进行精准的安全问题分析。

如果仅仅只是提取表名的话还比较简单(直接用字符串切分、匹配等处理就可以搞定),但如果还要提取列名(且SQL比较复杂)的话,就不那么容易搞定了,这时就需要对SQL进行语法分析以进行准确提取。问了一下别人,说 ANTLR4 可以做,于是在网上搜了一些资料,在此实践整理一下,方便以后参考。

正文:

参考解答:

这里先介绍一下如何通过字符串切分和正则的方式来提取SQL中的表名(列名提取比较复杂,不太好实现,后续找到好的方法了再补充进来)。

方法一(字符串切分):

def tables_in_query(sql):
     replace_list = ['\n', '(', ')', '*', '=', '`']
     for i in replace_list:
         sql = sql.replace(i, ' ')
     sql = sql.split()
     res = set()
     for i in range(1, len(sql)):
         if sql[i-1].lower() in ['from', 'join'] and sql[i].lower() != 'select': 
             res.add(sql[i])
     # print(res) # ['table1', 'table2', 'table3', 'table4', 'table5']
     return list(res)

方法二(正则提取):

import re

def tables_in_query2(sql_str):
    # remove the /* */ comments (先移除块注释)
    q = re.sub(r"/\*[^*]*\*+(?:[^*/][^*]*\*+)*/", "", sql_str)

    # remove whole line -- and # comments (再移除整行注释)
    lines = [line for line in q.splitlines() if not re.match("^\s*(--|#)", line)]

    # remove trailing -- and # comments (再移除行尾注释)
    q = " ".join([re.split("--|#", line)[0] for line in lines])

    # split on blanks, parens and semicolons (用「空白符、括号、分号」作为分隔符进行切分)
    tokens = re.split(r"[\s)(;`]+", q)

    # scan the tokens. if we see a FROM or JOIN, we set the get_next
    # flag, and grab the next one (unless it's SELECT).
    # 扫描经过上述分词的列表,当看到 FROM/JOIN 关键字时,获取下一个非「空/select」的作为表名

    result = set()
    get_next = False
    for tok in tokens:
        if get_next:
            if tok.lower() not in ["", "select"]:
                result.add(tok)
            get_next = False
        get_next = tok.lower() in ["from", "join"]

    return result

其实正则的这个方法考虑的挺完整的(比如在移除注释那里),但是上面的字符串版本会更简单而且效率更高(在线上MySQL操作中一般不会有注释),所以可以按需选择。下面提供一个测试效果用的SQL:

sql = """SELECT a.time_updated_server/1000,
content,
nick,
name
FROM table1 a
JOIN table2 b ON a.sender_id = b.user_id
JOIN table3 c ON a.channel_id = c.channel_id
JOIN table4 d ON c.store_id = d.store_id
WHERE sender_id NOT IN
(SELECT user_id
FROM table5
WHERE store_id IN ('agent_store:1', 'ask:1'))
AND to_timestamp(a.time_updated_server/1000)::date >= '2014-05-01'
GROUP BY 1,2,3,4
HAVING sum(1) > 500
ORDER BY 1 ASC
"""

提取列名的方式其实照着表名的方法改一下应该也OK,但这里暂时就不放出来了,有需要的可以按需进行修改。

参考链接:

=END=


《 “如何用Python从SQL中提取出涉及到的表名、列名?” 》 有 4 条评论

发表回复

您的电子邮箱地址不会被公开。 必填项已用 * 标注