Add post-processing SQL over semantic query results (#135)

* 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.

* Preserve root semantic rewrite when JOIN has subquery

When the root FROM references a semantic model and a JOIN contains a
subquery, the routing into _rewrite_with_ctes_or_subqueries must still
apply _rewrite_simple_query to the root SELECT so the explicit JOIN
guard and semantic rewriting are not bypassed.

* Keep CTE scope in root semantic rewrite and merge post_process CTEs

Fix two issues from PR review:

1. When the root SELECT references a semantic model and has user-defined
   CTEs, _rewrite_simple_query was discarding them. Now user CTEs are
   saved before rewriting and merged back into the generated SQL so
   filter references like IN (SELECT ... FROM user_cte) remain valid.

2. When post_process SQL has its own WITH clause and the inner semantic
   query also produces CTEs, the two WITH clauses are now merged into
   one instead of producing invalid double-WITH SQL.

* Preserve WITH RECURSIVE when merging root CTEs

Propagate the recursive flag from user-defined CTEs to the merged WITH
clause so self-referencing CTEs execute correctly.

* Handle CTE name collisions between user and generated CTEs

Two changes:

1. post_process path: remove CTE hoisting entirely. Inner SQL (with
   CTEs) is placed directly in the subquery position. CTEs inside
   subqueries are valid in all target databases and naturally scoped,
   so name collisions with post_process CTEs cannot occur.

2. Root semantic + user CTEs: detect name collisions between user CTEs
   and generated CTEs, raising a clear error instead of producing
   invalid SQL. Walk-based renaming was too aggressive (renamed user
   CTE references inside filter subqueries).

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,24 @@ 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])
+
+            # Inner SQL (including any CTEs) is placed directly in the
+            # subquery position. CTEs inside subqueries are valid SQL in
+            # all target databases and naturally scoped, avoiding name
+            # collisions with CTEs in the post_process SQL.
+            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,89 @@ 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)
+
+        # If the root SELECT itself references a semantic model, it must
+        # still go through _rewrite_simple_query (which enforces the
+        # explicit JOIN guard and performs semantic rewriting).
+        if self._references_semantic_model(parsed):
+            # Save user-defined CTEs before _rewrite_simple_query replaces
+            # the entire query with fresh generator output.
+            original_with = parsed.args.get("with")
+
+            rewritten_sql = self._rewrite_simple_query(parsed)
+
+            if original_with:
+                # Merge user CTEs into the generated SQL so references
+                # from filters/expressions (e.g. IN (SELECT ... FROM cte))
+                # remain valid.
+                rewritten = sqlglot.parse_one(rewritten_sql, dialect=self.dialect)
+                gen_with = rewritten.args.get("with")
+                if gen_with:
+                    # Check for CTE name collisions between user and generated CTEs
+                    user_names = {cte.alias for cte in original_with.expressions}
+                    for gen_cte in gen_with.expressions:
+                        if gen_cte.alias in user_names:
+                            raise ValueError(
+                                f"CTE name '{gen_cte.alias}' conflicts with an internally "
+                                f"generated name. Please choose a different CTE name."
+                            )
+
+                    user_ctes = [cte.copy() for cte in original_with.expressions]
+                    gen_with.set("expressions", user_ctes + list(gen_with.expressions))
+                    # Preserve WITH RECURSIVE from the original query
+                    if original_with.args.get("recursive"):
+                        gen_with.set("recursive", True)
+                else:
+                    rewritten.set("with", original_with.copy())
+                return rewritten.sql(dialect=self.dialect)
+
+            return rewritten_sql
+
+        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."""