OBSL-1027: SQL AST window functions + UNION-in-CTE

soda-core/src/soda_core/common/sql_ast.py

@@ -2,7 +2,7 @@
 
 import logging
 import weakref
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from numbers import Number
 from typing import Any, Optional
 
@@ -172,12 +172,12 @@
 class CTE(BaseSqlExpression):
     alias: str
     alias_columns: list[COLUMN] | None = None
-    cte_query: list | str | VALUES | None = None
+    cte_query: list | str | VALUES | UNION | None = None
 
     def __post_init__(self):
         super().__post_init__()
 
-    def AS(self, cte_query: list | str) -> CTE:
+    def AS(self, cte_query: list | str | VALUES | UNION) -> CTE:
         self.cte_query = cte_query
         self.handle_parent_node_update(cte_query)
         return self
@@ -401,6 +401,24 @@
         self.handle_parent_node_update(self.args)
 
 
+@dataclass
+class WINDOW_FUNCTION(SqlExpression):
+    """Window function: ``name(args) OVER ([PARTITION BY cols] [ORDER BY cols])``."""
+
+    name: str
+    args: list[SqlExpression | str] = field(default_factory=list)
+    partition_by: list[COLUMN] | None = None
+    order_by: list | None = None  # elements are ORDER_BY_ASC / ORDER_BY_DESC
+
+    def __post_init__(self):
+        super().__post_init__()
+        self.handle_parent_node_update(self.args)
+        if self.partition_by:
+            self.handle_parent_node_update(self.partition_by)
+        if self.order_by:
+            self.handle_parent_node_update(self.order_by)
+
+
 @dataclass
 class COALESCE(SqlExpression):
     args: list[SqlExpression | str]

soda-core/src/soda_core/common/sql_dialect.py

@@ -97,6 +97,7 @@
     VALUES,
     VALUES_ROW,
     WHERE,
+    WINDOW_FUNCTION,
     WITH,
     Operator,
     SqlExpression,
@@ -756,6 +757,8 @@ def _build_cte(self, cte: CTE) -> str:
             sql_str: str = self.build_select_sql(cte.cte_query)
         elif isinstance(cte.cte_query, VALUES):
             sql_str: str = self.build_cte_values_sql(values=cte.cte_query, alias_columns=cte.alias_columns)
+        elif isinstance(cte.cte_query, UNION):
+            sql_str: str = self.build_union_sql(cte.cte_query, add_semicolon=False)
         elif isinstance(cte.cte_query, str):
             sql_str: str = indent(cte.cte_query, "  ").strip()
         else:
@@ -842,6 +845,8 @@ def build_expression_sql(self, expression: SqlExpression | str | Number) -> str:
             return self._build_coalesce_sql(expression)
         elif isinstance(expression, CAST):
             return self._build_cast_sql(expression)
+        elif isinstance(expression, WINDOW_FUNCTION):
+            return self._build_window_function_sql(expression)
         elif isinstance(expression, FUNCTION):
             return self._build_function_sql(expression)
         elif isinstance(expression, DISTINCT):
@@ -1037,6 +1042,22 @@ def _build_function_sql(self, function: FUNCTION) -> str:
             args_list_sql: str = ", ".join(args_sqls)
             return f"{function.name}({args_list_sql})"
 
+    def _build_window_function_sql(self, wf: WINDOW_FUNCTION) -> str:
+        args: list[SqlExpression | str] = wf.args if isinstance(wf.args, list) else [wf.args]
+        args_list_sql: str = ", ".join(self.build_expression_sql(arg) for arg in args)
+        over_clauses: list[str] = []
+        if wf.partition_by:
+            partition_sqls = ", ".join(self._build_column_sql(col) for col in wf.partition_by)
+            over_clauses.append(f"PARTITION BY {partition_sqls}")
+        if wf.order_by:
+            order_by_parts: list[str] = []
+            for ob in wf.order_by:
+                direction: str = " ASC" if isinstance(ob, ORDER_BY_ASC) else " DESC"
+                order_by_parts.append(f"{self.build_expression_sql(ob.expression)}{direction}")
+            over_clauses.append(f"ORDER BY {', '.join(order_by_parts)}")
+        over_sql: str = " ".join(over_clauses)
+        return f"{wf.name}({args_list_sql}) OVER ({over_sql})"
+
     def _build_star_sql(self, star: STAR) -> str:
         if star.alias:
             return f"{self.quote_default(star.alias)}.*"

