fix(parser):adapt union sql has order/limit (#30)

* fix(parser):adapt union sql has order/limit

* fix(parser):adapt union sql has order/limit in oceanbase/mysql

Author: Ifffff

.gitignore

@@ -26,6 +26,7 @@ __pycache__/
 ######################
 *.sqlite
 *.log
+logs/
 
 # OS generated files #
 ######################

parser/mysql_parser/parser.py

@@ -281,14 +281,14 @@ def p_cursor_specification(p):
                          query.where,
                          query.group_by,
                          query.having,
-                         p[2],
-                         limit,
-                         offset,
+                         order_by=query.order_by if query.order_by else p[2],
+                         limit=query.limit if query.limit else limit,
+                         offset=query.offset if query.offset else offset,
                          for_update=query.for_update if query.for_update else for_update,
                          nowait_or_wait=query.nowait_or_wait if query.nowait_or_wait else nowait_or_wait
                      ),
-                     limit=limit,
-                     offset=offset
+                     limit=query.limit if query.limit else limit,
+                     offset=query.offset if query.offset else offset
                      )
     else:
         p[0] = Query(p.lineno(1), p.lexpos(1),
@@ -303,11 +303,11 @@ def p_subquery(p):
     if p_limit:
         offset = p_limit[0]
         limit = p_limit[1]
-
+    a = p[2]
     if isinstance(p[2], QuerySpecification):
-        p[2].limit = limit
-        p[2].offset = offset
-        p[2].order_by = p[3] or []
+        p[2].limit = p[2].limit if p[2].limit else limit
+        p[2].offset = p[2].offset if p[2].offset else offset
+        p[2].order_by = p[2].order_by if p[2].order_by else p[3] or []
     p[0] = SubqueryExpression(p.lineno(1), p.lexpos(1), query=p[2])
 
 
@@ -461,7 +461,7 @@ def _item_list(p):
 
 
 def p_query_spec(p):
-    r"""query_spec : SELECT select_items table_expression_opt"""
+    r"""query_spec : SELECT select_items table_expression_opt order_by_opt limit_opt"""
     select_items = p[2]
     table_expression_opt = p[3]
     from_relations = table_expression_opt.from_ if table_expression_opt else None
@@ -472,6 +472,13 @@ def p_query_spec(p):
     for_update = None
     nowait_or_wait = None
 
+    p_limit = p[5]
+    offset = 0
+    limit = 0
+    if p_limit:
+        offset = p_limit[0]
+        limit = p_limit[1]
+
     if p_for_update:
         for_update = p_for_update[0]
         nowait_or_wait = p_for_update[1]
@@ -490,7 +497,11 @@ def p_query_spec(p):
                               group_by=group_by,
                               having=having,
                               for_update=for_update,
-                              nowait_or_wait=nowait_or_wait)
+                              nowait_or_wait=nowait_or_wait,
+                              order_by=p[4],
+                              limit=limit,
+                              offset=offset
+                              )
 
 
 def p_where_opt(p):

parser/oceanbase_parser/parser.py

@@ -275,14 +275,14 @@ def p_cursor_specification(p):
                          query.where,
                          query.group_by,
                          query.having,
-                         p[2],
-                         limit,
-                         offset,
+                         order_by=query.order_by if query.order_by else p[2],
+                         limit=query.limit if query.limit else limit,
+                         offset=query.offset if query.offset else offset,
                          for_update=query.for_update if query.for_update else for_update,
                          nowait_or_wait=query.nowait_or_wait if query.nowait_or_wait else nowait_or_wait
                      ),
-                     limit=limit,
-                     offset=offset
+                     limit=query.limit if query.limit else limit,
+                     offset=query.offset if query.offset else offset
                      )
     else:
         p[0] = Query(p.lineno(1), p.lexpos(1),
@@ -299,9 +299,9 @@ def p_subquery(p):
         limit = p_limit[1]
 
     if isinstance(p[2], QuerySpecification):
-        p[2].limit = limit
-        p[2].offset = offset
-        p[2].order_by = p[3] or []
+        p[2].limit = p[2].limit if p[2].limit else limit
+        p[2].offset = p[2].offset if p[2].offset else offset
+        p[2].order_by = p[2].order_by if p[2].order_by else p[3] or []
     p[0] = SubqueryExpression(p.lineno(1), p.lexpos(1), query=p[2])
 
 
@@ -454,7 +454,7 @@ def _item_list(p):
 
 
 def p_query_spec(p):
-    r"""query_spec : SELECT select_items table_expression_opt"""
+    r"""query_spec : SELECT select_items table_expression_opt order_by_opt limit_opt"""
     select_items = p[2]
     table_expression_opt = p[3]
     from_relations = table_expression_opt.from_ if table_expression_opt else None
@@ -465,6 +465,13 @@ def p_query_spec(p):
     for_update = None
     nowait_or_wait = None
 
+    p_limit = p[5]
+    offset = 0
+    limit = 0
+    if p_limit:
+        offset = p_limit[0]
+        limit = p_limit[1]
+
     if p_for_update:
         for_update = p_for_update[0]
         nowait_or_wait = p_for_update[1]
@@ -483,7 +490,10 @@ def p_query_spec(p):
                               group_by=group_by,
                               having=having,
                               for_update=for_update,
-                              nowait_or_wait=nowait_or_wait)
+                              nowait_or_wait=nowait_or_wait,
+                              order_by=p[4],
+                              limit=limit,
+                              offset=offset)
 
 
 def p_where_opt(p):

test/parser/test_parser_dml.py

@@ -283,7 +283,6 @@ def test_limit_question_mark(self):
         result = oceanbase_parser.parse("""
         SELECT * FROM `antinvoice93`.einv_base_info WHERE einv_source = ? ORDER BY gmt_create DESC LIMIT ?
         """)
-        assert result.limit == '?'
         assert result.query_body.limit == '?'
 
     def test_subquery_limit(self):
@@ -449,6 +448,14 @@ def test_double_type(self):
         result = oceanbase_parser.parse(sql)
         assert isinstance(result, Statement)
 
+    def test_union_has_order_limit(self):
+        sql = """
+        ( select 球员id from 球员夺冠次数 order by 冠军次数 asc limit 3 ) union ( select 球员id from 球员夺冠次数 order by 亚军次数 desc limit 5 )
+        """
+        sql = Utils.remove_sql_text_affects_parser(sql)
+        result = oceanbase_parser.parse(sql)
+        assert isinstance(result, Statement)
+
 
 if __name__ == '__main__':
     unittest.main()