Add post-processing SQL over semantic query results

Support arbitrary SQL (CASE, window functions, arithmetic, etc.) on top
of semantic query results via subquery wrapping. The rewriter now walks
the query tree recursively so nested subqueries and JOIN subqueries that
reference semantic models are compiled correctly.

Also adds a post_process parameter to compile() and query() for the
Python API path, with automatic CTE hoisting.

Author: nicosuave

sidemantic/core/semantic_layer.py

@@ -435,6 +435,7 @@ def query(
         ungrouped: bool = False,
         parameters: dict[str, any] | None = None,
         use_preaggregations: bool | None = None,
+        post_process: str | None = None,
     ):
         """Execute a query against the semantic layer.
 
@@ -448,6 +449,9 @@ def query(
             ungrouped: If True, return raw rows without aggregation (no GROUP BY)
             parameters: Template parameters for Jinja2 rendering
             use_preaggregations: Override pre-aggregation routing setting for this query
+            post_process: Optional SQL to wrap around the semantic query result.
+                Use {inner} as a placeholder for the compiled semantic query, e.g.:
+                "SELECT *, revenue / count AS avg_value FROM ({inner})"
 
         Returns:
             DuckDB relation object (can convert to DataFrame with .df() or .to_df())
@@ -462,6 +466,7 @@ def query(
             ungrouped=ungrouped,
             parameters=parameters,
             use_preaggregations=use_preaggregations,
+            post_process=post_process,
         )
 
         return self.adapter.execute(sql)
@@ -479,6 +484,7 @@ def compile(
         ungrouped: bool = False,
         parameters: dict[str, any] | None = None,
         use_preaggregations: bool | None = None,
+        post_process: str | None = None,
     ) -> str:
         """Compile a query to SQL without executing.
 
