add format for sound_like window_function

-s

Author: qidi1

sqlgpt_parser/format/formatter.py

@@ -421,24 +421,31 @@ def visit_list_expression(self, node, unmangle_names):
         else:
             return "(%s)" % self._join_expressions(node.values, unmangle_names)
 
+    def visit_window_func(self, node, unmangle_names):
+        args = ", ".join([self.process(arg, unmangle_names) for arg in node.func_args])
+        ignore_null = f" {node.ignore_null} NULLS" if node.ignore_null else ""
+        window_spec = " OVER (" + self.process(node.window_spec, unmangle_names) + ")"
+        return f"{node.func_name.upper()}({args}){ignore_null}{window_spec}"
+
     def visit_window_spec(self, node, unmangle_names):
         parts = []
-
         if node.partition_by:
-            parts.append(
-                "PARTITION BY "
-                + self._join_expressions(node.partition_by, unmangle_names)
-            )
+            self.process(node.partition_by, unmangle_names)
         if node.order_by:
             parts.append("ORDER BY " + format_sort_items(node.order_by, unmangle_names))
-        if node.frame:
-            parts.append(self.process(node.frame, unmangle_names))
+        if node.frame_clause:
+            parts.append(self.process(node.frame_clause, unmangle_names))
 
-        return '(' + ' '.join(parts) + ')'
+        return ' '.join(parts)
 
-    def visit_window_frame(self, node, unmangle_names):
-        ret = node.type + " "
+    def visit_partition_by_clause(self, node, unmangle_names):
+        return "PARTITION BY " + self._join_expressions(node.items, unmangle_names)
+
+    def visit_frame_clause(self, node, unmangle_names):
+        return f"{node.type} {self.process(node.frame_range, unmangle_names)}"
 
