feat(wren): preserve SELECT * in CTE rewriter (#1536)

Author: goldmedal

wren/src/wren/mdl/cte_rewriter.py

@@ -119,9 +119,13 @@ def rewrite(self, sql: str) -> str:
 
     def _collect_model_columns(
         self, ast: exp.Expression, user_cte_names: set[str]
-    ) -> dict[str, list[str]]:
+    ) -> dict[str, list[str] | None]:
         """Return ``{model_name: [col1, col2, ...]}`` for all referenced models.
 
+        A value of ``None`` means the model was referenced via ``SELECT *``
+        and should be passed as-is to ``transform_sql`` so that wren-core
+        can apply column-level access control (CLAC).
+
         Uses sqlglot's ``qualify_columns`` to fully resolve all column
         references (including ``SELECT *`` expansion and
         correlated subquery outer references), then walks the qualified AST
@@ -131,6 +135,12 @@ def _collect_model_columns(
         copy = ast.copy()
         copy = qualify_tables(copy, dialect=self.dialect)
         copy = normalize_identifiers(copy, dialect=self.dialect)
+
+        # Detect models referenced via SELECT * BEFORE qualify_columns
+        # expands the star.  These will use SELECT * in transform_sql so
+        # that wren-core controls column visibility (CLAC).
+        star_models = self._detect_star_models(copy, user_cte_names)
+
         qualified = qualify_columns(
             copy,
             schema=self.schema,
@@ -161,17 +171,22 @@ def _collect_model_columns(
             if model_name:
                 used[model_name][col.name] = None
 
-        return {m: list(cols) for m, cols in used.items()}
+        return {m: None if m in star_models else list(cols) for m, cols in used.items()}
 
     # ------------------------------------------------------------------
     # CTE generation
     # ------------------------------------------------------------------
 
-    def _build_model_ctes(self, used_columns: dict[str, list[str]]) -> list[exp.CTE]:
+    def _build_model_ctes(
+        self, used_columns: dict[str, list[str] | None]
+    ) -> 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():
-            if columns:
+            if columns is None:
+                # SELECT * — let wren-core handle column visibility (CLAC)
+                col_list = "*"
+            elif columns:
                 orig = self._col_orig_name.get(model_name, {})
                 resolved = [orig.get(c, c) for c in columns]
                 col_list = ", ".join(f'"{model_name}"."{c}"' for c in resolved)
@@ -218,6 +233,43 @@ def _inject_ctes(self, ast: exp.Expression, model_ctes: list[exp.CTE]) -> None:
     # Helpers
     # ------------------------------------------------------------------
 
+    def _detect_star_models(
+        self, ast: exp.Expression, user_cte_names: set[str]
+    ) -> set[str]:
+        """Detect models selected via ``*`` before column qualification.
+
+        A bare ``SELECT *`` marks all models; ``SELECT t.*`` marks only
+        the referenced model.
+        """
+        star_models: set[str] = set()
+        select = ast.find(exp.Select)
+        if not select:
+            return star_models
+
+        # Build alias → model mapping from tables in FROM/JOIN
+        alias_to_model: dict[str, str] = {}
+        for table in ast.find_all(exp.Table):
+            name = table.name
+            if name not in self.model_dict or name in user_cte_names:
+                continue
+            alias = table.alias or name
+            alias_to_model[alias] = name
+            alias_to_model[name] = name
+
+        for sel_expr in select.expressions:
+            if isinstance(sel_expr, exp.Star):
+                # Bare * → all models
+                star_models.update(alias_to_model.values())
+            elif isinstance(sel_expr, exp.Column) and isinstance(
+                sel_expr.this, exp.Star
+            ):
+                # table.* → specific model
+                table_ref = sel_expr.table
+                if table_ref and table_ref in alias_to_model:
+                    star_models.add(alias_to_model[table_ref])
+
+        return star_models
+
     @staticmethod
     def _collect_user_cte_names(ast: exp.Expression) -> set[str]:
         """Collect all CTE names defined in the user's SQL (all scopes)."""

