feat(wren): CTE-based SQL planning with per-model expansion

wren/pyproject.toml

@@ -46,7 +46,7 @@ redshift    = ["redshift_connector"]
 spark       = ["pyspark>=3.5"]
 athena      = ["ibis-framework[athena]"]
 oracle      = ["ibis-framework[oracle]", "oracledb"]
-all         = ["wren[postgres,mysql,bigquery,snowflake,clickhouse,trino,mssql,databricks,redshift,athena,oracle]"]
+all         = ["wren-engine[postgres,mysql,bigquery,snowflake,clickhouse,trino,mssql,databricks,redshift,athena,oracle]"]
 dev         = ["pytest>=8", "ruff>=0.4", "orjson>=3", "testcontainers[postgres]>=4"]
 
 [project.scripts]

wren/src/wren/cli.py

@@ -147,13 +147,13 @@ def dry_run(
         engine.close()
 
 
-@app.command()
-def transpile(
+@app.command(name="dry-plan")
+def dry_plan(
     sql: SqlArg,
     datasource: DatasourceOpt,
     mdl: MdlOpt,
 ):
-    """Transform SQL through MDL and emit the target dialect SQL (no DB required)."""
+    """Plan SQL through MDL and print the expanded SQL (no DB required)."""
     from wren.engine import WrenEngine  # noqa: PLC0415
     from wren.model.data_source import DataSource  # noqa: PLC0415
 
@@ -164,10 +164,9 @@ def transpile(
         typer.echo(f"Error: unknown datasource '{datasource}'", err=True)
         raise typer.Exit(1)
 
-    # For transpile we don't need real connection_info — pass a dummy dict
     engine = WrenEngine(manifest_str=manifest_str, data_source=ds, connection_info={})
     try:
-        result = engine.transpile(sql)
+        result = engine.dry_plan(sql)
         typer.echo(result)
     except Exception as e:
         typer.echo(f"Error: {e}", err=True)

wren/src/wren/connector/duckdb.py

@@ -75,7 +75,7 @@ def __init__(self, connection_info):
     def query(self, sql: str, limit: int | None = None) -> pa.Table:
         if limit is not None:
             sql = f"SELECT * FROM ({sql}) AS _q LIMIT {int(limit)}"
-        return self.connection.execute(sql).fetch_arrow_table()
+        return self.connection.execute(sql).to_arrow_table()
 
     def dry_run(self, sql: str) -> None:
         self.connection.execute(f"EXPLAIN {sql}")

wren/src/wren/engine.py

@@ -11,8 +11,8 @@
         connection_info={"host": "localhost", "port": 5432, ...},
     )
 
-    # Transform only (no DB required)
-    planned_sql = engine.transpile("SELECT * FROM orders")
+    # Plan only (no DB required)
+    planned_sql = engine.dry_plan("SELECT * FROM orders")
 
     # Execute against the data source
     arrow_table = engine.query("SELECT * FROM orders", limit=100)
@@ -23,25 +23,12 @@
 from typing import Any
 
 import pyarrow as pa
-import sqlglot
 
 from wren.connector.factory import get_connector
 from wren.mdl import get_manifest_extractor, get_session_context, to_json_base64
+from wren.mdl.cte_rewriter import CTERewriter
 from wren.model.data_source import DataSource
-from wren.model.error import DIALECT_SQL, PLANNED_SQL, ErrorCode, ErrorPhase, WrenError
-
-
-def _get_write_dialect(data_source: DataSource) -> str:
-    if data_source == DataSource.canner:
-        return "trino"
-    if data_source in {
-        DataSource.local_file,
-        DataSource.s3_file,
-        DataSource.minio_file,
-        DataSource.gcs_file,
-    }:
-        return "duckdb"
-    return data_source.name
+from wren.model.error import DIALECT_SQL, ErrorCode, ErrorPhase, WrenError
 
 
 class WrenEngine:
@@ -66,13 +53,16 @@ def __init__(
         data_source: DataSource | str,
         connection_info: dict[str, Any] | object,
         function_path: str | None = None,
+        *,
+        fallback: bool = True,
     ):
         if isinstance(data_source, str):
             data_source = DataSource(data_source)
 
         self.manifest_str = manifest_str
         self.data_source = data_source
         self.function_path = function_path
+        self._fallback = fallback
 
         # Build typed ConnectionInfo if a raw dict was given.
         # An empty dict is allowed for transpile-only usage (no DB connection).
@@ -87,16 +77,8 @@ def __init__(
     # SQL transformation (no DB access)
     # ------------------------------------------------------------------
 
-    def transpile(self, sql: str, properties: dict | None = None) -> str:
-        """Transform SQL through MDL and transpile to the target dialect.
-
-        Returns the dialect SQL string without executing it.
-        """
-        planned = self._plan(sql, properties)
-        return self._transpile(planned)
-
     def dry_plan(self, sql: str, properties: dict | None = None) -> str:
-        """Return the wren-core planned SQL (DataFusion dialect, before transpile)."""
+        """Plan SQL through MDL and return the expanded SQL in the target dialect."""
         return self._plan(sql, properties)
 
     # ------------------------------------------------------------------
@@ -110,7 +92,7 @@ def query(
         properties: dict | None = None,
     ) -> pa.Table:
         """Transpile and execute SQL, return results as an Arrow table."""
-        dialect_sql = self.transpile(sql, properties)
+        dialect_sql = self.dry_plan(sql, properties)
         connector = self._get_connector()
         try:
             return connector.query(dialect_sql, limit)
@@ -126,7 +108,7 @@ def query(
 
     def dry_run(self, sql: str, properties: dict | None = None) -> None:
         """Transpile and dry-run SQL without returning results."""
-        dialect_sql = self.transpile(sql, properties)
+        dialect_sql = self.dry_plan(sql, properties)
         connector = self._get_connector()
         try:
             connector.dry_run(dialect_sql)
@@ -180,7 +162,13 @@ def _plan(self, sql: str, properties: dict | None) -> str:
                 processed,
                 self.data_source.name,
             )
-            return session.transform_sql(sql)
+            rewriter = CTERewriter(
+                effective_manifest,
+                session,
+                self.data_source,
+                fallback=self._fallback,
+            )
+            return rewriter.rewrite(sql)
         except Exception as e:
             raise WrenError(
                 ErrorCode.INVALID_SQL,
@@ -189,18 +177,6 @@ def _plan(self, sql: str, properties: dict | None) -> str:
                 metadata={DIALECT_SQL: sql},
             ) from e
 
-    def _transpile(self, planned_sql: str) -> str:
-        try:
-            write = _get_write_dialect(self.data_source)
-            return sqlglot.transpile(planned_sql, read="duckdb", write=write)[0]
-        except Exception as e:
-            raise WrenError(
-                ErrorCode.SQLGLOT_ERROR,
-                str(e),
-                phase=ErrorPhase.SQL_TRANSPILE,
-                metadata={PLANNED_SQL: planned_sql},
-            ) from e
-
     def _get_connector(self):
         if self._connector is None:
             self._connector = get_connector(self.data_source, self.connection_info)

wren/src/wren/mdl/cte_rewriter.py

@@ -0,0 +1,226 @@
+"""CTE-based SQL rewriter.
+
+Parses user SQL with sqlglot, uses ``qualify_columns`` to fully resolve
+all column references, then calls wren-core
+``transform_sql`` once per model with a simple ``SELECT col1, col2 FROM
+model`` and injects each expanded result as a CTE into the original query.
+"""
+
+from __future__ import annotations
+
+import base64
+import json
+from collections import defaultdict
+
+from sqlglot import exp, parse_one
+from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
+from sqlglot.optimizer.qualify_columns import qualify_columns
+from sqlglot.optimizer.qualify_tables import qualify_tables
+from sqlglot.schema import MappingSchema
+
+# Ensure the Wren dialect is registered with sqlglot on import.
+import wren.mdl.wren_dialect as _wren_dialect  # noqa: F401
+from wren.model.data_source import DataSource
+
+_SQLGLOT_DIALECT_MAP: dict[DataSource, str] = {
+    DataSource.canner: "trino",
+    DataSource.mssql: "tsql",
+    DataSource.local_file: "duckdb",
+    DataSource.s3_file: "duckdb",
+    DataSource.minio_file: "duckdb",
+    DataSource.gcs_file: "duckdb",
+}
+
+
+def get_sqlglot_dialect(data_source: DataSource) -> str:
+    """Map a DataSource to a valid sqlglot dialect name."""
+    return _SQLGLOT_DIALECT_MAP.get(data_source, data_source.name)
+
+
+class CTERewriter:
+    """Rewrite user SQL by expanding MDL model references into CTEs.
+
+    Parameters
+    ----------
+    manifest_str:
+        Base64-encoded MDL JSON string.
+    session_context:
+        A ``wren_core.SessionContext`` used to expand per-model SQL.
+    data_source:
+        The target data source (determines sqlglot dialect).
+    fallback:
+        When ``True`` (default), if no model references are detected in the
+        SQL, fall back to ``session_context.transform_sql()`` directly.
+        Set to ``False`` in tests to ensure the CTE path is always exercised
+        and silent fallbacks don't mask bugs.
+    """
+
+    def __init__(
+        self,
+        manifest_str: str,
+        session_context,
+        data_source: DataSource,
+        *,
+        fallback: bool = True,
+    ):
+        self.session_context = session_context
+        self.fallback = fallback
+        self.dialect = get_sqlglot_dialect(data_source)
+        self.manifest = json.loads(base64.b64decode(manifest_str))
+
+        self.model_dict: dict[str, dict] = {}
+        self.schema = MappingSchema(dialect=self.dialect)
+        # normalized column name → original manifest column name, per model
+        self._col_orig_name: dict[str, dict[str, str]] = {}
+
+        for model in self.manifest.get("models", []):
+            name = model["name"]
+            self.model_dict[name] = model
+            cols: dict[str, str] = {}
+            orig: dict[str, str] = {}
+            for col in model.get("columns", []):
+                if col.get("isHidden"):
+                    continue
+                if col.get("relationship"):
+                    continue
+                col_name = col["name"]
+                cols[col_name] = col.get("type", "TEXT")
+                orig[col_name.lower()] = col_name
+            self.schema.add_table(name, cols, dialect=self.dialect)
+            self._col_orig_name[name] = orig
+
+    def rewrite(self, sql: str) -> str:
+        """Rewrite *sql* by injecting model CTEs.
+
+        Returns the transformed SQL string in the target sqlglot dialect.
+        If no model references are found, falls back to
+        ``session_context.transform_sql(sql)`` directly.
+        """
+        ast = parse_one(sql, dialect=self.dialect)
+
+        user_cte_names = self._collect_user_cte_names(ast)
+        used_columns = self._collect_model_columns(ast, user_cte_names)
+
+        # No model references detected — either fall back to the legacy
+        # whole-query transform, or raise so tests can catch the miss.
+        if not used_columns:
+            if self.fallback:
+                return self.session_context.transform_sql(sql)
+            raise ValueError(f"No model references found in SQL: {sql}")
+
+        model_ctes = self._build_model_ctes(used_columns)
+        self._inject_ctes(ast, model_ctes)
+        return ast.sql(dialect=self.dialect)
+
+    # ------------------------------------------------------------------
+    # Column collection via qualify
+    # ------------------------------------------------------------------
+
+    def _collect_model_columns(
+        self, ast: exp.Expression, user_cte_names: set[str]
+    ) -> dict[str, set[str]]:
+        """Return ``{model_name: {col1, col2, ...}}`` for all referenced models.
+
+        Uses sqlglot's ``qualify_columns`` to fully resolve all column
+        references (including ``SELECT *`` expansion and
+        correlated subquery outer references), then walks the qualified AST
+        to collect model→column mappings.
+        """
+        copy = ast.copy()
+        copy = qualify_tables(copy, dialect=self.dialect)
+        copy = normalize_identifiers(copy, dialect=self.dialect)
+        qualified = qualify_columns(
+            copy,
+            schema=self.schema,
+            dialect=self.dialect,
+            allow_partial_qualification=True,
+        )
+
+        # Build alias → model name mapping from Table nodes
+        alias_to_model: dict[str, str] = {}
+        for table in qualified.find_all(exp.Table):
+            table_name = table.name
+            if table_name not in self.model_dict or table_name in user_cte_names:
+                continue
+            alias = table.alias
+            if alias:
+                alias_to_model[alias] = table_name
+            alias_to_model[table_name] = table_name
+
+        used: dict[str, set[str]] = defaultdict(set)
+        for col in qualified.find_all(exp.Column):
+            table_ref = col.table
+            if not table_ref:
+                continue
+            model_name = alias_to_model.get(table_ref)
+            if model_name:
+                used[model_name].add(col.name)
+
+        return dict(used)
+
+    # ------------------------------------------------------------------
+    # CTE generation
+    # ------------------------------------------------------------------
+
+    def _build_model_ctes(self, used_columns: dict[str, set[str]]) -> list[exp.CTE]:
+        """Generate one CTE per model via wren-core transform_sql."""
+        ctes: list[exp.CTE] = []
+        for model_name, columns in used_columns.items():
+            orig = self._col_orig_name.get(model_name, {})
+            resolved = [orig.get(c, c) for c in sorted(columns)]
+            col_list = ", ".join(f'"{model_name}"."{c}"' for c in resolved)
+            model_sql = f'SELECT {col_list} FROM "{model_name}"'
+            expanded = self.session_context.transform_sql(model_sql)
+
+            expanded_ast = parse_one(expanded, dialect="wren")
+            cte = exp.CTE(
+                this=expanded_ast,
+                alias=exp.TableAlias(this=exp.to_identifier(model_name, quoted=True)),
+            )
+            ctes.append(cte)
+        return ctes
+
+    # ------------------------------------------------------------------
+    # CTE injection
+    # ------------------------------------------------------------------
+
+    def _inject_ctes(self, ast: exp.Expression, model_ctes: list[exp.CTE]) -> None:
+        """Prepend *model_ctes* before any existing user CTEs in *ast*."""
+        if not model_ctes:
+            return
+
+        existing_with = ast.args.get("with_")
+
+        if existing_with:
+            # Prepend model CTEs before user CTEs
+            existing_ctes = list(existing_with.expressions)
+            all_ctes = model_ctes + existing_ctes
+            existing_with.set("expressions", all_ctes)
+        else:
+            with_clause = exp.With(expressions=model_ctes)
+            ast.set("with_", with_clause)
+
+        # Preserve RECURSIVE if the original WITH had it
+        final_with = ast.args.get("with_")
+        if existing_with and existing_with.args.get("recursive"):
+            final_with.set("recursive", True)
+
+    # ------------------------------------------------------------------
+    # Helpers
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def _collect_user_cte_names(ast: exp.Expression) -> set[str]:
+        """Collect all CTE names defined in the user's SQL."""
+        names: set[str] = set()
+        with_clause = ast.args.get("with_")
+        if with_clause:
+            for cte in with_clause.expressions:
+                alias = cte.args.get("alias")
+                if alias:
+                    names.add(
+                        alias.this.name
+                        if isinstance(alias.this, exp.Identifier)
+                        else str(alias.this)
+                    )
+        return names