侧边栏壁纸
博主头像
DA-技术分享-M版 博主等级

  • 累计撰写 148 篇文章
  • 累计创建 13 个标签
  • 累计收到 0 条评论

目 录CONTENT

文章目录

基于语法树的SQL自动改写工具开发系列（2）-使用PYTHON进行简单SQL改写的开发实战

DarkAthena
DarkAthena
2024-11-04 / 0 评论 / 0 点赞 / 14 阅读 / 0 字

一、前言

前面一篇写了如何搭建环境，本文接着讲怎么使用antlr4进行开发。

二、实战

根据上一篇，基于语法树的SQL自动改写工具开发系列（1）-离线安装语法树解析工具antlr4-DA-技术分享-M版,先在本地部署好开发环境。

DEMO 1

写一段期望改写的原始SQL，使用pygrun进行解析
比如在原生PG中不支持ORACLE中的table()函数，但有替代的unnest改写方式
ORACLE:

SELECT A FROM TABLE(VAR) T

PG:

SELECT A FROM (SELECT * FROM UNNEST(VAR) COLUMN_VALUE) T

要实现的改写规则为，找到table函数的整个表达式，并取出它的入参，将table函数表达式的节点替换成(SELECT * FROM UNNEST(入参) COLUMN_VALUE)
语法树：

C:\antlr>pygrun PlSql sql_script  --tree
SELECT A FROM TABLE(VAR) T;
^Z
(sql_script
   (unit_statement
      (data_manipulation_language_statements
         (select_statement
            (select_only_statement
               (subquery
                  (subquery_basic_elements
                     (query_block SELECT
                        (selected_list
                           (select_list_elements
                              (expression
                                 (logical_expression
                                    (unary_logical_expression
                                       (multiset_expression
                                          (relational_expression
                                             (compound_expression
                                                (concatenation
                                                   (model_expression
                                                      (unary_expression
                                                         (atom
                                                            (general_element
                                                               (general_element_part
                                                                  (id_expression
                                                                     (regular_id A))))))))))))))))
                        (from_clause FROM
                           (table_ref_list
                              (table_ref
                                 (table_ref_aux
                                    (table_ref_aux_internal
                                       (dml_table_expression_clause
                                          (table_collection_expression TABLE (
                                             (expression
                                                (logical_expression
                                                   (unary_logical_expression
                                                      (multiset_expression
                                                         (relational_expression
                                                            (compound_expression
                                                               (concatenation
                                                                  (model_expression
                                                                     (unary_expression
                                                                        (atom
                                                                           (general_element
                                                                              (general_element_part
                                                                                 (id_expression
                                                                                    (regular_id VAR)))))))))))))) ))))
                                 (table_alias
                                    (identifier
                                       (id_expression
                                          (regular_id T))))))))))))))) ; <EOF>)

C:\antlr>

以下为一个完整且可以运行的demo代码，用于说明如何进行语法改写

from antlr4 import FileStream, CommonTokenStream
from PlSqlLexer import PlSqlLexer
from PlSqlParser import PlSqlParser
from PlSqlParserVisitor import PlSqlParserVisitor
from antlr4.TokenStreamRewriter import TokenStreamRewriter

class SQLTransformer(PlSqlParserVisitor):
    def __init__(self, token_stream):
        super().__init__()
        self.rewriter = TokenStreamRewriter(token_stream)

    def visitDml_table_expression_clause(self, ctx: PlSqlParser.Dml_table_expression_clauseContext):
        if ctx.table_collection_expression() and ctx.table_collection_expression().getText().startswith('table'):
            argument = ctx.table_collection_expression().expression().getText()
            new_text = "(select * from unnest({}) column_value)".format(argument)
            self.rewriter.replace(TokenStreamRewriter.DEFAULT_PROGRAM_NAME, ctx.table_collection_expression().start.tokenIndex, ctx.table_collection_expression().stop.tokenIndex, new_text)
        return self.visitChildren(ctx)

def main(input_file, output_file):
    input_stream = FileStream(input_file, encoding='utf-8')
    lexer = PlSqlLexer(input_stream)
    stream = CommonTokenStream(lexer)
    parser = PlSqlParser(stream)
    tree = parser.sql_script()

    transformer = SQLTransformer(stream)
    transformer.visit(tree)
    output_text = transformer.rewriter.getDefaultText()
    with open(output_file, 'w', encoding='utf-8',newline='') as f:
        f.write(output_text)

if __name__ == '__main__':
    input_file = 'input.sql'
    output_file = 'output.sql'
    main(input_file, output_file)

