From 953685472e7a6c705cb17401e067ebf0b2b9ea9a Mon Sep 17 00:00:00 2001 From: Michael Van de Steene Date: Tue, 30 Jun 2026 22:42:21 +0200 Subject: [PATCH 1/3] feat(sql-ast): add WINDOW_FUNCTION node + OVER rendering, UNION-in-CTE (OBSL-1027) Co-Authored-By: Claude Opus 4.8 --- soda-core/src/soda_core/common/sql_ast.py | 24 +++- soda-core/src/soda_core/common/sql_dialect.py | 21 ++++ .../unit/test_sql_ast_window_functions.py | 119 ++++++++++++++++++ 3 files changed, 161 insertions(+), 3 deletions(-) create mode 100644 soda-tests/tests/unit/test_sql_ast_window_functions.py diff --git a/soda-core/src/soda_core/common/sql_ast.py b/soda-core/src/soda_core/common/sql_ast.py index e1381caf9..27bd16df2 100644 --- a/soda-core/src/soda_core/common/sql_ast.py +++ b/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 WITH(BaseSqlExpression): 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 @@ def __post_init__(self): 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] diff --git a/soda-core/src/soda_core/common/sql_dialect.py b/soda-core/src/soda_core/common/sql_dialect.py index 57169f320..ddc95bce7 100644 --- a/soda-core/src/soda_core/common/sql_dialect.py +++ b/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)}.*" diff --git a/soda-tests/tests/unit/test_sql_ast_window_functions.py b/soda-tests/tests/unit/test_sql_ast_window_functions.py new file mode 100644 index 000000000..142c80a53 --- /dev/null +++ b/soda-tests/tests/unit/test_sql_ast_window_functions.py @@ -0,0 +1,119 @@ +from __future__ import annotations + +from soda_core.common.sql_ast import ( + COLUMN, + CTE, + FROM, + ORDER_BY_ASC, + ORDER_BY_DESC, + SELECT, + UNION, + UNION_ALL, + WINDOW_FUNCTION, + WITH, +) +from soda_core.common.sql_dialect import SqlDialect + + +def dialect() -> SqlDialect: + return SqlDialect() + + +# --------------------------------------------------------------------------- +# P1 — 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") + ) + # Build expectation from actual output once node is implemented. + # The expression renders as () OVER () AS . + 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=["*"], + 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 + + +# --------------------------------------------------------------------------- +# P2 — 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 From 389ce83aef9a123128a39773e6674fbd37af085e Mon Sep 17 00:00:00 2001 From: Michael Van de Steene Date: Wed, 1 Jul 2026 23:30:47 +0200 Subject: [PATCH 2/3] fix(test): use STAR() not "*" string in window-function reference test (OBSL-1027) Co-Authored-By: Claude Fable 5 --- soda-tests/tests/unit/test_sql_ast_window_functions.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/soda-tests/tests/unit/test_sql_ast_window_functions.py b/soda-tests/tests/unit/test_sql_ast_window_functions.py index 142c80a53..92baae727 100644 --- a/soda-tests/tests/unit/test_sql_ast_window_functions.py +++ b/soda-tests/tests/unit/test_sql_ast_window_functions.py @@ -7,6 +7,7 @@ ORDER_BY_ASC, ORDER_BY_DESC, SELECT, + STAR, UNION, UNION_ALL, WINDOW_FUNCTION, @@ -61,11 +62,11 @@ def test_window_function_partition_only(): sql = dialect().build_expression_sql( WINDOW_FUNCTION( name="COUNT", - args=["*"], + args=[STAR()], partition_by=[COLUMN("dept")], ) ) - assert 'COUNT("*") OVER (PARTITION BY "dept")' == sql + assert 'COUNT(*) OVER (PARTITION BY "dept")' == sql def test_window_function_multiple_partition_cols(): From 5ab57404b243708e9b99a402f7954b5c5e9c64bb Mon Sep 17 00:00:00 2001 From: Michael Van de Steene Date: Sun, 5 Jul 2026 01:37:59 +0200 Subject: [PATCH 3/3] docs: tighten comments and docs (OBSL-1027) Co-Authored-By: Claude Fable 5 --- soda-tests/tests/unit/test_sql_ast_window_functions.py | 10 ++-------- 1 file changed, 2 insertions(+), 8 deletions(-) diff --git a/soda-tests/tests/unit/test_sql_ast_window_functions.py b/soda-tests/tests/unit/test_sql_ast_window_functions.py index 92baae727..2cedf73c8 100644 --- a/soda-tests/tests/unit/test_sql_ast_window_functions.py +++ b/soda-tests/tests/unit/test_sql_ast_window_functions.py @@ -20,9 +20,7 @@ def dialect() -> SqlDialect: return SqlDialect() -# --------------------------------------------------------------------------- -# P1 — WINDOW_FUNCTION rendering -# --------------------------------------------------------------------------- +# WINDOW_FUNCTION rendering def test_window_function_order_by_only(): @@ -33,8 +31,6 @@ def test_window_function_order_by_only(): order_by=[ORDER_BY_DESC("value_")], ).AS("rn") ) - # Build expectation from actual output once node is implemented. - # The expression renders as () OVER () AS . assert 'ROW_NUMBER() OVER (ORDER BY "value_" DESC)' in sql @@ -81,9 +77,7 @@ def test_window_function_multiple_partition_cols(): assert 'RANK() OVER (PARTITION BY "a", "b" ORDER BY "c" DESC)' == sql -# --------------------------------------------------------------------------- -# P2 — UNION inside a CTE body -# --------------------------------------------------------------------------- +# UNION inside a CTE body def test_cte_with_union_body():