wren/tests/unit/test_cte_rewriter.py

@@ -121,6 +121,19 @@ def _count_ctes(sql: str, dialect: str = "duckdb") -> int:
     return len(with_clause.expressions)
 
 
+def _cte_body_sql(sql: str, cte_name: str, dialect: str = "duckdb") -> str | None:
+    """Return the SQL body of a named CTE, or None if not found."""
+    ast = sqlglot.parse_one(sql, dialect=dialect)
+    with_clause = ast.args.get("with_")
+    if not with_clause:
+        return None
+    for cte in with_clause.expressions:
+        alias = cte.args.get("alias")
+        if alias and alias.this.name == cte_name:
+            return cte.this.sql(dialect=dialect)
+    return None
+
+
 # ---------------------------------------------------------------------------
 # Tests: basic model resolution
 # ---------------------------------------------------------------------------
@@ -134,10 +147,20 @@ def test_single_model_generates_cte(self):
         assert isinstance(result, str)
         assert len(result) > 0
 
-    def test_star_expansion(self):
+    def test_star_not_expanded(self):
+        """SELECT * should pass star to wren-core so all columns are included."""
         rw = _make_rewriter(_SINGLE_MODEL_MANIFEST)
         result = rw.rewrite('SELECT * FROM "orders"')
         assert _has_cte(result, "orders")
+        cte_body = _cte_body_sql(result, "orders")
+        assert cte_body is not None
+        # wren-core expands SELECT * into all model columns — verify every
+        # non-hidden, non-relationship column is present in the CTE body.
+        body_lower = cte_body.lower()
+        for col_name in ("o_orderkey", "o_custkey", "o_orderstatus", "order_cust_key"):
+            assert col_name.lower() in body_lower, (
+                f"Expected column {col_name!r} in CTE body: {cte_body}"
+            )
 
     def test_alias_resolution(self):
         rw = _make_rewriter(_SINGLE_MODEL_MANIFEST)
@@ -163,10 +186,39 @@ def test_join_generates_two_ctes(self):
         assert _has_cte(result, "customer")
         assert _count_ctes(result) >= 2
 
-    def test_qualified_star(self):
+    def test_qualified_star_not_expanded(self):
+        """SELECT table.* should pass star to wren-core so all columns appear."""
         rw = _make_rewriter(_MULTI_MODEL_MANIFEST)
         result = rw.rewrite('SELECT "orders".* FROM "orders"')
         assert _has_cte(result, "orders")
+        cte_body = _cte_body_sql(result, "orders")
+        assert cte_body is not None
+        body_lower = cte_body.lower()
+        for col_name in ("o_orderkey", "o_custkey", "o_orderstatus", "order_cust_key"):
+            assert col_name.lower() in body_lower, (
+                f"Expected column {col_name!r} in CTE body: {cte_body}"
+            )
+
+    def test_mixed_star_and_explicit_columns(self):
+        """orders.* should include all columns; customer CTE only referenced ones."""
+        rw = _make_rewriter(_MULTI_MODEL_MANIFEST)
+        result = rw.rewrite(
+            'SELECT "orders".*, c.c_name '
+            'FROM "orders" JOIN "customer" c ON "orders".o_custkey = c.c_custkey'
+        )
+        assert _has_cte(result, "orders")
+        assert _has_cte(result, "customer")
+        # orders used star → all columns present
+        orders_body = _cte_body_sql(result, "orders")
+        assert orders_body is not None
+        orders_lower = orders_body.lower()
+        for col_name in ("o_orderkey", "o_custkey", "o_orderstatus", "order_cust_key"):
+            assert col_name.lower() in orders_lower
+        # customer used explicit column → only c_name (not necessarily c_custkey
+        # in the select list, though it may appear in the subquery structure)
+        customer_body = _cte_body_sql(result, "customer")
+        assert customer_body is not None
+        assert "c_name" in customer_body.lower()
 
 
 # ---------------------------------------------------------------------------