@@ -493,6 +499,9 @@ def compile(
             dialect: SQL dialect override (defaults to layer's dialect)
             ungrouped: If True, return raw rows without aggregation (no GROUP BY)
             use_preaggregations: Override pre-aggregation routing setting for this query
+            post_process: Optional SQL to wrap around the semantic query result.
+                Use {inner} as a placeholder for the compiled semantic query, e.g.:
+                "SELECT *, revenue / count AS avg_value FROM ({inner})"
 
         Returns:
             SQL query string
@@ -520,7 +529,7 @@ def compile(
             preagg_schema=self.preagg_schema,
         )
 
-        return generator.generate(
+        inner_sql = generator.generate(
             metrics=metrics,
             dimensions=dimensions,
             filters=filters,
@@ -533,6 +542,34 @@ def compile(
             use_preaggregations=use_preaggs,
         )
 
+        if post_process is not None:
+            if "{inner}" not in post_process:
+                raise ValueError("post_process must contain a {inner} placeholder")
+
+            # Strip sidemantic instrumentation comment
+            stripped = inner_sql.rstrip()
+            last_line = stripped.split("\n")[-1].strip()
+            if last_line.startswith("-- sidemantic:"):
+                stripped = "\n".join(stripped.split("\n")[:-1])
+
+            # If inner SQL starts with WITH (CTEs), hoist them outside
+            # the subquery position so the SQL is valid.
+            if stripped.lstrip().upper().startswith("WITH "):
+                import sqlglot
+
+                target_dialect = dialect or self.dialect
+                parsed_inner = sqlglot.parse_one(stripped, dialect=target_dialect)
+                with_clause = parsed_inner.args.get("with")
+                if with_clause:
+                    parsed_inner.set("with", None)
+                    body = parsed_inner.sql(dialect=target_dialect)
+                    ctes = with_clause.sql(dialect=target_dialect)
+                    return ctes + "\n" + post_process.replace("{inner}", body)
+
+            return post_process.replace("{inner}", stripped)
+
+        return inner_sql
+
     def explain(
         self,
         metrics: list[str] | None = None,

sidemantic/sql/query_rewriter.py

@@ -118,8 +118,9 @@ def rewrite(self, sql: str, strict: bool = True) -> str:
         # Check if this is a CTE-based query or has subqueries
         has_ctes = parsed.args.get("with") is not None
         has_subquery_in_from = self._has_subquery_in_from(parsed)
+        has_subquery_in_joins = any(isinstance(join.this, exp.Subquery) for join in (parsed.args.get("joins") or []))
 
-        if has_ctes or has_subquery_in_from:
+        if has_ctes or has_subquery_in_from or has_subquery_in_joins:
             # Handle CTEs and subqueries
             return self._rewrite_with_ctes_or_subqueries(parsed)
 
@@ -1851,41 +1852,51 @@ def _has_subquery_in_from(self, select: exp.Select) -> bool:
     def _rewrite_with_ctes_or_subqueries(self, parsed: exp.Select) -> str:
         """Rewrite query that contains CTEs or subqueries.
 
-        Strategy:
-        1. Rewrite each CTE that references semantic models
-        2. Rewrite subqueries in FROM clause
-        3. Return the modified SQL
+        Recursively walks the query tree bottom-up, rewriting any
+        SELECT whose FROM target resolves to a semantic model.
+        Outer queries are left as plain SQL, so post-processing
+        (CASE, window functions, arithmetic, etc.) works naturally.
         """
-        # Handle CTEs
-        if parsed.args.get("with"):
-            with_clause = parsed.args["with"]
-            for cte in with_clause.expressions:
-                # Each CTE has a name (alias) and a query (this)
+        self._rewrite_select_tree(parsed)
+        return parsed.sql(dialect=self.dialect)
+
+    def _rewrite_select_tree(self, select: exp.Select):
+        """Recursively rewrite semantic subqueries and CTEs (bottom-up).
+
+        At each level: recurse into children first, then rewrite this
+        node if it directly references a semantic model.
+        """
+        # Recurse into CTEs
+        if select.args.get("with"):
+            for cte in select.args["with"].expressions:
                 cte_query = cte.this
                 if isinstance(cte_query, exp.Select):
-                    # Check if this CTE references a semantic model
+                    self._rewrite_select_tree(cte_query)
                     if self._references_semantic_model(cte_query):
-                        # Rewrite the CTE query
-                        rewritten_cte_sql = self._rewrite_simple_query(cte_query)
-                        # Parse the rewritten SQL and replace the CTE query
-                        rewritten_cte = sqlglot.parse_one(rewritten_cte_sql, dialect=self.dialect)
-                        cte.set("this", rewritten_cte)
-
-        # Handle subquery in FROM
-        from_clause = parsed.args.get("from")
+                        rewritten_sql = self._rewrite_simple_query(cte_query)
+                        cte.set("this", sqlglot.parse_one(rewritten_sql, dialect=self.dialect))
+
+        # Recurse into FROM subquery
+        from_clause = select.args.get("from")
         if from_clause and isinstance(from_clause.this, exp.Subquery):
             subquery = from_clause.this
             subquery_select = subquery.this
-            if isinstance(subquery_select, exp.Select) and self._references_semantic_model(subquery_select):
-                # Rewrite the subquery
-                rewritten_subquery_sql = self._rewrite_simple_query(subquery_select)
-                rewritten_subquery = sqlglot.parse_one(rewritten_subquery_sql, dialect=self.dialect)
-                subquery.set("this", rewritten_subquery)
-
-        # Return the modified SQL
-        # Note: Individual CTEs/subqueries are already instrumented by _rewrite_simple_query -> generator
-        # The outer query wrapper doesn't need separate instrumentation
-        return parsed.sql(dialect=self.dialect)
+            if isinstance(subquery_select, exp.Select):
+                self._rewrite_select_tree(subquery_select)
+                if self._references_semantic_model(subquery_select):
+                    rewritten_sql = self._rewrite_simple_query(subquery_select)
+                    subquery.set("this", sqlglot.parse_one(rewritten_sql, dialect=self.dialect))
+
+        # Recurse into JOIN subqueries
+        for join in select.args.get("joins") or []:
+            join_expr = join.this
+            if isinstance(join_expr, exp.Subquery):
+                join_select = join_expr.this
+                if isinstance(join_select, exp.Select):
+                    self._rewrite_select_tree(join_select)
+                    if self._references_semantic_model(join_select):
+                        rewritten_sql = self._rewrite_simple_query(join_select)
+                        join_expr.set("this", sqlglot.parse_one(rewritten_sql, dialect=self.dialect))
 
     def _references_semantic_model(self, select: exp.Select) -> bool:
         """Check if a SELECT statement references any semantic models."""