soda-tests/tests/unit/test_sql_ast_window_functions.py

@@ -0,0 +1,114 @@
+from __future__ import annotations
+
+from soda_core.common.sql_ast import (
+    COLUMN,
+    CTE,
+    FROM,
+    ORDER_BY_ASC,
+    ORDER_BY_DESC,
+    SELECT,
+    STAR,
+    UNION,
+    UNION_ALL,
+    WINDOW_FUNCTION,
+    WITH,
+)
+from soda_core.common.sql_dialect import SqlDialect
+
+
+def dialect() -> SqlDialect:
+    return SqlDialect()
+
+
+# WINDOW_FUNCTION rendering
+
+
+def test_window_function_order_by_only():
+    """ROW_NUMBER() OVER (ORDER BY value_ DESC)"""
+    sql = dialect().build_expression_sql(
+        WINDOW_FUNCTION(
+            name="ROW_NUMBER",
+            order_by=[ORDER_BY_DESC("value_")],
+        ).AS("rn")
+    )
+    assert 'ROW_NUMBER() OVER (ORDER BY "value_" DESC)' in sql
+
+
+def test_window_function_partition_and_order_by():
+    """SUM(amount) OVER (PARTITION BY category ORDER BY id ASC)"""
+    sql = dialect().build_expression_sql(
+        WINDOW_FUNCTION(
+            name="SUM",
+            args=["amount"],
+            partition_by=[COLUMN("category")],
+            order_by=[ORDER_BY_ASC("id")],
+        )
+    )
+    assert 'SUM("amount") OVER (PARTITION BY "category" ORDER BY "id" ASC)' == sql
+
+
+def test_window_function_no_partition_no_order_by():
+    """ROW_NUMBER() OVER ()"""
+    sql = dialect().build_expression_sql(WINDOW_FUNCTION(name="ROW_NUMBER"))
+    assert sql == "ROW_NUMBER() OVER ()"
+
+
+def test_window_function_partition_only():
+    """COUNT(*) OVER (PARTITION BY dept)"""
+    sql = dialect().build_expression_sql(
+        WINDOW_FUNCTION(
+            name="COUNT",
+            args=[STAR()],
+            partition_by=[COLUMN("dept")],
+        )
+    )
+    assert 'COUNT(*) OVER (PARTITION BY "dept")' == sql
+
+
+def test_window_function_multiple_partition_cols():
+    """RANK() OVER (PARTITION BY a, b ORDER BY c DESC)"""
+    sql = dialect().build_expression_sql(
+        WINDOW_FUNCTION(
+            name="RANK",
+            partition_by=[COLUMN("a"), COLUMN("b")],
+            order_by=[ORDER_BY_DESC("c")],
+        )
+    )
+    assert 'RANK() OVER (PARTITION BY "a", "b" ORDER BY "c" DESC)' == sql
+
+
+# UNION inside a CTE body
+
+
+def test_cte_with_union_body():
+    """WITH cte AS (SELECT ... UNION SELECT ...) SELECT * FROM cte"""
+    d = dialect()
+    select_a = [SELECT(COLUMN("id")), FROM("table_a")]
+    select_b = [SELECT(COLUMN("id")), FROM("table_b")]
+    sql = d.build_select_sql(
+        [
+            WITH([CTE("combined").AS(UNION([select_a, select_b]))]),
+            SELECT(COLUMN("id")),
+            FROM("combined"),
+        ]
+    )
+    assert '"combined" AS (' in sql
+    assert "UNION" in sql
+    assert 'FROM "table_a"' in sql
+    assert 'FROM "table_b"' in sql
+
+
+def test_cte_with_union_all_body():
+    """WITH cte AS (SELECT ... UNION ALL SELECT ...) SELECT * FROM cte"""
+    d = dialect()
+    select_a = [SELECT(COLUMN("val")), FROM("src_a")]
+    select_b = [SELECT(COLUMN("val")), FROM("src_b")]
+    sql = d.build_select_sql(
+        [
+            WITH([CTE("all_vals").AS(UNION_ALL([select_a, select_b]))]),
+            SELECT(COLUMN("val")),
+            FROM("all_vals"),
+        ]
+    )
+    assert '"all_vals" AS (' in sql
+    assert "UNION ALL" in sql