其中def main内的代码基本可以固定，我们直接看class SQLTransformer里的def visitDml_table_expression_clause。
visitDml_table_expression_clause其实是PlSqlParserVisitor.py里面的一个def，这个def的名称由visit加上节点名组成，也就是说，语法树中的每一个节点，都有一个对应的visit。而原本PlSqlParserVisitor.py里的每个visit里面都是空的，直接就return出去了：

def  visitDml_table_expression_clause(self, ctx:PlSqlParser.Dml_table_expression_clauseContext):

return  self.visitChildren(ctx)

我们自己写的这个visit就是实现了里面的具体内容。

如果我们需要修改某个语法，可以从语法树中，找到这个语法相关的上下文的最小节点，以本文前面输出的语法树为例，就应该是(table_collection_expression TABLE (，所以理论上，我们再写个visitTable_collection_expression就好了，但本文的demo代码是从visitDml_table_expression_clause开始，是为了说明如何引用当前节点的下级节点

if ctx.table_collection_expression() and ctx.table_collection_expression().getText().startswith('table')

这句是一个判断，作用是，判断当前节点下，是否存在table_collection_expression这个节点，我们可以去对比语法树，如果没有使用table函数，是不会有这个节点的；第二个条件就是，获取table_collection_expression这个节点的文本，判断它是不是使用 table开始。这里要注意，如果不写第一个条件，在没有table_collection_expression节点时，对它执行getText会报错。

argument = ctx.table_collection_expression().expression().getText()

这里是取出table_collection_expression的下一个叫expression的节点，对照语法树可以看到，虽然里面有很多层，但这个节点实际只包含VAR这个文本，因此这里就可以得到argument="VAR",即前文例子中，table函数的入参。

new_text = "(select * from unnest({}) column_value)".format(argument)

这一句很好理解，就是格式化一个字符串，把argument的值替换{}，得到(select * from unnest(VAR) column_value)

self.rewriter.replace(TokenStreamRewriter.DEFAULT_PROGRAM_NAME, ctx.table_collection_expression().start.tokenIndex, ctx.table_collection_expression().stop.tokenIndex, new_text)

这一句就是最关键的，self.rewriter是前面定义的TokenStreamRewriter(token_stream)，在TokenStreamRewriter里面，可以支持对节点的替换、删除、增加等操作。
TokenStreamRewriter.replace有4个入参,分别为程序名，开始位置、结束位置、需要替换成的文本。
程序名一般固定使用TokenStreamRewriter.DEFAULT_PROGRAM_NAME就行，如果期望一次解析，就能做多种替换，比如同时生成支持PG和MYSQL的两种语法，就可以在这里设置程序名,针对不同的程序名写不同的规则。
开始位置和结束位置，可以使用对应节点的start.tokenIndex和stop.tokenIndex。
至此，我们就完成了一个改写规则的开发。

DEMO 2

如果需要在一次语法树解析中就完成多种规则的执行，可以再添加几个def visit，比如我们再针对create type语句来进行改写。

ORACLE:

CREATE OR REPLACE TYPE TY_TEST AS OBJECT(COL1 INT,COL2 VARCHAR(20));

PG:

CREATE TYPE TY_TEST AS (COL1 INT,COL2 VARCHAR(20));

改写规则为，对于create_type的语法节点，将create or replace 改为create，并且删除object

语法树：

C:\antlr>pygrun PlSql sql_script  --tree
CREATE OR REPLACE TYPE TY_TEST AS OBJECT(COL1 INT,COL2 VARCHAR(20));
^Z
(sql_script
   (unit_statement
      (create_type CREATE OR REPLACE TYPE
         (type_definition
            (type_name
               (id_expression
                  (regular_id TY_TEST)))
            (object_type_def
               (object_as_part AS OBJECT) (
               (object_member_spec
                  (identifier
                     (id_expression
                        (regular_id COL1)))
                  (type_spec
                     (datatype
                        (native_datatype_element INT)))) ,
               (object_member_spec
                  (identifier
                     (id_expression
                        (regular_id COL2)))
                  (type_spec
                     (datatype
                        (native_datatype_element VARCHAR)
                        (precision_part (
                           (numeric 20) ))))) ))))) ; <EOF>)

改写代码

    def visitCreate_type(self, ctx:PlSqlParser.Create_typeContext):
        # 检查并修改 'create or replace' 为 'create'
        if ctx.getChild(0).getText() == 'create' and ctx.getChild(1).getText() == 'or' and ctx.getChild(2).getText() == 'replace':
           # print(f"Modifying: {ctx.getChild(0).getText()} {ctx.getChild(1).getText()} {ctx.getChild(2).getText()}")
            self.rewriter.replace(
                TokenStreamRewriter.DEFAULT_PROGRAM_NAME,
                ctx.getChild(0).symbol.tokenIndex,
                ctx.getChild(2).symbol.tokenIndex,
                'create'
            )   
        object_as_part = ctx.type_definition().object_type_def().object_as_part()
        #print(f"ctx.object_as_part.getText(): {object_type_def.object_as_part().getChild(1).getText()}")
        # 删除 object关键字
        if object_as_part.getChild(1).getText()=='OBJECT':
            self.rewriter.delete(
                TokenStreamRewriter.DEFAULT_PROGRAM_NAME,
                object_as_part.getChild(1).symbol.tokenIndex,
                object_as_part.getChild(1).symbol.tokenIndex
            )
        return self.visitChildren(ctx)

这里可以从(create_type CREATE OR REPLACE TYPE看到，CREATE OR REPLACE TYPE这一串都在create_type这个节点上，可以通过getChild(n)来取出中间的每一部分，而这每一部分的位置，则是通过symbol.tokenIndex获取。
然后从(object_as_part AS OBJECT) (中可以看到，我们期望删除的object是在object_as_part这个节点的第二个字符串，因此使用了getChild(1)。这里需要注意，由于object_as_part这个节点是在create_type这个节点的下面很多层，然后后面会要多次使用这个节点，所以可以定义一个object_as_part = ctx.type_definition().object_type_def().object_as_part()，减少冗余代码。
找到OBJECT后，执行self.rewriter.delete，就可以把OBJECT删掉了。

DEMO 3

在ORACLE的sql脚本中，create type应该以/结尾，但是在OG中，则不能有/，如果一个脚本文件里混合了多种语句，就会出现有的/要删，有的不能删，因此我们可以写一个规则，将所有create type语句后面的/删掉。

    def visitSql_script(self, ctx:PlSqlParser.Sql_scriptContext):
        for i in range(ctx.getChildCount() - 1):
            unit = ctx.getChild(i)
            if isinstance(unit, PlSqlParser.Unit_statementContext):
                create_type_stmt = unit.getChild(0)
                # 对于create type语句
                if isinstance(create_type_stmt, PlSqlParser.Create_typeContext):
                    # 遍历兄弟节点，找到 `/` 进行删除
                    for j in range(i + 1, ctx.getChildCount()):
                        sibling = ctx.getChild(j)
                        if sibling.getText() == '/':
                            #print(f"Deleting: {sibling.getText()}")
                            self.rewriter.delete(
                                TokenStreamRewriter.DEFAULT_PROGRAM_NAME,
                                sibling.symbol.tokenIndex,
                                sibling.symbol.tokenIndex
                            )
                            break
        return self.visitChildren(ctx)

这里没有使用visitCreate_type的原因是，在antlr4生成的语法树中，/这个节点并不在Create_type这个节点的内部，所以得对Sql_script这个根节点，找到所有的单条语句Unit_statement，然后判断里面是不是有Create_type，然后再回头删掉/

三、注意事项

antlr4目前的版本在sql语法解析规则存在一个我认为很严重的BUG，就是SQL中如果包含有pro或者rem，解析就会报错，参考这个issue
# [PlSql] “REM”, “REMARK”, “PRO”, “PROMPT” can not be a identifier #3817
其实规避手段也很简单，删除g4文件中的PROMPT_MESSAGE和REMARK_COMMENT这两个语法定义，然后重新执行antlr4vjava PlSqlParser.g4和antlr4vjava PlSqlLexer.g4即可。不过这样就将无法解析SQLPLUS中的prompt和remark命令了

四、总结

antlr4很强大,我们可以借助其语法树的能力，实现精准的语法改写，当然前提是我们要先想好怎么去改写，然后才能编写对应的规则。这种方式虽然比正则替换要慢很多，但是对于大量复杂的存储过程而言，使用语法树进行改写的方式无疑比正则替换更加可靠。

0

  1. 支付宝打赏

    qrcode alipay

  2. 微信打赏

    qrcode weixin
版权归属： DarkAthena
本文链接： https://www.darkathena.top/archives/antlr4-part2-dev-with-python
许可协议： 本文使用《署名-非商业性使用-相同方式共享 4.0 国际 (CC BY-NC-SA 4.0)》协议授权
博主关闭了所有页面的评论