+    def visit_window_frame(self, node, unmangle_names):
+        ret = ""
         if node.end:
             ret += "BETWEEN %s AND %s" % (
                 self.process(node.start, unmangle_names),
@@ -449,6 +456,19 @@ def visit_window_frame(self, node, unmangle_names):
 
         return ret
 
+    def visit_frame_bound(self, node, unmangle_names):
+        if node.type.upper() == "ROW":
+            return "CURRENT ROW"
+        expr = (
+            self.process(node.expr, unmangle_names)
+            if node.expr is not None
+            else "UNBOUNDED "
+        )
+        return f"{expr} {node.type.upper()}"
+
+    def visit_frame_expr(self, node, unmangle_names):
+        return self.process(node.value, unmangle_names)
+
     def visit_single_column(self, node, indent):
         format_expression(node.expression)
 
@@ -468,6 +488,9 @@ def visit_match_against_expression(self, node, unmangle_names):
         full_text_search_modifier = full_text_search_modifier.upper()
         return f"MATCH({columns}) AGAINST ({self.process(node.expr, unmangle_names)}{full_text_search_modifier})"
 
+    def visit_sound_like(self, node, unmangle_names):
+        return f"{self.process(node.arguments[0])} SOUNDS LIKE {self.process(node.arguments[1])}"
+
     def _format_binary_expression(self, operator, left, right, unmangle_names):
         return "%s %s %s" % (
             self.process(left, unmangle_names),
@@ -689,13 +712,14 @@ def visit_table_subquery(self, node, indent):
         return None
 
     def visit_union(self, node, indent):
-        all = node.all
         for i, relation in enumerate(node.relations):
             self._process_relation(relation, indent)
             self.builder.append("\n")
             if i != len(node.relations) - 1:
-                if all:
+                if node.all:
                     self._append(indent, "UNION ALL")
+                elif node.distinct:
+                    self._append(indent, "UNION DISTINCT")
                 else:
                     self._append(indent, "UNION")
                 self.builder.append("\n")
@@ -704,7 +728,12 @@ def visit_union(self, node, indent):
 
     def visit_except(self, node, indent):
         self._process_relation(node.left, indent)
-        self.builder.append("EXCEPT " + "ALL " if not node.distinct else "")
+        if node.all is not None:
+            self._append(indent, "EXCEPT ALL")
+        elif node.distinct is not None:
+            self._append(indent, "EXCEPT DISTINCT")
+        else:
+            self._append(indent, "EXCEPT")
         self._process_relation(node.right, indent)
 
         return None
@@ -756,7 +785,11 @@ def visit_intersect(self, node, indent):
         relations = [
             self._process_relation(relation, indent) for relation in node.relations
         ]
-        intersect = "INTERSECT " + "ALL " if not node.distinct else ""
+        intersect = "INTERSECT"
+        if node.all is not None:
+            intersect += " ALL"
+        elif node.distinct is not None:
+            intersect += " DISTINCT"
         self.builder.append(intersect.join(relations))
         return None
 

sqlgpt_parser/parser/mysql_parser/parser.py

@@ -1570,7 +1570,7 @@ def p_window_func_call(p):
     | ROW_NUMBER LPAREN RPAREN over_clause
     """
     length = len(p)
-    window_spec = p[-1]
+    window_spec = p[length-1]
     args = []
     ignore_null = None
 
@@ -1711,7 +1711,10 @@ def p_frame_start(p):
     | frame_expr PRECEDING
     | frame_expr FOLLOWING
     """
-    p[0] = FrameBound(p.lineno(1), p.lexpos(1), type=p[2], expr=p[1])
+    if p.slice[1].type == 'frame_expr':
+        p[0] = FrameBound(p.lineno(1), p.lexpos(1), type=p[2], expr=p[1])
+    else:
+        p[0] = FrameBound(p.lineno(1), p.lexpos(1), type=p[2], expr=None)
 
 
 def p_frame_end(p):
@@ -1726,13 +1729,8 @@ def p_frame_between(p):
 
 def p_frame_expr(p):
     r"""frame_expr : figure
-    | QM
-    | INTERVAL expression time_unit
-    |"""
-    if len(p) == 4:
-        p[0] = FrameExpr(p.lineno(1), p.lexpos(1), value=p[2], unit=p[3])
-    else:
-        p[0] = FrameExpr(p.lineno(1), p.lexpos(1), value=p[1])
+    | time_interval"""
+    p[0] = FrameExpr(p.lineno(1), p.lexpos(1), value=p[1])
 
 
 def p_lead_lag_info_opt(p):

sqlgpt_parser/parser/mysql_parser/parser_table.py

@@ -1649,7 +1649,7 @@ def p_window_func_call(p):
     | ROW_NUMBER LPAREN RPAREN over_clause
     """
     length = len(p)
-    window_spec = p[-1]
+    window_spec = p[length-1]
     args = []
     ignore_null = None
 
@@ -1790,7 +1790,10 @@ def p_frame_start(p):
     | frame_expr PRECEDING
     | frame_expr FOLLOWING
     """
-    p[0] = FrameBound(p.lineno(1), p.lexpos(1), type=p[2], expr=p[1])
+    if p.slice[1].type == 'frame_expr':
+        p[0] = FrameBound(p.lineno(1), p.lexpos(1), type=p[2], expr=p[1])
+    else:
+        p[0] = FrameBound(p.lineno(1), p.lexpos(1), type=p[2], expr=None)
 
 
 def p_frame_end(p):
@@ -1802,12 +1805,9 @@ def p_frame_between(p):
     r"""frame_between : BETWEEN frame_start AND frame_end"""
     p[0] = WindowFrame(p.lineno(1), p.lexpos(1), start=p[2], end=p[4])
 
-
 def p_frame_expr(p):
     r"""frame_expr : figure
-    | QM
-    | time_interval
-    |"""
+    | time_interval"""
     p[0] = FrameExpr(p.lineno(1), p.lexpos(1), value=p[1])
 
 

sqlgpt_parser/parser/oceanbase_parser/parser.py

@@ -1657,7 +1657,7 @@ def p_window_func_call(p):
     | ROW_NUMBER LPAREN RPAREN over_clause
     """
     length = len(p)
-    window_spec = p[-1]
+    window_spec = p[length-1]
     args = []
     ignore_null = None
 
@@ -1798,8 +1798,10 @@ def p_frame_start(p):
     | frame_expr PRECEDING
     | frame_expr FOLLOWING
     """
-    p[0] = FrameBound(p.lineno(1), p.lexpos(1), type=p[2], expr=p[1])
-
+    if p.slice[1].type == 'frame_expr':
+        p[0] = FrameBound(p.lineno(1), p.lexpos(1), type=p[2], expr=p[1])
+    else:
+        p[0] = FrameBound(p.lineno(1), p.lexpos(1), type=p[2], expr=None)
 
 def p_frame_end(p):
     r"""frame_end : frame_start"""
@@ -1810,12 +1812,9 @@ def p_frame_between(p):
     r"""frame_between : BETWEEN frame_start AND frame_end"""
     p[0] = WindowFrame(p.lineno(1), p.lexpos(1), start=p[2], end=p[4])
 
-
 def p_frame_expr(p):
     r"""frame_expr : figure
-    | QM
-    | time_interval
-    |"""
+    | time_interval"""
     p[0] = FrameExpr(p.lineno(1), p.lexpos(1), value=p[1])
 
 

sqlgpt_parser/parser/oceanbase_parser/parser_table.py

@@ -30,7 +30,7 @@ def __init__(
         self.window_spec = window_spec
 
     def accept(self, visitor, context):
-        return super().accept(visitor, context)
+        return visitor.visit_window_func(self, context)
 
 
 class WindowSpec(Node):
@@ -64,9 +64,8 @@ def accept(self, visitor, context):
 
 
 class WindowFrame(Node):
-    def __init__(self, line=None, pos=None, type=None, start=None, end=None):
+    def __init__(self, line=None, pos=None, start=None, end=None):
         super(WindowFrame, self).__init__(line, pos)
-        self.type = type
         self.start = start
         self.end = end
 
@@ -85,7 +84,9 @@ def accept(self, visitor, context):
 
 
 class FrameExpr(Node):
-    def __init__(self, line=None, pos=None, value=None, unit=None):
+    def __init__(self, line=None, pos=None, value=None):
         super(FrameExpr, self).__init__(line, pos)
         self.value = value
-        self.unit = unit
+
+    def accept(self, visitor, context):
+        return visitor.visit_frame_expr(self, context)

sqlgpt_parser/parser/odps_parser/parser.py

@@ -197,6 +197,23 @@ def test_time_interval(self):
             after_sql_rewrite_format = format_sql(statement, 0)
             assert after_sql_rewrite_format == except_sql
 
+    def test_windows_function(self):
+        test_sqls_except = {
+            "SELECT  first_value(value) OVER (PARTITION BY id ORDER BY date ROWS BETWEEN UNBOUNDED PRECEDING AND "
+            "UNBOUNDED FOLLOWING) AS first_val FROM  my_table": "SELECT\n"
+            "  FIRST_VALUE(value) OVER (ORDER BY date ASC ROWS BETWEEN UNBOUNDED  PRECEDING AND UNBOUNDED  FOLLOWING) AS first_val"
+            "\nFROM\n  my_table",
+            "SELECT first_value(value) OVER (PARTITION BY id ORDER BY date RANGE UNBOUNDED PRECEDING) "
+            "AS first_val FROM my_table": "SELECT"
+            "\n  FIRST_VALUE(value) OVER (ORDER BY date ASC RANGE UNBOUNDED  PRECEDING) AS first_val"
+            "\nFROM"
+            "\n  my_table",
+        }
+        for sql, except_sql in test_sqls_except.items():
+            statement = parser.parse(sql)
+            after_sql_rewrite_format = format_sql(statement, 0)
+            assert after_sql_rewrite_format == except_sql
+
 
 if __name__ == '__main__':
     unittest.main()

sqlgpt_parser/parser/odps_parser/parser_table.py

@@ -250,7 +250,7 @@ def test_with_operation(self):
                 result = oceanbase_parser.parse(sql)
                 assert isinstance(result, WithHasQuery)
 
-    def test_windows_func(self):
+    def test_window_func(self):
         test_sqls = [
             """
         SELECT