diff --git a/core/wren/pyproject.toml b/core/wren/pyproject.toml index 5cd14d261..8ccb9e139 100644 --- a/core/wren/pyproject.toml +++ b/core/wren/pyproject.toml @@ -22,7 +22,9 @@ classifiers = [ dependencies = [ "wren-core-py>=0.7.1", "duckdb>=1.5.0", - "sqlglot>=27", + # >=29 required: the CTE rewriter and policy read the ``with_`` / ``from_`` + # Expression arg keys, which sqlglot renamed from ``with`` / ``from`` in 28. + "sqlglot>=29", "typer>=0.12", "pydantic>=2", "pyarrow>=14", diff --git a/core/wren/src/wren/mdl/cte_rewriter.py b/core/wren/src/wren/mdl/cte_rewriter.py index 43367ba2a..f565cbdc2 100644 --- a/core/wren/src/wren/mdl/cte_rewriter.py +++ b/core/wren/src/wren/mdl/cte_rewriter.py @@ -18,6 +18,7 @@ import json import sqlglot from sqlglot import exp, parse_one +from sqlglot.dialects.dialect import Dialect, NormalizationStrategy from sqlglot.optimizer.normalize_identifiers import normalize_identifiers from sqlglot.optimizer.qualify_columns import qualify_columns from sqlglot.optimizer.qualify_tables import qualify_tables @@ -26,6 +27,7 @@ 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 +from wren.model.error import ErrorCode, ErrorPhase, WrenError from wren.policy import resolve_model_name _SQLGLOT_DIALECT_MAP: dict[DataSource, str] = { @@ -44,6 +46,24 @@ def get_sqlglot_dialect(data_source: DataSource) -> str: return _SQLGLOT_DIALECT_MAP.get(data_source, data_source.name) +# sqlglot dialects whose *physical column names* are case-sensitive — i.e. the +# backing database can hold two columns differing only in case (``Year`` and +# ``year``) and address them via quoting. Only these allow a model to declare +# case-distinct columns; everywhere else such a model is physically +# unrepresentable and is rejected at build time (see ``CTERewriter.__init__``). +# +# This is an explicit allow-list rather than a derivation from sqlglot's +# ``NORMALIZATION_STRATEGY``, because the strategy is about identifier *folding*, +# not column-name uniqueness, and the two disagree: MySQL/Doris report +# CASE_SENSITIVE yet their column names are case-insensitive, and Athena +# (LOWERCASE) lowercases columns in the Glue catalog. Postgres/Oracle/Snowflake +# fold unquoted identifiers but compare the stored (quoted) name +# case-sensitively; ClickHouse is case-sensitive throughout. +_CASE_SENSITIVE_COLUMN_DIALECTS: frozenset[str] = frozenset( + {"postgres", "oracle", "snowflake", "clickhouse"} +) + + class CTERewriter: """Rewrite user SQL by expanding MDL model references into CTEs. @@ -56,10 +76,12 @@ class CTERewriter: 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. + Controls SQL that references a table which is not an MDL model or + view. When ``True`` (default), fall back to + ``session_context.transform_sql()`` directly. Set to ``False`` in + tests so such a query raises instead of silently masking a rewriter + miss. (Pure scalar / TVF SQL with no base-table reference always + passes through, regardless of this flag.) """ def __init__( @@ -74,12 +96,44 @@ class CTERewriter: self.data_source = data_source self.fallback = fallback self.dialect = get_sqlglot_dialect(data_source) + # Upper-folding dialects (Oracle, Snowflake, …) uppercase every unquoted + # identifier, which would change result-set column names (aggregate + # aliases, cube columns, …). Render those with ``identify=True`` so the + # output is fully quoted and result casing stays as authored. Detected + # from the dialect's normalization strategy rather than hard-coded. + self._force_identify = ( + Dialect.get_or_raise(self.dialect).NORMALIZATION_STRATEGY + == NormalizationStrategy.UPPERCASE + ) self.manifest = json.loads(base64.b64decode(manifest_str)) + # A model may declare case-distinct columns (``Year`` and ``year``) only + # on dialects whose physical column names are case-sensitive AND only + # when the manifest actually contains such a collision. We engage the + # case-sensitive resolution path *only then*, so the well-trodden + # case-insensitive path (and its result-column casing) is unchanged for + # every existing manifest. On the case-sensitive path quoted refs match + # exactly and unquoted refs match exact-then-case-insensitively, erroring + # only when ambiguous (see ``_resolve_column`` / + # ``_normalize_model_column_case``). + has_case_distinct = self._manifest_has_case_distinct_columns() + if has_case_distinct and self.dialect not in _CASE_SENSITIVE_COLUMN_DIALECTS: + self._raise_case_collision() + self._case_sensitive_columns = has_case_distinct + self.model_dict: dict[str, dict] = {} - self.schema = MappingSchema(dialect=self.dialect) + # On the case-sensitive path the qualify schema must NOT fold identifiers + # (``normalize=False``) so a quoted ``"year"`` matches the stored ``year`` + # and not ``Year``; column keys are kept in manifest case. + self.schema = MappingSchema( + dialect=self.dialect, normalize=not self._case_sensitive_columns + ) # normalized column name → original manifest column name, per model + # (only used on the case-insensitive path). self._col_orig_name: dict[str, dict[str, str]] = {} + # manifest-case column names, per model (used on the case-sensitive path + # for exact-then-CI resolution). + self._model_cols: dict[str, list[str]] = {} for model in self.manifest.get("models", []): name = model["name"] @@ -92,19 +146,36 @@ class CTERewriter: if col.get("relationship"): continue col_name = col["name"] + # Case-only collisions were already vetted by the pre-scan: on + # case-insensitive-column dialects they raised INVALID_MDL; on + # case-sensitive-column dialects they are kept distinct here and + # resolved case-sensitively at query time. cols[col_name] = col.get("type", "TEXT") orig[col_name.lower()] = col_name - # ``qualify_columns`` runs against the post-``normalize_identifiers`` - # AST, so the schema must be keyed under the same normalized form - # of the model name. BigQuery / DuckDB lowercase, Oracle uppercases — - # registering the literal manifest name leaves a mismatch and the - # column qualification silently produces an empty CTE body. - schema_name = normalize_identifiers( - exp.to_identifier(name, quoted=True), dialect=self.dialect - ).name - self.schema.add_table(schema_name, cols, dialect=self.dialect) + self._model_cols[name] = list(cols) + if self._case_sensitive_columns: + # Keep manifest case as the schema key (``normalize=False``) so + # quoted refs resolve exactly and case-distinct columns coexist. + self.schema.add_table(name, cols) + else: + # ``qualify_columns`` runs against the post-``normalize_identifiers`` + # AST, so the schema must be keyed under the same normalized form + # of the model name. BigQuery / DuckDB lowercase, Oracle uppercases + # — registering the literal manifest name leaves a mismatch and the + # column qualification silently produces an empty CTE body. + schema_name = normalize_identifiers( + exp.to_identifier(name, quoted=True), dialect=self.dialect + ).name + self.schema.add_table(schema_name, cols, dialect=self.dialect) self._col_orig_name[name] = orig + # Flat union of every model's columns, for resolving *unqualified* + # column references on the case-sensitive path. Computed once here + # rather than per ``_normalize_model_column_case`` call. + self._all_model_cols: list[str] = [ + c for cols in self._model_cols.values() for c in cols + ] + # A view's ``statement`` is native-dialect SQL that references models. # It is NOT expanded by wren-core — it becomes a CTE kept verbatim, # preceded by model CTEs for the models it references. view_dict maps @@ -114,6 +185,54 @@ class CTERewriter: } self.view_names: set[str] = set(self.view_dict) + @staticmethod + def _iter_model_column_names(model: dict): + """Yield the visible column names of *model* (skips hidden / relationship). + + Mirrors the column filter used when populating the schema so the + case-collision pre-scan sees exactly the columns that get registered. + """ + for col in model.get("columns", []): + if col.get("isHidden") or col.get("relationship"): + continue + yield col["name"] + + def _manifest_has_case_distinct_columns(self) -> bool: + """True if any model has two visible columns differing only in case.""" + for model in self.manifest.get("models", []): + seen: set[str] = set() + for col_name in self._iter_model_column_names(model): + low = col_name.lower() + if low in seen: + return True + seen.add(low) + return False + + def _raise_case_collision(self) -> None: + """Raise ``INVALID_MDL`` for the first case-only column collision. + + Used on case-insensitive-column dialects, where Wren resolves column + names case-insensitively and two columns differing only in case would + silently collide — and the backing database cannot represent them. + """ + for model in self.manifest.get("models", []): + seen: dict[str, str] = {} + for col_name in self._iter_model_column_names(model): + low = col_name.lower() + if low in seen: + raise WrenError( + error_code=ErrorCode.INVALID_MDL, + message=( + f"Model '{model['name']}' has columns that differ " + f"only in case ('{seen[low]}' and '{col_name}'). Wren " + f"resolves column names case-insensitively on " + f"{self.dialect} and cannot distinguish them; rename " + "one of the columns." + ), + phase=ErrorPhase.MDL_EXTRACTION, + ) + seen[low] = col_name + def rewrite(self, sql: str) -> str: """Rewrite *sql* by injecting model and view CTEs. @@ -122,41 +241,82 @@ class CTERewriter: statement references are expanded as model CTEs placed before it. Returns the transformed SQL string in the target sqlglot dialect. - If no model or view references are found, falls back to - ``session_context.transform_sql(sql)`` directly (when ``fallback`` - is ``True``); otherwise raises ``ValueError``. + Pure scalar / TVF SQL (no model, view, or base-table reference, e.g. + ``SELECT 1``) is passed through transpiled to that dialect. A + reference to a table that is not an MDL model or view falls back to + ``session_context.transform_sql(sql)`` when ``fallback`` is ``True``; + otherwise it raises ``ValueError``. """ ast = parse_one(sql, dialect=self.dialect) user_cte_names = self._collect_user_cte_names(ast) - used_columns, user_table_refs = self._collect_model_columns(ast, user_cte_names) + + # Two situations need model-column references canonicalized to the + # manifest case before collection: + # * Upper-folding dialects (Oracle/Snowflake) render force-quoted + # (identify=True, below); a case-insensitive ref (``mixedcase`` for a + # manifest ``MixedCase``) would be quoted verbatim and never bind. + # * The case-sensitive path (manifest declares case-distinct columns) + # resolves quoted refs exactly and unquoted refs exact-then-CI, then + # force-quotes the resolved name so the dialect can't re-fold it. + if self._force_identify or self._case_sensitive_columns: + self._normalize_model_column_case(ast, user_cte_names) + + used_columns, user_table_refs, col_quoting = self._collect_model_columns( + ast, user_cte_names + ) view_refs = self._collect_view_refs(ast, user_cte_names) # A view's native-SQL statement references models; collect those so # they get model CTEs placed before the (verbatim) view CTEs. - self._collect_view_model_usage(view_refs, used_columns, user_table_refs) + self._collect_view_model_usage( + view_refs, used_columns, user_table_refs, col_quoting + ) + + # The user's SQL is never rewritten — we only append CTEs. To bind the + # user's column references to the injected CTE on dialects that + # case-fold unquoted identifiers, each model CTE *exposes* its columns + # with the same quoting the user wrote (see ``_collect_model_columns`` / + # ``_alias_projection_to_user_quoting``): both the user's reference and + # the CTE alias share the same quoting, so the dialect folds them + # identically and they match. + # + # Upper-folding dialects (Oracle, Snowflake, …) are the exception: they + # render with ``identify=True``. They uppercase every unquoted + # identifier, so without forced quoting the result-set column names + # (aggregate aliases like ``cnt``, cube columns like + # ``order_date__month``) would come back uppercased — a breaking change + # for callers that read columns by name. Forcing quoting keeps result + # casing stable; case-insensitive references there are accepted by the + # ``_normalize_model_column_case`` fold above. ``identify`` is purely an + # output-rendering flag — the input AST is otherwise left as authored. + identify = self._force_identify - # No model or view references detected — either fall back to the - # legacy whole-query transform, or raise so tests can catch the miss. if not used_columns and not view_refs: + # Nothing resolved to a model or view. If the query also has no + # base-table reference at all, it is pure scalar / TVF SQL + # (``SELECT 1`` or a standalone ``SELECT * FROM UNNEST([...])``) + # with nothing to expand — pass it through transpiled to the + # target dialect rather than rejecting it. + base_tables = [ + t + for t in ast.find_all(exp.Table) + if (t.name or "").lower() not in user_cte_names + ] + if not base_tables: + return ast.sql(dialect=self.dialect, identify=identify) + # Otherwise the query references a table that is not an MDL model + # or view. Fall back to the legacy whole-query transform (so a + # broken/stale reference still surfaces an error), or raise when + # ``fallback=False`` so tests catch a rewriter miss. if self.fallback: wren_sql = self.session_context.transform_sql(sql) return sqlglot.transpile(wren_sql, read="wren", write=self.dialect)[0] raise ValueError(f"No model or view references found in SQL: {sql}") - model_ctes = self._build_model_ctes(used_columns, user_table_refs) + model_ctes = self._build_model_ctes(used_columns, user_table_refs, col_quoting) view_ctes = self._build_view_ctes(view_refs) self._inject_ctes(ast, model_ctes + view_ctes) - # Oracle uppercases unquoted identifiers. Without forcing quoting - # on output, the user's ``SELECT o_orderkey FROM orders`` would - # land as ``SELECT O_ORDERKEY FROM ORDERS`` — both the table - # reference and the result column name. The injected CTE projects - # quoted lowercase columns, so the lookup misses (ORA-00904), and - # any caller asserting on result-column casing breaks. Forcing - # quoting on Oracle makes the dialect's output deterministic and - # matches the pre-fallback path where wren-core's whole-query - # transform had quoted everything implicitly. - identify = self.data_source == DataSource.oracle return ast.sql(dialect=self.dialect, identify=identify) # ------------------------------------------------------------------ @@ -165,8 +325,12 @@ class CTERewriter: def _collect_model_columns( self, ast: exp.Expression, user_cte_names: set[str] - ) -> tuple[dict[str, list[str] | None], dict[str, tuple[str, bool]]]: - """Return ``(used_columns, user_table_refs)`` for all referenced models. + ) -> tuple[ + dict[str, list[str] | None], + dict[str, tuple[str, bool]], + dict[str, dict[str, bool]], + ]: + """Return ``(used_columns, user_table_refs, col_quoting)``. ``used_columns``: ``{model_name: [col1, col2, ...]}``. A value of ``None`` means the model was referenced via ``SELECT *`` and should @@ -179,6 +343,16 @@ class CTERewriter: emitted CTE must fold to the same form) bind the user's outer reference to the injected CTE. + ``col_quoting``: ``{model_name: {column_lower: was_quoted}}`` capturing + whether the user quoted each column reference, scoped to the model the + reference resolved to. The model CTE then exposes that column with the + same quoting, so the user's (untouched) reference binds regardless of + how the dialect folds unquoted identifiers. Scoping per model keeps a + quoted reference to one source (a user CTE, or another model's + same-named column) from flipping an unrelated model's CTE alias. When a + model's column is referenced both quoted and unquoted, quoted wins + (preserves the manifest case). + 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 @@ -201,6 +375,65 @@ class CTERewriter: # that wren-core controls column visibility (CLAC). star_models = self._detect_star_models(copy, alias_to_model) + if self._case_sensitive_columns: + # Columns were already canonicalized + force-quoted to manifest case + # by ``_normalize_model_column_case``; resolve them case-sensitively + # against the ``normalize=False`` schema with no folding. + return self._collect_columns_case_sensitive( + copy, alias_to_model, user_table_refs, star_models + ) + + # Unquote model table refs on the collection-only copy so a quoted + # mixed-case model name (e.g. ``"WREN_AI_CaseTest"``) folds to the + # dialect's natural case during ``normalize_identifiers`` and matches + # the dialect-folded qualify schema key. This mirrors the column + # unquoting below: without it, the quoted reference preserves its case, + # never matches the schema, ``qualify_columns`` leaves the column + # unqualified, and the model CTE collapses to ``SELECT 1``. + # + # Only table refs that resolved to a model are unquoted (user CTEs are + # excluded by ``_build_alias_map``), so a user CTE's case matching + # between its definition and reference is untouched. ``user_table_refs`` + # has already captured the original quoting for CTE-alias mirroring, and + # this copy is never emitted — the user's SQL is not mutated. + model_refs = set(alias_to_model) + for tbl in copy.find_all(exp.Table): + ident = tbl.this + if ( + isinstance(ident, exp.Identifier) + and ident.quoted + and ident.name in model_refs + ): + ident.set("quoted", False) + + # Record each column reference's quoting (by node identity), then unquote + # it on the copy. The quoting drives CTE-alias mirroring; it is attributed + # to a resolved model *after* qualification (below) so a quoted reference + # to one source can't flip an unrelated model's CTE alias. + # - Recording quoting lets the model CTE expose the column with the same + # quoting the user wrote (mirroring), so the untouched user reference + # binds. + # - Unquoting on the copy makes a quoted mixed-case ref (e.g. ``"Year"``) + # fold to the dialect's natural case during ``normalize_identifiers`` + # so it matches the dialect-folded qualify schema. Wren resolves column + # names case-insensitively (see ``_col_orig_name``); the copy is only + # used for collection, so the fold is safe. Without it a quoted + # ``"Year"`` never binds and the model CTE collapses to ``SELECT 1``. + # + # Quoting is recorded from *every* column reference (SELECT, WHERE, JOIN, + # GROUP BY, ORDER BY, ...), not just the SELECT list — a column used only + # in WHERE still needs its CTE alias to mirror it. The node identities + # survive ``normalize_identifiers`` / ``qualify_columns``, so each + # reference can be matched back to its resolved model afterwards. + quoting_by_id: dict[int, bool] = {} + for col in copy.find_all(exp.Column): + ident = col.this + if not isinstance(ident, exp.Identifier): + continue + quoting_by_id[id(col)] = ident.quoted + if ident.quoted: + ident.set("quoted", False) + copy = normalize_identifiers(copy, dialect=self.dialect) qualified = qualify_columns( copy, @@ -217,17 +450,87 @@ class CTERewriter: # been normalized (lowercased) by qualify_columns even though we built # the alias map from the pre-normalize AST. alias_lookup = {k.lower(): v for k, v in alias_to_model.items()} + # Quoting recorded *per resolved model* (``{model: {col_lower: quoted}}``) + # so a quoted reference to one source can't flip another model's CTE + # alias. Quoted-wins within a model. A column referenced only via a + # non-model source (user CTE, external table) is never attributed here. + col_quoting: dict[str, dict[str, bool]] = {} for col in qualified.find_all(exp.Column): table_ref = col.table if not table_ref: continue model_name = alias_lookup.get(table_ref.lower()) + if not model_name: + continue + used[model_name][col.name] = None + quoted = quoting_by_id.get(id(col)) + if quoted is None: + continue + per = col_quoting.setdefault(model_name, {}) + low = col.name.lower() + if quoted: + per[low] = True + else: + per.setdefault(low, False) + + return ( + {m: None if m in star_models else list(cols) for m, cols in used.items()}, + user_table_refs, + col_quoting, + ) + + def _collect_columns_case_sensitive( + self, + copy: exp.Expression, + alias_to_model: dict[str, str], + user_table_refs: dict[str, tuple[str, bool]], + star_models: set[str], + ) -> tuple[ + dict[str, list[str] | None], + dict[str, tuple[str, bool]], + dict[str, dict[str, bool]], + ]: + """Collect model→columns on the case-sensitive path (no folding). + + The schema is keyed in manifest case (``normalize=False``) and column + refs were already canonicalized + force-quoted to manifest case, so + ``qualify_columns`` matches each ref exactly and ``SELECT *`` expands to + the manifest-case columns. Returns an empty ``col_quoting`` map: the + model CTE exposes *every* column quoted in manifest case (the default in + ``_alias_projection_to_user_quoting``), matching the force-quoted refs. + """ + alias_lower = {k.lower(): v for k, v in alias_to_model.items()} + + # Rewrite each model table ref to its manifest-case name (quoted) so it + # matches the manifest-case schema key without folding. Only the table + # *name* is touched; SQL aliases are left intact. + for tbl in copy.find_all(exp.Table): + ident = tbl.this + if isinstance(ident, exp.Identifier): + model_name = alias_lower.get(ident.name.lower()) + if model_name: + tbl.set("this", exp.to_identifier(model_name, quoted=True)) + + qualified = qualify_columns( + copy, + schema=self.schema, + dialect=self.dialect, + allow_partial_qualification=True, + ) + + used: dict[str, dict[str, None]] = {m: {} for m in alias_to_model.values()} + for col in qualified.find_all(exp.Column): + table_ref = col.table + if not table_ref: + continue + model_name = alias_lower.get(table_ref.lower()) if model_name: used[model_name][col.name] = None return ( {m: None if m in star_models else list(cols) for m, cols in used.items()}, user_table_refs, + {}, ) def _build_alias_map( @@ -265,6 +568,242 @@ class CTERewriter: user_table_refs.setdefault(model_name, (name, quoted)) return alias_to_model, user_table_refs + def _resolve_column( + self, candidates: list[str], name: str, quoted: bool, where: str + ) -> str | None: + """Resolve a column reference to its manifest-case name. + + SQL identifier semantics, mirroring ``resolve_model_name`` one level + down: a **quoted** reference must match a manifest column exactly + (case-sensitive); an **unquoted** reference prefers an exact match, then + falls back to a case-insensitive scan. Returns the manifest-case name, + or ``None`` if nothing matches (the ref then fails to bind, surfacing as + a "column not found" from wren-core). + + Raises ``INVALID_SQL`` only when an unquoted reference is genuinely + **ambiguous** — two or more candidates differ only in case and none + matches exactly (e.g. ``YEAR`` against ``Year`` and ``year``). + """ + # Single pass: an exact match wins immediately (quoted-exact or + # unquoted-exact); otherwise an unquoted ref collects case-insensitive + # candidates for the exact-then-CI fallback. + low = name.lower() + ci = [] + for c in candidates: + if c == name: + return name + if not quoted and c.lower() == low: + ci.append(c) + if quoted: # quoted is strict — no case-insensitive fallback + return None + if len(ci) == 1: + return ci[0] + if len(ci) > 1: + raise WrenError( + error_code=ErrorCode.INVALID_SQL, + message=( + f"Column reference '{name}' in {where} is ambiguous: it " + f"matches case-distinct columns {sorted(ci)}. Quote it to " + "select one exactly." + ), + phase=ErrorPhase.SQL_PARSING, + ) + return None + + @staticmethod + def _collect_output_alias_refs(ast: exp.Expression) -> set[int]: + """Return ``id()`` of column nodes that reference a SELECT output alias. + + An output alias (``SELECT x AS yr``) is referenceable by name only from + its own SELECT's ``ORDER BY`` / ``GROUP BY`` / ``HAVING`` / ``QUALIFY`` + clauses — not from WHERE/JOIN, and not from another (e.g. outer) scope. + Resolving this per scope (rather than with a query-wide alias-name set) + avoids mis-skipping an outer column ref whose name merely coincides with + a subquery's alias. Returns node identities so the caller can skip the + exact references without re-deriving scope. + """ + refs: set[int] = set() + for select in ast.find_all(exp.Select): + proj_aliases = { + proj.alias.lower() + for proj in select.expressions + if isinstance(proj, exp.Alias) and proj.alias + } + if not proj_aliases: + continue + for clause_key in ("order", "group", "having", "qualify"): + clause = select.args.get(clause_key) + if clause is None: + continue + for col in clause.find_all(exp.Column): + ident = col.this + if ( + isinstance(ident, exp.Identifier) + and not col.table + and ident.name.lower() in proj_aliases + ): + refs.add(id(col)) + return refs + + def _normalize_model_column_case( + self, ast: exp.Expression, user_cte_names: set[str] + ) -> None: + """Rewrite model-column references in *ast* to their manifest case. + + Invoked when the output is force-quoted (Oracle/Snowflake, ``identify= + True``) or the manifest declares case-distinct columns (the + case-sensitive path). In both, a user reference whose case differs from + the manifest would otherwise be emitted quoted-verbatim and never bind + the manifest-case CTE column. Canonicalizing it lets the input stay + case-insensitive (for unambiguous refs) while result-set casing is the + manifest case. + + On the **case-sensitive path** resolution is strict per SQL rules + (``_resolve_column``): quoted refs match exactly, unquoted refs match + exact-then-CI, ambiguous unquoted refs raise. The resolved identifier is + force-quoted so a folding dialect (Postgres) cannot re-fold it and so + two case-distinct columns stay distinct in the output. + + Scope guards keep this from rebinding non-model columns: + + - a column qualified by a user CTE (or any non-model alias) is left + untouched; + - an unqualified column is only rewritten when the query has no user + CTEs, so it cannot shadow a CTE-exposed column. + + A *qualified* reference is resolved against its specific model, so two + models whose columns share a lowercase form but differ in manifest case + (``A.year`` vs ``B.Year``) each canonicalize correctly. An unqualified + reference (only rewritten when there are no user CTEs) resolves against + the union of all models' columns. Mutates *ast* in place; only column + identifiers are touched (SELECT aliases, function names, etc. are left + as the user wrote them). + """ + if not self._model_cols: + return + + alias_to_model, _ = self._build_alias_map( + qualify_tables(ast.copy(), dialect=self.dialect), user_cte_names + ) + alias_to_model_lower = {a.lower(): m for a, m in alias_to_model.items()} + has_user_ctes = bool(user_cte_names) + cs = self._case_sensitive_columns + # A SELECT-list output alias (``SELECT x AS yr ... ORDER BY yr``) parses + # as a bare column but is not a model column — never canonicalize or + # error on it. It is only referenceable by name from its *own* SELECT's + # ORDER BY / GROUP BY / HAVING / QUALIFY clauses, so resolve this + # precisely per scope: collect the exact ``exp.Column`` *nodes* in those + # clauses that match a projection alias of the same SELECT. + alias_ref_cols = self._collect_output_alias_refs(ast) + + for col in ast.find_all(exp.Column): + ident = col.this + if not isinstance(ident, exp.Identifier): + continue + if id(col) in alias_ref_cols: + # references a SELECT output alias, not a model column; skip + continue + table = col.table + if table: + model_name = alias_to_model_lower.get(table.lower()) + if model_name is None: + # qualified by something that isn't a model (e.g. a user CTE) + continue + candidates = self._model_cols.get(model_name, []) + where = f"{table}" + elif has_user_ctes or not alias_to_model: + # Skip canonicalization (and the case-sensitive existence check) + # when the column can't be attributed to a model: + # * user CTEs present → it could be a CTE column, and we can't + # tell without full scope analysis; + # * no model table resolved at all → it belongs to an external + # table or a pure-TVF projection, which ``rewrite()`` handles + # via its passthrough/fallback branch. + # Known limitation: an unqualified model-column ref in a query + # that also defines CTEs is not force-quoted and bypasses the + # case-sensitive "column does not exist" check — it falls back to + # the old behavior (may collapse to ``SELECT 1``). Qualify the + # column (``model.col``) to get strict checking. + continue + else: + candidates = self._all_model_cols + where = "the query" + # On the case-sensitive path a quoted ref is strict (must match + # exactly). On the force-identify-only path (no case-distinct + # columns) preserve the existing lenient behavior: resolve every ref + # case-insensitively — the candidate list has no case collisions, so + # this lookup is unambiguous and never raises. + canonical = self._resolve_column( + candidates, ident.name, ident.quoted if cs else False, where + ) + if canonical is None: + if cs: + # Case-sensitive path: the ref is model-attributable but + # matches no manifest column — a quoted wrong-case ref + # (``"YEAR"`` vs ``Year``/``year``) or a genuine typo. Fail + # loudly here instead of dropping the column and emitting a + # ``SELECT 1`` CTE that explodes later as an opaque database + # "column does not exist" at execution. + quoted_hint = ( + " (quoted references are matched case-sensitively)" + if ident.quoted + else "" + ) + raise WrenError( + error_code=ErrorCode.INVALID_SQL, + message=( + f"Column '{ident.name}' in {where} does not exist" + f"{quoted_hint}. Available columns: {sorted(candidates)}." + ), + phase=ErrorPhase.SQL_PARSING, + ) + continue + if ident.name != canonical: + ident.set("this", canonical) + if cs: + # force-quote so a folding dialect cannot re-fold the resolved + # name and case-distinct columns stay distinct in the output + ident.set("quoted", True) + + @staticmethod + def _alias_projection_to_user_quoting( + expanded_ast: exp.Expression, col_quoting: dict[str, bool] + ) -> None: + """Alias the model CTE's outermost projection to mirror user quoting. + + wren-core projects manifest-case columns (e.g. ``"Year"``). We add an + explicit alias to each so the CTE *exposes* the column with the same + quoting the user wrote it with: + + - user wrote ``"Year"`` (quoted) → expose ``... AS "Year"`` (the + column stays case-sensitive ``Year``; the user's ``"Year"`` binds). + - user wrote ``Year`` (unquoted) → expose ``... AS Year`` (the dialect + folds it, e.g. Postgres → ``year``; the user's folded ``Year`` binds). + + Columns with no captured reference (``SELECT *`` / introspection) + default to quoted so the result preserves the manifest case. This keeps + the user's SQL untouched — only the CTE projection carries the mirror. + """ + # wren-core's single-model expansion is always a ``SELECT``; guard + # defensively so an unexpected shape can't crash the rewrite (it would + # just skip mirroring and expose manifest-case columns). + if not isinstance(expanded_ast, exp.Select): + return + new_exprs = [] + for proj in expanded_ast.expressions: + if isinstance(proj, exp.Alias): + name, inner = proj.alias, proj.this + elif isinstance(proj, exp.Column): + # Bare column with no alias wrapper. wren-core normally emits an + # Alias, so this is defensive — handled the same way. + name, inner = proj.name, proj + else: + new_exprs.append(proj) + continue + quoted = col_quoting.get(name.lower(), True) + new_exprs.append(exp.alias_(inner, exp.to_identifier(name, quoted=quoted))) + expanded_ast.set("expressions", new_exprs) + def _collect_view_refs( self, ast: exp.Expression, user_cte_names: set[str] ) -> dict[str, tuple[str, bool]]: @@ -306,6 +845,7 @@ class CTERewriter: self, used_columns: dict[str, list[str] | None], user_table_refs: dict[str, tuple[str, bool]], + col_quoting: dict[str, dict[str, bool]], ) -> list[exp.CTE]: """Generate one CTE per model via wren-core transform_sql.""" ctes: list[exp.CTE] = [] @@ -314,13 +854,19 @@ class CTERewriter: # SELECT * — let wren-core handle column visibility (CLAC) model_sql = f'SELECT * FROM "{model_name}"' elif columns: - # ``_col_orig_name`` is keyed by lowercase column names; the - # column refs come from the post-normalize AST whose case - # depends on the dialect (Oracle uppercases unquoted idents, - # Postgres lowercases them). Lower-case before lookup so the - # original manifest casing is restored either way. - orig = self._col_orig_name.get(model_name, {}) - resolved = [orig.get(c.lower(), c) for c in columns] + if self._case_sensitive_columns: + # Collected names are already in manifest case (the schema is + # not folded), and ``_col_orig_name`` would collide for + # case-distinct columns — use the names as collected. + resolved = columns + else: + # ``_col_orig_name`` is keyed by lowercase column names; the + # column refs come from the post-normalize AST whose case + # depends on the dialect (Oracle uppercases unquoted idents, + # Postgres lowercases them). Lower-case before lookup so the + # original manifest casing is restored either way. + orig = self._col_orig_name.get(model_name, {}) + resolved = [orig.get(c.lower(), c) for c in columns] col_list = ", ".join(f'"{model_name}"."{c}"' for c in resolved) model_sql = f'SELECT {col_list} FROM "{model_name}"' else: @@ -339,6 +885,13 @@ class CTERewriter: # the shadow chain breaks at the top scope. self._rename_outer_alias(expanded_ast, model_name) + # Expose each column with the quoting the user wrote it with, so the + # (untouched) user reference binds regardless of dialect folding. + # Use this model's own quoting map only. + self._alias_projection_to_user_quoting( + expanded_ast, col_quoting.get(model_name, {}) + ) + # Match the user's literal identifier (case + quoting) for the # CTE alias so dialects with case-folding still bind the user's # outer ``FROM `` to the CTE. Oracle uppercases unquoted @@ -362,9 +915,11 @@ class CTERewriter: view_refs: dict[str, tuple[str, bool]], used_columns: dict[str, list[str] | None], user_table_refs: dict[str, tuple[str, bool]], + col_quoting: dict[str, dict[str, bool]], ) -> None: """Merge the models each referenced view's statement uses into - *used_columns* / *user_table_refs*, so those models get CTEs. + *used_columns* / *user_table_refs* / *col_quoting*, so those models get + CTEs whose columns mirror how the view's statement quotes them. A view's statement is native SQL referencing models by name. Parsing it through the same model-collection path captures which model columns @@ -377,10 +932,30 @@ class CTERewriter: self.view_dict[view_name]["statement"], dialect=self.dialect ) view_cte_names = self._collect_user_cte_names(view_ast) - v_cols, v_refs = self._collect_model_columns(view_ast, view_cte_names) + # Resolve the view body's column refs the same way as the user query, + # so a case-distinct model referenced from a view statement gets the + # same strict treatment (quoted = exact, unquoted = exact-then-CI, + # ambiguous / not-found raise cleanly) rather than silently + # mis-resolving or collapsing to ``SELECT 1``. Mutating this AST is + # safe: it is used only to collect columns here; the view body is + # emitted verbatim from a fresh parse in ``_build_view_ctes``. + if self._force_identify or self._case_sensitive_columns: + self._normalize_model_column_case(view_ast, view_cte_names) + v_cols, v_refs, v_quoting = self._collect_model_columns( + view_ast, view_cte_names + ) self._merge_used_columns(used_columns, v_cols) for model_name, ref in v_refs.items(): user_table_refs.setdefault(model_name, ref) + for model_name, qmap in v_quoting.items(): + # merge per-model, quoted-wins (consistent with + # _collect_model_columns) + dest = col_quoting.setdefault(model_name, {}) + for low, quoted in qmap.items(): + if quoted: + dest[low] = True + else: + dest.setdefault(low, False) @staticmethod def _merge_used_columns( diff --git a/core/wren/src/wren/policy.py b/core/wren/src/wren/policy.py index 4fdc9810a..5ff4aed6c 100644 --- a/core/wren/src/wren/policy.py +++ b/core/wren/src/wren/policy.py @@ -6,13 +6,34 @@ manifest and does not use any denied functions. from __future__ import annotations +from functools import lru_cache from typing import Iterable -from sqlglot import exp +from sqlglot import exp, parse_one +from sqlglot.errors import SqlglotError from wren.config import WrenConfig from wren.model.error import ErrorCode, ErrorPhase, WrenError +# Dialects we probe when canonicalising the user's denylist. sqlglot can map +# the same function name (e.g. ``version()``) onto different concrete AST +# subclasses depending on the dialect — postgres/mysql/duckdb/trino/clickhouse +# normalise to ``CurrentVersion`` while tsql/oracle/bigquery/snowflake keep it +# as ``Anonymous``. Probing each dialect ensures the canonical class key is +# captured regardless of which one the user's SQL ends up parsed with. +_CANONICALISE_DIALECTS: tuple[str | None, ...] = ( + None, + "postgres", + "mysql", + "tsql", + "oracle", + "bigquery", + "snowflake", + "clickhouse", + "trino", + "duckdb", +) + def resolve_model_name( name: str, @@ -131,16 +152,46 @@ def _check_tables( ) +@lru_cache(maxsize=128) +def _canonical_denied(denied: frozenset[str]) -> frozenset[str]: + """Expand the user's denylist to also cover sqlglot's canonical keys. + + sqlglot >=29 maps several common functions onto concrete subclasses — + e.g. ``version()`` becomes ``exp.CurrentVersion`` in + postgres/mysql/duckdb/trino/clickhouse (with ``type(node).key == + "currentversion"``), while in tsql/oracle/bigquery/snowflake it stays + as ``exp.Anonymous(name="version")``. A user denylist entry of + ``"version"`` would only match the anonymous case without this + expansion. Probing each known dialect collects every class key the + name might land on at parse time and adds them all to the result. + """ + expanded: set[str] = {d.lower() for d in denied} + for name in list(expanded): + for dialect in _CANONICALISE_DIALECTS: + try: + ast = parse_one(f"SELECT {name}()", dialect=dialect) + except SqlglotError: + # A malformed denylist entry can fail tokenizing or parsing on + # some dialects; skip it for this dialect rather than crashing + # validation (SqlglotError covers ParseError and TokenError). + continue + first_func = next(ast.find_all(exp.Func), None) + if first_func is not None and not isinstance(first_func, exp.Anonymous): + expanded.add(type(first_func).key.lower()) + return frozenset(expanded) + + def _check_functions( ast: exp.Expression, denied: frozenset[str], ) -> None: + canonical = _canonical_denied(denied) for func in ast.find_all(exp.Func): if isinstance(func, exp.Anonymous): name = func.name else: name = type(func).key - if name.lower() in denied: + if name.lower() in canonical: raise WrenError( ErrorCode.BLOCKED_FUNCTION, f"Function '{name}' is not allowed. " diff --git a/core/wren/tests/unit/test_case_sensitivity.py b/core/wren/tests/unit/test_case_sensitivity.py index ddd573fa5..9d72fe0e9 100644 --- a/core/wren/tests/unit/test_case_sensitivity.py +++ b/core/wren/tests/unit/test_case_sensitivity.py @@ -183,7 +183,7 @@ def test_cte_rewriter_quoted_preserves_case( sql = f"SELECT u.id FROM {table_ref} AS u" ast = parse_one(sql, dialect="bigquery") user_ctes = rewriter._collect_user_cte_names(ast) - used, _refs = rewriter._collect_model_columns(ast, user_ctes) + used, _refs, _quoting = rewriter._collect_model_columns(ast, user_ctes) assert list(used.keys()) == [expected_model], ( f"expected the rewriter to bind {expected_model!r} for {table_ref}, " @@ -202,7 +202,7 @@ def test_cte_rewriter_unquoted_ci_fallback(dual_case_manifest_b64): # Unquoted ``Users`` matches the ``Users`` model (exact case) — not # ``users`` even though both exist. ast = parse_one("SELECT u.id FROM Users AS u", dialect="bigquery") - used, _refs = rewriter._collect_model_columns(ast, set()) + used, _refs, _quoting = rewriter._collect_model_columns(ast, set()) assert list(used.keys()) == ["Users"] diff --git a/core/wren/tests/unit/test_cte_rewriter.py b/core/wren/tests/unit/test_cte_rewriter.py index 6ea85583c..e801b9beb 100644 --- a/core/wren/tests/unit/test_cte_rewriter.py +++ b/core/wren/tests/unit/test_cte_rewriter.py @@ -11,6 +11,7 @@ import sqlglot from wren.mdl import get_session_context from wren.mdl.cte_rewriter import CTERewriter, get_sqlglot_dialect from wren.model.data_source import DataSource +from wren.model.error import ErrorCode, WrenError pytestmark = pytest.mark.unit @@ -316,12 +317,52 @@ class TestCTEEdgeCases: ast = sqlglot.parse_one(result, dialect="duckdb") assert ast.args["with_"].args.get("recursive") - def test_no_model_references_fallback(self): - """Query referencing no models falls back to direct transform_sql.""" - rw = _make_rewriter(_SINGLE_MODEL_MANIFEST, fallback=True) - # This should raise because 'unknown_table' is not in the manifest, - # and the fallback transform_sql will also fail. - with pytest.raises(Exception): + @pytest.mark.parametrize( + ("sql", "data_source"), + [ + ("SELECT 1", DataSource.postgres), + ("SELECT 1 + 1 AS two", DataSource.postgres), + ("SELECT CURRENT_DATE", DataSource.postgres), + # Oracle forces identifier quoting (identify=True) on output; the + # passthrough must still round-trip a scalar query cleanly. + ("SELECT 1", DataSource.oracle), + # A pure inline-row STRUCT spine — no model reference. + ( + "SELECT stage, sort_order FROM UNNEST([" + "STRUCT('Prospecting' AS stage, 1 AS sort_order), " + "('Qualification', 2), ('Proposal', 3)])", + DataSource.bigquery, + ), + ], + ) + def test_scalar_or_pure_tvf_passes_through(self, sql, data_source): + """SQL with no model/view reference is passed through, never rejected.""" + rw = _make_rewriter(_SINGLE_MODEL_MANIFEST, data_source, fallback=True) + out = rw.rewrite(sql) + + # Nothing to expand — no model CTE is injected. + assert "with" not in out.lower() + # It round-trips to the same AST in the target dialect. + dialect = get_sqlglot_dialect(data_source) + assert sqlglot.parse_one(out, dialect=dialect) == sqlglot.parse_one( + sql, dialect=dialect + ) + + def test_scalar_passes_through_without_fallback(self): + """Pure scalar SQL passes through even when fallback is disabled.""" + rw = _make_rewriter(_SINGLE_MODEL_MANIFEST, fallback=False) + out = rw.rewrite("SELECT 1") + assert "with" not in out.lower() + + def test_unresolved_table_raises_without_fallback(self): + """A non-model table reference raises when fallback is disabled. + + Only pure scalar / TVF SQL passes through unconditionally; a table + reference that resolves to no MDL model or view is still a rewriter + miss and must surface, not be silently passed on. + """ + rw = _make_rewriter(_SINGLE_MODEL_MANIFEST, fallback=False) + with pytest.raises(ValueError): rw.rewrite("SELECT * FROM unknown_table") @@ -431,3 +472,152 @@ class TestDialectMapping: def test_local_file_maps_to_duckdb(self): assert get_sqlglot_dialect(DataSource.local_file) == "duckdb" + + +# --------------------------------------------------------------------------- +# Tests: case-insensitive column & model binding +# --------------------------------------------------------------------------- + +_MIXED_CASE_COL_MANIFEST = { + "catalog": "wren", + "schema": "public", + "models": [ + { + "name": "events", + "tableReference": {"schema": "main", "table": "events"}, + "columns": [ + {"name": "id", "type": "integer"}, + {"name": "Amount", "type": "integer"}, + ], + "primaryKey": "id", + } + ], +} + +_MIXED_CASE_MODEL_MANIFEST = { + "catalog": "wren", + "schema": "public", + "models": [ + { + "name": "CaseModel", + "tableReference": {"schema": "main", "table": "case_model"}, + "columns": [{"name": "id", "type": "integer"}], + "primaryKey": "id", + } + ], +} + +_CASE_COLLISION_MANIFEST = { + "catalog": "wren", + "schema": "public", + "models": [ + { + "name": "clash", + "tableReference": {"schema": "main", "table": "clash"}, + "columns": [ + {"name": "Year", "type": "integer"}, + {"name": "year", "type": "integer"}, + ], + } + ], +} + + +class TestCaseSensitiveBinding: + def test_force_identify_detected_from_dialect(self): + """Upper-folding dialects force-quote output; others do not.""" + assert _make_rewriter(_SINGLE_MODEL_MANIFEST, DataSource.oracle)._force_identify + assert _make_rewriter( + _SINGLE_MODEL_MANIFEST, DataSource.snowflake + )._force_identify + assert not _make_rewriter( + _SINGLE_MODEL_MANIFEST, DataSource.postgres + )._force_identify + assert not _make_rewriter( + _SINGLE_MODEL_MANIFEST, DataSource.bigquery + )._force_identify + + def test_case_only_column_collision_rejected_on_ci_dialect(self): + """Case-distinct columns are rejected on a case-insensitive dialect.""" + with pytest.raises(WrenError) as exc: + _make_rewriter(_CASE_COLLISION_MANIFEST, DataSource.bigquery) + assert exc.value.error_code == ErrorCode.INVALID_MDL + + def test_case_distinct_columns_allowed_on_case_sensitive_dialect(self): + """Postgres can hold case-distinct columns — build succeeds, refs bind.""" + rw = _make_rewriter(_CASE_COLLISION_MANIFEST, DataSource.postgres) + assert rw._case_sensitive_columns + # Quoted refs select each case-distinct column exactly. + out = rw.rewrite('SELECT "Year", "year" FROM clash') + assert _has_cte(out, "clash", dialect="postgres") + assert "select 1" not in out.lower() + + def test_quoted_wrong_case_column_rejected_on_case_sensitive_dialect(self): + """A quoted ref to a non-existent case variant fails loudly.""" + rw = _make_rewriter(_CASE_COLLISION_MANIFEST, DataSource.postgres) + with pytest.raises(WrenError) as exc: + rw.rewrite('SELECT "YEAR" FROM clash') + assert exc.value.error_code == ErrorCode.INVALID_SQL + + def test_ambiguous_unquoted_ref_rejected_on_case_sensitive_dialect(self): + """An unquoted ref matching multiple case-distinct columns is ambiguous.""" + rw = _make_rewriter(_CASE_COLLISION_MANIFEST, DataSource.postgres) + with pytest.raises(WrenError) as exc: + rw.rewrite("SELECT YEAR FROM clash") + assert exc.value.error_code == ErrorCode.INVALID_SQL + + def test_quoted_mixed_case_column_binds(self): + """A quoted mixed-case column reference binds to the model CTE.""" + rw = _make_rewriter(_MIXED_CASE_COL_MANIFEST, DataSource.postgres) + out = rw.rewrite('SELECT "Amount" FROM events') + assert _has_cte(out, "events", dialect="postgres") + # The CTE exposes the column with the user's quoting and does not + # collapse to a column-less SELECT 1. + assert "Amount" in out + assert "select 1" not in out.lower() + + def test_unquoted_mixed_case_column_binds(self): + """An unquoted mixed-case column reference binds via dialect folding.""" + rw = _make_rewriter(_MIXED_CASE_COL_MANIFEST, DataSource.postgres) + out = rw.rewrite("SELECT Amount FROM events") + assert _has_cte(out, "events", dialect="postgres") + assert "select 1" not in out.lower() + + def test_quoted_mixed_case_model_binds(self): + """A quoted mixed-case model name binds to its CTE (not SELECT 1).""" + rw = _make_rewriter(_MIXED_CASE_MODEL_MANIFEST, DataSource.postgres) + out = rw.rewrite('SELECT id FROM "CaseModel"') + assert _has_cte(out, "CaseModel", dialect="postgres") + assert "select 1" not in out.lower() + + def test_quoting_does_not_leak_across_sources(self): + """A quoted same-name column elsewhere must not flip a model CTE alias. + + ``events`` has a mixed-case ``Amount`` and is referenced unquoted; a + user CTE also exposes ``amount`` referenced quoted. Quoting is scoped + per model, so the ``events`` CTE still mirrors the unquoted reference + (``AS Amount``) and the user's ``e.Amount`` binds on Postgres — the + CTE's quoted ``"amount"`` does not force-quote it. + """ + rw = _make_rewriter(_MIXED_CASE_COL_MANIFEST, DataSource.postgres) + out = rw.rewrite( + "WITH c AS (SELECT 1 AS amount) " + 'SELECT e.Amount, c."amount" FROM events e JOIN c ON 1=1' + ) + ast = sqlglot.parse_one(out, dialect="postgres") + events_cte = next( + cte for cte in ast.args["with_"].expressions if cte.alias == "events" + ) + proj = events_cte.this.expressions[0] + assert isinstance(proj, sqlglot.exp.Alias) + assert not proj.args["alias"].quoted, ( + f"events CTE alias should be unquoted (mirroring e.Amount), got: {out}" + ) + + def test_case_insensitive_ref_on_upper_folding_dialect(self): + """Oracle: a lower-case ref to a mixed-case column is canonicalized.""" + rw = _make_rewriter(_MIXED_CASE_COL_MANIFEST, DataSource.oracle) + out = rw.rewrite("SELECT amount FROM events") + # Output is force-quoted; the manifest-case column survives. + assert "Amount" in out + assert "select 1" not in out.lower() diff --git a/core/wren/tests/unit/test_policy.py b/core/wren/tests/unit/test_policy.py index 7ccafca0f..8ad0d4a37 100644 --- a/core/wren/tests/unit/test_policy.py +++ b/core/wren/tests/unit/test_policy.py @@ -6,7 +6,7 @@ These tests use sqlglot parsing only and do not require a database or wren-core. from __future__ import annotations import pytest -from sqlglot import parse_one +from sqlglot import exp, parse_one from wren.config import WrenConfig from wren.model.error import ErrorCode, WrenError @@ -109,6 +109,27 @@ def test_builtin_function_on_denied_list(): assert exc_info.value.error_code == ErrorCode.BLOCKED_FUNCTION +def test_denied_function_reclassified_by_sqlglot(): + """A denied name still matches when sqlglot maps it to a concrete subclass. + + sqlglot >=29 parses ``version()`` on postgres into ``exp.CurrentVersion`` + (``type(node).key == "currentversion"``), not ``exp.Anonymous``. Denying + ``"version"`` must still block it via the canonical-key expansion. + """ + ast = parse_one("SELECT version()", dialect="postgres") + # Guard the premise: the test only exercises the canonical-key expansion if + # sqlglot actually reclassified version() off exp.Anonymous. + func = next(ast.find_all(exp.Func)) + assert not isinstance(func, exp.Anonymous), ( + "version() should be reclassified to a concrete subclass in this " + "sqlglot version" + ) + config = WrenConfig(denied_functions=frozenset(["version"])) + with pytest.raises(WrenError) as exc_info: + validate_sql_policy(ast, _MODELS, config) + assert exc_info.value.error_code == ErrorCode.BLOCKED_FUNCTION + + def test_nested_denied_function(): sql = "SELECT * FROM (SELECT dblink('host=evil', 'SELECT 1') AS x) AS t" ast = parse_one(sql, dialect="postgres") diff --git a/core/wren/uv.lock b/core/wren/uv.lock index 0a403118b..6a88292b2 100644 --- a/core/wren/uv.lock +++ b/core/wren/uv.lock @@ -3449,7 +3449,7 @@ requires-dist = [ { name = "requests", specifier = ">=2.33.0" }, { name = "sentence-transformers", marker = "extra == 'memory'", specifier = ">=2.2" }, { name = "snowflake-connector-python", extras = ["pandas"], marker = "extra == 'snowflake'", specifier = ">=3.10" }, - { name = "sqlglot", specifier = ">=27" }, + { name = "sqlglot", specifier = ">=29" }, { name = "starlette", marker = "extra == 'ui'", specifier = ">=0.37" }, { name = "trino", marker = "extra == 'trino'", specifier = ">=0.333,<1" }, { name = "typer", specifier = ">=0.12" },