diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 28edec08..118a0707 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -4,7 +4,10 @@ on: push: branches: [main] pull_request: - branches: [main] + # Run CI for every PR regardless of base branch — stacked PRs + # targeting feature branches (e.g. `egor/dev-1450-…`) need the + # same lint + test gate as PRs to main. + branches: ['**'] jobs: lint-and-test: diff --git a/CLAUDE.md b/CLAUDE.md index 0196e51f..58e8fdc4 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -48,7 +48,7 @@ poetry run ruff check slayer/ tests/ - Use keyword arguments for functions with more than 1 parameter - Imports at the top of files - SQL generation uses sqlglot AST building (not string concatenation) -- **Two reference modes** (DEV-1369): SLayer has exactly two expression layers and the rules differ by design. **Mode A (SQL)** covers `Column.sql`, `Column.filter`, and `SlayerModel.filters` — sqlglot-parsed free SQL accepting any function call (`json_extract`, `coalesce`, `CASE WHEN`, …), bare names referencing the underlying table, and `__`-delimited join paths (`customers__regions.name` — `__` between hops, single dot before the leaf). **Mode B (DSL)** covers `ModelMeasure.formula`, `SlayerQuery.measures`, `SlayerQuery.filters`, and every other query field — Python-AST DSL accepting only `Column` / `ModelMeasure` references, single-dot dotted paths through joins, aggregation colon syntax (`revenue:sum`, `*:count`), transform calls, and arithmetic/boolean ops. DSL mode rejects raw SQL function calls, `__` in user input, and bare names that don't resolve. The internal carve-out: `Column.name` accepts `__` because `_query_as_model` flattens joined columns into virtual-model columns like `stores__name`. Single source of truth: [docs/concepts/references.md](docs/concepts/references.md). Predicate-promotion (DEV-1336) is removed — a query filter naming a `Column` whose `sql` contains a window function now raises with a suggestion to use a rank-family transform (`rank() <= N`, etc.). DEV-1378 closed two implementation gaps: (1) Mode A model filters now route through `parse_sql_predicate` at enrichment time (previously they hit the DSL parser and `Column.filter`/`SlayerModel.filters` containing arbitrary SQL functions raised at runtime); (2) `SlayerQuery.filters` (Mode B) accepts a small lowercase allowlist of string-hygiene scalars — `lower`, `upper`, `trim`, `replace`, `substr`, `instr`, `length`, `concat` — plus the SQL `||` concat operator (rewritten to `concat(...)`). +- **Two reference modes** (DEV-1369): SLayer has exactly two expression layers and the rules differ by design. **Mode A (SQL)** covers `Column.sql`, `Column.filter`, and `SlayerModel.filters` — sqlglot-parsed free SQL accepting any function call (`json_extract`, `coalesce`, `CASE WHEN`, …), bare names referencing the underlying table, and `__`-delimited join paths (`customers__regions.name` — `__` between hops, single dot before the leaf). **Mode B (DSL)** covers `ModelMeasure.formula`, `SlayerQuery.measures`, `SlayerQuery.filters`, and every other query field — Python-AST DSL accepting only `Column` / `ModelMeasure` references, single-dot dotted paths through joins, aggregation colon syntax (`revenue:sum`, `*:count`), transform calls, and arithmetic/boolean ops. DSL mode rejects raw SQL function calls, `__` in user input, and bare names that don't resolve. The internal carve-out: `Column.name` accepts `__` because `_query_as_model` flattens joined columns into virtual-model columns like `stores__name`. Single source of truth: [docs/concepts/references.md](docs/concepts/references.md). Predicate-promotion (DEV-1336) is removed — a query filter naming a `Column` whose `sql` contains a window function now raises with a suggestion to use a rank-family transform (`rank() <= N`, etc.). DEV-1378 closed two implementation gaps: (1) Mode A model filters now route through `parse_sql_predicate` at enrichment time (previously they hit the DSL parser and `Column.filter`/`SlayerModel.filters` containing arbitrary SQL functions raised at runtime); (2) `SlayerQuery.filters` (Mode B) accepts a small lowercase allowlist of string-hygiene scalars — `lower`, `upper`, `trim`, `replace`, `substr`, `instr`, `length`, `concat`, `like` — plus the SQL `||` concat operator. DEV-1484 wired up the last two (the reference docs had promised them but the typed pipeline rejected both): `||` is desugared to a `concat(...)` scalar call via a `|`/BitOr rewrite in `slayer/engine/syntax.py`, and `like(value, pattern)` is a 2-arg scalar that emits the SQL `LIKE` operator (sqlglot `exp.Like`) — wrap it in `not (...)` for `NOT LIKE`. - Models, columns, measures (formulas), and aggregations have an optional `meta: Dict[str, Any]` field for arbitrary user-defined JSON metadata. Persisted in storage, editable via MCP (`edit_model`), HTTP API, and CLI. `inspect_model` renders `meta` for any entity that has it set; the column is auto-pruned when no entity in the section uses meta. - **Schema versioning**: `SlayerModel`, `SlayerQuery`, and `DatasourceConfig` carry a `version: int` (currently `7` for `SlayerModel`, `3` for `SlayerQuery`, `1` for `DatasourceConfig`). On load, older versions are upgraded via the converter chain in `slayer/storage/migrations.py` before Pydantic validates the dict. Saves always emit the current version. The hook is on the Pydantic class itself (`@model_validator(mode="before")`), so every storage backend — YAML, SQLite, third-party backends, plus MCP/API/dbt entry points — gets migrations automatically without backend changes. The v1→v2 converter (in `slayer/storage/v2_migration.py`) merges v1 `dimensions`+`measures` into v2 `columns`, repurposes `measures` to hold `ModelMeasure` formulas, and renames `SlayerQuery.fields`→`measures`. The v2→v3 converter (in `slayer/storage/v3_migration.py`) drops the legacy `dry_run`/`explain` fields from `SlayerQuery` (they are now engine kwargs only — `engine.execute(query, dry_run=..., explain=...)`) and walks `SlayerModel.source_queries` entries through the SlayerQuery chain. The v3→v4 converter (in `slayer/storage/v4_migration.py`) requires non-empty `data_source` on table-backed SlayerModel dicts (`sql_table` or `sql` mode); query-backed models (`source_queries` set) are exempt because their `data_source` is filled by `engine._validate_and_populate_cache` from the resolved virtual model before save. The v4 converter also ships layout migrators that move legacy `models/.yaml` flat files into `models//.yaml` and rebuild the SQLite `models` table with a composite `(data_source, name)` primary key. The v4→v5 converter (in `slayer/storage/v5_migration.py`, DEV-1361) coarse-renames `Column.type` legacy values to the new sqlglot-aligned vocabulary (`string→TEXT`, `number→DOUBLE`, `time→TIMESTAMP`, etc.) and strips the dead aggregation pseudo-types (`count`/`sum`/...). The v5→v6 converter (in `slayer/storage/v6_migration.py`, DEV-1375) is a no-op forward — v6 introduces a single new optional field, `Column.sampled: Optional[str]`, that caches the per-column sample-value snapshot consumed by [`search`](docs/concepts/search.md) and `inspect_model`. The v6→v7 converter (in `slayer/storage/v7_migration.py`, DEV-1480) is also a no-op forward — v7 adds two more `Column` fields: `sampled_values: Optional[List[str]]` (the structured top-50-by-frequency list, paired with the existing text `sampled`) and `distinct_count: Optional[int]` (true cardinality at profile time, surfaced via a secondary `count_distinct` query on overflow). Storage backends additionally introspect each model's datasource on first load and refine `DOUBLE → INT` for base columns whose live SQL type is integer (`slayer/storage/type_refinement.py`); the refined model is written back so subsequent loads skip both steps. `SlayerQuery` v3 sets `extra="forbid"` so unknown fields raise. See [docs/concepts/models.md](docs/concepts/models.md#schema-versioning). - **DataType / CAST emission** (v5, DEV-1361): `DataType` values match sqlglot's `exp.DataType.Type` byte-for-byte: `TEXT`, `INT`, `DOUBLE`, `BOOLEAN`, `DATE`, `TIMESTAMP`. Auto-ingestion narrows integer DB types (INTEGER/BIGINT/SERIAL/INT8…/UINT8…) to `INT`, and floats/numerics (FLOAT/DOUBLE/DECIMAL with scale>0) to `DOUBLE`; NUMERIC(p,0) and DECIMAL(p,0) are integer-shaped and narrow to `INT`. The SQL generator wraps **non-bare** `Column.sql` (function calls, arithmetic, CASE WHEN) in `CAST(... AS )` driven by `Column.type`; bare identifiers and `sql=None` are emitted unchanged. `TEXT` is a no-op (skipped, since `CAST AS TEXT` is cosmetic and doesn't unwrap SQLite's JSON-quoted strings anyway). `ModelMeasure.type` (also `Optional[DataType]`) declares the formula's result type; when set, the aggregation expression is wrapped in an outer CAST. Lenient `before`-validators on `Column.type` and `ModelMeasure.type` absorb legacy lowercase agent input (`"string"` → `TEXT`, `"number"` → `DOUBLE`, etc.) and silently drop dropped pseudo-types. `slayer storage migrate-types --dry-run [--data-source X]` runs the storage refinement step explicitly for batch / inspectable usage. Schema-drift detection (`data_type_bucket`) keeps `INT` and `DOUBLE` in the same `"number"` bucket so a `DOUBLE`-typed persisted column does not flag as drift against an `INT` live column. @@ -69,7 +69,7 @@ poetry run ruff check slayer/ tests/ - **`apply_drift_deletes`** (DEV-1356): `await engine.apply_drift_deletes(deletes)` applies each entry via `engine.edit_model_remove` / `engine.delete_model_by_name` and returns `ApplyDriftResult(applied, errors, residual)`. Per-entry failures are captured; processing continues. Surfaced **only** via `slayer validate-models --force-clean [--yes]` — destructive auto-application is opt-in at the CLI layer. Not exposed via MCP or REST. - Queries support `measures` (renamed from `fields` in v2) — list of `{"formula": "...", "name": "...", "label": "..."}` parsed by `slayer/core/formula.py`. `label` is an optional human-readable display name (also supported on `ColumnRef` and `TimeDimension`). - **Dim-only queries deduplicate**: a `SlayerQuery` with no measures and at least one dimension/time-dimension auto-emits `GROUP BY ` and returns the distinct combinations (Cube.js semantics). The `GROUP BY` is applied **before** `LIMIT`, so a row cap can never silently drop unique tuples that surface past row N. The rule is unconditional — no opt-out flag. Implemented in `SQLGenerator._generate_base` as `dim_only_dedup = bool(group_by_columns) and not enriched.measures`, OR'd into `needs_group_by` alongside `has_aggregation` / `cross_model_measures` / `skip_isolated`. With aggregating measures present the aggregation `GROUP BY` is already mandatory, so the dim-only term is idempotent. -- **Result column naming**: `revenue:sum` → `orders.revenue_sum` (colon becomes underscore). `*:count` → `orders._count` — the `*` is dropped but the underscore is kept as a leading marker so the alias never collides with a user-defined column literally named `count`. When converting queries to models (`create_model_from_query`), the same colon-to-underscore mapping applies. An explicit `name` on the measure spec overrides the canonical form for **both** simple aggregations and arithmetic/transform formulas — `{"formula": "amount:sum", "name": "rev"}` surfaces as `orders.rev`. This matters most for inner stages of multi-stage `source_queries`, where downstream stages reference inner-stage outputs by the chosen name. **DEV-1443**: when a query measure is renamed via `name`, filters and ORDER BY entries in the same node may reference EITHER the raw colon form (`amount:sum > 100`) OR the user alias (`rev > 100`) — both resolve to the user alias, and a colon-form filter classifies as HAVING. Renaming never changes the legal filter/order form. Two enrichment-time validations on renames: (1) a query measure `name` colliding with a source column on the source model is rejected (the alias-form filter would otherwise silently bind to the source column instead of the aggregate); (2) a rename whose canonical alias (e.g. `revenue_sum`) literally shadows a source column on the same model is also rejected — the colon-form filter would otherwise be ambiguous between the renamed aggregate and the source column. **DEV-1448**: the user-supplied `name` on a *cross-model* aggregated measure (`{"formula": "customers.revenue:sum", "name": "cust_rev"}`) is now honored for the projection alias and the downstream-stage virtual model column. Only the canonical **leaf** of the dotted path is swapped to the user name; the hop path is preserved — same dot-syntax shape as every other multi-hop caller-facing key. So `{"formula": "customers.revenue:sum", "name": "cust_rev"}` surfaces as `orders.customers.cust_rev`, and `{"formula": "customers.regions.population:sum", "name": "region_pop"}` surfaces as `orders.customers.regions.region_pop`. Two caller-facing surfaces: (a) the top-level result-column key uses the hop-preserved form (`orders.customers.cust_rev`); (b) the downstream-stage virtual model column built by `_query_as_model` uses the BARE user name (`cust_rev`) — a special-case in the cross-model loop short-circuits the `__`-flattening of `_alias_to_short` whenever `cm.name` is a bare identifier, so a stage-2 reference like `cust_rev:max` resolves directly. The inner `CrossModelMeasure.measure` (used to build the CTE aggregate expression) intentionally retains the canonical form — only the outer handle is renamed. The canonical-collision guard from DEV-1443 was lifted into a pre-pass so it fires for both local AND cross-model renames symmetrically. Covers `customers.*:count` and `customers.col:count_distinct` cross-model variants. Same-stage **filter** referencing a renamed cross-model measure (`filters=["cust_rev > 100"]` or the colon form `"customers.revenue:sum > 100"`) is NOT auto-resolved — both raise at strict resolution; the SQL generator has no path to route the bare user alias to the cross-model CTE's output column for filter classification, so admitting the filter would emit broken SQL (`WHERE orders.cust_rev > 100` against a column that doesn't exist on the base table). Filter remap stays DEV-1445 territory. **ORDER BY** referencing the bare user alias (`order=[{"column": "cust_rev"}]`) DOES resolve via `SQLGenerator._resolve_order_column`'s `alias_lookup[cm.name] = cm.alias` mapping and emits `ORDER BY "orders.customers.cust_rev"` against the cross-model CTE's output column. Workaround until DEV-1445 lands: restructure as a multi-stage `source_queries` so the cross-model measure becomes a local measure in the downstream stage. See companion tickets DEV-1445 (cross-model filter remap) and DEV-1446 (transform-wrapped inner-ref dedup). +- **Result column naming**: `revenue:sum` → `orders.revenue_sum` (colon becomes underscore). `*:count` → `orders._count` — the `*` is dropped but the underscore is kept as a leading marker so the alias never collides with a user-defined column literally named `count`. When converting queries to models (`create_model_from_query`), the same colon-to-underscore mapping applies. An explicit `name` on the measure spec overrides the canonical form for **both** simple aggregations and arithmetic/transform formulas — `{"formula": "amount:sum", "name": "rev"}` surfaces as `orders.rev`. This matters most for inner stages of multi-stage `source_queries`, where downstream stages reference inner-stage outputs by the chosen name. **DEV-1443**: when a query measure is renamed via `name`, filters and ORDER BY entries in the same node may reference EITHER the raw colon form (`amount:sum > 100`) OR the user alias (`rev > 100`) — both resolve to the user alias, and a colon-form filter classifies as HAVING. Renaming never changes the legal filter/order form. Two enrichment-time validations on renames: (1) a query measure `name` colliding with a source column on the source model is rejected (the alias-form filter would otherwise silently bind to the source column instead of the aggregate); (2) a rename whose canonical alias (e.g. `revenue_sum`) literally shadows a source column on the same model is also rejected — the colon-form filter would otherwise be ambiguous between the renamed aggregate and the source column. **DEV-1448**: the user-supplied `name` on a *cross-model* aggregated measure (`{"formula": "customers.revenue:sum", "name": "cust_rev"}`) is now honored for the projection alias and the downstream-stage virtual model column. Only the canonical **leaf** of the dotted path is swapped to the user name; the hop path is preserved — same dot-syntax shape as every other multi-hop caller-facing key. So `{"formula": "customers.revenue:sum", "name": "cust_rev"}` surfaces as `orders.customers.cust_rev`, and `{"formula": "customers.regions.population:sum", "name": "region_pop"}` surfaces as `orders.customers.regions.region_pop`. Two caller-facing surfaces: (a) the top-level result-column key uses the hop-preserved form (`orders.customers.cust_rev`); (b) the downstream-stage virtual model column built by `_query_as_model` uses the BARE user name (`cust_rev`) — a special-case in the cross-model loop short-circuits the `__`-flattening of `_alias_to_short` whenever `cm.name` is a bare identifier, so a stage-2 reference like `cust_rev:max` resolves directly. The inner `CrossModelMeasure.measure` (used to build the CTE aggregate expression) intentionally retains the canonical form — only the outer handle is renamed. The canonical-collision guard from DEV-1443 was lifted into a pre-pass so it fires for both local AND cross-model renames symmetrically. Covers `customers.*:count` and `customers.col:count_distinct` cross-model variants. Same-stage **filter** referencing a renamed cross-model measure (`filters=["cust_rev > 100"]` or the colon form `"customers.revenue:sum > 100"`) is NOT auto-resolved — both raise at strict resolution; the SQL generator has no path to route the bare user alias to the cross-model CTE's output column for filter classification, so admitting the filter would emit broken SQL (`WHERE orders.cust_rev > 100` against a column that doesn't exist on the base table). Filter remap stays DEV-1445 territory. **ORDER BY** referencing the bare user alias (`order=[{"column": "cust_rev"}]`) DOES resolve via `SQLGenerator._resolve_order_column`'s `alias_lookup[cm.name] = cm.alias` mapping and emits `ORDER BY "orders.customers.cust_rev"` against the cross-model CTE's output column. **Cross-model parametric result keys** (DEV-1450): a cross-model parametric aggregate keeps its kwarg-signature suffix in the result key — `customers.revenue:percentile(p=0.5)` → `orders.customers.revenue_percentile_p_0_5`, and `p=0.5` / `p=0.95` on the same target column yield two DISTINCT keys. This diverges from legacy, which dropped the suffix so the two variants collided (a bug). Plain `customers.revenue:sum` → `orders.customers.revenue_sum` in both. Workaround until DEV-1445 lands: restructure as a multi-stage `source_queries` so the cross-model measure becomes a local measure in the downstream stage. See companion tickets DEV-1445 (cross-model filter remap) and DEV-1446 (transform-wrapped inner-ref dedup). - **Response attributes**: `SlayerResponse.attributes` is a `ResponseAttributes` with `.dimensions` and `.measures` dicts, each mapping column alias → `FieldMetadata(label, format)`. Split by type so consumers can distinguish dimension metadata from measure metadata. - Available formula transforms: cumsum, time_shift, change, change_pct, rank, percent_rank, dense_rank, ntile, first (FIRST_VALUE window ASC), last (FIRST_VALUE window DESC), lag, lead, consecutive_periods. time_shift uses a self-join CTE where the shifted sub-query has the time column expression offset by INTERVAL (calendar-based, gap-safe). change and change_pct are desugared at enrichment time into a hidden time_shift + arithmetic expression. lag/lead use LAG/LEAD window functions directly (more efficient but produce NULLs at edges). Non-transform SQL function calls (`nullif`, `coalesce`, `ln`, `sqrt`, etc.) may also wrap aggregated refs inside arithmetic expressions, e.g. `"*:count / nullif(revenue:max, 0)"` — the call passes through to emitted SQL while the inner refs resolve to their measure aliases - **Rank-family transforms** (DEV-1353): `rank`, `percent_rank`, `dense_rank`, and `ntile` are timeless window-function transforms emitted as `RANK() / PERCENT_RANK() / DENSE_RANK() / NTILE(n) OVER (... ORDER BY DESC)`. They default to **no `PARTITION BY`** (rank across the entire result set, unlike cumsum/lag/lead which auto-partition by query dimensions), and accept an optional `partition_by=col` or `partition_by=[col1, col2]` kwarg to opt into per-partition ranking; the columns referenced must be query dimensions or time dimensions. `ntile` additionally requires `n=`. Standard SQL across SQLite (≥3.25), Postgres, DuckDB, MySQL, and ClickHouse — no UDFs needed. @@ -78,6 +78,7 @@ poetry run ruff check slayer/ tests/ - **Bare same-model derived refs in `Column.sql`** (DEV-1410): A bare identifier inside a `Column.sql` that names a sibling **derived** column on the host model is inlined parenthesised, identical to the qualified form (`A.foo` and bare `foo` produce the same SQL). Inlining is scope-guarded — refs inside a nested scope (sub-query, set-op branch, CTE, `VALUES`) are left alone because they belong to the inner rowset. Cycles in the derived-ref graph (`c1.sql = "c2 + 1"`, `c2.sql = "c1 - 1"`) raise `slayer.core.errors.ColumnCycleError` at `storage.save_model` time, before the model reaches a query. `ColumnCycleError` subclasses both `SlayerError` and `ValueError` so existing `except ValueError` call sites keep working. `StorageBackend.save_model` is a template method that runs the cycle validator before delegating to the backend's `_save_model_impl`; the migration write-back at `_migrate_and_refine_on_load` passes `_validate=False` so legacy cyclic data remains loadable. Save-time validation stays within the model's `data_source`; unresolved join targets are silently skipped (best-effort) — the compile-time guard in `slayer/engine/column_expansion.py` remains authoritative. - Filters support `{variable}` placeholders substituted from `query.variables: Dict[str, Any]`. Values must be str/number, inserted as-is. `{{`/`}}` for literal braces. Undefined variables raise errors. - Models can have explicit `joins` to other models (LEFT JOINs). Cross-model measures use dotted syntax with colon aggregation (`customers.revenue:sum`) and multi-hop (`customers.regions.name`). Joins are auto-resolved by walking the join graph. Transforms work on cross-model measures (`cumsum(customers.revenue:sum)`) +- **Filtered-local isolation** (DEV-1503): a host aggregate whose `Column.filter` references a joined table (`loss_payment_amt:sum` where `loss_payment_amt` has `filter="loss_payment.has_flag = 1"`) is hoisted out of the host `_base` SELECT into its own `_cm_*` CTE — same isolation mechanism as forward cross-model aggregates, host-rooted variant. Without isolation, two such measures whose filter joins are different INNER joins on different tables would intersect in the host base to only the rows present in BOTH targets, silently corrupting both aggregates. The planner trigger is structural: cross-model planner fires whenever `AggregateKey.source.path` is non-empty OR `column_filter_key.referenced_join_paths` is non-empty. The plan carries `cte_root_model = host.name` as the host-rooted disambiguator (vs forward / re-rooted, which target-root). An AGGREGATE-phase host filter referencing the isolated aggregate (`loss_payment_amt:sum > 1000`) routes to an outer WHERE on the combined non-aggregating SELECT — not HAVING-into-the-CTE, which would surface host rows as NULL instead of dropping them. Mixed filters (`loss_payment_amt:sum > 1000 AND total_amount:sum > 10`) promote the non-isolated operand into `_base` as a hidden aux slot so the outer WHERE can reference both. See [docs/architecture/cross-model-aggregates.md](docs/architecture/cross-model-aggregates.md#strategy-3-filtered-local-isolation-dev-1503). - **Path-based table aliases**: Joined tables use `__`-delimited path aliases in SQL to disambiguate diamond joins. In queries, dots denote paths (`customers.regions.name`); in model SQL definitions, `__` denotes the table alias (`customers__regions.name`). For diamond joins (same table reached via different paths, e.g., `orders → customers → regions` AND `orders → warehouses → regions`), each path gets a unique alias (`customers__regions` vs `warehouses__regions`). Auto-ingestion creates only direct joins (one per FK on the source table); multi-hop paths are resolved at query time by walking each intermediate model's own joins - `SlayerQuery.source_model` accepts a model name, inline `SlayerModel`, or `ModelExtension` (extends a model with extra `columns`/`measures` formulas/`joins`). `create_model_from_query()` saves a query as a permanent model - Models can have `filters` (always-applied WHERE conditions, e.g., `"deleted_at IS NULL"`) diff --git a/docs/architecture/binding.md b/docs/architecture/binding.md new file mode 100644 index 00000000..1e7dd1f1 --- /dev/null +++ b/docs/architecture/binding.md @@ -0,0 +1,165 @@ +# Binding + +**Module:** `slayer/engine/binding.py` + +The binder takes a `ParsedExpr`, a scope (`ModelScope` or `StageSchema`), and a +`ResolvedSourceBundle`, and produces a `BoundExpr` whose leaves are resolved +`ValueKey`s. It is the stage that turns *names* into *structural identity* — and +it is a pure function of its inputs (P11). + +```mermaid +flowchart LR + parsed["ParsedExpr"] --> bind["_bind (recursive)"] + scope["scope: ModelScope | StageSchema"] --> bind + bundle["ResolvedSourceBundle"] --> bind + bind --> bk["BoundExpr(value_key: ValueKey)"] +``` + +## Output types + +- `BoundExpr(value_key)` — the whole expression's structural identity. `.phase` + is lifted from `value_key.phase`. +- `BoundFilter(value_key, phase, referenced_keys)` — adds the max phase any + referenced slot reaches (**P8**) and the full tuple of `ValueKey`s touched + anywhere in the tree (used by cross-model filter routing). + +Public entry points: `bind_expr`, `bind_filter`, `bind_time_dimension`, plus the +`walk_value_keys` traversal helper. + +## Resolving a reference + +`_resolve_ref` (bare) and `_resolve_dotted` (path) are where scope semantics +live (**P5**): + +**Against a `StageSchema`:** a bare name must be a column in the flat schema +(else `UnknownReferenceError`); a dotted ref raises `IllegalScopeReferenceError` +("downstream stages see a flat schema"). This is the DEV-1449 guard. + +**Against a `ModelScope`:** + +- A bare name resolves to a `ColumnKey(path=(), leaf=name)`, or to a + `ColumnSqlKey` if the column is derived (`col.sql` set and not a trivial + self-remap). If the name matches a `ModelMeasure` instead of a column, the + binder raises with a suggestion to expand it first — measure expansion is the + [parser stage](parsing.md)'s job, not the binder's. +- A `__`-bearing bare name is legal *only* if it exact-matches a literal column + name (the C11 carve-out); otherwise `IllegalScopeReferenceError`. +- A dotted ref walks the join graph: `parts[:-1]` are join targets, `parts[-1]` + is the leaf. Each hop must have a matching `ModelJoin` on the current model and + the target must be in the bundle; revisiting a model raises a legacy-compatible + `Circular join` error. The terminal column becomes a `ColumnKey` (or + `ColumnSqlKey`) carrying the hop path. + +### C14 — self-prefix stripping + +`_resolve_dotted` strips a leading segment equal to the host model's name before +walking: `orders.status` on an `orders`-rooted query → `status`; +`orders.customers.name` → `customers.name`. This is principle **C14**, +preserving the legacy convenience of qualifying with your own model name. + +## Binding aggregates + +`_bind_agg` builds an `AggregateKey`. The source is a `StarKey` for `*`, a +path-carrying `StarKey` for a cross-model star (`customers.*:count`, via +`_resolve_dotted_star`), or the bound column otherwise. Positional/kwarg +arguments bind through `_bind_agg_arg` — identifier args become `ColumnKey`, +literals normalize via `normalize_scalar`. Crucially, `_resolve_column_filter_key` +looks up the resolved source column's `Column.filter` and folds it into the key +as `column_filter_key` (a `SqlExprKey`) — so an aggregate over a filtered column +has a distinct identity (P3 / the `column_filter_key` invariant from +[Typed keys](typed-keys.md)). + +## Binding transforms (P9) + +`_bind_transform` produces a `TransformKey` whose `input` is the bound value to +transform — a `ValueSlotRef`, not a string. Two whitelists govern it: + +- `_TRANSFORM_KWARG_RULES` — per-op accepted kwargs. It is deliberately broader + than the legacy whitelist: every transform implicitly accepts `partition_by` + (so `change(measure, partition_by=…)` threads through to the desugared + `time_shift`, **C6**), and the rank family / `time_shift` / `lag` / `lead` / + `ntile` / `consecutive_periods` add their own. +- `_TRANSFORM_POSITIONAL_KWARGS` — the transforms whose documented DSL form + accepts positional params after the value (`time_shift(x, periods, + granularity)`, `lag(x, periods)`, `lead(x, periods)`). A name supplied both + positionally and as a kwarg is an error. + +`partition_by` values must bind to a `ColumnKey` / `ColumnSqlKey` (and become +`partition_keys`); other kwargs must fold to a scalar literal +(`_fold_to_scalar` also handles unary-minus over a numeric literal, the AST shape +of `periods=-1`). `_apply_transform_kwarg_defaults` validates required kwargs +(`ntile` needs a positive-integer `n`; `time_shift` needs `periods`) and applies +defaults (`lag`/`lead` default `periods=1`). Note the binder does **not** set +`time_key` — it lacks query context; the [stage planner](stage-planning.md) +attaches it after all binding completes. + +## Binding scalar calls (P1 / C12) + +`_bind_scalar` re-checks `SCALAR_FUNCTIONS` membership (defence-in-depth against +direct `ParsedExpr` construction that bypasses the parser) and builds a +`ScalarCallKey`. Arithmetic, comparison, boolean, and unary ops all become +`ArithmeticKey` with the operator string. + +## Filters and phase classification (P8) + +`bind_filter` binds the predicate, walks it via `walk_value_keys` to gather every +referenced `ValueKey`, takes `phase = max(referenced phases)`, and rejects raw +windows. The phase is computed entirely from the slots referenced — no text +analysis. `_reject_windowed_column_sql` raises `IllegalWindowInFilterError` if +any referenced `ColumnSqlKey` has a windowed `Column.sql` body — DEV-1369 +removed predicate promotion, so filtering on a window is an error, not an +auto-hoist. (Against a `StageSchema` this check is skipped — window detection +already happened when the upstream stage was bound.) + +### `alias_map` — filter/order refs by declared name (P4 / DEV-1445) + +`bind_filter` accepts an optional `alias_map: Dict[str, ValueKey]` mapping a +stage's declared-measure names (user `name`, declared name, canonical alias) to +their bound `ValueKey`. A bare ref matching an alias interns onto that exact slot +*before* any column lookup. This is what lets a filter reference a renamed measure +by alias: `filters=["rev >= 100"]` for a measure declared +`{"formula": "customers.revenue:sum", "name": "rev"}` binds `rev` onto the +cross-model aggregate slot — and because the dotted/colon form interns +structurally onto the *same* `AggregateKey`, both forms share one slot. Only +*measure* aliases enter the map (never dimension / time-dimension names), because +a time dimension's declared name is its raw column and a `created_at <= '…'` +filter must resolve to the raw column, not the truncated dimension slot. See +[Stage planning](stage-planning.md) for how the map is built. + +## `bind_time_dimension` + +Binds a `TimeDimension` into a `BoundExpr` carrying a `TimeTruncKey`. The +underlying column resolves against scope exactly like an identifier ref and must +have a temporal `Column.type` (`DATE` / `TIMESTAMP`). It may be a base +`ColumnKey` OR a **derived** `ColumnSqlKey` (DEV-1450 follow-up #4a): +`TimeTruncKey.column` is `Union[ColumnKey, ColumnSqlKey]`, and the generator +applies the `DATE_TRUNC` over the expanded `Column.sql` everywhere it would +over a bare column. + +> **Limitation.** Only `ModelScope` is accepted (a `StageSchema` raises — +> downstream stages already see the truncated column as a flat name). + +## `walk_value_keys` + +Yields every `ValueKey` reachable from a key (including the key itself), +recursing through `AggregateKey` source/args/kwargs, `TransformKey` +input/args/kwargs/partition_keys/time_key, `ArithmeticKey` operands, +`ScalarCallKey` args, and `BetweenKey` column/low/high. It is the typed +counterpart to the parser's `walk_parsed_refs`, used for phase computation and +cross-model filter routing. + +## Design rationale + +- **Why is the binder pure?** Everything it needs is in the bundle (P11). No + storage access, no `ContextVar`, no callback re-resolution — so binding a given + `(parsed, scope, bundle)` is deterministic and order-independent. This is the + single biggest simplification over the legacy enrichment closures. +- **Why fold `Column.filter` into the aggregate key rather than handle it at + render time?** Because two aggregates over the same column differ *as values* + when their filters differ; making that part of identity means the registry + interns correctly and the generator never has to reconcile two slots that are + "the same column but filtered differently". +- **Why does the binder leave `time_key` unset?** Resolving the active time + dimension needs query-level context (`main_time_dimension`, + `default_time_dimension`, the set of TDs in the query). The binder is + expression-local; the planner has the query, so it patches `time_key` there. diff --git a/docs/architecture/cross-model-aggregates.md b/docs/architecture/cross-model-aggregates.md new file mode 100644 index 00000000..151aed04 --- /dev/null +++ b/docs/architecture/cross-model-aggregates.md @@ -0,0 +1,225 @@ +# Cross-model aggregates + +**Modules:** `slayer/engine/cross_model_planner.py` (strategy + +`_maybe_reroot_cross_model_plan`), +`slayer/sql/generator.py` (`_render_with_cross_model_plans`, +`_render_rerooted_cross_model_cte`) + +A cross-model aggregate is `customers.revenue:sum` on an `orders`-rooted query — +an aggregate whose source carries a non-empty join path. Principle **P3** says it +shares the `AggregateKey` shape with a local aggregate (only `source.path` +differs), and that "base CTE vs cross-model CTE" is a *render strategy* decided +downstream, not a semantic split. The identity side of P3 holds cleanly. The +**render** side turned out to need two strategies — the most significant +deviation from the plan. + +## The identity is uniform; the rendering is not + +```mermaid +flowchart TB + agg["AggregateKey(source.path = ('customers',), agg='sum')"] + agg --> cmp["cross_model_planner.plan(...)"] + cmp --> plan["CrossModelAggregatePlan"] + plan -->|rerooted_plan is None| fwd["forward-path CTE
FROM bare target, GROUP BY forward dims"] + plan -->|rerooted_plan set| rr["re-rooted nested PlannedQuery
FROM target + target's joins"] +``` + +The planner detects the cross-model case structurally (`agg_path` non-empty in +`plan_query`) and invokes the strategy. The strategy is a substitutable +component — **I1**: `CrossModelPlanner` is a `Protocol`, +`IsolatedCteCrossModelPlanner` is the default — so the *shape* of the result +(`CrossModelAggregatePlan` in `planned.py`) is strategy-agnostic and only the +populating planner changes. + +## Strategy 1: `IsolatedCteCrossModelPlanner` (the plan's design) + +This is the planned design: one CTE per `(target_model, shared_grain)`. It walks +the join chain from host to target (`_walk_chain` → `JoinRequirement`s), groups +the aggregate at the first hop's target grain, and builds `join_back_pairs` so +the host LEFT JOINs the CTE back on the first-hop columns. `_make_cte_schema` +produces the CTE's typed projection. `_aggregate_alias` derives the output +column name via `canonical_agg_name`. + +### Host-filter routing (the `inherited_filter_policy` decision table) + +`classify_host_filter` is a pure classifier mapping each host filter to a +`FilterRoute`: + +| Filter references | Route | +| --- | --- | +| host-local row slot only | `DROP_HOST_LOCAL` (applied at host) | +| all on the joined-target path (row) | `PROPAGATE_WHERE` | +| cross-model agg-ref on the same target | `PROPAGATE_HAVING` | +| slots on a different joined branch | `DROP_UNREACHABLE` (+ warn) | +| mixed reachable + unreachable | `DROP_UNREACHABLE` (+ warn) | +| transform / POST phase | `STAY_AT_HOST_POST` | + +The planner threads each route into the explicit +`where_filter_ids` / `having_filter_ids` lists on `CrossModelAggregatePlan` so +the generator never re-classifies. The target model's own `SlayerModel.filters` +ride into `target_model_filters` (always-applied WHERE), and a `Column.filter` on +the aggregated column rides on the `AggregateKey` itself as a CASE-WHEN — neither +goes through host-filter classification. `shared_grain_slots` is the set of host +dimension/time-dimension slots reachable from the target, used to LEFT JOIN the +CTE back without changing cardinality. + +## Strategy 2: re-rooting (the deviation) + +`IsolatedCteCrossModelPlanner` alone is insufficient. When the host query carries +dimensions that are reachable from the target through the **target's own** join +graph (the legacy `_build_rerooted_enriched` case — e.g. +`policy_amount → policy → policy_number`), the forward-path CTE +("FROM bare target, GROUP BY forward-path dims only") collapses the host +dimension to a scalar `CROSS JOIN`: every host row gets the global aggregate +instead of a per-dimension value. + +`_maybe_reroot_cross_model_plan` detects this and attaches a nested re-rooted +plan. As of DEV-1450 follow-up #2 it lives in `cross_model_planner.py` and runs +**inside** `IsolatedCteCrossModelPlanner.plan` — the strategy owns the +forward-vs-re-rooted choice rather than `plan_query` patching the plan after the +fact: + +```mermaid +flowchart TB + detect["host dims/filters reachable from target
via the target's join graph?"] + detect -->|no| keep["keep forward-path plan"] + detect -->|yes| build["build a full SlayerQuery rooted at the target"] + build --> replan["subplan_builder(rerooted_query, rerooted_bundle)"] + replan --> attach["attach rerooted_plan / rerooted_grain_pairs / rerooted_agg_slot_id"] + attach --> gen["generator: _render_rerooted_cross_model_cte"] +``` + +It re-roots each host dimension/time-dimension/filter from the host's perspective +to the target's (`_reroot_ref`: host-local → `.`; on-target → bare; +through-target → strip the prefix), drops anything unreachable from the target +(matching legacy), reconstructs the local aggregate formula +(`_local_agg_formula`), builds a fresh `SlayerQuery` rooted at the target, and +compiles it via the injected `subplan_builder` callback (which `plan_query` +supplies as a `plan_query` recursion — keeping `cross_model_planner.py` free of a +`stage_planner` import). The sub-plan is rendered by +`_render_rerooted_cross_model_cte` as the `_cm_*` CTE (FROM target + the target's +joins, preserving host grain) and joined back on the re-rooted dimension via +`rerooted_grain_pairs`. + +### Why this is flagged as a deviation + +The plan envisioned the `inherited_filter_policy` decision table plus +`IsolatedCteCrossModelPlanner` as **the** cross-model mechanism. In practice +there are now **two** cross-model render strategies, both owned by the strategy +(`IsolatedCteCrossModelPlanner.plan` → `_maybe_reroot_cross_model_plan`), with +the re-rooting one bolted onto `CrossModelAggregatePlan` via `rerooted_plan` / +`rerooted_grain_pairs` / `rerooted_agg_slot_id`. P3's "one shape, render strategy +chosen downstream" holds for *identity* but not for *rendering* — and the +re-rooted path is, structurally, the legacy `_build_rerooted_enriched` shape +brought across to the typed plan. +This reintroduces (in a contained, typed form) the kind of "second resolution +path for a permutation" the redesign set out to eliminate. It works and is +tested, but it is the place a future reviewer should look first when reasoning +about cross-model behavior. + +## Strategy 3: filtered-local isolation (DEV-1503) + +A **cross-model-FILTERED local measure** is a host aggregate whose `Column.filter` +references a joined table — `loss_payment_amt:sum` where `loss_payment_amt` has +`filter="loss_payment.has_flag = 1"`. The aggregate's `source.path` is empty +(it's a local column), but its `column_filter_key.referenced_join_paths` is +non-empty, so emitting it inline in the host base SELECT would pull the +filter-target join into the host's FROM. With **two** such measures whose +filter targets are different INNER joins, the host base would intersect to +only the rows present in BOTH targets — silently corrupting both aggregates. + +The trigger predicate is structural: +`agg_path` non-empty (forward cross-model) **OR** +`column_filter_key.referenced_join_paths` non-empty (filtered-local). Both +route through `IsolatedCteCrossModelPlanner.plan`; the filtered-local branch +calls `_plan_filtered_local`, which builds a **host-rooted** nested +`PlannedQuery` (same `source_model`, same dims/TDs, only the filtered measure +as the single aggregate) and attaches it via the same +`rerooted_plan` / `rerooted_grain_pairs` / `rerooted_agg_slot_id` slots the +re-rooted path uses. The plan carries `cte_root_model = host_model.name` as +the disambiguator the renderer reads; `_render_rerooted_cross_model_cte` +short-circuits the source-model swap when `cte_root_model` is set. + +```mermaid +flowchart TB + detect["column_filter_key.referenced_join_paths non-empty?"] + detect -->|yes| build["_plan_filtered_local builds host-rooted SlayerQuery"] + build --> replan["subplan_builder(rerooted_query, bundle)"] + replan --> attach["attach with cte_root_model = host.name"] + attach --> gen["generator: _render_rerooted_cross_model_cte (host-rooted branch)"] +``` + +`subplan_builder` always passes `disable_dev1503_isolation=True` so the +recursive `plan_query` call inside the sub-plan does NOT re-trigger isolation +on the same measure. + +### Filter routing for filtered-local + +| Host filter phase | Route | +| --- | --- | +| ROW | propagate into the host-rooted sub-plan (so a non-dim filter like `status = 'active'` affects the isolated aggregate's rowset) | +| AGGREGATE | **outer combined-SELECT WHERE wrapper** (see below) | +| POST | stay at the existing host post-transform wrapper | + +### Outer combined-SELECT WHERE wrapper + +An AGGREGATE-phase host filter referencing an isolated aggregate +(`loss_payment_amt:sum > 1000`) cannot route as HAVING inside the `_cm_*` CTE: +the LEFT JOIN back to `_base` would surface host rows whose filtered +aggregate didn't meet the predicate with a NULL value instead of dropping +them. The renderer (`_render_with_cross_model_plans`) classifies each +AGGREGATE-phase filter; any that walks an `AggregateKey` matching an +isolated slot is routed to an outer WHERE on the **combined SELECT** (which +is non-aggregating — plain WHERE is legal). The renderer +(`_render_filter_for_outer_wrapper`) substitutes: + +- isolated `AggregateKey` → `.""` (the joined-back column), +- any other slot → `_base.""` (the host base's projection). + +Non-isolated aggregate operands of a mixed filter (`loss_payment_amt:sum > +1000 AND total_amount:sum > 10` where `total_amount:sum` isn't a public +measure) are promoted to hidden aux slots in `base_render_order` by the +existing `_add_local_aux_slots(aggregates_only=True)` pass — `_base` +materialises them so the outer WHERE can reference them, and the combined +public projection trims them out. + +## Generator side + +`generate_from_planned` delegates to `_render_with_cross_model_plans` when +`cross_model_aggregate_plans` is non-empty. Each plan renders as a `_cm_*` CTE +(forward-path or re-rooted), joined back to the host base. `Column.filter` on the +aggregated column renders as `SUM(CASE WHEN THEN END)`. See +[SQL generation](sql-generation.md). + +## Known limitations (documented, not blocking) + +- A host-local filter on a **no-dimension** cross-model-agg query is applied + nowhere (the empty `_base` placeholder doesn't filter; host-local filters are + excluded from the re-rooted CTE). Semantically ambiguous; rare. +- `time_shift` / `consecutive_periods` / `change` / `change_pct` over (or + alongside) a cross-model aggregate raise `NotImplementedError` — factor the + temporal transform into an earlier stage. +- Cross-model parametric-agg result keys diverge from legacy **by design**: + `customers.revenue:percentile(p=0.5)` → `…revenue_percentile_p_0_5` where + legacy dropped the kwarg suffix (`…revenue_percentile`). Legacy's drop was a + collision bug; the new path keeps the suffix. This violates **P10** for this + one combination and is tested structurally, not by parity. See + [the deviations list](index.md#deviations-from-the-plan). +- `weighted_avg(weight=qty)` / `corr(other=qty)` cross-model semantics (a + host-local weight column evaluated inside the target CTE) are not supported. + +## Design rationale + +- **Why a Protocol (I1)?** So the cross-model strategy is substitutable without + touching the plan shape or the generator. The re-rooting case shows the value: + it was added as a *second* population path for the same `CrossModelAggregatePlan` + struct, not as a new struct. +- **Why route filters in the planner, not the generator?** So the generator + renders each route mechanically. Classification needs the slot graph (which + slot is on which branch); putting it in the planner keeps the generator a + straight `WHERE`/`HAVING`/`CASE-WHEN` emitter. +- **Why re-root rather than emit a literal JOIN chain inside the CTE?** Parity + with legacy `_build_rerooted_enriched` for the grain-preserving case; emitting + the chain directly was the path not taken, and re-rooting reuses the whole + planner recursively, which is less code than a bespoke chain emitter — at the + cost of the second-strategy complexity above. diff --git a/docs/architecture/engine-orchestration.md b/docs/architecture/engine-orchestration.md new file mode 100644 index 00000000..f18c7780 --- /dev/null +++ b/docs/architecture/engine-orchestration.md @@ -0,0 +1,124 @@ +# Engine orchestration + +**Modules:** `slayer/engine/query_engine.py` (`_execute_pipeline`, +`save_model`), `slayer/engine/variables.py` + +`SlayerQueryEngine` is where the pipeline is wired into a runnable execution. It +also marks the boundary between the new typed pipeline and the legacy stack that +still co-exists. + +## `execute` → `_execute_pipeline` + +`execute(query, …)` dispatches over the input shape (str run-by-name, dict, list +DAG, `SlayerQuery`), then `_execute_pipeline` runs the linear pipeline: + +```mermaid +flowchart TB + pre["strip_source_model_prefix + snap_to_whole_periods"] + pre --> bundle["build_resolved_source_bundle (P11)"] + bundle --> qbexp["_expand_query_backed_model (LEGACY path)
source / referenced / stage-source models"] + qbexp --> norm["_normalize_stage → normalize_query (P0)"] + norm --> vars["apply_variables_to_query"] + vars --> plan["plan_stages (root last)"] + plan --> gen["generate_planned_stages → SQL"] + gen --> meta["build_response_metadata"] + meta --> exec["client.execute → SlayerResponse"] +``` + +The new typed pipeline is `build_resolved_source_bundle → _normalize_stage → +apply_variables_to_query → plan_stages → generate_planned_stages → +build_response_metadata`. Storage is consulted once, in +`build_resolved_source_bundle` (P11). + +`_normalize_stage` resolves each stage's source model from the bundle so +`MISPLACED_MEASURE` and custom-aggregation-aware `FUNC_STYLE_AGG` see the right +column / aggregation names; a sibling-sourced stage normalizes with `model=None`. +Slack warnings from every stage are collected and surface on +`SlayerResponse.warnings`. + +`_touched_models_for_plan` collects the model names a query-time DBAPI error +could be attributed to (bundle referenced models + cross-model targets + +query-backed base names) for schema-drift attribution. + +## Variables (`variables.py`) + +`merge_query_variables` collapses the four layers — **runtime > stage > outer > +model defaults** — into the effective dict that populates +`ResolvedSourceBundle.query_variables`. `apply_variables_to_query` returns a fresh +`SlayerQuery` with `{var}` substituted in `filters` (the only field legacy +substituted; formula text, `Column.sql`, `Column.filter`, `SlayerModel.filters` +are deliberately not substituted). `dry_run_placeholders=True` fills unresolved +valid placeholders with `"0"` (the legacy save-time dry-run behavior); invalid +names still raise. + +## `save_model` + +`save_model` runs `normalize_model` (the [slack layer](slack-normalization.md)) +so persisted formulas land canonical, then persists. For a **query-backed** model +it rejects user-supplied cache fields and calls `_validate_and_populate_cache`, +which renders the backing query and stores `columns` / `backing_query_sql` / +`data_source`. + +## Where the legacy stack still runs (the deviation) + +This is the most important orchestration fact, and the largest gap between the +plan and the implementation. The cutover routed **top-level query planning** +through the new pipeline, but **query-backed model expansion** still runs +entirely on the legacy stack — in production, not just tests: + +| Path | Renders the backing SQL via | +| --- | --- | +| `_execute_pipeline` → `_expand_query_backed_model` | `_query_as_model` → `enrich_query` → `SQLGenerator.generate` (legacy) | +| `save_model` → `_validate_and_populate_cache` | `_query_as_model` → … → `SQLGenerator.generate` (legacy) | + +`_expand_query_backed_model` turns a model's `source_queries` into a virtual +`sql`-mode model whose `.sql` is the rendered backing query — and it does so by +calling `_query_as_model`, which internally runs `_enrich` and the legacy +`SQLGenerator.generate`. The new pipeline then treats that virtual model as a +plain `sql`-mode model and plans/renders the **outer** query through the typed +path. + +`_execute_pipeline` invokes `_expand_query_backed_model` for the source model +(line ~581), for query-backed referenced (join/cross-model target) models +(~621), and for non-root stage sources (~640). So: + +- `enrichment.py` (~100 KB), `enriched.py` (`EnrichedQuery` / `EnrichedMeasure`), + `_query_as_model`, the legacy `SQLGenerator.generate`, + `_rewrite_funcstyle_aggregations`, and the `_forbidden_sibling_refs_var` / + `_join_target_resolving_var` `ContextVar`s **all still exist and still run**; +- the typed pipeline is the resolution path for the *outer* query and for the + four acceptance bugs; the legacy stack is load-bearing for query-backed inner + rendering and (via the synth adapter) dialect SQL emission. + +The plan's stage-7b bullet said the cutover would delete all of the above. In +practice every deletion is deferred to **DEV-1452**, which must first migrate +query-backed expansion onto the typed pipeline (it shares machinery with +cross-stage join rendering). Until then, do not read the legacy modules as dead +code — `_query_as_model` is on a hot path. + +## Pre-processing before the "single" slack pass + +`_execute_pipeline` runs `strip_source_model_prefix()` and (when +`whole_periods_only`) `snap_to_whole_periods()` *before* `_normalize_stage`. +These are query-shape transforms rather than slack-token rewrites, but they mean +the pipeline does not literally "begin with a single slack-normalization pass" +(**P0**). A minor deviation, noted for completeness in +[the deviations list](index.md#deviations-from-the-plan). + +## Design rationale + +- **Why split the cutover this way?** Bisectability. Flipping the outer query + while leaving query-backed expansion on legacy let the cutover land with all + non-integration tests green and integration green, without a single + thousand-line "delete everything" commit. The cost is the temporary + two-pipeline coexistence above. +- **Why does query-backed expansion produce a virtual `sql`-mode model rather + than planning the inner stages directly?** Because the outer typed pipeline + already knows how to consume a `sql`-mode model. Re-expressing a query-backed + model as `{sql_table: None, sql: }` lets the outer + planner stay oblivious to query-backedness — at the cost of rendering the inner + SQL through the legacy generator for now. DEV-1452's job is to make the inner + rendering go through `plan_stages` / `generate_planned_stages` too. +- **Why consult storage only in the bundle builder (P11)?** So everything after + it is pure and order-independent. The legacy `ContextVar` re-resolution exists + *only* on the query-backed/legacy path; the new path has none. diff --git a/docs/architecture/errors-and-warnings.md b/docs/architecture/errors-and-warnings.md new file mode 100644 index 00000000..8431e482 --- /dev/null +++ b/docs/architecture/errors-and-warnings.md @@ -0,0 +1,80 @@ +# Errors and warnings + +**Modules:** `slayer/core/errors.py`, `slayer/core/warnings.py` + +The redesign replaced anonymous `ValueError`s scattered through enrichment with a +typed error vocabulary. Each error carries the offending input, a scope summary, +and (where feasible) a did-you-mean suggestion, and renders with a **stable +`str()` format** so tests can snapshot it. + +## The stable message format + +`_format_error_message` builds every stage-5 error's message in one shape: + +```text +: + at + scope: + suggestion: +``` + +The first line always begins with the class name, so log greps and snapshot +tests bind to a stable prefix; the indented lines are optional. + +## The error classes + +| Class | Raised when | +| --- | --- | +| `UnknownReferenceError` | a bare or dotted ref doesn't resolve in scope | +| `AmbiguousReferenceError` | a ref matches multiple candidates in scope | +| `IllegalScopeReferenceError` | a dotted ref against a `StageSchema`, or `__` in a `ModelScope` without an exact match | +| `IllegalWindowInFilterError` | raw `OVER(...)` in a DSL filter, or a filter referencing a windowed `Column.sql` | +| `AggregationNotAllowedError` | type-bucket / PK / `allowed_aggregations` violation | +| `UnknownFunctionError` | a Mode-B call not in `SCALAR_FUNCTIONS` / transforms / aggregations | +| `MeasureRecursionLimitError` | named-measure expansion exceeded depth (32, env-configurable) | +| `MeasureCycleError` | a cycle in named-measure expansion | +| `DuplicateMeasureNameError` | two measures declare the same `name` | +| `MeasureNameCollidesWithColumnError` | a declared `name` matches a source column | +| `CanonicalAliasShadowsColumnError` | a formula's canonical alias shadows a source column | + +### `ValueError` multi-inheritance for back-compat + +`UnknownReferenceError`, `AmbiguousReferenceError`, `IllegalScopeReferenceError`, +and `IllegalWindowInFilterError` multi-inherit `ValueError` (alongside +`SlayerError`). This is deliberate: the cutover replaced legacy `ValueError` +resolution paths with these typed errors, and many pre-existing call sites and +tests catch `ValueError` (or use `pytest.raises(ValueError)`). Multi-inheriting +keeps them working unchanged. `ColumnCycleError` (DEV-1410) does the same. + +`SlayerError` is the base for SLayer's intentional failure modes, so callers can +distinguish them from unexpected `Exception` paths (driver errors, IO errors). + +## Warnings + +Two warning types are **not** exceptions: + +- `SlayerNormalizationWarning` (`core/warnings.py`) — a `UserWarning` carrying a + `NormalizationWarning` payload, emitted by the + [slack layer](slack-normalization.md) on every rewrite. Surfaced both via + `warnings.warn(...)` and on `SlayerResponse.warnings`. +- `UnreachableFilterDroppedWarning` (`core/errors.py`) — a `UserWarning` emitted + by the [cross-model planner](cross-model-aggregates.md) when a host filter + references slots unreachable from a CTE's root, so the filter is dropped from + the CTE (the host still applies it to its own rows). A visibility/debug + warning, not an error. + +## Design rationale + +- **Why typed errors over `ValueError`?** The legacy enrichment raised bare + `ValueError`s that callers couldn't distinguish, so error handling was + string-matching on messages. Typed classes let surfaces (REST/MCP/CLI) and + tests react to the *kind* of failure, and the `.name` / `.scope_summary` / + `.suggestion` attributes make programmatic remediation possible. +- **Why a stable `str()` format?** So snapshot tests can pin the message + (including suggestion text and scope summary) without brittle substring + matches, and so agents reading the error get a consistent, parseable shape. +- **Why keep `ValueError` in the MRO?** A clean break would have churned every + `except ValueError` call site and test in the same PR as the cutover. Multi- + inheriting defers that churn without weakening the new typed surface — callers + that want the specific class can catch it; callers that catch `ValueError` + still work. diff --git a/docs/architecture/index.md b/docs/architecture/index.md new file mode 100644 index 00000000..40c5af41 --- /dev/null +++ b/docs/architecture/index.md @@ -0,0 +1,252 @@ +# Architecture: the typed resolution pipeline + +This section documents how SLayer turns a `SlayerQuery` into SQL — the typed, +composable pipeline introduced by the DEV-1450 redesign. It describes the code +**as currently implemented**, not the original plan; where the implementation +diverges from the plan, [Deviations from the plan](#deviations-from-the-plan) +calls it out. + +The audience is contributors. If you only write queries, read +[Concepts](../concepts/queries.md) instead. + +## Why the redesign + +The expressive surface syntax (dotted joins, colon aggregation, transforms, +renamed measures, cross-model aggregates) used to be resolved by a single large +enrichment pass (`slayer/engine/enrichment.py`, ~2300 lines) that interleaved +string rewriting, alias remapping, a parallel `cross_model_measures` track, +virtual-model flattening, and implicit passthrough. Every new permutation of +"custom name × join × transform" added another resolution path, and the paths +interacted in ways that produced corner-case bugs (DEV-1445/1446/1448/1449). + +The redesign replaces that with a pipeline of small stages, each taking +well-typed input and producing a well-typed intermediate object that carries +everything the next stage needs. **Identity is structural, not textual** — the +single idea that makes the four bugs structurally impossible rather than +individually patched. + +## The pipeline at a glance + +```mermaid +flowchart LR + raw["SlayerQuery
(raw, slack-tolerant)"] --> norm["slack normalize
(normalization.py)"] + norm --> parse["parse_expr
(syntax.py)"] + parse --> expand["expand measures
(measure_expansion.py)"] + expand --> bind["bind_expr / bind_filter
(binding.py)"] + bind --> plan["plan_query / plan_stages
(stage_planner.py)"] + plan --> planned["PlannedQuery
(planned.py)"] + planned --> gen["generate_planned_stages
(generator.py)"] + gen --> sql["SQL string"] +``` + +Each arrow is the typed-object boundary from principle **P7** +(`raw → NormalizedInput → ParsedExpr → BoundExpr → ValueSlot → PlannedQuery → +SQL`). No stage string-rewrites the output of a previous stage after parsing. + +A more detailed view, showing the planner's internal sub-stages and the source +bundle that feeds resolution: + +```mermaid +flowchart TB + subgraph orchestrator["engine.execute → _execute_pipeline (query_engine.py)"] + bundle["build_resolved_source_bundle
(source_bundle.py) — storage read once (P11)"] + norm["_normalize_stage → normalize_query (P0)"] + vars["apply_variables_to_query (variables.py)"] + plan["plan_stages (stage_planner.py)"] + gen["generate_planned_stages (generator.py)"] + meta["build_response_metadata (response_meta.py)"] + end + + bundle --> norm --> vars --> plan --> gen --> meta + + subgraph perstage["plan_query — per stage"] + decl["parse + expand + bind
declared measures / dims / TDs"] + filt["bind_filter
filters (phase-classified, P8)"] + td["attach time_key → transforms"] + sugar["lower_sugar_transforms (change/change_pct)"] + proj["ProjectionPlanner + ValueRegistry
intern slots (P2/P4)"] + cm["cross_model_planner.plan
per cross-model aggregate (P3/I1)"] + end + + plan --> decl --> filt --> td --> sugar --> proj --> cm +``` + +## Principles, and where they live + +The redesign was specified as 12 principles (P0–P11). Each maps to a concrete +module: + +| Principle | Statement | Where | +| --- | --- | --- | +| **P0** | Pipeline begins with a single slack-normalization pass; rewrites are returned as typed warnings | [`normalization.py`](slack-normalization.md) | +| **P1** | Two surface languages (Mode A SQL / Mode B DSL), never mixed mid-expression; closed `SCALAR_FUNCTIONS` allowlist | [`syntax.py`](parsing.md), `keys.py` | +| **P2** | Identity is structural, not textual — two equal keys intern to one slot | [`keys.py`](typed-keys.md), `planning.py` | +| **P3** | Local and cross-model aggregates share one `AggregateKey` shape (path empty vs non-empty) | [`keys.py`](typed-keys.md), [`cross_model_planner.py`](cross-model-aggregates.md) | +| **P4** | Public names are a separate namespace; a slot has one declared name + many public aliases | [`planning.py`](planning.md) | +| **P5** | Scope determines what dots mean — `ModelScope` vs `StageSchema`, never confused | [`scope.py`](scopes-and-bundle.md), [`binding.py`](binding.md) | +| **P6** | Each stage emits an explicit `StageSchema`; stages compose only through schemas | [`scope.py`](scopes-and-bundle.md), [`stage_planner.py`](stage-planning.md) | +| **P7** | Typed pipeline; no string rewriting after parse | whole pipeline | +| **P8** | Phase (WHERE/HAVING/post) is a property of the slot, not the filter text | [`keys.py`](typed-keys.md) `Phase`, [`binding.py`](binding.md) | +| **P9** | Transforms are operators over slots, not over strings | [`planning.py`](planning.md), [`binding.py`](binding.md) | +| **P10** | Result-key contract preserved exactly | [`generator.py`](sql-generation.md), `response_meta.py` | +| **P11** | Resolution is pure — storage consulted once, no `ContextVar` re-resolution | [`source_bundle.py`](scopes-and-bundle.md) | + +## Module map + +The new pipeline modules, in dependency order: + +| Module | Role | Doc | +| --- | --- | --- | +| `slayer/core/keys.py` | The `ValueKey` family — structural identity primitives | [Typed keys](typed-keys.md) | +| `slayer/core/scope.py` | `ModelScope`, `StageSchema`, `StageColumn` | [Scopes & bundle](scopes-and-bundle.md) | +| `slayer/core/errors.py`, `warnings.py` | Typed errors + slack-warning carriers | [Errors & warnings](errors-and-warnings.md) | +| `slayer/engine/source_bundle.py` | `ResolvedSourceBundle` + eager builder (P11) | [Scopes & bundle](scopes-and-bundle.md) | +| `slayer/engine/normalization.py` | Slack-normalization layer (P0) | [Slack normalization](slack-normalization.md) | +| `slayer/engine/syntax.py` | Mode-B Python-AST parser → `ParsedExpr` | [Parsing](parsing.md) | +| `slayer/sql/sql_expr.py` | Mode-A sqlglot wrapper | [Parsing](parsing.md) | +| `slayer/engine/measure_expansion.py` | Pre-bind named-`ModelMeasure` expansion | [Parsing](parsing.md) | +| `slayer/engine/binding.py` | `ExpressionBinder` / `FilterBinder` → `BoundExpr` | [Binding](binding.md) | +| `slayer/engine/planning.py` | `ValueRegistry`, `ProjectionPlanner`, transform lowering | [Planning](planning.md) | +| `slayer/engine/cross_model_planner.py` | Cross-model aggregate strategy (I1) | [Cross-model aggregates](cross-model-aggregates.md) | +| `slayer/engine/planned.py` | `PlannedQuery` and its parts | [Planning](planning.md) | +| `slayer/engine/stage_planner.py` | `plan_query` / `plan_stages` orchestrators | [Stage planning](stage-planning.md) | +| `slayer/engine/variables.py` | `{var}` substitution + 4-layer merge | [Engine orchestration](engine-orchestration.md) | +| `slayer/sql/generator.py` | `generate_from_planned` / `generate_planned_stages` | [SQL generation](sql-generation.md) | +| `slayer/engine/response_meta.py` | `attributes` / `expected_columns` from the plan | [SQL generation](sql-generation.md) | +| `slayer/engine/query_engine.py` | `_execute_pipeline` orchestration + cutover | [Engine orchestration](engine-orchestration.md) | + +## The four bugs, made structurally impossible + +The acceptance criterion was that DEV-1445/1446/1448/1449 stop being reachable, +not that each gets a patch. How structural identity achieves that: + +- **DEV-1446** (transform-wrapped agg-ref of a renamed measure deduping): + `change(amount:sum)` and `amount:sum` share the same inner `AggregateKey` + instance, so the `ValueRegistry` interns one slot — `SUM(amount)` appears once. + See [Planning](planning.md). +- **DEV-1445** (cross-model renamed-measure filter by alias *or* dotted form): + `customers.revenue:sum` and the user alias `rev` both bind to one + `AggregateKey`; the filter's `rev` ref resolves through `alias_map` onto that + same slot. See [Binding](binding.md), [Stage planning](stage-planning.md). +- **DEV-1448** (user `name` on a join-traversed measure governs the stage column): + `StageColumn.name` is the declared name, flattened — downstream stages bind + against it. See [Stage planning](stage-planning.md). +- **DEV-1449** (downstream stages see upstream stages as flat schemas): + binding against a `StageSchema` rejects dotted refs with + `IllegalScopeReferenceError`; only flat `__` names resolve. See + [Scopes & bundle](scopes-and-bundle.md), [Binding](binding.md). + +Each has an `engine.execute`-level acceptance test in +`tests/test_dev1445_*.py` / `1446` / `1448` / `1449`. + +## Current state: two pipelines coexist + +The cutover (DEV-1450 stage 7b.15) routed **top-level query planning** through +the new pipeline. It deliberately did **not** delete the legacy stack — +`enrichment.py`, `enriched.py` (`EnrichedQuery` / `EnrichedMeasure`), +`_query_as_model`, the legacy `SQLGenerator.generate`, +`_rewrite_funcstyle_aggregations`, and the `ContextVar` machinery all still +exist and, in some paths, **still run in production** (not just in tests): + +- **Query-backed model expansion** (turning a model's `source_queries` into a + virtual `sql`-mode model whose `.sql` is the rendered backing query) runs + entirely on legacy `_query_as_model → enrich_query → SQLGenerator.generate`, + on **both** the execute path (`_expand_query_backed_model`) and the save path + (`_validate_and_populate_cache`). The new pipeline then plans/renders the + *outer* query against the resulting virtual model. +- `generate_from_planned` reuses the legacy dialect helpers via a synthetic + `EnrichedMeasure` adapter (see [SQL generation](sql-generation.md)). + +Deleting the legacy stack — and migrating query-backed expansion onto the typed +pipeline — is tracked as **DEV-1452**. Until then, treat the typed pipeline as +the resolution path for the *outer* query and for the four acceptance bugs, and +the legacy stack as still load-bearing for query-backed inner rendering and +dialect SQL emission. See [Engine orchestration](engine-orchestration.md) for +the exact call sites. + +## Deviations from the plan + +These are places where the implemented code departs from the DEV-1450 plan. +They are documented here so reviewers don't mistake them for the intended end +state. All are deliberate and tracked, but several reintroduce — temporarily — +the kind of multi-path coupling the redesign set out to remove. + +1. **Legacy stack still load-bearing for query-backed models** (above). The + plan's stage-7b bullet said the cutover would "delete `EnrichedQuery`, + `EnrichedMeasure`, … `_query_as_model`, … legacy `SQLGenerator.generate`". + In practice every deletion is deferred to DEV-1452, and the legacy stack is + not merely "kept reachable for tests" — it renders the backing SQL of every + query-backed model. This is the largest gap between plan and reality. + +2. **A second cross-model rendering path was needed** (re-rooting). The plan's + cross-model design was a single strategy: `IsolatedCteCrossModelPlanner` plus + the `inherited_filter_policy` decision table, producing one CTE per + `(target, grain)`. That proved insufficient: when host dimensions are + reachable from the target through the *target's own* join graph, the + forward-path CTE collapses the host grain to a scalar `CROSS JOIN`. The fix + (`_maybe_reroot_cross_model_plan` in `stage_planner.py`, rendered by + `_render_rerooted_cross_model_cte` in `generator.py`) builds a full nested + re-rooted `PlannedQuery` — mirroring legacy `_build_rerooted_enriched`. So + there are now **two** cross-model render strategies selected heuristically, + bolted onto `CrossModelAggregatePlan` via `rerooted_plan` / + `rerooted_grain_pairs` / `rerooted_agg_slot_id`. This is the most significant + architectural compromise — the "one shape, render strategy chosen downstream" + abstraction (P3) holds for identity but not for rendering. See + [Cross-model aggregates](cross-model-aggregates.md). + +3. **The new generator adapts back to `EnrichedMeasure`.** The plan said + "rewrite `generator.py` to consume `PlannedQuery`". The implemented + `generate_from_planned` consumes `PlannedQuery` at the top but synthesizes + `EnrichedMeasure` objects (`_synthesize_enriched_measure_from_planned`) to + reuse the legacy dialect helpers (`_build_agg`, `_build_percentile`, + `_build_stat_agg`, …). The new path is therefore coupled to a type DEV-1452 + wants to delete. See [SQL generation](sql-generation.md). + +4. **Derived-column parity with legacy restored (DEV-1450 follow-ups #4a / #4b).** + Two cases that legacy handled, and the typed pipeline briefly narrowed to + `NotImplementedError`, now work again: + - A `TimeDimension` over a derived (`Column.sql`) temporal column. + `TimeTruncKey.column` is `Union[ColumnKey, ColumnSqlKey]`; the binder + (`bind_time_dimension`) and every generator render site apply the + `DATE_TRUNC` over the EXPANDED derived expression (base SELECT, ORDER BY, + window/OVER transforms, the time_shift self-join CTE, `date_range` + BETWEEN, and the cross-model shared-grain CTE). See + [Typed keys](typed-keys.md) and [SQL generation](sql-generation.md). + - A `SlayerModel.filters` entry, OR a column-level `Column.filter` on an + aggregated measure, referencing a non-trivial derived column. + `_validate_model_filter` no longer rejects the model-filter form; the + generator inline-expands the predicate (the shared + `_render_mode_a_predicate`, used by both `_render_model_filter_sql` and the + `Column.filter` CASE-WHEN path) and pulls any join the expansion crosses + into the FROM. Inlining covers a bare derived ref (`is_eu` → + `customers.region`) **and** a dotted ref to a derived column on a joined + model (`loss_payment.has_flag` → its `sql`), matching the query-level + filter path so no dangling `.` (a non-physical column) + is emitted (DEV-1494). Join discovery for these Mode-A text filters + (`_filter_join_paths`) unions the paths of the **un-inlined** predicate + (so the dbt placeholder-join idiom — a constant `has_flag sql="1"` whose + only purpose is to force the join — keeps its alias) with the paths the + **inline-expanded** predicate crosses. The cross-model `_cm_*` CTE discovers + its OWN filter joins too — the target measure's `Column.filter` and the + target-model filters — and adds them to the CTE's FROM (each `_cm_*` CTE is + an isolated per-(target, grain) computation, so the join resolves the + filter's refs without affecting sibling measures). The windowed-`Column.sql` + and same-model `ModelMeasure`-ref rejects remain. + +5. **P10 is intentionally violated for cross-model parametric aggregates.** + Result keys for `customers.revenue:percentile(p=0.5)` now carry the kwarg + signature (`…revenue_percentile_p_0_5`) where legacy dropped it + (`…revenue_percentile`). Legacy's drop was a collision bug (two parametric + variants on one column produced the same alias); the new path fixes it but + at the cost of bit-identical parity. Tested structurally, not by parity. + +6. **Pre-processing runs before the "single" slack pass (P0).** + `_execute_pipeline` runs `strip_source_model_prefix()` and + `snap_to_whole_periods()` on the query *before* `_normalize_stage`. These are + query-shape transforms rather than slack-token rewrites, but they mean the + pipeline does not literally "begin with a single slack-normalization pass". + +Test-only deviations (parity oracle replacing the planned parity adapter; the +two retained `@pytest.mark.skip`s in `tests/test_filter_renamed_measure.py`; +production extractors using a scope-free `walk_parsed_refs` instead of binding) +are noted in the relevant component docs and are not architectural concerns. diff --git a/docs/architecture/parsing.md b/docs/architecture/parsing.md new file mode 100644 index 00000000..1129ab8e --- /dev/null +++ b/docs/architecture/parsing.md @@ -0,0 +1,158 @@ +# Parsing + +**Modules:** `slayer/engine/syntax.py` (Mode B), `slayer/sql/sql_expr.py` +(Mode A), `slayer/engine/measure_expansion.py` (pre-bind expansion) + +Parsing turns expression strings into typed `ParsedExpr` trees. It is **pure +syntax** — no scope resolution, no named-measure expansion, no function-style +rewriting. Those are separate stages (the [slack layer](slack-normalization.md) +does function-style → colon; the [binder](binding.md) does scope; expansion is +its own step). This separation is what keeps each stage small. + +```mermaid +flowchart LR + text["expr string
'change(amount:sum) > 0'"] --> pp["_preprocess_colons
amount:sum → placeholder"] + pp --> ast["ast.parse(mode='eval')"] + ast --> conv["_convert
AST → ParsedExpr"] + conv --> tree["ParsedExpr tree"] +``` + +## The `ParsedExpr` family + +Eleven frozen Pydantic node types with value-based equality (so tests assert via +`==`): + +| Node | Shape | +| --- | --- | +| `Ref` | `name` — a bare identifier | +| `DottedRef` | `parts` — a dotted path | +| `StarSource` | `*` | +| `Literal` | `value` (`Decimal` / `str` / `bool` / `None`) | +| `AggCall` | `source, agg, args, kwargs` — colon aggregation | +| `TransformCall` | `op, input, args, kwargs` | +| `ScalarCall` | `name, args` | +| `Arith` | `op, left, right` | +| `UnaryOp` | `op, operand` | +| `Cmp` | `op, left, right` | +| `BoolOp` | `op, operands` | + +## How `parse_expr` works + +Mode B is a *Python-AST* DSL — the grammar is a deliberate subset of Python +expression syntax, so the parser leans on `ast.parse(..., mode="eval")` rather +than a hand-rolled grammar. Two pre/post steps make the colon and `__` rules +work: + +1. **`_preprocess_colons`** replaces `:` with a placeholder + identifier (`__slayer_agg_N__`) before handing the text to Python's parser, + capturing the source kind (`*` / `Ref` / `DottedRef`) and agg name in a side + map. Any trailing `(args)` is left in place so Python parses it as a `Call` + naturally. String-literal spans are skipped so quoted contents aren't touched. +2. **`_reject_dunder_in_ast`** walks the parsed AST and rejects any user + identifier containing `__` (on `Name`, `Attribute.attr`, and `keyword.arg`), + unless `allow_dunder=True`. `__` is reserved for internal join-path aliases on + the SQL side; users write single-dot DSL paths. + +`_convert` then maps AST nodes to `ParsedExpr` nodes. A `Call` dispatches in a +fixed order: aggregation placeholder → transform (in `ALL_TRANSFORMS`, requires +≥1 positional) → scalar (in `SCALAR_FUNCTIONS`, rejects kwargs) → otherwise +`UnknownFunctionError`. List/tuple kwarg values (e.g. `partition_by=[a, b]`) +convert to a tuple of converted elements. + +### Rejections (P1) + +The parser is where the Mode-B contract is enforced: + +- a function call not in `SCALAR_FUNCTIONS` / `ALL_TRANSFORMS` / aggregations → + `UnknownFunctionError`; +- a raw `OVER(...)` clause anywhere in the text → `IllegalWindowInFilterError` + (checked by regex before AST parsing); +- `__` in a user identifier → `ValueError` (unless `allow_dunder`); +- chained comparisons (`1 < x < 10`) → `ValueError` (split into `1 < x and x < + 10`); the binder can't give a chained comparison a single phase. + +### `allow_dunder` — the StageSchema escape hatch (DEV-1449) + +`parse_expr(text, *, allow_dunder=False)` defaults to rejecting `__`. The +[stage planner](stage-planning.md) sets `allow_dunder=True` *only* when binding a +downstream stage against a flat `StageSchema`, whose columns **are** the +`__`-flattened multi-hop aliases of the upstream stage (`customers__region`). +Legality there is the binder's concern (the column must exist in the upstream +schema). This is the one place `__` is legal in a Mode-B ref, and it is exactly +what makes a downstream stage able to name an upstream joined dimension. + +### `parse_filter_expr` — SQL-operator leniency + +Filters historically accepted SQL operator spellings (`=`, `<>`, `NULL`, keyword +`AND`/`OR`/`NOT`/`IS`/`IN`) alongside Python ones. `parse_filter_expr` normalizes +those to Python equivalents (string-literal-aware) via +`_normalize_sql_filter_operators`, then delegates to `parse_expr`. Measures and +order parse with `parse_expr` directly; only filters get the leniency — matching +the legacy `parse_filter` contract. + +### `walk_parsed_refs` — scope-free reference extraction + +`walk_parsed_refs(parsed)` yields the reference-bearing leaves (`Ref`, +`DottedRef`, `AggCall`) of a tree without binding it. It is the scope-free +counterpart to the binder's `walk_value_keys`: production extractors that only +need the *names* a formula touches — schema-drift cascade attribution and memory +entity tagging — walk the parse tree directly instead of binding against a scope. +Its descent rules match the legacy `parse_formula` walk exactly (an `AggCall` is +yielded as a unit and its args/kwargs are *not* descended, so +`weighted_avg(weight=quantity)` surfaces `price`, never `quantity`). + +> **Deviation note.** The plan specified walking the typed-key `BoundExpr` via +> `walk_value_keys` for these extractors. That is infeasible: binding raises on +> bare named-measure refs (which need planner-side expansion) and resolves the +> very refs drift detection must find *pre*-resolution. `parse_expr` + +> `walk_parsed_refs` was the user-approved alternative. + +## Mode A — `sql_expr.py` + +Mode A (`Column.sql`, `Column.filter`, `SlayerModel.filters`) is sqlglot-native. +`parse_sql_expr` wraps the fragment as `SELECT () AS _` before sqlglot +parses it — necessary because sqlglot's SQLite/MySQL parser otherwise falls back +to a `Command` node for a top-level `replace(...)`. `has_window_function` is the +predicate the binder uses to reject filters that touch a windowed `Column.sql` +(DEV-1369). Mode A keeps full SQL expressiveness; the typed pipeline only needs +to detect windows and (in the slack layer) rewrite multi-dot paths. + +## Pre-bind measure expansion — `measure_expansion.py` + +The binder raises `UnknownReferenceError` for a bare *measure* name (measures +aren't columns). `expand_model_measures` runs *before* binding: it is an +AST → AST rewrite that replaces every `Ref(name=X)` whose `X` is a saved +`ModelMeasure` with the recursively-expanded `parse_expr(measure.formula)` tree, +turning measure refs into binder-resolvable column/aggregation nodes. + +```mermaid +flowchart LR + r["Ref('aov')"] -->|aov is a ModelMeasure| sub["parse_expr(aov.formula)"] + sub --> walk["recursively expand its refs"] + walk --> out["expanded ParsedExpr"] +``` + +Eligibility is principled: expansion fires at the root and in `Arith` / `UnaryOp` +/ `Cmp` / `BoolOp` operands, `ScalarCall.args`, and `TransformCall.input` / args +/ kwarg values. It does **not** fire on `DottedRef` segments (those resolve +through joins), `AggCall` in any position (sources/args/kwargs are column-level +by contract), or function-name slots. Recursion is bounded: depth limit 32 +(configurable via `SLAYER_MEASURE_EXPANSION_DEPTH`) raising +`MeasureRecursionLimitError`, plus per-chain cycle detection raising +`MeasureCycleError` with the offending chain. A `parse_cache` memoizes each +measure's parse. The node-type tuple is derived from the `ParsedExpr` union via +`get_args`, so a new node type added to `syntax.py` is automatically walked. + +## Design rationale + +- **Why reuse Python's AST for Mode B?** The DSL was always a Python-expression + subset; `ast.parse` gives precedence, grouping, and operator handling for free, + and the conversion layer stays small. The colon preprocessor is the only piece + that bridges the one construct Python doesn't have. +- **Why is parsing pure (no scope)?** So the same parser serves the binder, the + measure expander, and the scope-free extractors. Mixing in resolution would + re-couple parsing to the model graph — the coupling the redesign removes. +- **Why a separate expansion pass instead of expanding in the binder?** Expansion + is an AST → AST rewrite with its own recursion/cycle concerns; keeping it + before binding means the binder only ever sees column/aggregation refs and can + stay a straight scope lookup. diff --git a/docs/architecture/planning.md b/docs/architecture/planning.md new file mode 100644 index 00000000..5937d598 --- /dev/null +++ b/docs/architecture/planning.md @@ -0,0 +1,170 @@ +# Planning: interning, projection, and the plan shape + +**Modules:** `slayer/engine/planning.py` (ValueRegistry, ProjectionPlanner, +transform lowering), `slayer/engine/planned.py` (the `PlannedQuery` types) + +Planning turns bound expressions into a `PlannedQuery` — the fully resolved, +render-ready target. Three composable concerns live in `planning.py`; the typed +result types live in `planned.py`. The [stage planner](stage-planning.md) +composes them. + +## `ValueRegistry` — interning by structural identity (P2 / P4) + +The registry maps `ValueKey → ValueSlot`. `intern(...)` either returns the +existing slot for a structurally-equal key or allocates a fresh one. This is the +mechanism behind **P2**: `change(amount:sum)` and a filter `amount:sum` build the +same inner `AggregateKey`, so they intern to one slot and `SUM(amount)` is +emitted once (DEV-1446). + +```mermaid +flowchart TB + k1["AggregateKey(amount, sum)
from measure"] --> reg + k2["AggregateKey(amount, sum)
inner of change()"] --> reg + k3["AggregateKey(amount, sum)
from filter"] --> reg + reg["ValueRegistry.intern"] --> slot["one ValueSlot
id=s1"] +``` + +### Public names are a separate namespace (P4 / C13) + +A slot has at most one *declared name* but can accumulate **multiple** +`public_aliases` — if the same structural key is declared with two different +user `name`s, both aliases appear in the projection pointing at one slot +(`_merge_into_existing`). A filter/order expression may reference a declared name +as an alias for the slot, but cannot *synthesize* a new slot from a +canonical-looking bare name when no corresponding measure was declared. + +### Alias-collision validations (DEV-1443) + +`intern` enforces three validations against the host model's column names: + +- a declared `public_name` colliding with a source column → + `MeasureNameCollidesWithColumnError`; +- a `canonical_alias` (e.g. `amount_sum`) shadowing a source column → + `CanonicalAliasShadowsColumnError`; +- two different keys declaring the same `public_name` → + `DuplicateMeasureNameError`. + +With carefully chosen exemptions: a *self-named dimension* (a `ColumnKey` / +`ColumnSqlKey` / `TimeTruncKey` whose public name **is** its own column name) is +the column, not a rename, so the collision check is skipped; and an unnamed +`*:` re-aggregation (whose canonical `_count` is a structural marker, not a +column ref) is exempt. + +## Transform lowering (P9 / C6) + +`change` and `change_pct` are sugar. `desugar_change` rewrites +`change(x)` → `x - time_shift(x, periods=-1)` and `desugar_change_pct` +→ `(x - time_shift(x, -1)) / time_shift(x, -1)`. The inner `x` is the **same +`ValueKey` instance** across the arithmetic and the time_shift, so a downstream +registry interns it once — this is the identity-preservation that makes DEV-1446 +hold even through desugaring. `partition_by` and `time_key` thread through to the +underlying `time_shift` (**C6**). + +`lower_sugar_transforms(key)` is the recursive walker that applies the desugar +functions anywhere in a `ValueKey` tree (`TransformKey` / `ArithmeticKey` / +`ScalarCallKey` / `BetweenKey`), rebuilding only the path that contains a +change/change_pct so identity is preserved elsewhere. The stage planner runs it +*after* `time_key` patching so the desugared `time_shift` inherits the patched +key. + +## `ProjectionPlanner` — declared + hidden slots + +`ProjectionPlanner.plan(...)` interns each declared measure (in dim → time-dim → +measure order) into the registry, builds the public projection, and then +materializes **hidden** slots for any value referenced only in filters / order +or as an auxiliary dependency of a declared measure. + +The dependency-selection rule is `_iter_slot_deps`, and it encodes which keys +need a materialised slot versus which the generator inlines: + +| Key | Slotted? | +| --- | --- | +| `ColumnKey` / `ColumnSqlKey` / `TimeTruncKey` | yes (row slot) | +| `AggregateKey` | yes (stops — its inner source materializes inside the aggregate) | +| `TransformKey` | yes, **and** recurse into `input`, `partition_keys`, `time_key` | +| `ArithmeticKey` / `ScalarCallKey` | no — recurse into operands/args; the op/call is inlined | +| `BetweenKey` | no — inlined into WHERE; recurse into column/low/high | +| `LiteralKey` / `StarKey` | never slottable alone | + +So `ORDER BY revenue:sum DESC LIMIT 10` with no declared `revenue:sum` measure +interns the aggregate as a `hidden=True` slot: the base CTE materializes it, the +outer SELECT trims it from the public projection, and `StageSchema.columns` +excludes it (downstream stages see no extra column). The same rule covers +filter-only refs. The no-transform "plain" path follows the same pattern via a +conditional outer-trim wrapper (DEV-1501): the wrap fires only when the base +materialises a hidden slot, so simple flat queries stay flat. Hidden parametric +aggregates (`revenue:last(created_at)` vs `revenue:last(updated_at)`, +`revenue:percentile(p=0.5)` vs `…(p=0.95)`) route their declared name through +`canonical_agg_name` so the args/kwargs surface in the materialised alias — +two distinct hidden parametric aggregates get distinct base-CTE aliases instead +of colliding on `revenue_last` / `revenue_percentile`. + +`filter_referenced_slot_ids(bound_filter, registry)` walks the predicate via +`_iter_slot_deps` and looks each dep up in the registry, returning `set[SlotId]` +— the input the [cross-model planner](cross-model-aggregates.md) needs for filter +routing (it gets slot ids, not pre-interning `ValueKey`s, and it sees +composite-predicate leaves rather than just the top-level key). + +## The `PlannedQuery` shape (planned.py) + +`PlannedQuery` is the typed target the [SQL generator](sql-generation.md) +consumes. It carries everything needed to emit SQL without re-walking the model +graph (**P7**): + +| Field | Role | +| --- | --- | +| `source_relation` | the FROM relation name (model name or stage CTE) | +| `join_plan` | `JoinRequirement` hops | +| `row_slots` / `aggregate_slots` / `combined_expression_slots` | slots bucketed by phase | +| `cross_model_aggregate_plans` | one `CrossModelAggregatePlan` per cross-model aggregate | +| `transform_layers` | one `TransformLayer` per transform slot, in dependency order | +| `filters_by_phase` | `FilterPhase` entries (WHERE / HAVING / post) | +| `projection` / `order` / `limit` / `offset` | output shape | +| `stage_schema` | the projection downstream stages bind against (P6) | +| `active_time_dimension_slot_id` | the TD slot used for OVER `ORDER BY` | +| `render_source_model` | the concrete `SlayerModel` this stage renders against | + +A `ValueSlot` carries `id`, `key`, `declared_name`, `public_name`, +`public_aliases`, `hidden`, `phase`, `label`, `type`, and `expression` +(a `BoundExpr`). A model-validator enforces the hidden invariant: a hidden slot +must have `public_name=None` and `public_aliases=[]`, so the generator can never +accidentally emit it in the public projection. + +`FilterPhase` has two mutually-exclusive carrier modes: a typed `expression` +(`BoundExpr`, for Mode-B DSL filters and the planner-emitted `BetweenKey` +date_range) or `text` + `text_columns` (a Mode-A SQL fragment, for +`SlayerModel.filters` — the renderer qualifies the named columns and emits the +text verbatim). + +### `BoundExpr` unification + +`planned.py` re-exports `binding.BoundExpr` as the canonical class. Earlier the +planned side had a separate `BoundExpr` with an optional `sql_text` cache; that +was folded into the binder's `BoundExpr(value_key=ValueKey)` so +`ValueSlot.expression` and `FilterPhase.expression` store binder output directly. +There is no cached SQL string — the generator renders from the typed `value_key` +against the slot registry. + +### `CrossModelAggregatePlan` — the re-rooting fields + +The struct carries the route-explicit filter ids +(`where_filter_ids` / `having_filter_ids` / `target_model_filters`) plus, for the +re-rooting case, `rerooted_plan` (a nested `PlannedQuery`), `rerooted_grain_pairs`, +and `rerooted_agg_slot_id`. See [Cross-model aggregates](cross-model-aggregates.md) +— the re-rooting fields are the largest deviation from the plan's single-strategy +design. + +## Design rationale + +- **Why intern at all?** Because the four bugs are all "the same value got two + slots" or "two values shared one". Structural interning is the single mechanism + that resolves both directions, instead of per-permutation alias bookkeeping. +- **Why hidden slots rather than special-casing order/filter refs?** A hidden + slot is materialised in the base CTE like any other, then trimmed from the + public projection — so the generator has one uniform notion of "a value to + compute" and the projection logic decides visibility. Order-only and + filter-only aggregates fall out of this for free. +- **Why does `_iter_slot_deps` inline `ArithmeticKey` / `ScalarCallKey`?** + Because they have no independent column to materialise — they're operators over + their operands. Slotting them would create spurious hidden columns; the + generator inlines the operator into the SELECT/WHERE and slots only the leaves. diff --git a/docs/architecture/scopes-and-bundle.md b/docs/architecture/scopes-and-bundle.md new file mode 100644 index 00000000..28b0e377 --- /dev/null +++ b/docs/architecture/scopes-and-bundle.md @@ -0,0 +1,138 @@ +# Scopes and the source bundle + +**Modules:** `slayer/core/scope.py`, `slayer/engine/source_bundle.py` + +Binding needs two things the keys don't carry: *what a name resolves against* +(the scope) and *the resolved models it resolves through* (the bundle). These +two modules supply them, and together they implement principles **P5**, **P6**, +and **P11**. + +## Two scope kinds (P5) + +A reference like `customers.regions.name` means different things depending on +where it is being resolved. The redesign makes that explicit with two scope +types that are never confused: + +```mermaid +flowchart TB + subgraph ModelScope + direction TB + ms["source_model: SlayerModel"] + msd["dotted refs walk the join graph
customers.regions.name → hop, hop, leaf"] + msu["__ in a ref → IllegalScopeReferenceError
(unless it exact-matches a literal column name)"] + end + subgraph StageSchema + direction TB + ss["columns: List[StageColumn] (flat)"] + ssd["dotted refs → IllegalScopeReferenceError"] + ssu["__-bearing names are legal flat names
customers__regions__name"] + end +``` + +- **`ModelScope`** — joins exist. Dotted refs walk the join graph rooted at + `source_model`. A `__` in a Mode-B ref is illegal *unless* it exact-matches a + column literally named that way (the C11 carve-out for legacy persisted + query-backed columns). +- **`StageSchema`** — a flat namespace. Dots are *not* join syntax (a dotted ref + raises); `__`-bearing identifiers are ordinary flat names. This is what a + downstream stage binds against — and exactly why DEV-1449 is impossible: a + downstream ref to an upstream multi-hop dimension must use the flat form + (`robot_details__modelseriesval`), and the dotted form raises. + +`ModelScope.source_model` is `Optional` from day one (the **I2** extension +point). The binder asserts `source_model is not None` at use sites today; a +future anchor-less mode would set it `None` and take a different binder branch. +Keeping the type optional avoids a breaking change later. + +## `StageSchema` and `StageColumn` (P6) + +`StageSchema` is the typed projection a stage emits. Downstream stages bind only +against it — they never re-walk the upstream join graph. The fields that make +DEV-1448/1449 work are on `StageColumn`: + +| Field | Meaning | +| --- | --- | +| `name` | the downstream **bind** name — flat (`customers__revenue_sum`, or a user `rev`) | +| `sql_alias` | the identifier emitted in the stage's SELECT projection | +| `public_alias` | the result-key piece returned to the user (dotted form, or the user name) | +| `type, label, format, hidden, description, meta, sampled, provenance` | per-column metadata carried downstream | + +The split between `name`, `sql_alias`, and `public_alias` is what lets the +planner reserve a hidden or alias-bearing form without coupling the downstream +bind name to the public result key. `StageSchema` supports `__getitem__`, `get`, +and `__contains__` for name lookup; `relation_name` is the CTE name / subquery +alias used when a downstream stage references it. + +## `ResolvedSourceBundle` and "storage consulted once" (P11) + +`ResolvedSourceBundle` is the eagerly-resolved set of everything the binder +might need, built **once** at the top of execution by +`build_resolved_source_bundle`. After that, the binder is provably scope-only — +no `ContextVar`, no callback that re-enters storage. This is **P11**, the +principle that most directly kills the old tangle: the legacy enrichment path +re-resolved models lazily through `ContextVar`-threaded callbacks, which is why +concurrent and nested resolution was so hard to reason about. + +The bundle carries: + +| Field | Contents | +| --- | --- | +| `source_model` | the host of the query (the real base the root chain bottoms out at) | +| `referenced_models` | transitive join-graph walk + each sibling stage's base; host first | +| `inline_extensions` | a root `ModelExtension` overlay over a non-sibling base, re-applied after query-backed expansion | +| `named_queries` | the raw sibling `SlayerQuery`s of a multi-stage DAG | +| `stage_source_models` | per-named-stage resolved source model (for heterogeneous DAGs) | +| `query_variables` | merged variables (runtime > stage > outer > model defaults) | +| `datasource_hint` | the `data_source=` kwarg that wins over the priority list | + +### How the builder resolves the source + +`build_resolved_source_bundle` (`source_bundle.py:189`) handles every input +shape the public API accepts — stored-model name, inline `SlayerModel`, +`ModelExtension` overlay, and the dict forms of both — via `_resolve_source_spec`. + +Two subtleties worth knowing: + +- **Sibling-chain following.** A root whose `source_model` points at a named + sibling is followed down (`_follow_sibling_chain`) to the real base it + ultimately reads from, so the bundle's `source_model` is always a concrete + base, not a sibling name. A cycle raises (mirrors the legacy circular-reference + guard). +- **Root overlay preservation.** When the root source is a `ModelExtension` over + a non-sibling base, the overlay is recorded in `inline_extensions` so the + engine can re-apply it *after* a query-backed base expands (expansion derives + columns from the backing query and would otherwise drop the overlay's extra + columns). + +The join-graph walk (`_collect_referenced_models`) is a best-effort BFS scoped +to the source model's own `data_source` — joins never cross datasource +boundaries. Absent join targets are skipped silently (matching the legacy +`_expand_join_graph`). The source model is returned first so +`get_referenced_model` finds the host before any same-named join target. + +### Synthetic models for sibling stages + +When a downstream stage joins to — or cross-model-references — a *sibling* stage +(materialised elsewhere as a CTE), the planner needs a `SlayerModel` to resolve +against. `synthetic_model_from_stage_schema` builds a stand-in whose +`sql_table` is the stage's CTE name and whose columns are the stage's flat +output columns. `stage_bundle_with_siblings` threads these synthetic models into +a per-stage bundle so a join/cross-model ref to a sibling resolves to its CTE +relation. These two helpers are what let the [stage planner](stage-planning.md) +treat sibling stages uniformly with stored models. + +## Design rationale + +- **Why a bundle at all, rather than passing `storage` to the binder?** Purity. + If the binder can reach storage, it can re-resolve, and re-resolution is where + the old order-dependence and `ContextVar` machinery came from. Resolving + everything up front makes the binder a pure function of `(parsed, scope, + bundle)`. +- **Why optional `source_model` everywhere (I2)?** So a future + "resolve against a whole datasource, no anchor model" mode is a type-additive + change, not a breaking one. The cost today is a handful of + `assert source_model is not None` lines. +- **Why two scope classes instead of a flag?** A flag (`is_stage: bool`) would + re-merge the two resolution rules into one function with internal branching — + precisely the shape the redesign is removing. Distinct types force the binder + to dispatch, and force callers to be explicit about which world they're in. diff --git a/docs/architecture/slack-normalization.md b/docs/architecture/slack-normalization.md new file mode 100644 index 00000000..210a50ac --- /dev/null +++ b/docs/architecture/slack-normalization.md @@ -0,0 +1,144 @@ +# Slack normalization + +**Module:** `slayer/engine/normalization.py` (warning types in +`slayer/core/warnings.py`) + +The pipeline begins (principle **P0**) with a single pass that rewrites +*slack-but-unambiguous* agent input into canonical form, so every downstream +stage sees only the canonical shape. Each rewrite is returned as a typed +`NormalizationWarning` and surfaced two ways at once. + +This is how SLayer stays tolerant of the natural things agents type +(`sum(revenue)`, a bare column listed under `measures`) without letting that +tolerance leak into the resolution logic — the parser, binder, and planner never +have to know that `sum(revenue)` is even a thing. + +## The three rules + +```mermaid +flowchart LR + subgraph "Mode B (DSL fields)" + f["FUNC_STYLE_AGG
sum(revenue) → revenue:sum
count(*) → *:count"] + end + subgraph "Query shape" + m["MISPLACED_MEASURE
bare column in query.measures
→ moved to query.dimensions"] + end + subgraph "Mode A (SQL fields)" + d["DOT_PATH_IN_SQL
customers.regions.name
→ customers__regions.name"] + end +``` + +| Rule | Mode | Detects | Rewrites to | +| --- | --- | --- | --- | +| `FUNC_STYLE_AGG` | Mode B | `sum(col)`, `count(*)`, `percentile(amount, p=0.5)` | colon form (`col:sum`, `*:count`, `amount:percentile(p=0.5)`) | +| `MISPLACED_MEASURE` | query shape | a bare (no colon, no call) entry in `query.measures` that names a column | moved to `query.dimensions` | +| `DOT_PATH_IN_SQL` | Mode A | a root-scope dotted ref whose leading segment is a known join target | the `__` alias form | + +### `FUNC_STYLE_AGG` + +Applies to Mode-B fields (`ModelMeasure.formula`, `SlayerQuery.measures[].formula`, +`SlayerQuery.filters`). It scans for `(` where `` is a builtin or +custom aggregation name (and not already preceded by `:`), finds the balanced +close paren (string-literal-aware), and rewrites the first argument into colon +form, keeping any remaining args as the parametric tail. `first` / `last` are +also transform names, so the rewrite skips them when the inner is already a +colon-form aggregate (`_AMBIGUOUS_AGG_TRANSFORMS`). Custom aggregation names are +threaded in via `custom_agg_names` so model-defined aggregations are recognized. + +`func_style_agg_to_colon` is the **quiet** variant for read-only consumers +(schema-drift attribution, memory entity tagging) that need the rewrite but must +not re-surface slack advice to the user — it suppresses the warning. + +### `MISPLACED_MEASURE` + +Mirrors the legacy `_auto_move_fields_to_dimensions` heuristic but emits a +structured warning. A bare token in `measures` that names a known `ModelMeasure` +stays a measure; one that names a column moves to `dimensions`; an unknown token +is left for the downstream resolver to error on. It is a no-op when the stage has +no resolved model (a sibling-sourced stage), because column classification needs +the model's column names. + +### `DOT_PATH_IN_SQL` + +The subtle one. It rewrites `customers.regions.name` → `customers__regions.name` +in Mode-A SQL (`Column.sql`, `Column.filter`, `SlayerModel.filters`), but only +when the leading segment is a real join target on the host model — and it is +**AST-based and scope-aware**, not a regex: + +- It parses with sqlglot and identifies the **root-scope** `Column` nodes by + walking lexical ancestors (`_dot_path_root_scope_analysis`), *not* by trusting + `Scope.columns` (which would pull in correlated subquery refs). Refs inside + subqueries, CTE bodies, and set-op branches are left alone. +- It collects shadow names — CTE definitions, explicit `AS` aliases, + Subquery/CTE sources, and schema/catalog qualifiers on FROM tables — and a ref + whose leading segment matches both a join target and a shadow is flagged + *ambiguous*: no rewrite, a warning carrying + `normalized="(ambiguous: …)"`. + +The scope-guard reuses the `column_expansion.py` precedent from DEV-1410. Why +AST and not the old construction-time regex: the rewrite needs the model's join +graph to know whether the first segment is a join target (vs. a +catalog/schema-qualified name like `mydb.customers.x`), which a `Column.sql` +field validator has no access to. So multi-dot normalization is **boundary-only, +by design** — it runs in the slack pass at `engine.execute` / `engine.save_model`, +not at Pydantic construction. A consequence to state honestly: a `SlayerModel` +built in memory and read back without crossing execute/save shows the raw +multi-dot form; `save_model` canonicalizes before persisting. + +## Warning shape and dual surfacing + +```python +class NormalizationWarning(BaseModel): # slayer/core/warnings.py + rule_id: str # "FUNC_STYLE_AGG" + original: str # "sum(revenue)" + normalized: str # "revenue:sum" + location: str # "measures[2].formula" + rule_doc_url: Optional[str] # "docs/agent_input_slack.md#func-style-agg" + +class SlayerNormalizationWarning(UserWarning): + """Carrier UserWarning around a NormalizationWarning payload.""" +``` + +Every rewrite is surfaced **both** as a Python warning +(`warnings.warn(SlayerNormalizationWarning(payload))`, so +`warnings.catch_warnings()` callers see it) **and** appended to +`SlayerResponse.warnings: List[NormalizationWarning]` (so REST/MCP/CLI consumers +get the structured payload alongside the result). One source of truth, two +surfaces. The payload Pydantic type lives in `slayer.core.warnings` rather than +in the engine module so storage/REST schemas can reference it without importing +engine code. + +## Entry points and boundaries + +- `normalize_query(query, *, model, custom_agg_names)` — runs `FUNC_STYLE_AGG` + over Mode-B fields and `MISPLACED_MEASURE` over the query shape, returning a + `NormalizationResult(query, warnings)`. (The query-side `DOT_PATH_IN_SQL` + wiring is present but a no-op — Mode-A on a query is rare; most Mode-A lives on + the model.) +- `normalize_model(model)` — runs `FUNC_STYLE_AGG` over `ModelMeasure.formula` + and `DOT_PATH_IN_SQL` over `Column.sql` / `Column.filter` / + `SlayerModel.filters`, returning `NormalizationResult(model, warnings)`. + +These are invoked at the engine boundaries: `engine.execute` (per stage, via +`_normalize_stage`) and `engine.save_model`. CLI / REST / MCP go through those +entry points automatically. See [Engine orchestration](engine-orchestration.md) +for the call sites. + +## Design rationale + +- **Why normalize before parsing rather than teaching the parser to accept slack + forms?** Keeping the slack rules in one pass means the rest of the pipeline has + exactly one shape to reason about. If `parse_expr` accepted `sum(revenue)`, the + binder and planner would each have to handle both spellings. +- **Why typed warnings rather than logging?** Agents (and the REST/MCP consumers + driving them) need to *see* that their input was rewritten, structurally, so + they can learn the canonical form. A log line is invisible to them; the + `rule_doc_url` points at the canonical-form documentation. +- **Why AST for `DOT_PATH_IN_SQL`?** The old construction-time regex blindly + rewrote any `a.b.c`, including `mydb.customers.x` — a latent bug. Being + scope-aware and join-graph-aware is only possible at a boundary that has the + model in hand. + +The reference page for the rules (with the `#func-style-agg` / +`#dot-path-in-sql` / `#misplaced-measure` anchors that `rule_doc_url` points at) +is `docs/agent_input_slack.md`, authored as part of the user-facing docs update. diff --git a/docs/architecture/sql-generation.md b/docs/architecture/sql-generation.md new file mode 100644 index 00000000..2de7f6ff --- /dev/null +++ b/docs/architecture/sql-generation.md @@ -0,0 +1,215 @@ +# SQL generation + +**Modules:** `slayer/sql/generator.py` (the planned-consuming path), +`slayer/engine/response_meta.py` (response metadata) + +The generator renders a `PlannedQuery` (or a list of them) to a SQL string. It +preserves the result-key contract exactly (**P10**) and emits SQL via sqlglot +AST building, not string concatenation. + +## Entry points + +```mermaid +flowchart TB + gps["generate_planned_stages(planned_list, bundle, dialect)"] + gps -->|single stage| gfp["generate_from_planned(planned, bundle, dialect)"] + gps -->|multi-stage| loop["render each stage → CTE; root = outer SELECT"] + loop --> gfp + gfp --> inst["SQLGenerator(dialect).generate_from_planned"] + inst -->|cross-model| cm["_render_with_cross_model_plans"] + inst -->|transforms| tl["WITH base, step CTEs, outer wrap"] + inst -->|plain| base["single SELECT"] +``` + +- `generate_from_planned(planned_query, *, bundle, dialect)` — module-level + entry that constructs an `SQLGenerator` and delegates to the instance method. + Renders **one** stage. +- `generate_planned_stages(planned_queries, *, bundle, dialect)` — renders a + multi-stage DAG to one SQL string. Each non-root stage becomes a CTE; the root + is the outer SELECT. + +## `generate_from_planned` (instance method) + +Reads from typed `PlannedQuery` fields (`row_slots` / `aggregate_slots` / +`filters_by_phase` / `order` / `transform_layers`) and dispatches: + +- `cross_model_aggregate_plans` non-empty → `_render_with_cross_model_plans`; +- `transform_layers` present → `WITH base AS (...)`, Kahn-batched step CTEs + carrying the window functions, an outer wrap projecting in user-spec order; + POST-phase filters that reference transform slots wrap as `SELECT * FROM (...) + AS _filtered WHERE …`; `time_shift` / `consecutive_periods` emit dedicated + self-join CTE pairs; +- otherwise → a single base SELECT with WHERE/HAVING, GROUP BY, ORDER BY, LIMIT. + When the base CTE materialises any hidden aggregate (an aggregate referenced + ONLY by ORDER BY or a filter, never declared as a measure), a conditional + outer-trim wrapper projects exactly the public projection — same shape as the + transform path's outer wrap, minus the step CTEs — so the hidden alias does + not leak into the result columns (DEV-1501). + +It builds its own `slot_id_by_key` map (the `PlannedQuery` doesn't carry the +registry), materializes hidden aux slots referenced as transform inputs / +partition keys / time keys / POST-filter operands, and renders. + +### The synthetic-`EnrichedMeasure` adapter (deviation) + +To render aggregations identically to legacy across all dialects, the new path +**reuses the legacy dialect helpers** (`_build_agg`, `_build_percentile`, +`_build_stat_agg`, `_wrap_cast_for_type`, `_resolve_sql`, `_build_date_trunc`). +It does so by synthesizing `EnrichedMeasure` objects from planned slots +(`_synthesize_enriched_measure_from_planned`) and feeding them to those helpers. + +This is a real coupling: `generate_from_planned` consumes `PlannedQuery` at the +top but adapts back to `EnrichedMeasure` — a type DEV-1452 wants to delete — to +emit aggregate SQL. The plan said "rewrite `generator.py` to consume +`PlannedQuery`"; the implemented path is a hybrid. It is flagged in +[the deviations list](index.md#deviations-from-the-plan). The upside is that +dialect-specific behavior (SQLite UDFs, ClickHouse `quantile`, the MySQL +`median` `NotImplementedError`, etc.) is rendered by exactly one code path, +shared with legacy — so the two pipelines can't drift on dialect SQL while both +exist. + +## Multi-stage chaining (`generate_planned_stages`) + +Each non-root stage renders independently (against a per-stage bundle from +`_bundle_for_stage`) and is wrapped by `_stage_rename_wrapper` so its output +columns become the flat names downstream stages bound against +(`orders.customers.region` → `customers__region`). The wrapper derives those from +the *actual* rendered `named_selects` (robust to the cross-model renderer +emitting columns out of `public_projection` order) and asserts they match the +stage's `StageSchema` — a planner/generator divergence fails here rather than as +a confusing downstream bind miss. Stage CTEs are prepended before any CTEs the +root already emits (the root reads `FROM `). + +`_bundle_for_stage` picks the host model the stage renders against from the +planner's `render_source_model` (the stage's own source / overlay / +synthetic-over-sibling), falling back to a synthetic model over the upstream CTE +for a `StageSchema` chain stage — so the generator's FROM/joins bind against +exactly what the binder used. + +## Cross-model rendering + +`_render_with_cross_model_plans` emits one `_cm_*` CTE per +`CrossModelAggregatePlan` joined back to the host base. When `plan.rerooted_plan` +is set, `_render_rerooted_cross_model_cte` renders the nested re-rooted plan +(FROM target + the target's joins) preserving host grain; otherwise the +forward-path CTE renders (FROM bare target, grouped at the forward dims). +`Column.filter` on the aggregated column renders as +`SUM(CASE WHEN THEN END)`. See +[Cross-model aggregates](cross-model-aggregates.md). + +## Mode-A filter inlining and join discovery (DEV-1494) + +A column-level `Column.filter` on an aggregated measure becomes a CASE-WHEN +wrapper (`SUM(CASE WHEN THEN END)`), and a `SlayerModel.filters` +entry becomes a WHERE term. Both are Mode-A SQL and share one renderer, +`_render_mode_a_predicate`, which inline-expands references to derived columns — +bare (`is_eu` → its `CASE WHEN customers.region …`) or dotted to a derived +column on a joined model (`loss_payment.has_flag` → its `sql`) — so the emitted +predicate is runnable and never references a non-physical `.`. +A predicate with only base refs takes the cheap qualify path +(`_qualify_mode_a_sql_filter` regex for model filters, `_qualify_column_filter_sql` +AST for column filters), byte-identical to before. On sqlglot parse failure the +predicate falls through to the qualify path unchanged. + +Join discovery for these text filters (`_filter_join_paths`) is the **union** of +the join paths in the **un-inlined** predicate and those in the **inline-expanded** +predicate. Both are needed: the dbt placeholder-join idiom — a constant derived +column such as `has_flag sql="1"` whose only purpose is to force the (inner) +join — keeps its alias only in the un-inlined form (it inlines to the constant +`(1)`), while a derived ref's *crossed* joins (`is_eu` → `customers`; +`loss_payment.deep_flag` → `loss_payment__claim`) appear only after expansion. +Discovery for column filters in the base SELECT is restricted to **local** +aggregate sources (empty `AggregateKey.path`); a cross-model aggregate's filter +joins are discovered inside its `_cm_*` CTE instead — `_render_cross_model_cte` +collects the join paths of the target measure's `Column.filter` and the +target-model filters and adds them to the CTE's own FROM. Because each `_cm_*` +CTE is an isolated per-(target, grain) computation, adding the join resolves the +filter's refs without affecting sibling measures. Discovery is root-scope-only, +so a correlated ref inside an `EXISTS (...)` subquery does not pull an outer join. + +## Host-base join discovery (the three symmetric sources) + +The host base FROM at `_build_base_select_for_planned` pulls in `LEFT JOIN`s from +three symmetric sources, each handled by a dedicated collector wired in the same +call chain just before `_build_from_and_joins`: + +1. **Dimension / time-dimension `Column.sql`** (DEV-1484): `_expand_derived_row_dims` + pre-expands derived ROW slots (`ColumnSqlKey` dims and `TimeTruncKey` columns + that are themselves derived) and scans the expansion through + `_joined_paths_in_sql`, appending crossed paths to `needed_join_paths`. +2. **Aggregated-measure `Column.filter`** (DEV-1494): `_collect_column_filter_join_paths` + recurses through AGGREGATE-phase composite keys (`ArithmeticKey` / + `ScalarCallKey`) and, for each `AggregateKey` with a `column_filter_key`, + collects the paths the predicate touches via `_filter_join_paths` (the union + of un-inlined and inline-expanded predicate paths, per the section above). +3. **Aggregate-source `Column.sql`** (DEV-1502): `_collect_aggregate_source_join_paths` + mirrors the filter helper — recurses through the same composite keys, and for + each `AggregateKey` whose `source` is a `ColumnSqlKey` with `path == ()`, + expands the column via `_expand_derived_column_sql` and scans the result + through `_joined_paths_in_sql`. The render-time expansion in + `_build_agg_render_spec_from_planned` already produces `SUM()` SQL; + this collector closes the join-discovery loop so a measure source like + `customers__regions.population` emits both `LEFT JOIN`s. + +All three collectors restrict to **local** aggregate sources (empty +`AggregateKey.source.path`); cross-model aggregates own their own join +discovery inside the per-plan `_cm_*` CTE for the `Column.filter` side +(DEV-1494 / DEV-1503). The symmetric source-`Column.sql` discovery inside +the `_cm_*` CTE is a known gap (DEV-1526) — a cross-model aggregate whose +target column's `Column.sql` crosses a further join does not yet have that +join pulled into the CTE FROM. All three host-side collectors feed the +shared `needed_join_paths` list, so repeated paths surfaced by different +sources dedupe naturally via `_build_from_and_joins`'s `emitted_aliases` +guard. + +## Result-key contract (P10) + +The generator preserves the result keys byte-for-byte: `orders.revenue_sum`, +`orders._count` (the `*` dropped, the leading `_` kept), joined dimensions as the +full dotted path `orders.customers.regions.name`, and renamed measures as +`orders.`. `_full_alias_for_slot` derives these from the slot's key / +public aliases. Two documented exceptions, both routed through the same +`canonical_agg_name` helper: cross-model parametric aggregates carry the kwarg +suffix legacy dropped, and hidden parametric `first`/`last` (DEV-1501) carry the +explicit time-arg suffix so distinct time-column specs get distinct +materialised aliases (`orders.revenue_last_created_at`, +`orders.revenue_last_updated_at`). + +## Response metadata (`response_meta.py`) + +The legacy engine derived `SlayerResponse.attributes` and `expected_columns` from +an `EnrichedQuery`. The typed pipeline has none, so `build_response_metadata` +rebuilds both from the root `PlannedQuery` plus the rendered SQL: + +- **`expected_columns`** comes from the final SQL's `named_selects` — the literal + result-key columns rows come back under. Reading them from the SQL (rather than + re-deriving from slots) is bulletproof: it is exactly the outer SELECT the + generator emitted. +- **`attributes`** (`ResponseAttributes.dimensions` / `.measures`) come from the + root plan's public `ValueSlot`s, classified dimension (ROW phase) vs measure + (everything else), with each public result key mapped to its + `FieldMetadata(label, format)`. `_slot_result_keys` mirrors + `_full_alias_for_slot` so the keys line up with the rendered projection; only + keys actually present in the rendered SQL are surfaced (a guard against + divergence). Aggregate formats come from `_infer_aggregated_format` (INTEGER + for count/star, FLOAT for avg-family, source-column format for sum/min/max). + +`FieldMetadata` / `ResponseAttributes` / `_infer_aggregated_format` live here (not +in `query_engine`) so the module imports nothing from the engine; +`query_engine` re-exports them, keeping the public import path unchanged. + +## Design rationale + +- **Why reuse legacy dialect helpers instead of reimplementing aggregation SQL?** + Dialect coverage (SQLite UDFs, ClickHouse parametric quantiles, MySQL's + unsupported-function `NotImplementedError`, the `log10`/`log2` literal + preservation, JSON-extract rewriting) is large and well-tested. Sharing one + emitter keeps the two pipelines from drifting on dialect SQL while both exist — + at the cost of the `EnrichedMeasure` coupling, which DEV-1452 removes. +- **Why derive `expected_columns` from the SQL?** Because the SQL is the ground + truth for what rows come back keyed by. Re-deriving from slots risks a subtle + mismatch; reading `named_selects` cannot. +- **Why assert in `_stage_rename_wrapper`?** A leaked hidden column or a C13 + over-projection would otherwise surface as a downstream "column not found" + deep in the next stage's binding. Asserting at the boundary turns a confusing + failure into a precise one. diff --git a/docs/architecture/stage-planning.md b/docs/architecture/stage-planning.md new file mode 100644 index 00000000..0bb6081f --- /dev/null +++ b/docs/architecture/stage-planning.md @@ -0,0 +1,164 @@ +# Stage planning + +**Module:** `slayer/engine/stage_planner.py` + +The stage planner is the orchestrator that turns `SlayerQuery` stages into +`PlannedQuery`s. `plan_query` compiles one stage; `plan_stages` compiles a +multi-stage DAG. It composes the [binder](binding.md), the +[planning](planning.md) primitives, and the +[cross-model planner](cross-model-aggregates.md), and emits each stage's +`StageSchema` (**P6**). + +## `plan_query` — one stage end to end + +```mermaid +flowchart TB + q["SlayerQuery + scope + bundle"] --> dm["_declared_measures_from_query
parse + expand + bind dims/TDs/measures"] + dm --> filt["bind filters (date_range, model.filters, user filters)"] + filt --> ord["resolve ORDER BY (alias map → bind)"] + ord --> td["resolve active TD → _attach_time_keys"] + td --> sugar["lower_sugar_transforms (change/change_pct)"] + sugar --> tval["validate: every time-needing transform has a time_key"] + tval --> proj["ProjectionPlanner.plan → registry + projection"] + proj --> cmp["per cross-model aggregate: cross_model_planner.plan + maybe re-root"] + cmp --> emit["emit transform_layers, filters_by_phase, stage_schema"] + emit --> pq["PlannedQuery"] +``` + +### Declared measures, in projection order + +`_declared_measures_from_query` builds the `DeclaredMeasure` list in **dim → +time-dim → measure** order (matching the legacy `user_projection` order). Each +dimension/time-dimension binds and is declared under its flattened `__` name +(`stores.opened_at` → `stores__opened_at`) — that flat name is the +`StageColumn.name` a downstream stage binds against (DEV-1448/1449). Measures +run through `expand_model_measures` first (against a `ModelScope` only — +downstream `StageSchema` stages don't expose saved measures), then bind, then get +a canonical alias via `_canonical_alias_for_formula`. + +`_canonical_alias_for_formula` routes any aggregate-rooted formula (including +parametric `revenue:percentile(p=0.5)`) through `canonical_agg_name` so kwargs are +sanitized consistently (`p=0.5` → `_p_0_5`); a cross-model star keeps its +`customers.` prefix. (This is also where cross-model parametric aliases keep the +kwarg suffix legacy dropped — the documented P10 divergence.) + +### Filters, in legacy WHERE order + +The planner constructs filters in the exact order the legacy generator emitted +them, so SQL stays parity-stable: + +1. `date_range` filters — one per time dimension with a 2-element range, built by + `_build_date_range_filter` as a `BetweenKey` over the **bare** underlying + column (a `ColumnKey`, or a `ColumnSqlKey` for a derived temporal column — + DEV-1450 #4a — which the generator renders as ` BETWEEN …`), + not the `TimeTruncKey`, so the self-join CTE path can read raw data while the + filter applies to the outer projection. +2. `SlayerModel.filters` — Mode-A SQL, validated by `_validate_model_filter` + (rejects DSL constructs, raw windows, measure refs, and windowed columns). + A reference to a non-trivial derived column is accepted (DEV-1450 #4b): the + generator's `_render_model_filter_sql` inline-expands the predicate at render + time. These are text-only `FilterPhase` entries with no typed value-key. +3. user query filters — Mode-B DSL, bound with the `filter_alias_map` so renamed + measures resolve by alias (DEV-1445). Two filter strings that bind to the same + structural key are deduped (P2) so a HAVING isn't duplicated. + +The `filter_alias_map` is built from **measure** aliases only (the tail of +`declared_measures` past the dim/time-dim prefix) — never dimension/time-dimension +names, because a time dimension's declared name is its raw column and a +`created_at <= '…'` filter must resolve to the raw column. + +### ORDER BY resolution + +A user order column may name a declared measure by user `name`, declared name, +canonical alias, the flattened dotted form (a joined dimension), or the `_count` +form of `*:count`. `plan_query` checks `declared_alias_to_bound` in that order +before falling back to `bind_expr` on the preserved `raw_formula` — so an +aggregate alias like `amount_sum` (not a column on the model) interns onto the +projection slot rather than raising. + +### Time-key attachment + +`_attach_time_keys` walks every measure/filter/order value-key and, for each +time-needing `TransformKey` (`cumsum` / `change` / `time_shift` / `first` / +`last` / `lag` / `lead` / `consecutive_periods`) whose `time_key` is `None`, +patches in the active TD's key. The active TD is resolved by +`_resolve_main_time_dimension` (0 TDs → none; 1 TD → that one; +2+ → `main_time_dimension` by full-name then leaf, else +`model.default_time_dimension` host-local). After patching, +`_find_unresolved_time_needing_op` validates that no time-needing transform was +left without a TD, raising the legacy error phrase. Sugar lowering runs *after* +this so the desugared `time_shift` inherits the patched `time_key`. + +### Emitting the plan + +`ProjectionPlanner.plan` builds the registry; `_bucket_slots` splits slots into +row / aggregate / combined by phase; the cross-model loop runs the strategy once +per cross-model aggregate — passing `host_query` / `public_projection` / a +`subplan_builder` callback so the strategy itself decides forward-vs-re-rooted +(DEV-1450 #2; the re-root logic moved into `cross_model_planner.py`); +`_emit_transform_layers` +emits one `TransformLayer` per transform slot in Kahn-topological dependency +order (so `cumsum(change(...))` renders inner before outer); and +`_emit_stage_schema` builds the `StageSchema` from public slots. + +## `_emit_stage_schema` — the downstream contract (P6) + +Only public (non-hidden) slots appear, one column per `public_projection` +occurrence (so a C13 multi-alias slot emits one column per alias). Each column's +downstream `name` and `sql_alias` are the `__`-flattened alias +(`customers.revenue_sum` → `customers__revenue_sum`), while `public_alias` keeps +the dotted result-key form. Two distinct public columns that flatten to the same +downstream name raise a collision error rather than silently binding the first +match. This flat-name schema is precisely what a downstream stage binds against — +the reason DEV-1449's dotted-ref-downstream raises. + +## `plan_stages` — the multi-stage DAG + +```mermaid +flowchart LR + qs["queries: List[SlayerQuery]"] --> topo["_topo_sort (Kahn, root last)"] + topo --> loop["for each stage in order"] + loop --> sb["_stage_scope_and_bundle
resolve scope + per-stage bundle"] + sb --> pq["plan_query"] + pq --> sch["record stage_schema by name"] + sch --> loop +``` + +`_topo_sort` orders stages so each appears after the siblings it references via +`source_model` (Kahn's algorithm; rejects duplicate names and cycles; unnamed +stages — typically the root — go last). For each stage, +`_stage_scope_and_bundle` resolves the right `(scope, bundle)`: + +- a `ModelExtension`/dict **over a sibling** → overlay the extra columns onto a + synthetic model of the sibling CTE and bind `ModelScope`-style; +- a **bare-string sibling** source (a chain) → bind against the upstream flat + `StageSchema` (P6 / DEV-1449), with the synthetic upstream model as the host + for cross-model/generation consistency; +- otherwise **model-scoped** → the stage's own resolved source (the root uses the + bundle's `source_model`; a named sibling uses its pre-resolved + `stage_source_models` entry). + +Each per-stage bundle threads in synthetic models for the *other* already-planned +siblings (`stage_bundle_with_siblings`), so a join/cross-model ref to a sibling +resolves to its CTE. After planning a named stage, its `StageSchema` is recorded +so later stages can bind against it. + +## Design rationale + +- **Why build declared measures in dim/time-dim/measure order?** So the public + projection order matches legacy exactly (P10), and so the `filter_alias_map` + can slice off the measure tail without tracking indices separately. +- **Why attach `time_key` in the planner rather than the binder?** Only the + planner has the query (the set of TDs, `main_time_dimension`, + `default_time_dimension`). The binder is expression-local. See + [Binding](binding.md). +- **Why emit `StageSchema` per stage rather than let downstream re-walk joins?** + Because re-walking is the legacy tangle. A stage composes with the next *only* + through its schema (P6); the downstream binder sees a flat namespace and + literally cannot reach the upstream join graph — which is what makes DEV-1449 + a structural impossibility rather than a guarded special case. +- **Why topo-sort here when the engine also sorts the runtime list?** The engine + sorts the user-submitted list (and validates root-as-sink) before planning; + `_topo_sort` re-establishes the planning order from `source_model` references so + `plan_stages` is correct regardless of how it's called (it shares the algorithm + but is the planner's own guarantee). diff --git a/docs/architecture/typed-keys.md b/docs/architecture/typed-keys.md new file mode 100644 index 00000000..c73c9bc9 --- /dev/null +++ b/docs/architecture/typed-keys.md @@ -0,0 +1,163 @@ +# Typed keys — structural identity + +**Module:** `slayer/core/keys.py` + +The `ValueKey` family is the foundation of the whole pipeline. A key answers +exactly one question — *"are these two expression occurrences the same value?"* +— and carries nothing else. Rendering state (SQL text, alias, projection +position, hidden-ness) lives on `ValueSlot` (in `planned.py`), never on the key. + +This separation is principle **P2**: identity is structural, not textual. +`revenue:sum`, the inner `revenue:sum` in `change(revenue:sum)`, and a filter +occurrence of `revenue:sum` all build the same `AggregateKey`, so the +`ValueRegistry` interns them to one slot. That is what makes the dedup bugs +(DEV-1446) structurally impossible rather than patched. + +## The family + +```mermaid +classDiagram + class ValueKey { + <> + } + ValueKey <|-- ColumnKey + ValueKey <|-- ColumnSqlKey + ValueKey <|-- TimeTruncKey + ValueKey <|-- StarKey + ValueKey <|-- LiteralKey + ValueKey <|-- SqlExprKey + ValueKey <|-- AggregateKey + ValueKey <|-- TransformKey + ValueKey <|-- ArithmeticKey + ValueKey <|-- ScalarCallKey + ValueKey <|-- BetweenKey +``` + +| Key | Phase | Identifies | +| --- | --- | --- | +| `ColumnKey(path, leaf)` | ROW | a base column; `path` empty for local, non-empty for joined | +| `ColumnSqlKey(path, model, column_name)` | ROW | a derived column (`Column.sql` set) | +| `TimeTruncKey(column, granularity)` | ROW | a time-truncated column at one grain | +| `StarKey(path)` | ROW | the `*` source for `*:count` (local or cross-model) | +| `LiteralKey(value)` | ROW | a literal operand inside an expression tree | +| `SqlExprKey(canonical_sql)` | ROW | a Mode-A SQL fragment (a `Column.filter`) | +| `AggregateKey(source, agg, args, kwargs, column_filter_key)` | AGGREGATE | one aggregation slot | +| `TransformKey(op, input, args, kwargs, partition_keys, time_key)` | POST | a window/temporal transform over a value | +| `ArithmeticKey(op, operands)` | max(operands) | arithmetic / comparison / boolean | +| `ScalarCallKey(name, args)` | max(args) | a closed-allowlist scalar function call | +| `BetweenKey(column, low, high)` | ROW | a `BETWEEN` predicate (today only `date_range`) | + +All are frozen Pydantic models (`_FrozenKey` sets `frozen=True`), so they are +hashable and immutable — usable directly as `dict` keys in the `ValueRegistry`. + +## Phase + +`Phase` is an `IntEnum` (`ROW=0 < AGGREGATE=1 < POST=2`). It is the engine of +filter routing (**P8**): a composite key's phase is the **max** of its operands' +phases, and a filter routes to `WHERE` / `HAVING` / post-filter by the highest +phase it reaches. Phase is computed, not stored — `ArithmeticKey.phase` is +`max(o.phase for o in operands)`, `ScalarCallKey.phase` is the max over args +that carry a phase, and the leaf keys hard-code their level. Keeping phase a +*property of the key* means no separate "is this a HAVING filter?" text analysis +exists anywhere. + +## Design choices + +### Local and cross-model share one shape (P3) + +`ColumnKey`, `ColumnSqlKey`, and `StarKey` all carry a `path: Tuple[str, ...]`. +Empty path = local; non-empty = a join walk from the query's source model +(`("customers",)`, `("customers", "regions")`). `AggregateKey` inherits this via +its `source`. There is **no separate `cross_model_measures` track** in the +intermediate representation — `path == ()` is the only thing distinguishing a +local aggregate from a cross-model one, and "base CTE vs cross-model CTE" is a +*render* decision made downstream by the planner, not a semantic split baked +into the key. + +### Structural identity has to survive Python's scalar coercion + +Python collapses `True == 1 == Decimal("1")` (and `False == 0`). A naive +tuple-of-bare-values hash would intern `args=(True,)` with `args=(Decimal("1"),)` +— wrong. `AggregateKey`, `TransformKey`, `ScalarCallKey`, and `LiteralKey` +therefore override `__hash__` / `__eq__` to wrap each scalar leaf in a +`(type_tag, value)` pair via `_typed_leaf` (`"__bool__"` / `"__num__"` / +`"__str__"` / `"__none__"`). This restores the type distinction at hash/eq time +without changing the stored representation users see via `key.args[0]`. This was +a review fix (the original keys interned distinct values together). + +### Kwargs are canonicalized to sorted order + +`AggregateKey.kwargs` and `TransformKey.kwargs` run through a `before`-validator +(`_sort_kwargs_tuple`) that sorts by key name, so `weighted_avg(weight=qty)` +interns regardless of input order. Numeric scalars are expected to already be +normalized to `Decimal` (via `normalize_scalar`) so `percentile(p=0.5)` and +`percentile(p=0.50)` are the same key; identifier kwargs arrive as `ColumnKey` +so `weight=quantity` and `weight=quantity_v2` differ. + +### `normalize_scalar` is the one place ints/floats become `Decimal` + +`int → Decimal(value)`; `float → Decimal(str(value))` (via `str`, so floats land +on their displayed decimal form, not their binary approximation); `bool` / `str` +/ `None` / `Decimal` pass through. Booleans are checked *before* int (because +`bool` is-a `int` in Python). Anything else raises `TypeError`. + +### `column_filter_key` folds `Column.filter` into aggregate identity + +A column's `Column.filter` (a Mode-A CASE-WHEN applied at aggregation time) +becomes part of the `AggregateKey` via `column_filter_key: Optional[SqlExprKey]`. +Two aggregates over the same column with different attached filters are therefore +different slots; same-filter ones intern. `*:count` (a `StarKey` source) has no +column, so `column_filter_key` stays `None`. + +### `TimeTruncKey` is a distinct key, not a slot flag + +A time dimension is identified by `(column, granularity)`. Month, day, and +raw uses of the same column are distinct slots automatically, with no +special-casing in the registry. The granularity is stored as a plain `str` +(the `TimeGranularity` member's value) so the key stays a pure-data frozen model +without an enum import. The underlying column is recoverable, so a +`date_range` filter can bind against the raw column independently of the +truncation. (Codex weighed three encodings — a new key, slot metadata, or +`TransformKey(op="date_trunc")` — and the new key won for keeping the registry +uniform.) `TimeTruncKey.column` is `Union[ColumnKey, ColumnSqlKey]` (DEV-1450 +follow-up #4a): a derived (`Column.sql`) temporal column is a first-class time +dimension — the generator applies the `DATE_TRUNC` over the expanded +expression. The kind-agnostic helpers `column_leaf` / `column_path` (in +`keys.py`) unwrap either form. See [Binding](binding.md). + +### `BetweenKey` exists only for legacy SQL parity + +A `col BETWEEN low AND high` and the Mode-B compound `col >= low and col <= high` +render to different SQL text. The legacy generator emits `BETWEEN` for +`date_range`, so `BetweenKey` marks exactly that spot. The Mode-B parser never +produces it — a user-written `col >= a and col <= b` stays an `ArithmeticKey`, +preserving its parity with legacy (which keeps the AND form verbatim). This is a +deliberately narrow key whose only job is "don't drift from legacy on this one +construct". + +## The closed scalar allowlist (P1 / C12) + +`SCALAR_FUNCTIONS` is a `frozenset` living here (not in `formula.py`) so the keys +module is the single source of truth for what counts as a structurally-keyed +scalar call: + +```python +SCALAR_FUNCTIONS = frozenset({ + # null handling + "nullif", "coalesce", "ifnull", + # math + "ln", "log10", "log2", "log", "exp", "sqrt", "pow", "power", + "abs", "floor", "ceil", "round", + # string hygiene (was DEV-1378's STRING_HYGIENE_OPS) + "lower", "upper", "trim", "replace", "substr", "instr", "length", "concat", +}) +``` + +Anything outside this set (plus the transform and aggregation registries) raises +`UnknownFunctionError` in Mode B. This replaces the deleted +`MixedArithmeticField` implicit passthrough: arbitrary dialect-specific +functions (`regexp_match`, `date_part`, JSON ops) belong in Mode A — the user +moves them into a derived `Column.sql`. The `keys.py` constant is imported by +both the parser (`syntax.py`) and the binder (`binding.py`); the binder +re-checks membership as defence-in-depth against direct `ParsedExpr` +construction that bypasses the parser. diff --git a/docs/concepts/queries.md b/docs/concepts/queries.md index 02eac171..ad6212e1 100644 --- a/docs/concepts/queries.md +++ b/docs/concepts/queries.md @@ -119,11 +119,16 @@ Filter formulas define conditions for the query. They go in the `filters` parame | `<` | `"amount < 1000"` | | `<=` | `"amount <= 1000"` | | `in` | `"status in ('active', 'pending')"` | +| `not in` | `"status not in ('cancelled', 'expired')"` | | `IS NULL` | `"discount IS NULL"` | | `IS NOT NULL` | `"discount IS NOT NULL"` | | `like` | `"name like '%acme%'"` | | `not like` | `"name not like '%test%'"` | +The right-hand side of `in` / `not in` must be a non-empty tuple of literal +values (strings, numbers, or booleans) — references and expressions on the +RHS are not supported. Both `(...)` and `[...]` syntax are accepted. + ### Boolean Logic Use `and`, `or`, `not` within a single filter string: @@ -394,6 +399,21 @@ When models have [joins](models.md#joins), you can reference measures from joine This generates a sub-query for the joined measure, scoped to shared dimensions, and LEFT JOINs it to the main query — avoiding aggregation errors from row multiplication. +A cross-model **parametric** aggregate keeps its kwarg signature in the result key, so two variants on the same target column do not collide: + +```json +{ + "source_model": "orders", + "dimensions": ["customers.region"], + "measures": [ + "customers.revenue:percentile(p=0.5)", + "customers.revenue:percentile(p=0.95)" + ] +} +``` + +surfaces two distinct result keys — `orders.customers.revenue_percentile_p_0_5` and `orders.customers.revenue_percentile_p_0_95`. (A non-parametric `customers.revenue:sum` surfaces as `orders.customers.revenue_sum`.) + ### Query lists Pass a list of queries to `execute()`. Earlier queries are named sub-queries, the last is the main query. Named queries can be referenced by `source_model` name or joined via `joins`: diff --git a/docs/concepts/references.md b/docs/concepts/references.md index 2fb593d8..2c60d8e0 100644 --- a/docs/concepts/references.md +++ b/docs/concepts/references.md @@ -7,7 +7,7 @@ SLayer has two distinct expression layers and the rules for what each one accept | Mode | Fields | Parser | Accepts | Rejects | |---|---|---|---|---| | **A — SQL** | `Column.sql`, `Column.filter`, each entry of `SlayerModel.filters` | sqlglot | Any valid SQL expression for the underlying dialect — function calls (`json_extract`, `coalesce`, `nullif`, `lower`, `length`, …), arithmetic, `CASE WHEN`, string literals, comparison and boolean operators in SQL spelling (`=`, `<>`, `IS NULL`, `AND`, `OR`, `NOT`, `IN`, `LIKE`). Bare names and `__`-delimited join paths. | Aggregation colon syntax (`revenue:sum`); SLayer transform calls (`cumsum`, `change`, `rank`, …); references to `ModelMeasure` formulas; raw `OVER (...)` window functions inside `Column.filter` / `SlayerModel.filters` (allowed only in `Column.sql`). | -| **B — DSL** | `ModelMeasure.formula`, `SlayerQuery.measures`, `SlayerQuery.filters`, `SlayerQuery.dimensions`, `SlayerQuery.time_dimensions`, `SlayerQuery.order`, `SlayerQuery.main_time_dimension` | Python AST formula parser | Bare names that resolve to a `Column` or `ModelMeasure` on the model; single-dot dotted paths through joins (`customers.regions.name`, `customers.revenue:sum`); aggregation colon syntax (`:`, `*:count`, parametric forms); transform calls (`cumsum(revenue:sum)`, `rank(revenue:sum, partition_by=region)`); arithmetic / boolean / comparison operators; `LIKE` / `NOT LIKE`; the SQL `\|\|` concat operator (folded into `concat(...)`); a small allowlist of lowercase string-hygiene scalars in `SlayerQuery.filters` only — `lower`, `upper`, `trim`, `replace`, `substr`, `instr`, `length`, `concat`; `{variable}` placeholders (filters only). | `__`-delimited tokens in user input; raw SQL function calls outside the string-hygiene allowlist (`json_extract`, `coalesce`, …); raw `OVER (...)`; bare names that don't resolve to a Column / ModelMeasure / custom aggregation / query alias; **uppercase** spellings of the string-hygiene functions (`LOWER`, `TRIM`, …) — DSL is case-sensitive. | +| **B — DSL** | `ModelMeasure.formula`, `SlayerQuery.measures`, `SlayerQuery.filters`, `SlayerQuery.dimensions`, `SlayerQuery.time_dimensions`, `SlayerQuery.order`, `SlayerQuery.main_time_dimension` | Python AST formula parser | Bare names that resolve to a `Column` or `ModelMeasure` on the model; single-dot dotted paths through joins (`customers.regions.name`, `customers.revenue:sum`); aggregation colon syntax (`:`, `*:count`, parametric forms); transform calls (`cumsum(revenue:sum)`, `rank(revenue:sum, partition_by=region)`); arithmetic / boolean / comparison operators; the SQL `\|\|` concat operator (folded into `concat(...)`); pattern matching via the `like(value, pattern)` scalar (emits the SQL `LIKE` operator — wrap in `not (...)` for `NOT LIKE`); a small allowlist of lowercase string-hygiene scalars in `SlayerQuery.filters` only — `lower`, `upper`, `trim`, `replace`, `substr`, `instr`, `length`, `concat`, `like`; `{variable}` placeholders (filters only). | `__`-delimited tokens in user input; raw SQL function calls outside the string-hygiene allowlist (`json_extract`, `coalesce`, …); raw `OVER (...)`; bare names that don't resolve to a Column / ModelMeasure / custom aggregation / query alias; **uppercase** spellings of the string-hygiene functions (`LOWER`, `TRIM`, …) — DSL is case-sensitive. | ## Identifier resolution diff --git a/docs/examples/02_sql_vs_dsl/sql_vs_dsl_nb.ipynb b/docs/examples/02_sql_vs_dsl/sql_vs_dsl_nb.ipynb index 34beafc2..fb254fa4 100644 --- a/docs/examples/02_sql_vs_dsl/sql_vs_dsl_nb.ipynb +++ b/docs/examples/02_sql_vs_dsl/sql_vs_dsl_nb.ipynb @@ -25,10 +25,10 @@ "id": "8785c43a", "metadata": { "execution": { - "iopub.execute_input": "2026-05-05T20:19:03.971500Z", - "iopub.status.busy": "2026-05-05T20:19:03.971386Z", - "iopub.status.idle": "2026-05-05T20:19:04.151515Z", - "shell.execute_reply": "2026-05-05T20:19:04.151069Z" + "iopub.execute_input": "2026-05-27T14:24:25.372595Z", + "iopub.status.busy": "2026-05-27T14:24:25.371226Z", + "iopub.status.idle": "2026-05-27T14:24:30.107987Z", + "shell.execute_reply": "2026-05-27T14:24:30.099565Z" } }, "outputs": [], @@ -62,10 +62,10 @@ "id": "4bbdb270", "metadata": { "execution": { - "iopub.execute_input": "2026-05-05T20:19:04.152763Z", - "iopub.status.busy": "2026-05-05T20:19:04.152695Z", - "iopub.status.idle": "2026-05-05T20:19:04.154781Z", - "shell.execute_reply": "2026-05-05T20:19:04.154539Z" + "iopub.execute_input": "2026-05-27T14:24:30.127282Z", + "iopub.status.busy": "2026-05-27T14:24:30.125235Z", + "iopub.status.idle": "2026-05-27T14:24:30.161721Z", + "shell.execute_reply": "2026-05-27T14:24:30.154353Z" } }, "outputs": [ @@ -76,13 +76,13 @@ "orders model has 7 columns:\n", "\n", "Columns (name -> SQL expression):\n", - " id -> sql: 'id' type: string [PK]\n", - " customer_id -> sql: 'customer_id' type: string\n", - " ordered_at -> sql: 'ordered_at' type: date\n", - " store_id -> sql: 'store_id' type: string\n", - " subtotal -> sql: 'subtotal' type: number\n", - " tax_paid -> sql: 'tax_paid' type: number\n", - " order_total -> sql: 'order_total' type: number\n" + " id -> sql: 'id' type: TEXT [PK]\n", + " customer_id -> sql: 'customer_id' type: TEXT\n", + " ordered_at -> sql: 'ordered_at' type: DATE\n", + " store_id -> sql: 'store_id' type: TEXT\n", + " subtotal -> sql: 'subtotal' type: DOUBLE\n", + " tax_paid -> sql: 'tax_paid' type: DOUBLE\n", + " order_total -> sql: 'order_total' type: DOUBLE\n" ] } ], @@ -132,10 +132,10 @@ "id": "b6a98979", "metadata": { "execution": { - "iopub.execute_input": "2026-05-05T20:19:04.155914Z", - "iopub.status.busy": "2026-05-05T20:19:04.155859Z", - "iopub.status.idle": "2026-05-05T20:19:04.210358Z", - "shell.execute_reply": "2026-05-05T20:19:04.209996Z" + "iopub.execute_input": "2026-05-27T14:24:30.176710Z", + "iopub.status.busy": "2026-05-27T14:24:30.175123Z", + "iopub.status.idle": "2026-05-27T14:24:32.225452Z", + "shell.execute_reply": "2026-05-27T14:24:32.216376Z" } }, "outputs": [ @@ -145,14 +145,19 @@ "text": [ "Month Revenue MoM Change\n", "--------------------------------------\n", - "2024-05 $ 118,106.62 $+17,512\n", - "2024-06 $ 154,699.94 $+36,593\n", - "2024-07 $ 168,960.69 $+14,261\n", - "2024-08 $ 180,589.12 $+11,628\n", - "2024-09 $ 177,773.88 $-2,815\n", - "2024-10 $ 187,103.17 $+9,329\n", + "2024-05 $ 114,871.03 $+16,044\n", + "2024-06 $ 134,270.36 $+19,399\n", + "2024-07 $ 166,394.63 $+32,124\n", + "2024-08 $ 180,528.01 $+14,133\n", + "2024-09 $ 178,288.75 $-2,239\n", + "2024-10 $ 187,832.31 $+9,544\n", "\n", - " SQL of the query WITH base AS (\n", + " SQL of the query SELECT\n", + " \"orders.ordered_at\",\n", + " \"orders.order_total_sum\",\n", + " \"orders.mom_change\"\n", + "FROM (\n", + "WITH base AS (\n", "SELECT\n", " DATE_TRUNC('MONTH', orders.ordered_at) AS \"orders.ordered_at\",\n", " SUM(orders.order_total) AS \"orders.order_total_sum\"\n", @@ -160,34 +165,35 @@ "GROUP BY\n", " DATE_TRUNC('MONTH', orders.ordered_at)\n", "),\n", - "shifted__ts_mom_change AS (\n", + "shifted__time_shift_inner AS (\n", "SELECT\n", - " DATE_TRUNC('MONTH', orders.ordered_at + INTERVAL '1' MONTH) AS \"orders.ordered_at\",\n", + " DATE_TRUNC('MONTH', orders.ordered_at + INTERVAL 1 MONTH) AS \"orders.ordered_at\",\n", " SUM(orders.order_total) AS \"orders.order_total_sum\"\n", "FROM orders AS orders\n", "GROUP BY\n", - " DATE_TRUNC('MONTH', orders.ordered_at + INTERVAL '1' MONTH)\n", + " DATE_TRUNC('MONTH', orders.ordered_at + INTERVAL 1 MONTH)\n", "),\n", - "sjoin__ts_mom_change AS (\n", - "SELECT base.\"orders.order_total_sum\", base.\"orders.ordered_at\", shifted__ts_mom_change.\"orders.order_total_sum\" AS \"orders._ts_mom_change\"\n", + "sjoin__time_shift_inner AS (\n", + "SELECT base.\"orders.order_total_sum\", base.\"orders.ordered_at\", shifted__time_shift_inner.\"orders.order_total_sum\" AS \"orders._time_shift_inner\"\n", "FROM base\n", - "LEFT JOIN shifted__ts_mom_change\n", - " ON base.\"orders.ordered_at\" = shifted__ts_mom_change.\"orders.ordered_at\"\n", + "LEFT JOIN shifted__time_shift_inner\n", + " ON base.\"orders.ordered_at\" = shifted__time_shift_inner.\"orders.ordered_at\"\n", "),\n", - "step2 AS (\n", + "step1 AS (\n", "SELECT\n", - " \"orders._ts_mom_change\",\n", + " \"orders._time_shift_inner\",\n", " \"orders.order_total_sum\",\n", " \"orders.ordered_at\",\n", - " \"orders.order_total_sum\" - \"orders._ts_mom_change\" AS \"orders.mom_change\"\n", - "FROM sjoin__ts_mom_change\n", + " \"orders.order_total_sum\" - \"orders._time_shift_inner\" AS \"orders.mom_change\"\n", + "FROM sjoin__time_shift_inner\n", ")\n", "SELECT\n", - " \"orders._ts_mom_change\",\n", + " \"orders._time_shift_inner\",\n", " \"orders.mom_change\",\n", " \"orders.order_total_sum\",\n", " \"orders.ordered_at\"\n", - "FROM step2\n", + "FROM step1\n", + ") AS _outer\n", "ORDER BY \"orders.ordered_at\" ASC\n", "LIMIT 6\n", "OFFSET 12\n" @@ -245,10 +251,10 @@ "id": "56baba38", "metadata": { "execution": { - "iopub.execute_input": "2026-05-05T20:19:04.211664Z", - "iopub.status.busy": "2026-05-05T20:19:04.211590Z", - "iopub.status.idle": "2026-05-05T20:19:04.250403Z", - "shell.execute_reply": "2026-05-05T20:19:04.250015Z" + "iopub.execute_input": "2026-05-27T14:24:32.250993Z", + "iopub.status.busy": "2026-05-27T14:24:32.246549Z", + "iopub.status.idle": "2026-05-27T14:24:33.078293Z", + "shell.execute_reply": "2026-05-27T14:24:33.071928Z" } }, "outputs": [ @@ -256,8 +262,8 @@ "name": "stdout", "output_type": "stream", "text": [ - " Brooklyn: 256,532 orders, $2,793,658.90\n", - " Philadelphia: 192,450 orders, $2,099,292.44\n", + " Brooklyn: 252,910 orders, $2,838,779.79\n", + " Philadelphia: 193,786 orders, $2,175,586.96\n", "\n", "The filter 'stores.name' resolved to SQL:\n", " WHERE\n", @@ -304,10 +310,10 @@ "id": "e3947bac", "metadata": { "execution": { - "iopub.execute_input": "2026-05-05T20:19:04.251470Z", - "iopub.status.busy": "2026-05-05T20:19:04.251398Z", - "iopub.status.idle": "2026-05-05T20:19:04.266120Z", - "shell.execute_reply": "2026-05-05T20:19:04.265743Z" + "iopub.execute_input": "2026-05-27T14:24:33.100876Z", + "iopub.status.busy": "2026-05-27T14:24:33.097649Z", + "iopub.status.idle": "2026-05-27T14:24:33.399629Z", + "shell.execute_reply": "2026-05-27T14:24:33.391348Z" } }, "outputs": [ @@ -316,11 +322,11 @@ "output_type": "stream", "text": [ "Orders where subtotal > 5x tax (raw SQL filter via ModelExtension):\n", - " Brooklyn: 256,532 orders, $2,793,658.90\n", - " Philadelphia: 192,450 orders, $2,099,292.44\n", - " Chicago: 106,290 orders, $1,183,765.61\n", - " San Francisco: 90,401 orders, $1,059,640.66\n", - " New Orleans: 15,473 orders, $180,782.69\n" + " Brooklyn: 252,910 orders, $2,838,779.79\n", + " Philadelphia: 193,786 orders, $2,175,586.96\n", + " Chicago: 106,715 orders, $1,146,106.66\n", + " San Francisco: 92,703 orders, $1,026,546.09\n", + " New Orleans: 15,398 orders, $170,325.62\n" ] } ], @@ -398,7 +404,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.11.14" + "version": "3.11.11" } }, "nbformat": 4, diff --git a/mkdocs.yml b/mkdocs.yml index e435c389..e418201a 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -69,6 +69,19 @@ nav: - Configuration: - Datasources: configuration/datasources.md - Storage Backends: configuration/storage.md + - Architecture: + - Overview: architecture/index.md + - Typed keys: architecture/typed-keys.md + - Scopes & source bundle: architecture/scopes-and-bundle.md + - Slack normalization: architecture/slack-normalization.md + - Parsing: architecture/parsing.md + - Binding: architecture/binding.md + - Planning: architecture/planning.md + - Cross-model aggregates: architecture/cross-model-aggregates.md + - Stage planning: architecture/stage-planning.md + - SQL generation: architecture/sql-generation.md + - Errors & warnings: architecture/errors-and-warnings.md + - Engine orchestration: architecture/engine-orchestration.md - Development: development.md plugins: diff --git a/slayer/core/errors.py b/slayer/core/errors.py index 22aabac8..4a7cfb92 100644 --- a/slayer/core/errors.py +++ b/slayer/core/errors.py @@ -114,3 +114,307 @@ def __init__(self, cycle: List[Tuple[str, str]]) -> None: self.cycle: List[Tuple[str, str]] = list(cycle) chain = " → ".join(f"{m}.{c}" for m, c in self.cycle) super().__init__(f"Circular column reference detected: {chain}") + + +# --------------------------------------------------------------------------- +# DEV-1450 stage-5 errors — typed, stable str() format. +# +# All error classes below build their message via ``_format_error_message`` +# so ``str(error)`` follows the documented snapshot-friendly shape:: +# +# : +# at +# scope: +# suggestion: +# +# Each indented line is optional. The first line ALWAYS starts with the +# class name so log greps and snapshot tests bind to a stable prefix. +# --------------------------------------------------------------------------- + + +def _format_error_message( + *, + cls_name: str, + summary: str, + location: str | None = None, + scope: str | None = None, + suggestion: str | None = None, + extras: List[Tuple[str, str]] | None = None, +) -> str: + """Build the stable error-message string used by stage-5 error classes. + + ``extras`` lets a class add bespoke key/value rows after the summary + while keeping the leading ``ClassName:`` token intact. + """ + lines = [f"{cls_name}: {summary}"] + if location: + lines.append(f" at {location}") + if scope: + lines.append(f" scope: {scope}") + for k, v in (extras or []): + lines.append(f" {k}: {v}") + if suggestion: + lines.append(f" suggestion: {suggestion}") + return "\n".join(lines) + + +class UnknownReferenceError(SlayerError, ValueError): + """A bare or dotted reference cannot be resolved in the current scope. + + Multi-inherits ``ValueError`` (like :class:`ColumnCycleError`) so the + pre-existing call sites and tests that catch ``ValueError`` for a failed + reference / model resolution keep working after the DEV-1450 cutover + replaced the legacy ``ValueError`` resolution paths with this typed error. + """ + + def __init__( + self, + name: str, + scope_kind: str, + scope_summary: str, + suggestion: str | None = None, + ) -> None: + self.name = name + self.scope_kind = scope_kind + self.scope_summary = scope_summary + self.suggestion = suggestion + super().__init__(_format_error_message( + cls_name=type(self).__name__, + summary=f"Cannot resolve reference {name!r}.", + scope=f"{scope_kind}: {scope_summary}", + suggestion=suggestion, + )) + + +class AmbiguousReferenceError(SlayerError, ValueError): + """A reference matches multiple candidates in scope and can't pick one. + + Multi-inherits ``ValueError`` for back-compat (see + :class:`UnknownReferenceError`). + """ + + def __init__(self, name: str, candidates: List[str]) -> None: + self.name = name + self.candidates = sorted(candidates) + super().__init__(_format_error_message( + cls_name=type(self).__name__, + summary=f"Reference {name!r} has multiple candidates.", + extras=[("candidates", repr(self.candidates))], + )) + + +class IllegalScopeReferenceError(SlayerError, ValueError): + """A reference is syntactically rejected by the current scope kind. + + Examples: ``__`` in a Mode-B ``ModelScope`` ref (use the dotted form); + a dotted ref against a ``StageSchema`` (downstream stages see a flat + namespace, no join syntax). + + Multi-inherits ``ValueError`` for back-compat (see + :class:`UnknownReferenceError`). + """ + + def __init__(self, name: str, scope_kind: str, reason: str) -> None: + self.name = name + self.scope_kind = scope_kind + self.reason = reason + super().__init__(_format_error_message( + cls_name=type(self).__name__, + summary=f"Reference {name!r} is not legal in this scope.", + scope=scope_kind, + extras=[("reason", reason)], + )) + + +class IllegalWindowInFilterError(SlayerError, ValueError): + """A filter contains a raw ``OVER(...)`` window expression, or refers + to a ``Column.sql`` whose body contains a window function (DEV-1369 / + DEV-1336 — predicate promotion was removed). Use a rank-family + transform instead. + + Multi-inherits ``ValueError`` (like :class:`UnknownReferenceError`) so + the pre-existing call sites and tests that catch ``ValueError`` for the + legacy windowed-filter rejection keep working after the cutover. + """ + + def __init__( + self, + filter_expr: str, + source: str, + suggestion: str = "use a rank-family transform (e.g. `rank() <= N`).", + ) -> None: + self.filter_expr = filter_expr + self.source = source + self.suggestion = suggestion + super().__init__(_format_error_message( + cls_name=type(self).__name__, + summary="Window expressions are not allowed in filters.", + extras=[ + ("expr", repr(filter_expr)), + ("source", source), + ], + suggestion=suggestion, + )) + + +class AggregationNotAllowedError(SlayerError, ValueError): + """An aggregation cannot apply to a column. + + Covers type-bucket violations (``sum`` on TEXT), primary-key + restrictions (only ``count`` / ``count_distinct``), and explicit + ``Column.allowed_aggregations`` whitelist violations. + + Subclasses ``ValueError`` (like the other resolution-time errors in + this module) so callers wrapping the engine in ``except ValueError`` + keep catching aggregation-gating failures — the legacy enrichment + pipeline raised a bare ``ValueError`` here. + """ + + def __init__(self, column: str, agg: str, reason: str) -> None: + self.column = column + self.agg = agg + self.reason = reason + super().__init__(_format_error_message( + cls_name=type(self).__name__, + summary=f"Aggregation {agg!r} is not allowed on column {column!r}.", + extras=[("reason", reason)], + )) + + +class UnknownFunctionError(SlayerError, ValueError): + """A function call in Mode B is not in the ``SCALAR_FUNCTIONS`` allowlist, + the transform registry, or the model's aggregation set (C12). + + Subclasses ``ValueError`` (like the other binding-time errors here) so + the REST ``ValueError -> 400`` mapping and ``except ValueError`` callers + keep catching it — the legacy enrichment pipeline raised a bare + ``ValueError`` for this case. + """ + + _DEFAULT_SUGGESTION = "move the call to a derived Column.sql (Mode A)." + + def __init__( + self, + name: str, + location: str, + suggestion: str | None = None, + ) -> None: + self.name = name + self.location = location + self.suggestion = suggestion or self._DEFAULT_SUGGESTION + super().__init__(_format_error_message( + cls_name=type(self).__name__, + summary=f"Function {name!r} is not allowed in Mode B.", + location=location, + suggestion=self.suggestion, + )) + + +class MeasureRecursionLimitError(SlayerError, ValueError): + """Named-measure expansion exceeded the configurable depth limit + (default 32; ``SLAYER_MEASURE_EXPANSION_DEPTH``). + + ValueError-derived for REST/caller parity with the other binding-time + errors (the legacy pipeline raised a bare ``ValueError``). + """ + + def __init__(self, chain: List[str], limit: int = 32) -> None: + self.chain = list(chain) + self.limit = limit + super().__init__(_format_error_message( + cls_name=type(self).__name__, + summary=f"Named-measure expansion exceeded depth (limit={limit}).", + extras=[("chain", " → ".join(self.chain))], + )) + + +class MeasureCycleError(SlayerError, ValueError): + """Named-measure expansion encountered a cycle. + + ValueError-derived for REST/caller parity with the other binding-time + errors (the legacy pipeline raised a bare ``ValueError``). + """ + + def __init__(self, chain: List[str]) -> None: + self.chain = list(chain) + super().__init__(_format_error_message( + cls_name=type(self).__name__, + summary="Cyclic reference in named-measure expansion.", + extras=[("chain", " → ".join(self.chain))], + )) + + +class DuplicateMeasureNameError(SlayerError, ValueError): + """Two measures in the same query declare the same explicit ``name`` + (DEV-1443). + + ValueError-derived for REST/caller parity with the other binding-time + errors (the legacy pipeline raised a bare ``ValueError``). + """ + + def __init__(self, name: str, occurrences: List[str]) -> None: + self.name = name + self.occurrences = list(occurrences) + super().__init__(_format_error_message( + cls_name=type(self).__name__, + summary=f"Measure name {name!r} is declared more than once.", + extras=[("occurrences", repr(self.occurrences))], + )) + + +class MeasureNameCollidesWithColumnError(SlayerError, ValueError): + """A declared measure ``name`` matches a source column on the model + (DEV-1443) — the alias-form filter would silently bind to the source + column instead of the aggregate. + + ValueError-derived for REST/caller parity with the other binding-time + errors (the legacy pipeline raised a bare ``ValueError``). + """ + + def __init__(self, name: str, model: str) -> None: + self.name = name + self.model = model + super().__init__(_format_error_message( + cls_name=type(self).__name__, + summary=( + f"Declared measure name {name!r} matches a source column on " + f"model {model!r}." + ), + )) + + +class CanonicalAliasShadowsColumnError(SlayerError, ValueError): + """A formula's canonical alias (e.g., ``amount_sum`` for ``amount:sum``) + shadows a source column on the same model (DEV-1443). + + ValueError-derived for REST/caller parity with the other binding-time + errors (the legacy pipeline raised a bare ``ValueError``). + """ + + def __init__(self, formula: str, canonical: str, model: str) -> None: + self.formula = formula + self.canonical = canonical + self.model = model + super().__init__(_format_error_message( + cls_name=type(self).__name__, + summary=( + f"Canonical alias {canonical!r} for formula {formula!r} " + f"shadows a source column on model {model!r}." + ), + )) + + +class UnreachableFilterDroppedWarning(UserWarning): + """A host filter referenced slots that aren't reachable from a + cross-model CTE's root, so the filter was dropped from the CTE. + The host query still applies the filter to its own rows; this is a + visibility/debug warning, not an error. + """ + + def __init__(self, filter_text: str, reason: str) -> None: + self.filter_text = filter_text + self.reason = reason + super().__init__( + f"Filter {filter_text!r} dropped from cross-model CTE " + f"(unreachable from CTE root): {reason}" + ) diff --git a/slayer/core/keys.py b/slayer/core/keys.py new file mode 100644 index 00000000..df57c4ee --- /dev/null +++ b/slayer/core/keys.py @@ -0,0 +1,662 @@ +"""Stage 1 (DEV-1450) — typed identity primitives for the new resolution +pipeline. + +Identity is structural (P2 of the DEV-1450 spec). Two expression occurrences +with the same key intern to the same slot — whether the occurrence is a +declared measure, an inner reference inside a transform, or a filter +predicate. + +Rendering state (SQL text, public alias, projection position, hidden-ness) +does not live here. Those decisions belong to the planner and the SQL +generator. The keys carry only what's needed to decide "are these the same +slot?". + +Public types: ``ValueKey`` (Union alias), ``Phase`` (IntEnum), ``ColumnKey``, +``ColumnSqlKey``, ``StarKey``, ``SqlExprKey``, ``AggregateKey``, +``TransformKey``, ``ArithmeticKey``, ``ScalarCallKey``. Helpers: +``normalize_scalar``, ``SCALAR_FUNCTIONS``. + +These types are dormant in stage 1 — no engine code routes through them. +Stages 7a and 7b wire them up. +""" + +from __future__ import annotations + +from decimal import Decimal +from enum import IntEnum +from typing import Optional, Tuple, Union + +from pydantic import BaseModel, ConfigDict, field_validator + + +# --------------------------------------------------------------------------- +# Closed scalar-function allowlist (C12). +# --------------------------------------------------------------------------- + +# Anything outside this set in Mode B raises ``UnknownFunctionError`` at +# binding time. Lives here (not in formula.py) so the keys module is the +# single source of truth for what counts as a structurally-keyed scalar +# call. The binder (stage 7a) imports from here. +SCALAR_FUNCTIONS: frozenset[str] = frozenset({ + # Null handling + "nullif", "coalesce", "ifnull", + # Math + "ln", "log10", "log2", "log", "exp", "sqrt", "pow", "power", + "abs", "floor", "ceil", "round", + # String hygiene (was DEV-1378's STRING_HYGIENE_OPS) + "lower", "upper", "trim", "replace", "substr", "instr", "length", "concat", + # Pattern match — ``like(value, pattern)`` emits the SQL ``LIKE`` operator + # (sqlglot ``exp.Like``); see SQLGenerator scalar-call rendering. + "like", +}) + + +# --------------------------------------------------------------------------- +# Phase +# --------------------------------------------------------------------------- + + +class Phase(IntEnum): + """Resolution phase of a ValueKey (P8). + + Filters and arithmetic compose by taking the maximum phase of their + operands; the filter's phase then routes it to WHERE (ROW), HAVING + (AGGREGATE), or post-filter on the outer SELECT (POST). + """ + + ROW = 0 + AGGREGATE = 1 + POST = 2 + + +# --------------------------------------------------------------------------- +# Scalar +# --------------------------------------------------------------------------- + +Scalar = Union[Decimal, str, bool, None] + + +def normalize_scalar(value): + """Canonicalize a raw scalar before keying. + + - Booleans pass through unchanged (checked BEFORE int because bool + is-a int in Python). + - ``None`` passes through unchanged. + - ``Decimal`` passes through unchanged. + - ``int`` becomes ``Decimal(value)``. + - ``float`` becomes ``Decimal(str(value))`` — via ``str`` so floats + land on their displayed decimal form, not their binary + approximation (``Decimal(0.5)`` differs from ``Decimal("0.5")``). + - ``str`` passes through unchanged. + + Raises ``TypeError`` for anything else (lists, dicts, custom objects). + Caller-side conversion of identifiers to ``ColumnKey`` happens in the + binder; this helper does not touch ColumnKey. + """ + if isinstance(value, bool): + return value + if value is None: + return None + if isinstance(value, Decimal): + return value + if isinstance(value, int): + return Decimal(value) + if isinstance(value, float): + return Decimal(str(value)) + if isinstance(value, str): + return value + raise TypeError( + f"Cannot normalize scalar of type {type(value).__name__!r}: " + f"only int/float/Decimal/str/bool/None are accepted (got {value!r})." + ) + + +# --------------------------------------------------------------------------- +# Base +# --------------------------------------------------------------------------- + + +class _FrozenKey(BaseModel): + """Common config for the typed-key family: frozen (hashable, immutable).""" + + model_config = ConfigDict(frozen=True) + + +def _typed_leaf(v): + """Return a hash- and equality-friendly representation of a scalar + leaf that does NOT conflate numerically-equal values of different + types. + + Python collapses ``True == 1 == Decimal("1")`` (and the same for + ``False`` / ``0``), so a key built from ``args=(True,)`` would + intern with one built from ``args=(Decimal("1"),)`` if the + container's hash/eq blindly delegate to tuple-of-bare-values. + Wrapping the leaf in a ``(type_tag, value)`` pair at hash/eq time + restores the type distinction without changing the stored + representation users see via ``key.args[0]``. + + ``ValueKey`` leaves (ColumnKey, AggregateKey, ...) are themselves + frozen Pydantic models with value-based equality — they ride in the + generic ``("__key__", v)`` slot. Every branch returns a uniform + ``(tag, value)`` pair so callers never have to special-case the + container shape. + """ + if isinstance(v, bool): + return ("__bool__", v) + if v is None: + return ("__none__", None) + if isinstance(v, Decimal): + return ("__num__", v) + if isinstance(v, str): + return ("__str__", v) + return ("__key__", v) + + +def _typed_args(args): + return tuple(_typed_leaf(a) for a in args) + + +def _typed_kwargs(kwargs): + return tuple((k, _typed_leaf(v)) for k, v in kwargs) + + +# --------------------------------------------------------------------------- +# Row-phase keys +# --------------------------------------------------------------------------- + + +class ColumnKey(_FrozenKey): + """Row-level reference to a base column on a model. + + ``path`` is the join walk from the query's source model to the + terminal model — empty for local refs, non-empty for joined refs + (``("customers",)``, ``("customers", "regions")``, …). ``leaf`` is + the column name on the terminal model. + + Local and cross-model references share this shape (P3) — the only + difference is whether ``path`` is empty. The planner uses + ``path == ()`` to decide whether to materialize the value in the + base CTE or in a cross-model sub-query. + """ + + path: Tuple[str, ...] = () + leaf: str + + @property + def phase(self) -> Phase: + return Phase.ROW + + +class ColumnSqlKey(_FrozenKey): + """Reference to a derived column (one whose ``Column.sql`` is set). + + The expansion AST is recovered from the model definition at binding + time — the key only carries identity. Two references to the same + derived column on the same model intern to one slot. + + ``path`` is the join walk from the query's source model to the + model that owns the derived column — empty for local references, + non-empty for joined ones (``("customers",)``, + ``("customers", "regions")``, …). Cross-model planners use + ``path`` the same way they use ``ColumnKey.path``. + """ + + path: Tuple[str, ...] = () + model: str + column_name: str + + @property + def phase(self) -> Phase: + return Phase.ROW + + +class TimeTruncKey(_FrozenKey): + """Row-level reference to a time-truncated column (DEV-1450 stage 7b.3). + + Identifies a time dimension by ``(column, granularity)``. The + underlying column is recoverable via ``column`` so date-range filters + can bind against the raw column independently of the truncation. + + Identity is structural: two ``TimeTruncKey``s with the same + ``column`` and the same ``granularity`` intern to the same slot; + different granularities on the same column are distinct slots. This + lets the ``ValueRegistry`` keep month / day / raw uses of the same + column as separate materialised values without special-casing. + + ``column`` is a ``ColumnKey`` (base temporal column) or a + ``ColumnSqlKey`` (DEV-1450 follow-up #4a — a DERIVED temporal column + whose ``Column.sql`` is set). The SQL generator applies the + ``DATE_TRUNC`` over the bare identifier (``ColumnKey``) or over the + expanded derived expression (``ColumnSqlKey``). + + ``granularity`` is the string value of a ``TimeGranularity`` member + (``"day"`` / ``"month"`` / ...). Stored as ``str`` so the key stays + a pure-data frozen Pydantic model without an enum import here. + """ + + column: Union["ColumnKey", "ColumnSqlKey"] + granularity: str + + @property + def phase(self) -> Phase: + return Phase.ROW + + +def column_leaf(col: Union["ColumnKey", "ColumnSqlKey"]) -> str: + """The leaf column name of a ``TimeTruncKey.column`` regardless of kind. + + ``ColumnKey`` carries ``leaf``; ``ColumnSqlKey`` carries + ``column_name``. Using this helper everywhere a ``TimeTruncKey``'s + column is unwrapped avoids ``leaf`` / ``column_name`` drift. + """ + return getattr(col, "leaf", None) or getattr(col, "column_name") + + +def column_path(col: Union["ColumnKey", "ColumnSqlKey"]) -> Tuple[str, ...]: + """The join path of a ``TimeTruncKey.column`` regardless of kind. + + Both ``ColumnKey`` and ``ColumnSqlKey`` carry ``.path``. + """ + return col.path + + +class StarKey(_FrozenKey): + """Sentinel source for ``*:count`` aggregations. + + ``path`` is empty for the local star (``*:count`` over the host) and + non-empty for a cross-model star (``customers.*:count`` → + ``path=("customers",)``), mirroring ``ColumnKey.path`` so the + cross-model planner can route a star aggregate through the join graph + (P3). Two stars with the same path intern; the default empty path + keeps the local-star identity bit-identical to before. + """ + + path: Tuple[str, ...] = () + + @property + def phase(self) -> Phase: + return Phase.ROW + + +class LiteralKey(_FrozenKey): + """Identity for a literal value inside an expression tree. + + Used wherever an ``ArithmeticKey``, ``TransformKey``, or other + composite key needs a literal operand (``revenue:sum + 1`` — the + ``1`` is a ``LiteralKey``). Carries phase ROW so it doesn't + artificially elevate the phase of expressions it appears in. + + Scalar normalization (int → Decimal, float → Decimal via str) + happens at the call site via ``normalize_scalar`` so equality + is type-stable (``LiteralKey(Decimal(1))`` and + ``LiteralKey(True)`` are distinct). + """ + + value: Union[Decimal, str, bool, None] = None + + @property + def phase(self) -> Phase: + return Phase.ROW + + def __hash__(self) -> int: + return hash(("LiteralKey", _typed_leaf(self.value))) + + def __eq__(self, other: object) -> bool: + if not isinstance(other, LiteralKey): + return NotImplemented + return _typed_leaf(self.value) == _typed_leaf(other.value) + + +class SqlExprKey(_FrozenKey): + """Identity for a Mode-A SQL fragment. + + Currently used as ``AggregateKey.column_filter_key`` so a + ``Column.filter`` wired in at aggregation time becomes part of the + aggregate's structural identity. Two aggregates over the same column + differ when their attached ``Column.filter`` differs; same-filter + ones intern. + + ``canonical_sql`` is a sqlglot-normalized string (the binder is + responsible for normalization — the key trusts the form it receives). + + DEV-1503 — ``referenced_join_paths`` is the typed SET (semantically; + stored as an ordered tuple for hashability) of non-anchor join-path + prefixes the filter touches after derived-ref expansion. Computed + once at bind time via + ``slayer.engine.column_filter_paths.compute_column_filter_join_paths``; + the planner reads it to decide whether a filtered-local measure must + isolate (the DEV-1503 trigger predicate). ``()`` for same-model + filters; non-empty for cross-model column filters. The field + participates in structural identity — two filters with the same + canonical SQL but different referenced paths would be a bug, so + folding it into the key catches that invariant violation by + comparison. + + The ``before``-validator canonicalises the input to a sorted, + de-duplicated tuple of tuples — so callers can pass any iterable + (list, set, generator) and order doesn't affect identity (otherwise + two semantically-equal SqlExprKeys built with paths in different + order would intern as different keys; CodeRabbit nitpick). + """ + + canonical_sql: str + referenced_join_paths: Tuple[Tuple[str, ...], ...] = () + + @field_validator("referenced_join_paths", mode="before") + @classmethod + def _canonicalize_referenced_join_paths(cls, v): + if not v: + return () + return tuple(sorted({tuple(p) for p in v})) + + @property + def phase(self) -> Phase: + return Phase.ROW + + def __hash__(self) -> int: + return hash(("SqlExprKey", self.canonical_sql, self.referenced_join_paths)) + + def __eq__(self, other: object) -> bool: + if not isinstance(other, SqlExprKey): + return NotImplemented + return ( + self.canonical_sql == other.canonical_sql + and self.referenced_join_paths == other.referenced_join_paths + ) + + +# --------------------------------------------------------------------------- +# Aggregate / Transform / Arithmetic / ScalarCall +# --------------------------------------------------------------------------- + + +_AggregateSource = Union[ColumnKey, ColumnSqlKey, StarKey] +# Positional and keyword arg values accept the same union — both +# `last(created_at)` (positional ColumnKey time arg) and +# `weighted_avg(weight=qty)` (kwarg ColumnKey) bind to identifier columns +# via `_bind_agg_arg`. Reusing one alias for both keeps the surface tight. +_AggregateArgValue = Union[ColumnKey, ColumnSqlKey, Decimal, str, bool, None] +_AggregateKwargValue = _AggregateArgValue + + +def _sort_kwargs_tuple(v): + """Validator helper: canonicalize a kwargs tuple to sorted order by key.""" + if v is None: + return () + return tuple(sorted(v, key=lambda kv: kv[0])) + + +class AggregateKey(_FrozenKey): + """Identity for an aggregation slot (P3). + + Local and cross-model aggregates share this shape: ``source.path`` + is empty for local, non-empty for joined. The render strategy + (base CTE vs cross-model CTE) is decided downstream by the planner. + + ``args`` and ``kwargs`` carry the aggregation's parameters. Numeric + scalars must already be normalized to ``Decimal`` (use + ``normalize_scalar``). Identifier kwargs (``weighted_avg(weight=quantity)``) + arrive as ``ColumnKey`` / ``ColumnSqlKey``. ``kwargs`` is canonicalized + to sorted-by-key order by the validator so input order does not affect + identity. + + ``column_filter_key`` is the ``Column.filter`` attached to the + aggregated column, if any — pulled into the structural key so two + aggregates with different attached filters do not collide. + """ + + source: _AggregateSource + agg: str + args: Tuple[_AggregateArgValue, ...] = () + kwargs: Tuple[Tuple[str, _AggregateKwargValue], ...] = () + column_filter_key: Optional[SqlExprKey] = None + + @field_validator("kwargs", mode="before") + @classmethod + def _canonicalize_kwargs(cls, v): + return _sort_kwargs_tuple(v) + + @property + def phase(self) -> Phase: + return Phase.AGGREGATE + + def __hash__(self) -> int: + return hash(( + "AggregateKey", + self.source, + self.agg, + _typed_args(self.args), + _typed_kwargs(self.kwargs), + self.column_filter_key, + )) + + def __eq__(self, other: object) -> bool: + if not isinstance(other, AggregateKey): + return NotImplemented + return ( + self.source == other.source + and self.agg == other.agg + and _typed_args(self.args) == _typed_args(other.args) + and _typed_kwargs(self.kwargs) == _typed_kwargs(other.kwargs) + and self.column_filter_key == other.column_filter_key + ) + + +class TransformKey(_FrozenKey): + """Identity for a transform slot (window / temporal operator over a value). + + The ``input`` is the value the transform operates on — typically an + aggregate or another transform, occasionally a row-level column. + + ``partition_keys`` is a frozenset (order-independent); ``time_key`` is + addressed separately as the sort dimension for time-ordered transforms. + """ + + op: str + input: "ValueKey" + args: Tuple[Scalar, ...] = () + kwargs: Tuple[Tuple[str, Scalar], ...] = () + partition_keys: frozenset["ValueKey"] = frozenset() + time_key: Optional["ValueKey"] = None + + @field_validator("kwargs", mode="before") + @classmethod + def _canonicalize_kwargs(cls, v): + return _sort_kwargs_tuple(v) + + @property + def phase(self) -> Phase: + return Phase.POST + + def __hash__(self) -> int: + return hash(( + "TransformKey", + self.op, + self.input, + _typed_args(self.args), + _typed_kwargs(self.kwargs), + self.partition_keys, + self.time_key, + )) + + def __eq__(self, other: object) -> bool: + if not isinstance(other, TransformKey): + return NotImplemented + return ( + self.op == other.op + and self.input == other.input + and _typed_args(self.args) == _typed_args(other.args) + and _typed_kwargs(self.kwargs) == _typed_kwargs(other.kwargs) + and self.partition_keys == other.partition_keys + and self.time_key == other.time_key + ) + + +class ArithmeticKey(_FrozenKey): + """Identity for an arithmetic / comparison / boolean expression. + + ``op`` is the operator symbol (``+``, ``-``, ``*``, ``/``, ``<``, + ``<=``, ``and``, ``or``, …). Operand order matters — subtraction + and division are non-commutative, comparisons have a fixed LHS/RHS, + and even commutative ops keep their textual order for deterministic + SQL emission. + + Phase is the maximum of operand phases (P8). + """ + + op: str + operands: Tuple["ValueKey", ...] + + @property + def phase(self) -> Phase: + return max((o.phase for o in self.operands), default=Phase.ROW) + + +_ScalarCallArg = Union["ValueKey", Decimal, str, bool, None] + + +def _arg_phase(arg) -> Optional[Phase]: + """Return ``arg.phase`` for ValueKey args, ``None`` for pure scalars.""" + return getattr(arg, "phase", None) + + +class ScalarCallKey(_FrozenKey): + """Identity for a closed-allowlist scalar function call (C12). + + ``name`` must be a member of ``SCALAR_FUNCTIONS``. The key constructor + does NOT validate this — the binder rejects unknown names with + ``UnknownFunctionError``. Keeping validation out of the key keeps + identity construction cheap on the hot path. + + Phase is the maximum of arg phases over the args that carry a phase + (i.e., ``ValueKey``s); pure-scalar args contribute the ROW floor. + """ + + name: str + args: Tuple[_ScalarCallArg, ...] = () + + @property + def phase(self) -> Phase: + phases = [p for a in self.args if (p := _arg_phase(a)) is not None] + return max(phases) if phases else Phase.ROW + + def __hash__(self) -> int: + return hash(("ScalarCallKey", self.name, _typed_args(self.args))) + + def __eq__(self, other: object) -> bool: + if not isinstance(other, ScalarCallKey): + return NotImplemented + return ( + self.name == other.name + and _typed_args(self.args) == _typed_args(other.args) + ) + + +# --------------------------------------------------------------------------- +# BetweenKey — DEV-1450 stage 7b.9 +# --------------------------------------------------------------------------- + + +class BetweenKey(_FrozenKey): + """Typed identity for a ``col BETWEEN low AND high`` predicate. + + Closed-form Mode-A SQL constructs (``BETWEEN``) and equivalent + Mode-B compound forms (``col >= low and col <= high``) render to + different SQL text. The planner uses ``BetweenKey`` to mark the + spots where ``BETWEEN`` is the right legacy-parity rendering — today + only ``TimeDimension.date_range`` produces them. User-written DSL + filters never produce ``BetweenKey``: the syntax parser doesn't + have a ``between`` construct, and a user-written ``col >= a and + col <= b`` stays as ``ArithmeticKey(and, [GE, LE])`` so its parity + with the legacy generator (which keeps the AND form verbatim) is + preserved. + + Phase is always ROW — ``BetweenKey`` predicates filter row-level + columns. The renderer emits ``exp.Between``. + """ + + column: "ValueKey" + low: "ValueKey" + high: "ValueKey" + + @property + def phase(self) -> Phase: + return Phase.ROW + + +# --------------------------------------------------------------------------- +# InKey — DEV-1475 +# --------------------------------------------------------------------------- + + +class InKey(_FrozenKey): + """Typed identity for a ``col IN (lit, lit, …)`` / ``NOT IN`` predicate. + + Modelled on ``BetweenKey``: a closed-form SQL predicate with a column + LHS and a fixed tuple of literal-valued RHS operands. Two ``InKey``s + with the same column and the same set of values (in the same order) + intern; ``negated`` flips IN vs NOT IN without doubling the class + count. + + ``values`` is a tuple of ``LiteralKey`` (not bare scalars) so equality + rides through ``LiteralKey``'s type-stable ``_typed_leaf`` machinery + — ``InKey(values=(LiteralKey(True),))`` does not collide with + ``InKey(values=(LiteralKey(Decimal(1)),))``. + + Phase is always ROW; the renderer emits ``exp.In`` (wrapped in + ``exp.Not`` when ``negated``). + """ + + column: "ValueKey" + values: Tuple[LiteralKey, ...] + negated: bool = False + + @field_validator("values") + @classmethod + def _reject_empty_values( + cls, v: Tuple[LiteralKey, ...], + ) -> Tuple[LiteralKey, ...]: + # Defense in depth (Codex review): the parser's ``ast.Compare`` + # branch already rejects empty RHS, but direct construction can + # bypass it and reach the SQL generator, which would emit + # ``col IN ()`` — invalid in every supported dialect. + if not v: + raise ValueError( + "InKey requires a non-empty ``values`` tuple; ``col IN " + "()`` is invalid SQL across every supported dialect.", + ) + return v + + @property + def phase(self) -> Phase: + return Phase.ROW + + +# --------------------------------------------------------------------------- +# Union alias + rebuild for forward refs +# --------------------------------------------------------------------------- + + +ValueKey = Union[ + ColumnKey, + ColumnSqlKey, + TimeTruncKey, + StarKey, + LiteralKey, + AggregateKey, + TransformKey, + ArithmeticKey, + ScalarCallKey, + BetweenKey, + InKey, +] + + +# Resolve the recursive forward references on the keys that take ValueKey. +TransformKey.model_rebuild() +ArithmeticKey.model_rebuild() +ScalarCallKey.model_rebuild() +BetweenKey.model_rebuild() +InKey.model_rebuild() +# TimeTruncKey.column is a Union[ColumnKey, ColumnSqlKey] (DEV-1450 #4a). +TimeTruncKey.model_rebuild() diff --git a/slayer/core/models.py b/slayer/core/models.py index 5b4cfb90..6382ab9f 100644 --- a/slayer/core/models.py +++ b/slayer/core/models.py @@ -23,9 +23,6 @@ logger = logging.getLogger(__name__) -_MULTIDOT_COLUMN_RE = re.compile(r'\b([a-zA-Z_]\w*(?:\.[a-zA-Z_]\w*){2,})\b') -_STRING_LITERAL_RE = re.compile(r"'[^']*'") - class _SubstringRule: """Single source of truth for a forbidden substring inside a name. @@ -115,41 +112,6 @@ def _validate_column_name(name: str, context: str) -> str: return name -def _convert_multidot_ref(match: re.Match) -> str: - """Convert a multi-dot reference like ``a.b.c`` to ``a__b.c``.""" - ref = match.group(1) - parts = ref.split(".") - return "__".join(parts[:-1]) + "." + parts[-1] - - -def _fix_multidot_sql(sql: str, context: str) -> str: - """Auto-convert multi-dot references in a SQL snippet to __ alias syntax. - - Single-dot references (``table.column``) are left as-is. - Multi-dot references (``a.b.c``) are converted to ``a__b.c`` with a warning. - String literals are skipped. - """ - # Build a map of string-literal spans to skip - literal_spans = [m.span() for m in _STRING_LITERAL_RE.finditer(sql)] - - def _in_literal(start: int) -> bool: - return any(s <= start < e for s, e in literal_spans) - - result = sql - for match in list(_MULTIDOT_COLUMN_RE.finditer(sql)): - if _in_literal(match.start()): - continue - ref = match.group(1) - fixed = _convert_multidot_ref(match) - logger.warning( - "%s: auto-converting multi-dot reference '%s' to '%s'. " - "Use '__' for join paths in SQL snippets (e.g., '%s').", - context, ref, fixed, fixed, - ) - result = result.replace(ref, fixed) - return result - - class Column(BaseModel): """A row-level column on a model. @@ -193,18 +155,10 @@ def _coerce_legacy_type(cls, data: Any) -> Any: def _validate_name(cls, v: str) -> str: return _validate_column_name(v, "Column") - @field_validator("sql") - @classmethod - def _fix_multidot_sql(cls, v: Optional[str]) -> Optional[str]: - if v is not None: - v = _fix_multidot_sql(v, context="Column sql") - return v - @field_validator("filter") @classmethod - def _fix_multidot_filter(cls, v: Optional[str]) -> Optional[str]: + def _validate_filter_predicate(cls, v: Optional[str]) -> Optional[str]: if v is not None: - v = _fix_multidot_sql(v, context="Column filter") # DEV-1369: Column.filter is SQL-mode — validate at construction # time so DSL constructs (aggregation colon, transform calls) are # caught early. Result is discarded; we only care about the @@ -495,23 +449,18 @@ def _require_data_source_unless_query_backed(self) -> "SlayerModel": @field_validator("filters") @classmethod - def _fix_multidot_filters(cls, v: List[str]) -> List[str]: - """Auto-convert multi-dot column references in model filters and - validate each entry as a SQL-mode predicate (DEV-1369). + def _validate_filter_predicates(cls, v: List[str]) -> List[str]: + """Validate each model filter as a SQL-mode predicate (DEV-1369). Model filters are SQL snippets: joined column references use the - ``__`` alias syntax (``customers__regions.name``), not the - multi-dot query syntax (``customers.regions.name``). Single-dot - references like ``customers.name`` (table.column) are left as-is. - - After the multi-dot rewrite each entry is parsed with - :func:`parse_sql_predicate` so DSL constructs (aggregation colon, - transform calls, raw OVER) are caught at construction time. + ``__`` alias syntax (``customers__regions.name``). Each entry is + parsed with :func:`parse_sql_predicate` so DSL constructs + (aggregation colon, transform calls, raw OVER) are caught at + construction time. """ - rewritten = [_fix_multidot_sql(f, context="Model filter") for f in v] - for f in rewritten: + for f in v: parse_sql_predicate(f) - return rewritten + return v @model_validator(mode="after") def _validate_column_measure_disjoint(self) -> "SlayerModel": diff --git a/slayer/core/refs.py b/slayer/core/refs.py index 01aa8fb8..dadf648d 100644 --- a/slayer/core/refs.py +++ b/slayer/core/refs.py @@ -6,14 +6,20 @@ Keeping a single source of truth prevents the four copies from drifting out of sync. -This module is intentionally side-effect-free and depends on nothing from -``slayer.core.models`` / ``slayer.core.query`` so it can be imported from -those modules' validators without circular import risk. +This module is intentionally side-effect-free and depends only on +``slayer.core.keys`` (for the ``ColumnKey`` shape ``agg_kwarg_canonical_str`` +canonicalises). It does NOT import ``slayer.core.models`` / +``slayer.core.query`` so it can be imported from those modules' +validators without circular import risk. ``slayer.core.keys`` is itself +free of ``slayer`` imports. """ from __future__ import annotations import re -from typing import List, Optional, Tuple +from decimal import Decimal +from typing import Any, List, Optional, Tuple + +from slayer.core.keys import ColumnKey, ColumnSqlKey # --------------------------------------------------------------------------- # Identifier shapes @@ -85,6 +91,104 @@ def agg_signature_suffix( return "_" + "_".join(parts) if parts else "" +def _decimal_to_plain_str(value: Decimal) -> str: + """Return ``value`` as a plain-decimal string with no scientific notation. + + ``str(Decimal("1E-7"))`` yields ``"1E-7"``, which the generator's + ``_SAFE_AGG_PARAM_RE`` SQL-injection allowlist rejects. ``f"{x:f}"`` + forces plain notation but pads short fractional values with extra + zeros (``f"{Decimal('0.5'):f}"`` -> ``"0.5"`` is fine, but + ``f"{0.5:f}"`` on a float yields ``"0.500000"``). To preserve + short forms while expanding exponents, normalize via the Decimal + layer's own ``normalize()`` + a fix-up for the + ``Decimal('-0E+1')`` "-0" exponent quirk. + """ + # Trip the exponent down so ``Decimal("1E-7")`` becomes + # ``Decimal("0.0000001")`` (and ``Decimal("1.0E+3")`` becomes + # ``Decimal("1000")``). For sign normalization use the standard + # ``f"{value:f}"`` then trim trailing zeros after a decimal point. + s = f"{value:f}" + if "." in s: + s = s.rstrip("0").rstrip(".") + return s or "0" + + +def agg_kwarg_canonical_str(value: Any) -> str: + """Canonicalize an AggregateKey kwarg / arg value to SQL-string form. + + DEV-1450 stage 7b.13: ``EnrichedMeasure.agg_kwargs`` is typed as + ``Dict[str, str]`` and every value flows through + ``_validate_agg_param_value`` (``slayer/sql/generator.py:172``) which + accepts only identifiers, qualified names, or numeric literals. + Sites that build the synth ``EnrichedMeasure`` from a typed + ``AggregateKey`` -- AND the two canonical-alias renderers that + previously called ``str(v)`` directly (``slayer/sql/generator.py:3753`` + and ``slayer/engine/cross_model_planner.py:286``) -- route every + kwarg value through this helper instead, so a ``ColumnKey`` never + surfaces as Pydantic-repr noise. + + Conversion rules: + + * ``bool`` / ``None`` -> ``TypeError`` (legacy never accepted these; + ``AggregateKey``'s structural-key normalisation at + ``slayer/core/keys.py:139-142`` keeps them distinct from numerics + precisely so they fail loudly here). + * ``Decimal`` -> ``str(value)`` (Decimal's ``__str__`` matches + ``_SAFE_AGG_PARAM_RE`` for the planner's normalised forms; ``0.5`` + / ``0.95`` / ``100``). + * ``int`` / ``float`` -> ``str(value)`` (planner-side normalisation + already routes literals through ``Decimal``, but the helper stays + total for direct callers). + * ``str`` -> returned unchanged. Callers writing strings into the + key are responsible for safety; downstream validation catches + malformed input at the generator boundary. + * ``ColumnKey(path=(), leaf=L)`` -> ``L``. + * ``ColumnKey(path=P, leaf=L)`` -> ``".".join(P) + "." + L``. + * ``ColumnSqlKey`` (a derived-column arg/kwarg, e.g. + ``corr(other=derived_col)`` — DEV-1450 #4a/#4b) -> the same + ``[path.]column_name`` form so a parametric agg over a derived + column canonicalizes instead of raising. + + Anything else raises ``TypeError`` -- the AggregateKey key shape is + closed over these branches. + """ + if isinstance(value, bool): + # bool is-a int, must check first. + raise TypeError( + f"AggregateKey kwarg cannot be bool: {value!r}", + ) + if value is None: + raise TypeError("AggregateKey kwarg cannot be None") + if isinstance(value, Decimal): + # Route through ``_decimal_to_plain_str`` to force plain + # decimal notation: ``str(Decimal("1E-7"))`` yields ``"1E-7"``, + # which the generator's ``_SAFE_AGG_PARAM_RE`` rejects (no + # scientific notation in the SQL-injection allowlist). + return _decimal_to_plain_str(value) + if isinstance(value, int): + return str(value) + if isinstance(value, float): + # Route floats through Decimal(str(float)) so the + # human-readable decimal text is preserved (matches the + # planner's ``normalize_scalar`` recipe at + # ``slayer/core/keys.py:102``). + return _decimal_to_plain_str(Decimal(str(value))) + if isinstance(value, str): + return value + if isinstance(value, ColumnKey): + if value.path: + return ".".join(value.path) + "." + value.leaf + return value.leaf + if isinstance(value, ColumnSqlKey): + if value.path: + return ".".join(value.path) + "." + value.column_name + return value.column_name + raise TypeError( + f"AggregateKey kwarg value of type {type(value).__name__!r} " + f"is not supported: {value!r}", + ) + + def canonical_agg_name( measure_name: str, aggregation_name: str, diff --git a/slayer/core/scope.py b/slayer/core/scope.py new file mode 100644 index 00000000..cc25ad07 --- /dev/null +++ b/slayer/core/scope.py @@ -0,0 +1,117 @@ +"""Stage 2 (DEV-1450) — typed scope and stage-schema for the new pipeline. + +Two scope kinds, never confused (P5): + +- ``ModelScope``: joins exist; dotted refs walk the join graph rooted at + ``source_model``. ``__`` in a Mode-B ref is an error unless it exact- + matches a column literally named that way (legacy persisted query-backed + columns). +- ``StageSchema``: flat namespace; dots are not join syntax; + ``__``-bearing identifiers are flat names. + +``StageColumn`` is the typed projection element (P6): explicit ``name`` +(downstream bind name), ``sql_alias`` (emitted SQL identifier), +``public_alias`` (result-key piece), plus the per-column metadata that +downstream stages need. + +Per I2 of the DEV-1450 execution plan, ``ModelScope.source_model`` is +``Optional`` from day one so a future anchor-less mode is a type-additive +change. DEV-1450's binder will assert ``source_model is not None`` at +use sites — the type-level optionality is the extension point. +""" + +from __future__ import annotations + +from typing import Any, Dict, List, Optional + +from pydantic import BaseModel, ConfigDict + +from slayer.core.enums import DataType +from slayer.core.format import NumberFormat +from slayer.core.models import SlayerModel + + +class StageColumn(BaseModel): + """Typed projection element for one stage (P6). + + ``name`` is the downstream bind name — flat (e.g. + ``robot_details__modelseriesval`` or ``rev``). ``sql_alias`` is the + identifier emitted in the stage's SELECT projection (usually equal + to ``name``, but the typed split lets the planner reserve hidden + or alias-bearing forms without coupling them). ``public_alias`` is + the result-key piece returned to the user — set only for non-hidden + columns. + + ``format`` (DEV-1452 Stage B decision #8) is the typed ``NumberFormat`` + inherited from the source ``ModelMeasure`` / ``Column`` or computed + by ``_infer_aggregated_format``. ``description`` propagates the source + column's documentation through the typed plan. + """ + + model_config = ConfigDict(frozen=True) + + name: str + sql_alias: str + public_alias: Optional[str] = None + type: Optional[DataType] = None + label: Optional[str] = None + format: Optional[NumberFormat] = None + hidden: bool = False + description: Optional[str] = None + meta: Optional[Dict[str, Any]] = None + sampled: Optional[str] = None + provenance: Optional[str] = None + + +class StageSchema(BaseModel): + """The typed projection of one query stage (P6). + + Downstream stages bind against this as a flat namespace (P5). They + never re-walk the upstream join graph through a StageSchema — the + only legal refs are entries in ``columns``. + + ``relation_name`` is the SQL identifier used when this stage is + referenced from a downstream stage (CTE name or subquery alias). + ``sql`` is the emitted text of the stage's SELECT — populated by the + planner; left ``None`` until rendering. + """ + + model_config = ConfigDict(arbitrary_types_allowed=True) + + relation_name: str + sql: Optional[str] = None + columns: List[StageColumn] + + def __getitem__(self, name: str) -> StageColumn: + for c in self.columns: + if c.name == name: + return c + raise KeyError( + f"No column named {name!r} in stage {self.relation_name!r}." + ) + + def get(self, name: str) -> Optional[StageColumn]: + for c in self.columns: + if c.name == name: + return c + return None + + def __contains__(self, name: object) -> bool: + return isinstance(name, str) and self.get(name) is not None + + +class ModelScope(BaseModel): + """Scope for binding Mode-B refs against a model with joins (P5). + + Dotted refs walk the join graph rooted at ``source_model``; + ``__``-bearing refs are flat-only and reject unless they exact-match + a column literally named that way on the model. + + I2: ``source_model`` is ``Optional`` from day one. DEV-1450's binder + asserts ``source_model is not None`` at use sites so behavior is + unchanged. A future anchor-less mode uses ``source_model=None`` and + a different binder branch (DatasourceScope-style binding). Keeping + the type optional avoids a breaking change later. + """ + + source_model: Optional[SlayerModel] = None diff --git a/slayer/core/warnings.py b/slayer/core/warnings.py new file mode 100644 index 00000000..0f5a1659 --- /dev/null +++ b/slayer/core/warnings.py @@ -0,0 +1,55 @@ +"""Stage 5 (DEV-1450) — slack-normalization warning types. + +The slack-normalization layer (stage 6) rewrites tolerant-but-unambiguous +agent input to canonical form before the typed pipeline sees it, and +emits one ``NormalizationWarning`` payload per rewrite. The payload is +surfaced two ways: + +- Emitted as a Python warning via ``warnings.warn(SlayerNormalizationWarning(payload), ...)`` + so callers using ``warnings.catch_warnings()`` see the rewrite. +- Appended to ``SlayerResponse.warnings: List[NormalizationWarning]`` so + REST/MCP/CLI consumers get the structured payload alongside the result. + +Living in ``slayer.core.warnings`` (not ``slayer.engine.normalization``) +lets memory/storage/REST schemas reference the Pydantic payload without +pulling in engine code. +""" + +from __future__ import annotations + +from typing import Optional + +from pydantic import BaseModel + + +class NormalizationWarning(BaseModel): + """Structured payload describing one slack-normalization rewrite. + + ``rule_id`` identifies the rule that fired (``FUNC_STYLE_AGG``, + ``DOT_PATH_IN_SQL``, ``MISPLACED_MEASURE``). ``location`` is a + human-readable pointer into the query input (e.g. + ``measures[2].formula``). ``rule_doc_url`` is an optional anchor + into ``docs/agent_input_slack.md``. + """ + + rule_id: str + original: str + normalized: str + location: str + rule_doc_url: Optional[str] = None + + +class SlayerNormalizationWarning(UserWarning): + """Carrier ``UserWarning`` for a ``NormalizationWarning`` payload. + + Lets callers route both via ``warnings.catch_warnings(...)`` and + via the structured ``SlayerResponse.warnings`` list — same data, + two surfaces, one source of truth. + """ + + def __init__(self, payload: NormalizationWarning) -> None: + self.payload = payload + super().__init__( + f"[{payload.rule_id}] {payload.original!s} → {payload.normalized!s} " + f"(at {payload.location})" + ) diff --git a/slayer/engine/agg_registry.py b/slayer/engine/agg_registry.py new file mode 100644 index 00000000..4b7c224a --- /dev/null +++ b/slayer/engine/agg_registry.py @@ -0,0 +1,152 @@ +"""Stage 4 (DEV-1450) — aggregation registry helpers. + +Lifts the agg-name collection BFS from ``enrichment.py`` and the +parameter-resolution helpers from ``sql/generator.py`` so the new binder +modules don't have to reach into those tangles. The helpers are pure: +given a model + a resolve_join_target callback, they produce structured +results without touching storage or spawning side maps. + +Public surface: +- ``collect_reachable_agg_names`` — BFS the join graph for custom + aggregation names. +- ``resolve_aggregation`` — find an ``Aggregation`` definition by name. +- ``is_known_aggregation_name`` — built-in or in the custom set. +- ``required_params_for`` — required built-in params (e.g., + ``weighted_avg`` requires ``weight``). +- ``merge_agg_params`` — defaults from the agg-def overridden by + query-time kwargs. + +These are dormant in stage 4 — the existing call sites still inline +their own logic. Stages 7a/7b switch them over. +""" + +from __future__ import annotations + +from typing import Any, Awaitable, Callable, Dict, FrozenSet, List, Optional, Tuple + +from slayer.core.enums import ( + BUILTIN_AGGREGATION_REQUIRED_PARAMS, + BUILTIN_AGGREGATIONS, +) +from slayer.core.models import Aggregation, SlayerModel + + +# --------------------------------------------------------------------------- +# Agg-name collection +# --------------------------------------------------------------------------- + + +ResolveJoinTarget = Callable[..., Awaitable[Optional[Tuple[Any, SlayerModel]]]] + + +async def collect_reachable_agg_names( + source_model: SlayerModel, + resolve_join_target: ResolveJoinTarget, + named_queries: Optional[Dict] = None, +) -> Optional[FrozenSet[str]]: + """Collect custom aggregation names from ``source_model`` and every + join-reachable model. + + BFS bounded only by the visited set (no fixed depth cap — dotted-path + resolution supports arbitrary depth, so the agg-name rewrite must too). + Returns ``None`` when no custom aggregations exist anywhere in the + reachable subgraph. + + ``resolve_join_target`` is the existing engine callback whose return + shape is ``(target_sql, target_model) | None``; this helper only uses + the ``target_model`` element. + """ + names: set[str] = set() + visited: set[str] = set() + queue: List[SlayerModel] = [source_model] + + while queue: + current = queue.pop(0) + if current.name in visited: + continue + visited.add(current.name) + + if current.aggregations: + names.update(a.name for a in current.aggregations) + + for join in current.joins: + if join.target_model in visited: + continue + target_info = await resolve_join_target( + target_model_name=join.target_model, + named_queries=named_queries or {}, + ) + if target_info: + _, target_model_obj = target_info + if target_model_obj is not None: + queue.append(target_model_obj) + + return frozenset(names) if names else None + + +# --------------------------------------------------------------------------- +# Name-based lookups +# --------------------------------------------------------------------------- + + +def is_known_aggregation_name( + name: str, + custom_names: Optional[FrozenSet[str]], +) -> bool: + """``True`` if ``name`` is a built-in aggregation or appears in the + model-collected custom set. + """ + if name in BUILTIN_AGGREGATIONS: + return True + return bool(custom_names) and name in custom_names + + +def resolve_aggregation( + name: str, + available_aggs: List[Aggregation], +) -> Optional[Aggregation]: + """Return the ``Aggregation`` definition for ``name`` if one is + declared in ``available_aggs``, else ``None``. + + A model-level entry whose ``name`` matches a built-in is treated as + an override and returned. ``None`` for a built-in name with no + override is the signal to use the default built-in formula. + """ + for agg in available_aggs: + if agg.name == name: + return agg + return None + + +# --------------------------------------------------------------------------- +# Parameter resolution +# --------------------------------------------------------------------------- + + +def required_params_for(agg_name: str) -> Tuple[str, ...]: + """Required parameter names for a built-in aggregation (e.g., + ``weighted_avg`` requires ``weight``). + + Custom aggregations declare their parameter shape via + ``Aggregation.params``; this helper only knows about the built-in + table in ``slayer.core.enums``. Returns an empty tuple for unknown + names so callers can branch on emptiness rather than ``KeyError``. + """ + return tuple(BUILTIN_AGGREGATION_REQUIRED_PARAMS.get(agg_name, [])) + + +def merge_agg_params( + agg_def: Optional[Aggregation], + query_kwargs: Dict[str, Any], +) -> Dict[str, Any]: + """Combine ``Aggregation.params`` defaults with query-time kwargs. + + Query-time kwargs override defaults. Kwargs not declared by the + ``agg_def`` pass through unchanged — validation of param names + (e.g., rejecting unknown ones) is the binder's responsibility, + not this helper's. + """ + if agg_def is None: + return dict(query_kwargs) + defaults = {p.name: p.sql for p in agg_def.params} + return {**defaults, **query_kwargs} diff --git a/slayer/engine/binding.py b/slayer/engine/binding.py new file mode 100644 index 00000000..a75b0d76 --- /dev/null +++ b/slayer/engine/binding.py @@ -0,0 +1,1309 @@ +"""Stage 7a.5 (DEV-1450) — ExpressionBinder + FilterBinder. + +The binder consumes a ``ParsedExpr`` (from ``slayer/engine/syntax.py``) +plus a scope (``ModelScope`` or ``StageSchema``) and a +``ResolvedSourceBundle`` (for join resolution). It produces a typed +``BoundExpr`` whose leaves are resolved ``ValueKey``s. + +Public surface: + +* ``bind_expr(parsed, *, scope, bundle) -> BoundExpr`` +* ``bind_filter(parsed, *, scope, bundle) -> BoundFilter`` + +Two scope kinds (P5): + +* ``ModelScope``: joins exist; dotted refs walk the join graph rooted + at ``source_model``. ``__``-bearing refs raise + ``IllegalScopeReferenceError`` unless they exact-match a column on + the model. I2: ``source_model is not None`` is asserted. +* ``StageSchema``: flat namespace; dotted refs raise + ``IllegalScopeReferenceError``; flat names with ``__`` are legal. + +C14: same-model self-prefix in Mode-B (`orders.status` over an +``orders``-rooted query) is stripped before the join walk. + +FilterBinder layers on top: ``bind_expr`` + phase classification +(``Phase.ROW`` / ``AGGREGATE`` / ``POST`` = the max phase of any +referenced slot) + walk for the referenced ``ValueKey``s + reject +filters that touch a windowed ``Column.sql``. + +Dormant in 7a — no engine wiring. The planner (7a.6) is the first +consumer. +""" + +from __future__ import annotations + +from typing import Dict, List, Optional, Tuple, Union + +from pydantic import BaseModel, ConfigDict, Field + +from slayer.core.errors import ( + AggregationNotAllowedError, + IllegalScopeReferenceError, + IllegalWindowInFilterError, + UnknownFunctionError, + UnknownReferenceError, +) +from slayer.core.enums import ( + DEFAULT_AGGREGATIONS_BY_TYPE, + PRIMARY_KEY_AGGREGATIONS, + DataType, +) +from slayer.core.keys import ( + SCALAR_FUNCTIONS, + AggregateKey, + ArithmeticKey, + BetweenKey, + ColumnKey, + ColumnSqlKey, + InKey, + LiteralKey, + Phase, + ScalarCallKey, + SqlExprKey, + StarKey, + TimeTruncKey, + TransformKey, + ValueKey, + column_leaf, + column_path, + normalize_scalar, +) +from slayer.core.models import SlayerModel +from slayer.core.query import TimeDimension +from slayer.core.scope import ModelScope, StageSchema +from slayer.engine.column_filter_paths import compute_column_filter_join_paths +from slayer.engine.source_bundle import ResolvedSourceBundle +from slayer.engine.syntax import ( + AggCall, + Arith, + BoolOp, + Cmp, + DottedRef, + Literal, + ParsedExpr, + Ref, + ScalarCall, + StarSource, + TransformCall, + TupleLit, + UnaryOp, +) +from slayer.sql.sql_expr import has_window_function + +__all__ = [ + "BoundExpr", + "BoundFilter", + "bind_expr", + "bind_filter", + "bind_time_dimension", + "walk_value_keys", +] + + +_TEMPORAL_TYPES = frozenset({DataType.DATE, DataType.TIMESTAMP}) + + +# --------------------------------------------------------------------------- +# BoundExpr / BoundFilter +# --------------------------------------------------------------------------- + + +class BoundExpr(BaseModel): + """A bound expression — its leaves are resolved ``ValueKey``s. + + ``value_key`` is the structural identity of the entire expression. + ``phase`` is the property of ``value_key.phase`` (lifted for + convenience). + """ + + model_config = ConfigDict(arbitrary_types_allowed=True, frozen=True) + + value_key: ValueKey + + @property + def phase(self) -> Phase: + return self.value_key.phase + + +class BoundFilter(BaseModel): + """A bound filter predicate. + + The same ``value_key`` shape as ``BoundExpr`` (boolean ops and + comparisons are encoded as ``ArithmeticKey`` with the corresponding + op string), plus: + + * ``phase`` — the maximum phase any referenced slot reaches. + * ``referenced_keys`` — every ``ValueKey`` touched anywhere in the + bound tree (used by the cross-model planner's filter routing). + """ + + model_config = ConfigDict(arbitrary_types_allowed=True, frozen=True) + + value_key: ValueKey + phase: Phase + referenced_keys: Tuple[ValueKey, ...] = Field(default_factory=tuple) + + +# --------------------------------------------------------------------------- +# Public entry points +# --------------------------------------------------------------------------- + + +def bind_expr( + parsed: ParsedExpr, + *, + scope: Union[ModelScope, StageSchema], + bundle: ResolvedSourceBundle, +) -> BoundExpr: + """Bind a parsed expression against a scope. + + Returns a ``BoundExpr`` carrying the structural identity of the + entire expression. Raises ``UnknownReferenceError`` if a ref doesn't + resolve; ``IllegalScopeReferenceError`` if a dotted ref is used + against a ``StageSchema`` (or vice versa for ``__`` against a + ``ModelScope``). + """ + value_key = _bind(parsed, scope=scope, bundle=bundle, in_filter=False) + return BoundExpr(value_key=value_key) + + +def bind_time_dimension( + td: TimeDimension, + *, + scope: Union[ModelScope, StageSchema], + bundle: ResolvedSourceBundle, +) -> BoundExpr: + """Bind a ``TimeDimension`` into a ``BoundExpr`` carrying a + ``TimeTruncKey``. + + The underlying column is resolved against ``scope`` exactly like a + Mode-B identifier ref (local name or dotted-join path); the bound + column must be a plain ``ColumnKey`` whose ``Column.type`` is in the + temporal bucket (``DATE`` / ``TIMESTAMP``). + + Stage 7b.3b limitations: + + * Only ``ModelScope`` with a non-None ``source_model`` is accepted. + Downstream stages bind upstream-emitted truncated columns by flat + name through ``bind_expr``; they do not re-truncate at a different + grain through this entry point. Passing a ``StageSchema`` raises + ``IllegalScopeReferenceError``. + * Derived (``Column.sql`` is set) temporal columns route through + ``ColumnSqlKey`` rather than ``ColumnKey``, and ``TimeTruncKey`` + is typed as ``column: ColumnKey``. Rather than silently widen the + typed key, this stage rejects derived-TD columns with + ``NotImplementedError`` and a clear message. + """ + if isinstance(scope, StageSchema): + raise IllegalScopeReferenceError( + name=td.dimension.full_name, + scope_kind="StageSchema", + reason=( + "time dimensions only bind against a ModelScope; downstream " + "stages already see the truncated column as a flat name " + "from the upstream stage's schema." + ), + ) + + assert isinstance(scope, ModelScope) + if scope.source_model is None: + raise UnknownReferenceError( + name=td.dimension.full_name, + scope_kind="ModelScope", + scope_summary="(no source_model anchor; anchor-less mode not implemented)", + suggestion=None, + ) + + full = td.dimension.full_name + if "." in full: + parts = tuple(full.split(".")) + bound_col = _resolve_dotted(parts, scope=scope, bundle=bundle) + else: + bound_col = _resolve_ref(full, scope=scope, bundle=bundle) + + if not isinstance(bound_col, (ColumnKey, ColumnSqlKey)): + # Defensive — the binder should never produce a non-column key + # for an identifier ref against a ModelScope. + raise ValueError( + f"TimeDimension {full!r} did not resolve to a column " + f"reference (got {type(bound_col).__name__})." + ) + + # DEV-1450 follow-up #4a: a derived (Column.sql) temporal column routes + # through ColumnSqlKey; TimeTruncKey.column accepts both kinds, so the + # leaf / path are read via the kind-agnostic helpers. + terminal_model = _terminal_model_for_path( + path=column_path(bound_col), + scope=scope, + bundle=bundle, + ) + if terminal_model is None: + # Shouldn't be reachable: _resolve_ref / _resolve_dotted would + # already have raised. Defensive only. + raise UnknownReferenceError( + name=full, + scope_kind="ModelScope", + scope_summary=f"could not resolve terminal model for {full!r}", + suggestion=None, + ) + col = next( + (c for c in terminal_model.columns if c.name == column_leaf(bound_col)), + None, + ) + if col is None or col.type not in _TEMPORAL_TYPES: + observed = col.type if col is not None else "" + raise ValueError( + f"TimeDimension {full!r} must reference a temporal column " + f"(DATE / TIMESTAMP); got column type {observed!r}." + ) + + return BoundExpr( + value_key=TimeTruncKey( + column=bound_col, granularity=str(td.granularity.value), + ), + ) + + +def _terminal_model_for_path( + *, + path: Tuple[str, ...], + scope: ModelScope, + bundle: ResolvedSourceBundle, +) -> Optional[SlayerModel]: + """Walk ``path`` from ``scope.source_model`` and return the terminal + model. Returns the host when ``path`` is empty. + """ + current = scope.source_model + if current is None: + return None + for hop in path: + nxt = bundle.get_referenced_model(hop) + if nxt is None: + return None + current = nxt + return current + + +def bind_filter( + parsed: ParsedExpr, + *, + scope: Union[ModelScope, StageSchema], + bundle: ResolvedSourceBundle, + alias_map: Optional[Dict[str, "ValueKey"]] = None, +) -> BoundFilter: + """Bind a parsed filter predicate + classify its phase. + + Walks the bound tree to gather every referenced ``ValueKey`` and + raises ``IllegalWindowInFilterError`` if any referenced + ``Column.sql`` contains a window function (DEV-1369: no + auto-promotion). + + ``alias_map`` maps a stage's declared-measure names (user ``name``, + canonical alias, declared name) to their bound ``ValueKey`` so a + filter may reference a declared measure by alias (P4 / DEV-1445: + ``filters=["rev >= 100"]`` for a measure declared ``name="rev"``). + A bare ref that matches an alias interns onto that exact slot rather + than resolving against the model columns — so the colon form and the + alias form share one slot. + """ + value_key = _bind( + parsed, scope=scope, bundle=bundle, in_filter=True, alias_map=alias_map, + ) + refs = tuple(walk_value_keys(value_key)) + phase = max( + (k.phase for k in refs), + default=value_key.phase, + ) + _reject_windowed_column_sql(refs, scope=scope, bundle=bundle, parsed=parsed) + return BoundFilter( + value_key=value_key, phase=phase, referenced_keys=refs, + ) + + +# --------------------------------------------------------------------------- +# Walk helper +# --------------------------------------------------------------------------- + + +_VALUE_KEY_TYPES = ( + ColumnKey, ColumnSqlKey, StarKey, LiteralKey, + AggregateKey, TransformKey, ArithmeticKey, ScalarCallKey, + BetweenKey, InKey, TimeTruncKey, +) + + +def walk_value_keys(key: ValueKey): + """Yield every ``ValueKey`` reachable from ``key``, including ``key``.""" + yield key + if isinstance(key, AggregateKey): + if isinstance(key.source, _VALUE_KEY_TYPES): + yield from walk_value_keys(key.source) + for a in key.args: + if isinstance(a, _VALUE_KEY_TYPES): + yield from walk_value_keys(a) + for _, v in key.kwargs: + if isinstance(v, _VALUE_KEY_TYPES): + yield from walk_value_keys(v) + elif isinstance(key, TransformKey): + if isinstance(key.input, _VALUE_KEY_TYPES): + yield from walk_value_keys(key.input) + for a in key.args: + if isinstance(a, _VALUE_KEY_TYPES): + yield from walk_value_keys(a) + for _, v in key.kwargs: + if isinstance(v, _VALUE_KEY_TYPES): + yield from walk_value_keys(v) + for pk in key.partition_keys: + yield from walk_value_keys(pk) + if key.time_key is not None: + yield from walk_value_keys(key.time_key) + elif isinstance(key, ArithmeticKey): + for op in key.operands: + yield from walk_value_keys(op) + elif isinstance(key, ScalarCallKey): + for arg in key.args: + if isinstance(arg, _VALUE_KEY_TYPES): + yield from walk_value_keys(arg) + elif isinstance(key, BetweenKey): + yield from walk_value_keys(key.column) + yield from walk_value_keys(key.low) + yield from walk_value_keys(key.high) + elif isinstance(key, InKey): + # DEV-1475: walk the column LHS and every literal RHS so the + # cross-model filter router and the windowed-column rejection + # check both see InKey-rooted predicates the same way they see + # BetweenKey ones. + yield from walk_value_keys(key.column) + for v in key.values: + yield from walk_value_keys(v) + + +# --------------------------------------------------------------------------- +# Internals +# --------------------------------------------------------------------------- + + +def _bind( + parsed: ParsedExpr, + *, + scope: Union[ModelScope, StageSchema], + bundle: ResolvedSourceBundle, + in_filter: bool, + alias_map: Optional[Dict[str, "ValueKey"]] = None, +) -> ValueKey: + if isinstance(parsed, Literal): + return LiteralKey(value=normalize_scalar(parsed.value)) + + if isinstance(parsed, Ref): + return _resolve_ref( + parsed.name, scope=scope, bundle=bundle, alias_map=alias_map, + ) + + if isinstance(parsed, DottedRef): + return _resolve_dotted(parsed.parts, scope=scope, bundle=bundle) + + if isinstance(parsed, StarSource): + return StarKey() + + if isinstance(parsed, AggCall): + return _bind_agg(parsed, scope=scope, bundle=bundle) + + if isinstance(parsed, TransformCall): + return _bind_transform( + parsed, scope=scope, bundle=bundle, alias_map=alias_map, + ) + + if isinstance(parsed, ScalarCall): + return _bind_scalar( + parsed, scope=scope, bundle=bundle, in_filter=in_filter, + alias_map=alias_map, + ) + + if isinstance(parsed, Arith): + return ArithmeticKey( + op=parsed.op, + operands=( + _bind(parsed.left, scope=scope, bundle=bundle, in_filter=in_filter, alias_map=alias_map), + _bind(parsed.right, scope=scope, bundle=bundle, in_filter=in_filter, alias_map=alias_map), + ), + ) + + if isinstance(parsed, UnaryOp): + return ArithmeticKey( + op=parsed.op, + operands=(_bind(parsed.operand, scope=scope, bundle=bundle, in_filter=in_filter, alias_map=alias_map),), + ) + + if isinstance(parsed, Cmp): + # DEV-1475: ``IN`` / ``NOT IN`` predicates fold into a single + # ``InKey`` rather than an ``ArithmeticKey`` so the SQL generator + # has a structured handle on the column + literal-tuple shape. + # The parser already validated that ``parsed.right`` is a + # ``TupleLit`` of ``Literal`` elements for these ops. + if parsed.op in ("in", "not in"): + return _bind_in( + parsed, + scope=scope, bundle=bundle, in_filter=in_filter, + alias_map=alias_map, + ) + return ArithmeticKey( + op=parsed.op, + operands=( + _bind(parsed.left, scope=scope, bundle=bundle, in_filter=in_filter, alias_map=alias_map), + _bind(parsed.right, scope=scope, bundle=bundle, in_filter=in_filter, alias_map=alias_map), + ), + ) + + if isinstance(parsed, BoolOp): + operands = tuple( + _bind(v, scope=scope, bundle=bundle, in_filter=in_filter, alias_map=alias_map) + for v in parsed.operands + ) + return ArithmeticKey(op=parsed.op, operands=operands) + + raise ValueError( + f"Unsupported ParsedExpr node: {type(parsed).__name__}" + ) + + +def _bind_in( + parsed: Cmp, + *, + scope: Union[ModelScope, StageSchema], + bundle: ResolvedSourceBundle, + in_filter: bool, + alias_map: Optional[Dict[str, "ValueKey"]] = None, +) -> InKey: + """Bind an ``IN`` / ``NOT IN`` predicate into an ``InKey`` (DEV-1475). + + The LHS is bound through the normal column-resolution path + (``ColumnKey`` for a bare ref, ``ColumnKey`` with a non-empty + ``path`` for a dotted join ref, ``ColumnSqlKey`` for a derived + column, or an alias-map hit for a declared-measure name). + + The RHS is a ``TupleLit`` of ``Literal`` nodes (the parser already + enforced that shape); every element binds to a ``LiteralKey`` after + scalar normalization. + """ + if not isinstance(parsed.right, TupleLit): + # Defensive — the parser's ``ast.Compare`` branch guarantees a + # ``TupleLit`` on the RHS for ``in`` / ``not in``. Surface a + # clear runtime error if a future caller bypasses the parser. + raise ValueError( + f"_bind_in: expected TupleLit on RHS of {parsed.op!r}, got " + f"{type(parsed.right).__name__}." + ) + column = _bind( + parsed.left, + scope=scope, bundle=bundle, in_filter=in_filter, alias_map=alias_map, + ) + values = tuple( + LiteralKey(value=normalize_scalar(elt.value)) + for elt in parsed.right.elements + ) + return InKey( + column=column, + values=values, + negated=(parsed.op == "not in"), + ) + + +def _resolve_ref( + name: str, + *, + scope: Union[ModelScope, StageSchema], + bundle: ResolvedSourceBundle, + alias_map: Optional[Dict[str, "ValueKey"]] = None, +) -> ValueKey: + """Resolve a bare identifier against the scope. + + A name present in ``alias_map`` (a stage's declared-measure aliases, + supplied only on the filter/order path) interns onto that declared + slot's ``ValueKey`` before any column lookup — so a filter referencing + a measure by its user ``name`` shares the measure's slot (P4). + """ + if alias_map and name in alias_map: + return alias_map[name] + + if isinstance(scope, StageSchema): + col = scope.get(name) + if col is None: + raise UnknownReferenceError( + name=name, + scope_kind="StageSchema", + scope_summary=( + f"stage {scope.relation_name!r} columns: " + f"{[c.name for c in scope.columns]}" + ), + suggestion=None, + ) + return ColumnKey(path=(), leaf=name) + + assert isinstance(scope, ModelScope) + if scope.source_model is None: + raise UnknownReferenceError( + name=name, + scope_kind="ModelScope", + scope_summary="(no source_model anchor; anchor-less mode not implemented)", + suggestion=None, + ) + model = scope.source_model + + if "__" in name: + # The Mode-B parser already rejects `__` for user input; this + # branch is reached only via direct ParsedExpr.Ref construction + # (e.g., downstream binders for StageSchema flat columns). The + # `__` is legal iff it exact-matches a column literally named + # that way on the model (legacy persisted query-backed columns). + if any(c.name == name for c in model.columns): + return ColumnKey(path=(), leaf=name) + raise IllegalScopeReferenceError( + name=name, + scope_kind="ModelScope", + reason=( + "`__` is reserved for internal join-path aliases. " + "Use single-dot DSL paths in queries." + ), + ) + + col = next((c for c in model.columns if c.name == name), None) + if col is None: + # Try ModelMeasure as a fallback for bare measure refs. + mm = next((m for m in model.measures if m.name == name), None) + if mm is not None: + # ModelMeasure expansion lives in the planner; the binder + # raises here so callers know expansion is required. + raise UnknownReferenceError( + name=name, + scope_kind="ModelScope", + scope_summary=f"model {model.name!r}", + suggestion=( + f"{name!r} is a saved measure on {model.name!r}; " + f"expand via ModelMeasure expansion before binding." + ), + ) + raise UnknownReferenceError( + name=name, + scope_kind="ModelScope", + scope_summary=( + f"model {model.name!r} columns: " + f"{[c.name for c in model.columns]}" + ), + suggestion=None, + ) + + if col.sql is not None and col.sql.strip() != name: + return ColumnSqlKey(path=(), model=model.name, column_name=col.name) + return ColumnKey(path=(), leaf=col.name) + + +def _resolve_dotted( + parts: Tuple[str, ...], + *, + scope: Union[ModelScope, StageSchema], + bundle: ResolvedSourceBundle, +) -> ValueKey: + """Resolve a dotted ref against the scope.""" + if isinstance(scope, StageSchema): + raise IllegalScopeReferenceError( + name=".".join(parts), + scope_kind="StageSchema", + reason=( + "downstream stages see a flat schema — dotted refs are " + "not legal. Use the flat column name." + ), + ) + + assert isinstance(scope, ModelScope) + if scope.source_model is None: + raise UnknownReferenceError( + name=".".join(parts), + scope_kind="ModelScope", + scope_summary="(no source_model anchor; anchor-less mode not implemented)", + suggestion=None, + ) + + # C14: strip same-model self-prefix. + host = scope.source_model + if parts and parts[0] == host.name: + parts = parts[1:] + if not parts: + raise UnknownReferenceError( + name=host.name, + scope_kind="ModelScope", + scope_summary=f"model {host.name!r}", + suggestion="self-prefix only — expected a column or join target.", + ) + if len(parts) == 1: + return _resolve_ref(parts[0], scope=scope, bundle=bundle) + + # parts now has the join walk to perform. + if len(parts) == 1: + # Single-segment after possible stripping — already a local ref. + return _resolve_ref(parts[0], scope=scope, bundle=bundle) + + # Walk join chain. parts[:-1] are join targets; parts[-1] is the leaf column. + hop_path = parts[:-1] + leaf = parts[-1] + current = host + visited_models = {host.name} + for hop in hop_path: + join = next( + (j for j in current.joins if j.target_model == hop), None, + ) + if join is None: + raise UnknownReferenceError( + name=".".join(parts), + scope_kind="ModelScope", + scope_summary=( + f"model {current.name!r} joins: " + f"{[j.target_model for j in current.joins]}" + ), + suggestion=f"model {current.name!r} has no join to {hop!r}.", + ) + nxt = bundle.get_referenced_model(hop) + if nxt is None: + raise UnknownReferenceError( + name=".".join(parts), + scope_kind="ModelScope", + scope_summary=f"target {hop!r} not in source bundle", + suggestion=None, + ) + # A dotted path that walks back to an already-visited model is a + # circular join (e.g. ``a -> b -> a``); the leaf can never resolve + # and an unguarded walk would otherwise just fail confusingly on the + # leaf. Raise the legacy-compatible ``Circular join`` ValueError. + if nxt.name in visited_models: + raise ValueError( + f"Circular join detected resolving {'.'.join(parts)!r}: " + f"revisits model {nxt.name!r}." + ) + visited_models.add(nxt.name) + current = nxt + + # `current` is the terminal model; `leaf` is the column on it. + col = next((c for c in current.columns if c.name == leaf), None) + if col is None: + raise UnknownReferenceError( + name=".".join(parts), + scope_kind="ModelScope", + scope_summary=( + f"model {current.name!r} columns: " + f"{[c.name for c in current.columns]}" + ), + suggestion=None, + ) + + if col.sql is not None and col.sql.strip() != leaf: + # Derived column on a joined model. The path is part of the key + # so the cross-model planner can route via the join graph. + return ColumnSqlKey( + path=tuple(hop_path), model=current.name, column_name=leaf, + ) + return ColumnKey(path=tuple(hop_path), leaf=leaf) + + +def _resolve_dotted_star( + parts: Tuple[str, ...], + *, + scope: Union[ModelScope, StageSchema], + bundle: ResolvedSourceBundle, +) -> StarKey: + """Resolve a dotted star (``customers.*``, trailing ``*``) to a StarKey. + + Mirrors ``_resolve_dotted``'s self-prefix strip (C14) and join-chain + validation, but the leaf is ``*`` (no terminal column) so the result + is a ``StarKey`` whose ``path`` is the validated hop chain. An empty + path after stripping is the local star (``orders.*`` on ``orders``). + """ + assert parts and parts[-1] == "*" + if isinstance(scope, StageSchema): + raise IllegalScopeReferenceError( + name=".".join(parts), + scope_kind="StageSchema", + reason=( + "downstream stages see a flat schema — dotted refs are " + "not legal. Use the flat column name." + ), + ) + assert isinstance(scope, ModelScope) + host = scope.source_model + if host is None: + raise UnknownReferenceError( + name=".".join(parts), + scope_kind="ModelScope", + scope_summary="(no source_model anchor; anchor-less mode not implemented)", + suggestion=None, + ) + hop_path = parts[:-1] + # C14: strip same-model self-prefix (``orders.*`` on ``orders``). + if hop_path and hop_path[0] == host.name: + hop_path = hop_path[1:] + current = host + visited_models = {host.name} + for hop in hop_path: + join = next((j for j in current.joins if j.target_model == hop), None) + if join is None: + raise UnknownReferenceError( + name=".".join(parts), + scope_kind="ModelScope", + scope_summary=( + f"model {current.name!r} joins: " + f"{[j.target_model for j in current.joins]}" + ), + suggestion=f"no join from {current.name!r} to {hop!r}.", + ) + nxt = bundle.get_referenced_model(hop) + if nxt is None: + raise UnknownReferenceError( + name=".".join(parts), + scope_kind="ModelScope", + scope_summary=f"target {hop!r} not in source bundle", + suggestion=None, + ) + # A dotted star that revisits a model is a circular join (``a.b.a.*``) + # — reject it the same way ``_resolve_dotted`` rejects ``a.b.a.col`` + # so the two stay consistent (CR). + if nxt.name in visited_models: + raise ValueError( + f"Circular join detected resolving {'.'.join(parts)!r}: " + f"revisits model {nxt.name!r}." + ) + visited_models.add(nxt.name) + current = nxt + return StarKey(path=tuple(hop_path)) + + +def _bind_agg( + parsed: AggCall, *, + scope: Union[ModelScope, StageSchema], + bundle: ResolvedSourceBundle, +) -> AggregateKey: + if isinstance(parsed.source, StarSource): + source = StarKey() + elif ( + isinstance(parsed.source, DottedRef) + and parsed.source.parts + and parsed.source.parts[-1] == "*" + ): + # Cross-model star: ``customers.*:count`` → a StarKey carrying the + # join path so the cross-model planner routes COUNT(*) through the + # join graph, exactly like ``customers.revenue:sum`` (P3). Parity + # with the legacy dotted-star path. + source = _resolve_dotted_star( + parsed.source.parts, scope=scope, bundle=bundle, + ) + else: + bound_source = _bind( + parsed.source, scope=scope, bundle=bundle, in_filter=False, + ) + if not isinstance(bound_source, (ColumnKey, ColumnSqlKey, StarKey)): + raise ValueError( + f"Aggregation source must resolve to a column / star, " + f"got {type(bound_source).__name__}." + ) + source = bound_source + + # Bind args / kwargs. For aggregations, identifier args/kwargs become + # ColumnKey via the binder; scalars normalise. + args = tuple( + _bind_agg_arg(a, scope=scope, bundle=bundle) for a in parsed.args + ) + kwargs = tuple( + (k, _bind_agg_arg(v, scope=scope, bundle=bundle)) + for k, v in parsed.kwargs + ) + # DEV-1450 stage 7b.12: propagate ``Column.filter`` into the + # AggregateKey's structural identity. The resolved source's column + # may carry a Mode-A SQL fragment (``filter="status = 'paid'"``) + # that wraps the aggregate argument as ``SUM(CASE WHEN ... THEN col + # END)``. Two aggregates over the same column with different + # ``Column.filter`` therefore differ at the key level; same-filter + # ones intern (legacy CASE-WHEN-at-agg-time semantics, preserved by + # the spec's C5 + ``column_filter_key`` invariants). + column_filter_key = _resolve_column_filter_key( + source=source, bundle=bundle, + ) + # Codex review: enforce per-column aggregation eligibility gates + # the legacy enrichment site at ``enrichment.py:401-417`` enforced. + # Without this, ``id:sum`` (a PK) or ``status:avg`` (text) compile + # silently in the typed pipeline. The check is best-effort against + # the bundle — sources whose target model can't be resolved (e.g. + # an unreferenced join target) skip the check. + _validate_agg_eligibility( + source=source, agg=parsed.agg, bundle=bundle, + ) + return AggregateKey( + source=source, + agg=parsed.agg, + args=args, + kwargs=kwargs, + column_filter_key=column_filter_key, + ) + + +def _resolve_column_filter_key( + *, source, bundle: ResolvedSourceBundle, +) -> Optional[SqlExprKey]: + """Look up the resolved source's ``Column.filter`` and convert it + to a ``SqlExprKey``. + + Returns ``None`` for ``StarKey`` sources (``*:count`` has no column + to attach a filter to) and for any column whose ``filter`` is + unset. For ``ColumnKey`` / ``ColumnSqlKey`` sources the resolver + walks ``source.path`` through the bundle and reads the target + model's column entry. Models the planner doesn't have access to + (e.g. an unresolved join target) are tolerated — no exception is + raised; the key just stays ``None`` (the compile-time validator + in path resolution would have caught a genuinely missing model). + """ + if isinstance(source, StarKey): + return None + path = getattr(source, "path", ()) + leaf = getattr(source, "leaf", None) or getattr(source, "column_name", None) + if leaf is None: + return None + host = bundle.source_model + if host is None: + return None + current: SlayerModel = host + for hop in path: + nxt = bundle.get_referenced_model(hop) + if nxt is None: + return None + current = nxt + col = next((c for c in current.columns if c.name == leaf), None) + if col is None or not col.filter: + return None + # DEV-1503 — stamp the typed non-anchor join paths on the SqlExprKey so + # the planner's isolation trigger reads typed data, not parsed SQL. The + # anchor is the model the filter is bound against — the joined target for + # cross-model aggregates, the host for filtered-local. The anchor relation + # uses the ``__``-canonical path alias when the anchor is a joined model. + anchor_relation = "__".join(path) if path else current.name + paths = compute_column_filter_join_paths( + canonical_sql=col.filter, + anchor_model=current, + anchor_relation=anchor_relation, + bundle=bundle, + ) + return SqlExprKey(canonical_sql=col.filter, referenced_join_paths=paths) + + +def _validate_agg_eligibility( + *, source, agg: str, bundle: ResolvedSourceBundle, +) -> None: + """Enforce per-column aggregation eligibility gates. + + Mirrors the legacy ``enrichment.py:401-417`` v2 contract: + + 1. Primary-key columns are always restricted to ``count`` / + ``count_distinct`` regardless of type or explicit whitelist. + 2. An explicit ``Column.allowed_aggregations`` whitelist overrides + type defaults. + 3. Otherwise, built-in aggregations are gated by + ``DEFAULT_AGGREGATIONS_BY_TYPE``; model-custom aggregations + (registered in ``SlayerModel.aggregations``) are exempt. + + ``StarKey`` sources (``*:count`` / ``customers.*:count``) have no + column to attach a whitelist to and pass through. Cross-model and + derived (``ColumnSqlKey``) sources are best-effort: if the target + model can't be resolved through the bundle (an unresolved join + target) the gate is skipped — the compile-time path validator would + have raised earlier on a truly broken ref. + """ + if isinstance(source, StarKey): + return + path = tuple(getattr(source, "path", ())) + leaf = getattr(source, "leaf", None) or getattr(source, "column_name", None) + if leaf is None: + return + host = bundle.source_model + if host is None: + return + current: SlayerModel = host + for hop in path: + nxt = bundle.get_referenced_model(hop) + if nxt is None: + return + current = nxt + col = next((c for c in current.columns if c.name == leaf), None) + if col is None: + return + if col.primary_key: + if agg not in PRIMARY_KEY_AGGREGATIONS: + raise AggregationNotAllowedError( + column=leaf, + agg=agg, + reason=( + f"primary-key column {leaf!r} restricted to " + f"{sorted(PRIMARY_KEY_AGGREGATIONS)}; got {agg!r}." + ), + ) + return + if col.allowed_aggregations is not None: + if agg not in col.allowed_aggregations: + raise AggregationNotAllowedError( + column=leaf, + agg=agg, + reason=( + f"column {leaf!r} restricts allowed_aggregations to " + f"{sorted(col.allowed_aggregations)}; got {agg!r}." + ), + ) + return + # Model-custom aggregations are exempt from the type-default gate. + if any(a.name == agg for a in (current.aggregations or [])): + return + allowed = DEFAULT_AGGREGATIONS_BY_TYPE.get(col.type, frozenset()) + if agg not in allowed: + raise AggregationNotAllowedError( + column=leaf, + agg=agg, + reason=( + f"aggregation {agg!r} is not applicable to " + f"{col.type} column {leaf!r}; default aggregations are " + f"{sorted(allowed)}." + ), + ) + + +def _bind_agg_arg( + parsed: ParsedExpr, *, + scope: Union[ModelScope, StageSchema], + bundle: ResolvedSourceBundle, +): + """Bind one positional / kwarg argument of an aggregation. + + The AggregateKey shape stores Scalars inline (not as LiteralKey) + so identity matches the spec — see ``slayer/core/keys.py``. + Identifier args become ``ColumnKey`` / ``ColumnSqlKey``; literal + args normalise via ``normalize_scalar``. + """ + if isinstance(parsed, Literal): + return normalize_scalar(parsed.value) + if isinstance(parsed, (Ref, DottedRef)): + return _bind(parsed, scope=scope, bundle=bundle, in_filter=False) + raise ValueError( + f"Aggregation argument of kind {type(parsed).__name__} is not " + f"supported. Pass a column reference or a scalar." + ) + + +_NOT_SCALAR = object() # sentinel returned by _fold_to_scalar when the input isn't a literal-resolvable scalar + + +def _fold_to_scalar(parsed: ParsedExpr): + """Resolve a parsed expression to a scalar literal if possible. + + Folds ``Literal`` directly, and unary ``-`` over a numeric ``Literal`` + (the AST shape Python emits for ``periods=-1``) into the negated + literal value. Returns ``_NOT_SCALAR`` for anything that doesn't + reduce — transform kwargs are typed as ``Scalar``, so a non-scalar + expression is a binding error. + """ + if isinstance(parsed, Literal): + return normalize_scalar(parsed.value) + if ( + isinstance(parsed, UnaryOp) + and parsed.op == "-" + and isinstance(parsed.operand, Literal) + ): + from decimal import Decimal + + inner = parsed.operand.value + if isinstance(inner, bool): + # Reject explicitly — ``-True`` is nonsense and bool is an + # int subclass that would otherwise pass the next branch. + return _NOT_SCALAR + if isinstance(inner, (int, float, Decimal)): + return normalize_scalar(-inner) + return _NOT_SCALAR + + +# Per-op kwarg whitelist for the typed pipeline. Broader than the legacy +# ``slayer.core.formula._ALLOWED_TRANSFORM_KWARGS`` because the new +# pipeline allows ``partition_by`` on more than just the rank family +# (DEV-1450 C6: ``change(measure, partition_by=...)`` threads through to +# the desugared time_shift). Every transform also implicitly accepts +# ``partition_by`` — that branch is handled before the whitelist check. +_TRANSFORM_KWARG_RULES: dict = { + "cumsum": frozenset(), + "change": frozenset(), + "change_pct": frozenset(), + "first": frozenset(), + "last": frozenset(), + "time_shift": frozenset({"periods", "granularity"}), + "lag": frozenset({"periods"}), + "lead": frozenset({"periods"}), + "rank": frozenset(), + "percent_rank": frozenset(), + "dense_rank": frozenset(), + "ntile": frozenset({"n"}), + "consecutive_periods": frozenset({"period"}), +} + +# Positional-parameter signature (after the value) for the transforms whose +# documented DSL form accepts positional args: ``time_shift(x, periods, +# granularity)``, ``lag(x, periods)``, ``lead(x, periods)``. Each name maps the +# i-th positional onto the matching kwarg. Transforms absent here are +# keyword-only after the value. +_TRANSFORM_POSITIONAL_KWARGS: dict = { + "time_shift": ("periods", "granularity"), + "lag": ("periods",), + "lead": ("periods",), +} + + +def _bind_transform( + parsed: TransformCall, *, + scope: Union[ModelScope, StageSchema], + bundle: ResolvedSourceBundle, + alias_map: Optional[Dict[str, "ValueKey"]] = None, +) -> TransformKey: + # ``alias_map`` lets a transform input reference a declared-measure + # alias inside a filter (``change(rev) > 0``); partition_by must still + # be a real column, so it is bound without the alias map. + inp = _bind( + parsed.input, scope=scope, bundle=bundle, in_filter=False, + alias_map=alias_map, + ) + # The value to transform is the first positional (``parsed.input``). + # A few transforms accept further POSITIONAL params per the documented + # DSL surface (``time_shift(x, periods, granularity)``, + # ``lag(x, periods)``, ``lead(x, periods)``); map those onto their kwarg + # names. Every other transform (rank family, cumsum, change, + # consecutive_periods, ...) stays keyword-only after the value. + positional_pairs: List = [] + pos_names = _TRANSFORM_POSITIONAL_KWARGS.get(parsed.op) + if parsed.args: + if pos_names is None: + raise ValueError( + f"Transform {parsed.op!r} accepts exactly one positional " + f"argument (the value to transform); pass any offset, " + f"partition, or other settings as keyword arguments " + f"(e.g. ``{parsed.op}(value, partition_by=...)``)." + ) + if len(parsed.args) > len(pos_names): + raise ValueError( + f"Transform {parsed.op!r} accepts at most {len(pos_names)} " + f"positional argument(s) after the value " + f"({', '.join(pos_names)}); got {len(parsed.args)}." + ) + positional_pairs = list(zip(pos_names, parsed.args)) + args: List = [] + kwargs: List = [] + partition_keys: List = [] + allowed_kwargs = _TRANSFORM_KWARG_RULES.get(parsed.op, frozenset()) + seen_kwargs: set = set() + # Positional params first, then explicit kwargs; a name supplied BOTH + # ways is an error (ambiguous, e.g. ``time_shift(x, -1, periods=-2)``). + _explicit_kw_names = {k for k, _ in parsed.kwargs} + for k, _ in positional_pairs: + if k in _explicit_kw_names: + raise ValueError( + f"Transform {parsed.op!r} got {k!r} both positionally and " + f"as a keyword argument." + ) + for k, v in [*positional_pairs, *parsed.kwargs]: + if k == "partition_by": + # ``partition_by`` accepts a single column ref OR a tuple/list of + # them (Codex review): ``rank(x, partition_by=[region, channel])``. + # ``_convert_kwarg_value`` returns a Python tuple for the list + # form; bind each element independently and accumulate into + # ``partition_keys`` so the SQL gen emits a multi-column OVER + # (PARTITION BY ...). A single ref still flows through the + # scalar branch. + elements = v if isinstance(v, tuple) else (v,) + for elem in elements: + bound_elem = _bind( + elem, scope=scope, bundle=bundle, in_filter=False, + ) + if isinstance(bound_elem, (ColumnKey, ColumnSqlKey)): + partition_keys.append(bound_elem) + else: + raise ValueError( + f"transform {parsed.op!r} partition_by must resolve " + f"to a column reference; got " + f"{type(bound_elem).__name__}." + ) + continue + if k not in allowed_kwargs: + raise ValueError( + f"Transform {parsed.op!r} does not accept keyword " + f"argument {k!r}. Accepted: " + f"{sorted(allowed_kwargs | {'partition_by'})}." + ) + seen_kwargs.add(k) + scalar = _fold_to_scalar(v) + if scalar is _NOT_SCALAR: + raise ValueError( + f"Transform {parsed.op!r} keyword {k!r} must be a " + f"scalar literal; got expression of kind " + f"{type(v).__name__}." + ) + kwargs.append((k, scalar)) + # Per-op required-kwarg validation + defaults. + kwargs = _apply_transform_kwarg_defaults( + op=parsed.op, kwargs=kwargs, seen=seen_kwargs, + ) + return TransformKey( + op=parsed.op, + input=inp, + args=tuple(args), + kwargs=tuple(kwargs), + partition_keys=frozenset(partition_keys), + ) + + +def _apply_transform_kwarg_defaults( + *, op: str, kwargs: list, seen: set, +) -> list: + """Validate required kwargs and apply per-op defaults for the typed + TransformKey. + + Validation: + * ``ntile`` requires ``n``; ``n`` must be a positive integer + (``bool`` rejected — it's an ``int`` subclass in Python but a + boolean ``True``/``False`` is never a sensible bucket count). + * ``time_shift`` requires ``periods`` (integer; may be negative). + + Defaults: + * ``lag`` / ``lead`` default ``periods=1`` when missing so the + typed TransformKey carries the resolved kwarg list; the SQL + generator can render PARTITION/ORDER without re-applying defaults. + + ``normalize_scalar`` wraps numeric literals in ``Decimal``, so the + integer checks accept ``Decimal`` whose value is integral as well + as plain ``int``. + """ + from decimal import Decimal + + def _ensure_positive_integer(value: object, *, kw: str) -> None: + if isinstance(value, bool): + raise ValueError( + f"Transform {op!r} keyword {kw} must be a positive " + f"integer; got {value!r}." + ) + if isinstance(value, int): + ival = value + elif isinstance(value, Decimal): + if value != value.to_integral_value(): + raise ValueError( + f"Transform {op!r} keyword {kw} must be a positive " + f"integer; got {value!r}." + ) + ival = int(value) + else: + raise ValueError( + f"Transform {op!r} keyword {kw} must be a positive " + f"integer; got {value!r}." + ) + if ival <= 0: + raise ValueError( + f"Transform {op!r} keyword {kw} must be a positive " + f"integer; got {value!r}." + ) + + if op == "ntile": + if "n" not in seen: + raise ValueError( + "Transform 'ntile' requires keyword argument n (the " + "number of buckets, a positive integer)." + ) + n_value = next(v for k, v in kwargs if k == "n") + _ensure_positive_integer(n_value, kw="n") + if op == "time_shift" and "periods" not in seen: + raise ValueError( + "Transform 'time_shift' requires keyword argument periods " + "(the integer offset, negative for a backward shift)." + ) + if op in ("lag", "lead") and "periods" not in seen: + kwargs.append(("periods", normalize_scalar(1))) + return kwargs + + +def _bind_scalar( + parsed: ScalarCall, *, + scope: Union[ModelScope, StageSchema], + bundle: ResolvedSourceBundle, + in_filter: bool, + alias_map: Optional[Dict[str, "ValueKey"]] = None, +) -> ScalarCallKey: + if parsed.name not in SCALAR_FUNCTIONS: + # Defence in depth: the parser already enforces the allowlist, + # but direct ParsedExpr construction can bypass the parser. + # Re-check here so the typed key family is always sound. + raise UnknownFunctionError( + name=parsed.name, + location="(binder)", + suggestion=( + f"Mode-B scalar calls are restricted to " + f"{sorted(SCALAR_FUNCTIONS)}." + ), + ) + if parsed.name == "like" and len(parsed.args) != 2: + raise ValueError( + f"Scalar function 'like' takes exactly 2 arguments " + f"(value, pattern); got {len(parsed.args)}." + ) + args = tuple( + _bind(a, scope=scope, bundle=bundle, in_filter=in_filter, alias_map=alias_map) + for a in parsed.args + ) + return ScalarCallKey(name=parsed.name, args=args) + + +def _reject_windowed_column_sql( + refs: Tuple[ValueKey, ...], + *, + scope: Union[ModelScope, StageSchema], + bundle: ResolvedSourceBundle, + parsed: ParsedExpr, +) -> None: + """Raise ``IllegalWindowInFilterError`` if any referenced + ``ColumnSqlKey`` has a windowed ``Column.sql`` body. + + DEV-1369 removed predicate-promotion; filters touching a windowed + column SQL now raise. + """ + if isinstance(scope, StageSchema): + # StageSchema columns don't carry Column.sql in the bundle; + # window detection is handled when the upstream stage was bound. + return + for k in refs: + if not isinstance(k, ColumnSqlKey): + continue + model = _lookup_model(name=k.model, scope=scope, bundle=bundle) + if model is None: + continue + col = next((c for c in model.columns if c.name == k.column_name), None) + if col is None or col.sql is None: + continue + if has_window_function(col.sql): + raise IllegalWindowInFilterError( + filter_expr=str(parsed), + source=( + f"filter references column {k.column_name!r} on model " + f"{k.model!r} whose Column.sql contains a window " + f"function" + ), + suggestion=( + "use a rank-family transform (rank, percent_rank, " + "dense_rank, ntile) in the formula instead, or " + "compute the windowed value in an earlier stage." + ), + ) + + +def _lookup_model( + *, + name: str, + scope: Union[ModelScope, StageSchema], + bundle: ResolvedSourceBundle, +) -> Optional[SlayerModel]: + if isinstance(scope, ModelScope) and scope.source_model is not None: + if scope.source_model.name == name: + return scope.source_model + return bundle.get_referenced_model(name) diff --git a/slayer/engine/column_expansion.py b/slayer/engine/column_expansion.py index 3653e4ef..e928d9b9 100644 --- a/slayer/engine/column_expansion.py +++ b/slayer/engine/column_expansion.py @@ -17,7 +17,7 @@ """ from __future__ import annotations -from typing import Any, Awaitable, Callable, Dict, Optional, Set, Tuple +from typing import Any, Awaitable, Callable, Dict, List, Optional, Protocol, Set, Tuple import sqlglot from sqlglot import exp @@ -95,6 +95,91 @@ def _root_scope_column_ids(*, parsed: exp.Expression) -> Set[int]: return root_ids +class _SyncBundle(Protocol): + """Minimal contract for ``collect_root_scope_joined_paths``'s bundle — + matches ``ResolvedSourceBundle.get_referenced_model``. Declared inline so + this helper stays import-free of the engine layer. + """ + + def get_referenced_model(self, name: str) -> Optional[SlayerModel]: ... + + +def _resolve_alias_to_join_segments( + *, + alias: str, + source_model: SlayerModel, + bundle: _SyncBundle, +) -> Optional[Tuple[str, ...]]: + """Walk the ``__``-segmented join alias against ``source_model``'s joins. + + Returns the segments tuple when EVERY hop resolves (so a prefix walk + can emit join-path tuples), or ``None`` when any hop fails — that + aborts the caller's emission for this column (CTE / subquery alias + or a spurious dotted ref). + """ + segments = tuple(alias.split("__")) + current = source_model + for seg in segments: + join = next( + (j for j in current.joins if j.target_model == seg), None, + ) + if join is None: + return None + nxt = bundle.get_referenced_model(seg) + if nxt is None: + return None + current = nxt + return segments + + +def collect_root_scope_joined_paths( + *, + parsed: exp.Expression, + source_model: SlayerModel, + source_relation: str, + bundle: _SyncBundle, +) -> List[Tuple[str, ...]]: + """Collect the ordered de-duplicated list of join-path prefixes a parsed + SQL fragment references in its root scope. + + Each ROOT-scope ``.`` whose ``alias`` is not the source + relation and fully resolves as a join walk on ``source_model`` + contributes its prefixes (``a__b`` yields ``("a",)`` AND ``("a", "b")``). + Aliases that don't resolve as a join walk (CTE / subquery aliases, + spurious dotted refs) are skipped, as are columns inside a nested + scope (subquery / set-op branch) — those belong to the inner rowset. + + Shared by the SQL generator (``SQLGenerator._joined_paths_in_sql``) and + the planner-side column filter discovery + (``slayer.engine.column_filter_paths._walk_root_scope_paths``) so the + two surfaces agree on what counts as "crosses a join." + """ + root_ids = _root_scope_column_ids(parsed=parsed) + seen: Set[Tuple[str, ...]] = set() + ordered: List[Tuple[str, ...]] = [] + anchor_aliases = (source_relation, source_model.name) + for col in parsed.find_all(exp.Column): + tbl = col.args.get("table") + if tbl is None or col.args.get("db") or col.args.get("catalog"): + continue + if id(col) not in root_ids: + continue + alias = tbl.name + if alias in anchor_aliases: + continue + segments = _resolve_alias_to_join_segments( + alias=alias, source_model=source_model, bundle=bundle, + ) + if segments is None: + continue + for i in range(1, len(segments) + 1): + prefix = segments[:i] + if prefix not in seen: + seen.add(prefix) + ordered.append(prefix) + return ordered + + async def _walk_path_to_target( *, source_model: SlayerModel, @@ -300,3 +385,134 @@ async def expand_derived_refs( ) return parsed.sql(dialect=dialect) + + +# --------------------------------------------------------------------------- +# Synchronous expansion (DEV-1450 typed pipeline) +# --------------------------------------------------------------------------- +# +# The async functions above resolve join targets through ``storage.get_model`` +# (genuinely async). The DEV-1450 generator runs synchronously over a +# ``ResolvedSourceBundle`` that has already loaded every referenced model +# (P11: storage consulted once, up front), so it expands derived refs through +# a *sync* model resolver. These twins mirror the async control flow exactly, +# reusing the shared (sync) ``_is_trivial_base`` / ``_root_scope_column_ids`` +# helpers. The async path is removed with the rest of enrichment in DEV-1452. + +SyncResolveModel = Callable[[str], Optional[SlayerModel]] + + +def _walk_path_to_target_sync( + *, + source_model: SlayerModel, + source_alias: str, + table_alias: str, + resolve_model: SyncResolveModel, + is_root: bool, +) -> Tuple[Optional[SlayerModel], Optional[str]]: + """Sync mirror of :func:`_walk_path_to_target`.""" + if table_alias == source_alias or table_alias == source_model.name: + return source_model, source_alias + parts = table_alias.split("__") if "__" in table_alias else [table_alias] + current = source_model + for hop in parts: + join = next((j for j in current.joins if j.target_model == hop), None) + if join is None: + return None, None + nxt = resolve_model(hop) + if nxt is None: + return None, None + current = nxt + walked = "__".join(parts) + canonical = walked if is_root else f"{source_alias}__{walked}" + return current, canonical + + +def _process_column_node_sync( + *, + col: exp.Column, + model: SlayerModel, + alias_path: str, + resolve_model: SyncResolveModel, + dialect: str, + visited: Tuple[Tuple[str, str], ...], + is_root: bool, + root_scope_ids: Set[int], +) -> None: + """Sync mirror of :func:`_process_column_node`.""" + if col.args.get("db") or col.args.get("catalog"): + return + table_id = col.args.get("table") + col_name = col.name + table_alias = table_id.name if table_id is not None else alias_path + target_model, canonical_alias = _walk_path_to_target_sync( + source_model=model, + source_alias=alias_path, + table_alias=table_alias, + resolve_model=resolve_model, + is_root=is_root, + ) + if target_model is None or canonical_alias is None: + return + target_col = target_model.get_column(col_name) + if target_col is None or _is_trivial_base(column=target_col): + col.set("table", exp.to_identifier(canonical_alias)) + return + if id(col) not in root_scope_ids: + return + next_is_root = is_root and (target_model is model) + key = (target_model.name, col_name) + if key in visited: + cycle_start = visited.index(key) + cycle = (*visited[cycle_start:], key) + raise ColumnCycleError(cycle=list(cycle)) + expanded_sql = expand_derived_refs_sync( + sql=target_col.sql, + model=target_model, + alias_path=canonical_alias, + resolve_model=resolve_model, + dialect=dialect, + visited=(*visited, key), + is_root=next_is_root, + ) + if expanded_sql is None: + return + expanded_ast = sqlglot.parse_one(expanded_sql, dialect=dialect) + col.replace(exp.Paren(this=expanded_ast)) + + +def expand_derived_refs_sync( + *, + sql: Optional[str], + model: SlayerModel, + alias_path: str, + resolve_model: SyncResolveModel, + dialect: str, + visited: Optional[Tuple[Tuple[str, str], ...]] = None, + is_root: bool = True, +) -> Optional[str]: + """Sync mirror of :func:`expand_derived_refs` for the DEV-1450 pipeline. + + ``resolve_model`` is a plain ``name -> Optional[SlayerModel]`` lookup + (typically ``bundle.get_referenced_model``); there is no + ``named_queries`` parameter because the bundle has already resolved the + full referenced-model set. + """ + if not sql: + return sql + visited = visited or () + parsed = sqlglot.parse_one(sql, dialect=dialect) + column_nodes = list(parsed.find_all(exp.Column)) + root_scope_ids = _root_scope_column_ids(parsed=parsed) + for col in column_nodes: + _process_column_node_sync( + col=col, + model=model, + alias_path=alias_path, + resolve_model=resolve_model, + dialect=dialect, + visited=visited, + is_root=is_root, + root_scope_ids=root_scope_ids, + ) + return parsed.sql(dialect=dialect) diff --git a/slayer/engine/column_filter_paths.py b/slayer/engine/column_filter_paths.py new file mode 100644 index 00000000..2aa2278e --- /dev/null +++ b/slayer/engine/column_filter_paths.py @@ -0,0 +1,286 @@ +"""DEV-1503 — planner-facing helper for ``Column.filter`` join-path discovery. + +When the binder constructs a ``SqlExprKey`` for an ``AggregateKey.column_filter_key``, +it computes the typed set of non-anchor join paths the filter touches and stamps +them on the key as ``SqlExprKey.referenced_join_paths``. The DEV-1503 trigger +predicate then reads this typed field instead of re-parsing SQL text at plan +time — preserving the typed pipeline's "no string rewriting after parse" P7 +boundary at the planner layer. + +The actual rewriting for RENDERING (inlining derived refs inside the +``SUM(CASE WHEN THEN col END)`` wrapper) still lives in +``slayer/sql/generator.py`` (DEV-1494's ``_expand_column_filter_sql``). The +helper here is a parallel structural-analysis pass that returns paths only, +shared by the binder. + +The discovery rules match the generator-side ones byte-for-byte: + +* Same-model bare refs ("status") qualify to the anchor relation — no path. +* Cross-model dotted refs ("loss_payment.has_flag") contribute ``("loss_payment",)``. +* Multi-hop ``__``-delimited refs ("loss_payment__claim.state") contribute + ``("loss_payment",)`` AND ``("loss_payment", "claim")``. +* Bare refs that name a non-trivial DERIVED column whose own ``Column.sql`` + crosses a join (``is_eu`` reaching ``customers.region``) expand to the + derived sql and contribute the expansion's paths. +* Refs inside nested scopes (subqueries, set-op branches) are ignored — + they belong to the inner rowset. +* Aliases that don't resolve as a join walk on the anchor model are + silently skipped — they may be CTE / subquery aliases out of scope. +""" + +from __future__ import annotations + +from typing import Optional, Tuple + +import sqlglot +from sqlglot import exp + +from slayer.core.models import Column, SlayerModel +from slayer.engine.column_expansion import ( + _is_trivial_base, + collect_root_scope_joined_paths, + expand_derived_refs_sync, +) +from slayer.engine.source_bundle import ResolvedSourceBundle + + +# Fallback dialect chain (Codex round 7) — the planner doesn't carry the +# datasource's dialect, so a user-configured backend with dialect-specific +# syntax in ``Column.filter`` (MySQL backticks, T-SQL square brackets, +# ClickHouse-specific functions) could fail the Postgres parse and +# silently return no referenced join paths. The DEV-1503 trigger would +# then miss and the generator would render the filter inline in ``_base``, +# pulling the cross-model join into the host rowset. Try each dialect in +# order; only return ``()`` if ALL fail. The path discovery itself is +# dialect-agnostic (AST walking, no SQL emission); the dialect-aware +# re-parse for emission still happens on the generator side. +_PLANNER_PARSE_DIALECT_CHAIN: Tuple[Optional[str], ...] = ( + "postgres", + None, # sqlglot's permissive default — accepts ANSI SQL broadly + "mysql", + "clickhouse", # ClickHouse-only constructs (countIf, distinct identifier escapes) — CR PR #153 r3350000228 + "bigquery", + "tsql", +) + + +def _parse_filter_sql_any_dialect(sql: str) -> Optional[exp.Expression]: + """Parse ``sql`` trying each dialect in the fallback chain. + + Returns the first successful parse, or ``None`` when every dialect + rejects the input — that fall-through preserves the catch-all the + original ``except Exception`` provided for malicious / unparseable + payloads (the generator's dialect-aware emission is the + authoritative gate). + """ + for dialect in _PLANNER_PARSE_DIALECT_CHAIN: + try: + return sqlglot.parse_one(sql, dialect=dialect) + except Exception: + continue + return None + + +def _expand_derived_refs_any_dialect( + *, + sql: str, + model: SlayerModel, + alias_path: str, + bundle: ResolvedSourceBundle, +) -> Optional[str]: + """Run ``expand_derived_refs_sync`` against each dialect in the chain. + + ``expand_derived_refs_sync`` parses ``sql`` (and the derived columns' + own ``sql`` fields) internally with the supplied dialect. If a + derived column whose ``Column.sql`` uses dialect-specific syntax + (MySQL backticks, BigQuery struct literals, etc.) tips over the + Postgres parser, the expansion would silently drop join paths the + DEV-1503 trigger needs (Codex round 9). Try the same chain + ``_parse_filter_sql_any_dialect`` uses; return the first + successful expansion or ``None`` if every dialect fails. + """ + for dialect in _PLANNER_PARSE_DIALECT_CHAIN: + if dialect is None: + # ``expand_derived_refs_sync`` requires a dialect string. + continue + try: + expanded = expand_derived_refs_sync( + sql=sql, + model=model, + alias_path=alias_path, + resolve_model=bundle.get_referenced_model, + dialect=dialect, + ) + except Exception: + continue + if expanded: + return expanded + return None + + +def _is_nontrivial_derived(model: SlayerModel, name: str) -> bool: + """True iff ``name`` is a column on ``model`` whose ``Column.sql`` is a + non-trivial expression (set, and not just a bare-identifier remap). + + Mirrors ``SQLGenerator._is_nontrivial_derived``; duplicated here to keep + the planner-facing helper free of generator imports. + """ + col: Optional[Column] = next( + (c for c in model.columns if c.name == name), None, + ) + return col is not None and col.sql is not None and not _is_trivial_base( + column=col, + ) + + +def _is_anchor_local_col_ref( + col: exp.Column, *, anchor_aliases: set, +) -> bool: + """A column ref counts as "on the anchor" when it is bare (no ``table``) + OR self-qualified to the anchor (``orders.is_eu`` where ``orders`` is + the anchor relation / model name). + """ + if not isinstance(col.this, exp.Identifier): + return False + tbl = col.args.get("table") + if tbl is None: + return True + return tbl.name in anchor_aliases + + +def _expand_filter_sql_if_anchor_derived( + *, + parsed: exp.Expression, + canonical_sql: str, + anchor_model: SlayerModel, + anchor_relation: str, + bundle: ResolvedSourceBundle, +) -> Optional[exp.Expression]: + """Expand any non-trivial derived anchor-local refs in ``parsed``, + re-parsing the expanded SQL. Returns the new AST, the original + ``parsed`` when no derived ref was present, or ``None`` when the + expansion produced unparseable output (caller should bail to ``()``). + + Mirrors the generator's ``_expand_column_filter_sql`` gate so the + planner and the renderer surface the same set of crossed joins. + """ + anchor_aliases = {anchor_relation, anchor_model.name} + anchor_local_names = { + col.this.name + for col in parsed.find_all(exp.Column) + if _is_anchor_local_col_ref(col, anchor_aliases=anchor_aliases) + } + has_derived = any( + _is_nontrivial_derived(anchor_model, n) for n in anchor_local_names + ) + if not has_derived: + return parsed + + # Degenerate: the whole predicate IS a single derived column ref + # (``filter="is_eu"`` or self-qualified ``filter="orders.is_eu"``). + # ``expand_derived_refs_sync`` rewrites refs via in-place + # ``col.replace`` which is a no-op on the AST root, so expand the + # column's sql directly. + if ( + isinstance(parsed, exp.Column) + and _is_anchor_local_col_ref(parsed, anchor_aliases=anchor_aliases) + and _is_nontrivial_derived(anchor_model, parsed.name) + ): + col = next( + (c for c in anchor_model.columns if c.name == parsed.name), None, + ) + if col is None or not col.sql: + return None + sql_to_expand = col.sql + else: + sql_to_expand = canonical_sql + + expanded = _expand_derived_refs_any_dialect( + sql=sql_to_expand, + model=anchor_model, + alias_path=anchor_relation, + bundle=bundle, + ) + if not expanded: + return parsed + return _parse_filter_sql_any_dialect(expanded) + + +def compute_column_filter_join_paths( + *, + canonical_sql: Optional[str], + anchor_model: SlayerModel, + anchor_relation: str, + bundle: ResolvedSourceBundle, +) -> Tuple[Tuple[str, ...], ...]: + """Return the ordered tuple of non-anchor join-path prefixes a + ``Column.filter`` predicate touches after derived-ref expansion. + + ``anchor_model`` is the model the filter is bound against — for a + filtered-local measure on the host (``AggregateKey.source.path == ()``), + the anchor is the host; for a cross-model aggregate on a target column + (``source.path == ("customers",)``), the anchor is ``customers``. + + Returns ``()`` for same-model filters (no cross-anchor joins), for an + empty / unparseable canonical_sql, and as a defensive fallback if join + resolution fails partway through. + + Multi-hop alias paths emit each prefix once (``loss_payment__claim`` + yields ``("loss_payment",)`` and ``("loss_payment", "claim")``). + """ + if not canonical_sql: + return () + parsed = _parse_filter_sql_any_dialect(canonical_sql) + if parsed is None: + return () + + expanded = _expand_filter_sql_if_anchor_derived( + parsed=parsed, + canonical_sql=canonical_sql, + anchor_model=anchor_model, + anchor_relation=anchor_relation, + bundle=bundle, + ) + if expanded is None: + return () + parsed = expanded + + # ``_walk_root_scope_paths`` exercises sqlglot's scope analyser, which + # can raise ``TypeError`` etc. on unusual / malicious payloads (SQL + # injection attempts like ``status = 'x' UNION SELECT * FROM users``). + # The dialect-aware generator path is the authoritative gate for those — + # it raises ``ParseError`` / ``ValueError`` at SQL emission. The planner + # discovery here is a best-effort structural pass; swallow any internal + # parser failure so it can't shadow the generator's rejection. + try: + return _walk_root_scope_paths( + parsed=parsed, + anchor_model=anchor_model, + anchor_relation=anchor_relation, + bundle=bundle, + ) + except Exception: + return () + + +def _walk_root_scope_paths( + *, + parsed: exp.Expression, + anchor_model: SlayerModel, + anchor_relation: str, + bundle: ResolvedSourceBundle, +) -> Tuple[Tuple[str, ...], ...]: + """Collect every root-scope ``.`` whose ``alias`` resolves as + a join walk on ``anchor_model``, returning the ordered tuple of path + prefixes (de-duplicated). + + Thin shim over the shared ``collect_root_scope_joined_paths`` helper so + the planner and ``SQLGenerator._joined_paths_in_sql`` agree on what + counts as "crosses a join." + """ + return tuple(collect_root_scope_joined_paths( + parsed=parsed, + source_model=anchor_model, + source_relation=anchor_relation, + bundle=bundle, + )) diff --git a/slayer/engine/cross_model_planner.py b/slayer/engine/cross_model_planner.py new file mode 100644 index 00000000..b8277db9 --- /dev/null +++ b/slayer/engine/cross_model_planner.py @@ -0,0 +1,1147 @@ +"""Stage 7a.2 (DEV-1450) — CrossModelPlanner Protocol + IsolatedCte impl (I1). + +The cross-model aggregate strategy is a substitutable component (I1): +``CrossModelPlanner`` is a Protocol; ``IsolatedCteCrossModelPlanner`` is +the default impl encoding today's "one CTE per (target_model, +shared_grain)" pattern and the ``inherited_filter_policy`` decision +table from the DEV-1450 spec. + +The Protocol's ``plan(...)`` consumes: + +* the aggregate slot id + ``AggregateKey`` (whose ``source.path`` + identifies the cross-model target), +* a ``ResolvedSourceBundle`` (the eagerly-resolved model graph), +* ``host_slots`` (every ``ValueSlot`` on the host query — used to + classify filter routing and compute shared grain), +* ``host_filters`` as ``HostFilterRouting`` records (filter id + + phase + referenced slot ids). + +It produces a ``CrossModelAggregatePlan`` (in ``planned.py``) with +explicit ``where_filter_ids`` / ``having_filter_ids`` / +``target_model_filters`` routes so the SQL generator (stage 7b) doesn't +re-classify. + +Decision table (host filter routing only): + +| Filter references | Route | +| -------------------------------------------- | ---------------------- | +| Host-local row slot only | DROP_HOST_LOCAL | +| All on joined-target path (row) | PROPAGATE_WHERE | +| Cross-model agg-ref on same target | PROPAGATE_HAVING | +| Slots on a different joined branch | DROP_UNREACHABLE | +| Mixed reachable + unreachable | DROP_UNREACHABLE | +| Transform / POST phase | STAY_AT_HOST_POST | + +Target model's own ``SlayerModel.filters`` and ``Column.filter`` on the +aggregated column are intrinsic — they ride on the target / the +``AggregateKey`` itself and don't go through host-filter classification. + +Dormant in 7a — no engine code calls these yet. ProjectionPlanner +(stage 7a.6) is the first consumer. +""" + +from __future__ import annotations + +from enum import Enum +from typing import Callable, List, Optional, Protocol, Tuple + +from pydantic import BaseModel, ConfigDict, Field + +from slayer.core.enums import DataType +from slayer.core.errors import ( + AmbiguousReferenceError, + IllegalScopeReferenceError, + UnknownReferenceError, + UnreachableFilterDroppedWarning, +) +from slayer.core.keys import ( + AggregateKey, + ColumnKey, + ColumnSqlKey, + Phase, + StarKey, + TimeTruncKey, + ValueKey, + column_path, +) +from slayer.core.models import ModelMeasure, SlayerModel +from slayer.core.query import ColumnRef, SlayerQuery, TimeDimension +from slayer.core.refs import agg_kwarg_canonical_str, canonical_agg_name +from slayer.core.scope import ModelScope, StageColumn, StageSchema +from slayer.engine.binding import ( + bind_expr, + bind_filter, + bind_time_dimension, + walk_value_keys, +) +from slayer.engine.planned import ( + BoundFilterId, + CrossModelAggregatePlan, + JoinRequirement, + PlannedQuery, + SlotId, + ValueSlot, +) +from slayer.engine.source_bundle import ResolvedSourceBundle +from slayer.engine.syntax import parse_expr, parse_filter_expr + + +# --------------------------------------------------------------------------- +# Public types +# --------------------------------------------------------------------------- + + +class FilterRoute(str, Enum): + """Routing decision for one host filter on a cross-model CTE.""" + + DROP_HOST_LOCAL = "drop_host_local" + PROPAGATE_WHERE = "propagate_where" + PROPAGATE_HAVING = "propagate_having" + DROP_UNREACHABLE = "drop_unreachable" + STAY_AT_HOST_POST = "stay_at_host_post" + + +class HostFilterRouting(BaseModel): + """A host filter + the slot ids it references. + + The planner consumes a list of these; each is classified per + ``classify_host_filter`` and routed into the resulting + ``CrossModelAggregatePlan``. + """ + + model_config = ConfigDict(arbitrary_types_allowed=True) + + filter_id: BoundFilterId + phase: Phase + referenced_slot_ids: List[SlotId] = Field(default_factory=list) + text: Optional[str] = None + + +# --------------------------------------------------------------------------- +# Classifier +# --------------------------------------------------------------------------- + + +def classify_host_filter( + *, + host_filter: HostFilterRouting, + host_slots: List[ValueSlot], + target_path: Tuple[str, ...], + host_model_name: Optional[str] = None, +) -> FilterRoute: + """Classify one host filter for cross-model CTE propagation. + + See the module docstring for the decision table. The classifier is + pure: same inputs → same output, no side effects. + + ``host_model_name`` is used to route ``ColumnSqlKey`` refs (derived + columns): the key carries its host model name but not a path, so we + compare it against the host model name and the path to decide + reachable / local / unreachable. When ``host_model_name`` is None, + ColumnSqlKey refs default to local — conservative for callers that + don't have the host model in scope. + """ + if host_filter.phase == Phase.POST: + return FilterRoute.STAY_AT_HOST_POST + if not host_filter.referenced_slot_ids: + # No referenced slots — nothing to route into the CTE. + return FilterRoute.STAY_AT_HOST_POST + + by_id = {s.id: s for s in host_slots} + + local_row: List[SlotId] = [] + reachable_path: List[SlotId] = [] + unreachable: List[SlotId] = [] + aggregate_on_target: List[SlotId] = [] + aggregate_other: List[SlotId] = [] + + for sid in host_filter.referenced_slot_ids: + s = by_id.get(sid) + if s is None: + # Unknown slot id — be conservative, treat as unreachable. + unreachable.append(sid) + continue + if isinstance(s.key, AggregateKey): + agg_source = s.key.source + agg_path = getattr(agg_source, "path", ()) + if agg_path == target_path: + aggregate_on_target.append(sid) + else: + aggregate_other.append(sid) + elif isinstance(s.key, ColumnKey): + if not s.key.path: + local_row.append(sid) + elif s.key.path == target_path[: len(s.key.path)]: + reachable_path.append(sid) + else: + unreachable.append(sid) + elif isinstance(s.key, ColumnSqlKey): + # Derived column. Route by its host model: on host → local; + # on any model in target_path → reachable; otherwise → + # unreachable. + cm = s.key.model + if host_model_name is not None and cm == host_model_name: + local_row.append(sid) + elif cm in target_path: + reachable_path.append(sid) + elif host_model_name is None: + # No host name to compare against — conservative default. + local_row.append(sid) + else: + unreachable.append(sid) + else: + # Transform / Arithmetic / ScalarCall: phase already decided. + # POST was checked above; ROW/AGGREGATE land here from + # arithmetic / scalar calls. Treat as local for routing. + local_row.append(sid) + + if unreachable or aggregate_other: + # Any unreachable ref → drop + warn (covers pure-unreachable AND + # mixed-with-reachable cases per decision table rows 6/7). + return FilterRoute.DROP_UNREACHABLE + if local_row and not (aggregate_on_target or reachable_path): + return FilterRoute.DROP_HOST_LOCAL + if local_row: + # Mixed local + (target-path / target-agg). The local refs can't + # be evaluated in the CTE, so the filter stays at host. + return FilterRoute.DROP_HOST_LOCAL + if aggregate_on_target: + return FilterRoute.PROPAGATE_HAVING + if reachable_path: + return FilterRoute.PROPAGATE_WHERE + return FilterRoute.STAY_AT_HOST_POST + + +# --------------------------------------------------------------------------- +# Protocol +# --------------------------------------------------------------------------- + + +class CrossModelPlanner(Protocol): + """Strategy for compiling one cross-model aggregate slot. + + DEV-1450 follow-up #2: re-rooting is owned by the strategy, not a + post-hoc mutation in ``plan_query``. When the host carries dimensions / + filters reachable from the target only by walking the TARGET's own join + graph (off the host→target forward path), the strategy may build a nested + re-rooted ``PlannedQuery`` and attach it to the returned plan. To do so it + needs the host query, its public projection, and a callback that compiles + a sub-query — all keyword-only and defaulting to ``None`` so direct + callers (and test doubles) that don't re-root keep working unchanged. + """ + + def plan( + self, + *, + aggregate_slot_id: SlotId, + aggregate_key: AggregateKey, + bundle: ResolvedSourceBundle, + host_slots: List[ValueSlot], + host_filters: List[HostFilterRouting], + public_alias: Optional[str] = None, + hidden: bool = False, + host_query: Optional[SlayerQuery] = None, + public_projection: Optional[List[SlotId]] = None, + subplan_builder: Optional[ + Callable[[SlayerQuery, ResolvedSourceBundle], PlannedQuery] + ] = None, + ) -> CrossModelAggregatePlan: + ... + + +# --------------------------------------------------------------------------- +# Default impl +# --------------------------------------------------------------------------- + + +def _walk_chain( + *, + host_model: SlayerModel, + hops: Tuple[str, ...], + bundle: ResolvedSourceBundle, +) -> Tuple[SlayerModel, List[JoinRequirement]]: + """Walk the join graph from ``host_model`` through ``hops``. + + Returns ``(terminal_model, [JoinRequirement, ...])``. Raises + ``ValueError`` if a hop has no matching join on the current model + or the referenced model isn't in ``bundle.referenced_models``. + + The walker is sync — the bundle holds eagerly-resolved models, so + no async I/O is needed (P11). + """ + current = host_model + chain: List[JoinRequirement] = [] + for hop in hops: + join = next( + (j for j in current.joins if j.target_model == hop), None, + ) + if join is None: + raise ValueError( + f"Model {current.name!r} has no join to {hop!r}. " + f"Available joins: {[j.target_model for j in current.joins]}" + ) + nxt = bundle.get_referenced_model(hop) + if nxt is None: + raise ValueError( + f"Join target {hop!r} from {current.name!r} not found in " + f"resolved source bundle." + ) + chain.append(JoinRequirement( + source_model=current.name, + target_model=hop, + join_pairs=[list(p) for p in join.join_pairs], + join_type=join.join_type, + )) + current = nxt + return current, chain + + +def _aggregate_alias(*, key: AggregateKey) -> str: + """Canonical alias for the aggregate's output column in the CTE. + + Mirrors the result-key contract: ``leaf`` + ``_`` + ``agg`` plus an + args/kwargs signature suffix that disambiguates parameterised + aggregates (``revenue:percentile(p=0.5)`` vs ``p=0.95``). The + ``*:count`` star form collapses to ``_count``. + + Built on ``slayer.core.refs.canonical_agg_name`` so the signature + suffix matches the rest of the engine (legacy enrichment, search, + DBT converter). + """ + # ColumnKey -> leaf, ColumnSqlKey (derived agg source) -> column_name, + # StarKey -> "*" (CR / Codex: a derived source must alias as + # ``net_sum``, not ``_sum``). + measure_name = ( + getattr(key.source, "leaf", None) + or getattr(key.source, "column_name", None) + or "*" + ) + # AggregateKey.args / kwargs are normalised tuples of scalars / + # ColumnKey-shaped values; convert to the (List[str], + # Dict[str, Any]) shape ``canonical_agg_name`` expects. DEV-1450 + # stage 7b.13: route through ``agg_kwarg_canonical_str`` so a + # ColumnKey kwarg renders as ``leaf`` (or ``path.leaf`` for joined + # paths) instead of Pydantic-repr noise from naive ``str(v)``. + # The kwarg suffix is preserved here -- the legacy enrichment at + # ``query_engine.py:2160`` drops it, causing two parametric aggs + # with different kwargs to collide on CTE alias (legacy bug). + # The 7b.5 fix added kwarg-aware aliases here as a correctness + # improvement -- ``test_cross_model_planner_wiring.py:: + # test_parameterized_aggregates_get_distinct_cte_aliases`` pins + # this. Parity with legacy for cross-model parametric aggs is + # not achievable on this axis. + args_list = [agg_kwarg_canonical_str(a) for a in key.args] + kwargs_dict = {k: agg_kwarg_canonical_str(v) for k, v in key.kwargs} + return canonical_agg_name( + measure_name=measure_name, + aggregation_name=key.agg, + agg_args=args_list or None, + agg_kwargs=kwargs_dict or None, + ) + + +def _make_cte_schema( + *, + aggregate_owner: SlayerModel, + join_back_target_owner: SlayerModel, + aggregate_key: AggregateKey, + join_back_pairs: List[Tuple], +) -> StageSchema: + """Build the typed projection schema for the CTE. + + The CTE walks the join chain inside its body but groups at the + FIRST hop's target grain — so the projection's join-back keys are + columns on ``join_back_target_owner`` (the first hop's target model), + while the aggregate output column's type comes from + ``aggregate_owner`` (the terminal/aggregated model). + + For single-hop plans the two owners are the same model. For multi- + hop (``orders → customers → regions``), ``aggregate_owner`` is + ``regions`` and ``join_back_target_owner`` is ``customers``. + + Stage 7b's SQL generator consumes the schema when emitting the CTE + body and the join-back ON clause. + """ + columns: List[StageColumn] = [] + agg_alias = _aggregate_alias(key=aggregate_key) + src_leaf = ( + getattr(aggregate_key.source, "leaf", None) + or getattr(aggregate_key.source, "column_name", None) + ) + agg_type: Optional[DataType] = None + if src_leaf and hasattr(aggregate_owner, "get_column"): + src_col = aggregate_owner.get_column(src_leaf) + if src_col is not None: + agg_type = src_col.type + columns.append(StageColumn( + name=agg_alias, + sql_alias=agg_alias, + public_alias=None, + hidden=True, + type=agg_type or DataType.DOUBLE, + provenance=f"agg:{aggregate_key.agg}", + )) + for _, target_key in join_back_pairs: + leaf = getattr(target_key, "leaf", None) + if leaf is None: + continue + if any(c.name == leaf for c in columns): + continue + target_col = ( + join_back_target_owner.get_column(leaf) + if hasattr(join_back_target_owner, "get_column") else None + ) + col_type = target_col.type if target_col is not None else None + columns.append(StageColumn( + name=leaf, + sql_alias=leaf, + public_alias=None, + hidden=True, + type=col_type, + provenance="join_back_key", + )) + return StageSchema( + relation_name=f"cm_{aggregate_owner.name}", + columns=columns, + ) + + +def _match_filtered_local_grain_pairs( + *, + host_slots: List[ValueSlot], + public_projection: List[SlotId], + sub_plan: PlannedQuery, +) -> List[Tuple[SlotId, SlotId]]: + """Pair each host dimension / time-dimension slot with the sub-plan's + corresponding row slot for the LEFT JOIN ON clause. + + Both plans bind against the SAME underlying column on the host model, + so slot identity (the ValueKey) matches across plans. + """ + sub_row_by_key = {s.key: s.id for s in sub_plan.row_slots} + grain_pairs: List[Tuple[SlotId, SlotId]] = [] + for host_sid in public_projection: + host_slot = next( + (s for s in host_slots if s.id == host_sid), None, + ) + if host_slot is None: + continue + sub_sid = sub_row_by_key.get(host_slot.key) + if sub_sid is not None: + grain_pairs.append((host_sid, sub_sid)) + return grain_pairs + + +def _find_filtered_local_sub_agg_slot( + *, + sub_plan: PlannedQuery, + formula: str, + host_model: SlayerModel, +) -> SlotId: + """Locate the sub-plan's single local aggregate slot. + + Recursion suppression guarantees no nested cross-model plans so the + sub-plan has exactly one local aggregate — the filtered measure being + isolated. + """ + for s in sub_plan.aggregate_slots: + if isinstance(s.key, AggregateKey) and not getattr( + s.key.source, "path", (), + ): + return s.id + raise ValueError( + "DEV-1503 sub-plan produced no local aggregate slot for " + f"{formula!r} on {host_model.name!r} — planner bug." + ) + + +def _build_filtered_local_cte_schema( + *, + aggregate_key: AggregateKey, + host_model: SlayerModel, +) -> StageSchema: + """Build the minimal CTE schema for a filtered-local plan. + + The actual CTE columns are derived from the sub-plan's stage_schema / + projection at render time; this entry exists so external consumers see + a schema shape that matches the existing CrossModelAggregatePlan + contract. + """ + agg_alias = _aggregate_alias(key=aggregate_key) + leaf = getattr(aggregate_key.source, "leaf", None) or getattr( + aggregate_key.source, "column_name", None, + ) + agg_type: Optional[DataType] = None + if leaf is not None and hasattr(host_model, "get_column"): + col = host_model.get_column(leaf) + if col is not None: + agg_type = col.type + return StageSchema( + relation_name=f"cm_{host_model.name}", + columns=[StageColumn( + name=agg_alias, + sql_alias=agg_alias, + public_alias=None, + hidden=True, + type=agg_type or DataType.DOUBLE, + provenance=f"agg:{aggregate_key.agg}", + )], + ) + + +def _classify_subplan_filters( + *, + host_filters: List[HostFilterRouting], +) -> Optional[List[str]]: + """Decide which host-query filters propagate into the DEV-1503 sub-plan. + + ROW: pass through — the sub-plan applies them to the aggregate's rowset + (otherwise a non-dim filter like ``status = 'active'`` has no effect on + the join-back aggregate value). + AGGREGATE (any slot ref): skip. Pushing such a filter into the sub-plan + as HAVING would drop CTE rows where the aggregate fails the test; the + outer LEFT JOIN then surfaces the host row with a NULL aggregate + instead of dropping it — wrong semantics. The generator's outer-WHERE + wrapper applies the filter on the joined-back column so the row is + actually dropped (DEV-1503 spec). + POST: skip — stays at the existing host post-transform wrapper. + + Consume ``routing.text`` directly — it carries the original user-filter + string for user-filter routings (None for date_range bounds) and is + populated by ``stage_planner`` from the deduped ``bound_filters`` list, + so it stays in lock-step with ``host_filters`` even after Mode-B + dedup-by-bound-key collapses two textually-different filter spellings + onto one routing (CR PR #153 thread r3350000254). Slicing + ``host_query.filters`` here would mis-pair phases when a ``date_range``- + bearing time_dimension is present (Codex review) OR when dedup drops + user-filter entries. + """ + sub_filter_texts: List[str] = [] + for routing in host_filters: + if routing.text is None: + # date_range bound — not a user filter, do not propagate. + continue + if routing.phase in (Phase.POST, Phase.AGGREGATE): + continue + # ROW phase — propagate. + sub_filter_texts.append(routing.text) + return sub_filter_texts or None + + +class IsolatedCteCrossModelPlanner: + """Default impl — one CTE per (target_model, shared_grain) tuple. + + Encodes the ``inherited_filter_policy`` decision table from the + DEV-1450 spec via ``classify_host_filter`` for host filters; pulls + target ``SlayerModel.filters`` automatically. + """ + + def plan( + self, + *, + aggregate_slot_id: SlotId, + aggregate_key: AggregateKey, + bundle: ResolvedSourceBundle, + host_slots: List[ValueSlot], + host_filters: List[HostFilterRouting], + public_alias: Optional[str] = None, + hidden: bool = False, + host_query: Optional[SlayerQuery] = None, + public_projection: Optional[List[SlotId]] = None, + subplan_builder: Optional[ + Callable[[SlayerQuery, ResolvedSourceBundle], PlannedQuery] + ] = None, + ) -> CrossModelAggregatePlan: + host_model = bundle.source_model + if host_model is None: + raise ValueError( + "ResolvedSourceBundle.source_model is None — " + "IsolatedCteCrossModelPlanner needs a host model anchor " + "(I2 anchor-less mode is not yet implemented)." + ) + + agg_source = aggregate_key.source + path = getattr(agg_source, "path", ()) + if not path: + return self._dispatch_filtered_local( + aggregate_slot_id=aggregate_slot_id, + aggregate_key=aggregate_key, + bundle=bundle, + host_model=host_model, + host_slots=host_slots, + host_filters=host_filters, + public_alias=public_alias, + hidden=hidden, + host_query=host_query, + public_projection=public_projection, + subplan_builder=subplan_builder, + ) + + terminal_model, join_chain = _walk_chain( + host_model=host_model, hops=path, bundle=bundle, + ) + + # Build join_back_pairs from the FIRST hop's join_pairs. The CTE + # is grouped at the first hop's target columns; the host joins + # back on those. + join_back_pairs: List[Tuple] = [] + if join_chain: + first_hop = join_chain[0] + for pair in first_hop.join_pairs: + host_col, target_col = pair + join_back_pairs.append(( + ColumnKey(path=(), leaf=host_col), + ColumnKey(path=(), leaf=target_col), + )) + + target_path = path + applied: List[BoundFilterId] = [] + where_ids: List[BoundFilterId] = [] + having_ids: List[BoundFilterId] = [] + dropped: List[UnreachableFilterDroppedWarning] = [] + for hf in host_filters: + route = classify_host_filter( + host_filter=hf, + host_slots=host_slots, + target_path=target_path, + host_model_name=host_model.name, + ) + if route is FilterRoute.PROPAGATE_WHERE: + where_ids.append(hf.filter_id) + applied.append(hf.filter_id) + elif route is FilterRoute.PROPAGATE_HAVING: + having_ids.append(hf.filter_id) + applied.append(hf.filter_id) + elif route is FilterRoute.DROP_UNREACHABLE: + dropped.append(UnreachableFilterDroppedWarning( + filter_text=hf.text or hf.filter_id, + reason=( + f"filter {hf.filter_id!r} references slot(s) outside " + f"the join path to {terminal_model.name!r}; " + f"unreachable filters are dropped." + ), + )) + # DROP_HOST_LOCAL and STAY_AT_HOST_POST: not propagated, not warned. + + target_model_filters = list(terminal_model.filters or []) + + # Shared grain: local ROW slots on host (dimensions) flow through. + # Cross-branch ROW slots and aggregate / transform slots do not. + # DEV-1450 stage 7b.12: ``TimeTruncKey`` slots count as grain + # candidates too — a joined TD (``customers.created_at`` MONTH) + # whose column path lies on the target's join chain is shared + # between the host base and the cross-model CTE, so legacy + # ``LEFT JOIN`` on the truncated alias replaces the global + # ``CROSS JOIN``. + shared_grain: List[SlotId] = [] + for s in host_slots: + if isinstance(s.key, ColumnKey): + p = s.key.path + if not p or p == target_path[: len(p)]: + shared_grain.append(s.id) + elif isinstance(s.key, TimeTruncKey): + td_path = column_path(s.key.column) + if not td_path or td_path == target_path[: len(td_path)]: + shared_grain.append(s.id) + + first_hop = join_chain[0] + first_hop_target = ( + bundle.get_referenced_model(first_hop.target_model) + or terminal_model + ) + cte_schema = _make_cte_schema( + aggregate_owner=terminal_model, + join_back_target_owner=first_hop_target, + aggregate_key=aggregate_key, + join_back_pairs=join_back_pairs, + ) + + forward_plan = CrossModelAggregatePlan( + aggregate_slot_id=aggregate_slot_id, + target_model=terminal_model.name, + datasource=host_model.data_source, + join_chain=join_chain, + join_back_pairs=join_back_pairs, + cte_stage_schema=cte_schema, + shared_grain_slots=shared_grain, + applied_filter_ids=applied, + where_filter_ids=where_ids, + having_filter_ids=having_ids, + target_model_filters=target_model_filters, + dropped_filter_warnings=dropped, + hidden=hidden, + public_alias=public_alias, + ) + + # DEV-1450 #2: re-rooting is the strategy's call. When the caller + # supplies the host query + a sub-plan builder, decide forward-plan + # vs re-rooted-plan here; without them (direct ``plan(...)`` callers / + # test doubles) return the forward plan unchanged. + if subplan_builder is not None and host_query is not None: + return _maybe_reroot_cross_model_plan( + plan=forward_plan, + query=host_query, + agg_key=aggregate_key, + bundle=bundle, + host_model=host_model, + public_projection=public_projection or [], + subplan_builder=subplan_builder, + ) + return forward_plan + + # ---------------------------------------------------------------------- + # DEV-1503 — filtered-local isolation + # ---------------------------------------------------------------------- + + def _dispatch_filtered_local( + self, + *, + aggregate_slot_id: SlotId, + aggregate_key: AggregateKey, + bundle: ResolvedSourceBundle, + host_model: SlayerModel, + host_slots: List[ValueSlot], + host_filters: List[HostFilterRouting], + public_alias: Optional[str], + hidden: bool, + host_query: Optional[SlayerQuery], + public_projection: Optional[List[SlotId]], + subplan_builder: Optional[ + Callable[[SlayerQuery, ResolvedSourceBundle], PlannedQuery] + ], + ) -> CrossModelAggregatePlan: + """Validate the filtered-local trigger preconditions and dispatch + into ``_plan_filtered_local`` — the aggregate is on a HOST column + but its ``Column.filter`` crosses a join, so a host-rooted nested + sub-plan owns the aggregation and the host base LEFT JOINs back. + """ + agg_source = aggregate_key.source + cfk = aggregate_key.column_filter_key + if cfk is None or not cfk.referenced_join_paths: + raise ValueError( + f"AggregateKey on {agg_source!r} has empty source.path " + f"AND no cross-model column_filter_key — this is a plain " + f"local aggregate. The cross-model planner should not " + f"have been invoked." + ) + if subplan_builder is None or host_query is None: + # The DEV-1503 strategy requires a sub-plan builder + the host + # query for grain-pair matching. Direct callers without these + # (legacy test doubles) can't trigger filtered-local — raise + # loudly so the call site is fixed rather than emitting + # silently wrong SQL. + raise ValueError( + "DEV-1503 filtered-local isolation requires host_query " + "and subplan_builder; received None for one or both. " + "Confirm the stage_planner is wired to pass them." + ) + return self._plan_filtered_local( + aggregate_slot_id=aggregate_slot_id, + aggregate_key=aggregate_key, + bundle=bundle, + host_model=host_model, + host_slots=host_slots, + host_filters=host_filters, + host_query=host_query, + public_alias=public_alias, + public_projection=public_projection or [], + hidden=hidden, + subplan_builder=subplan_builder, + ) + + def _plan_filtered_local( + self, + *, + aggregate_slot_id: SlotId, + aggregate_key: AggregateKey, + bundle: ResolvedSourceBundle, + host_model: SlayerModel, + host_slots: List[ValueSlot], + host_filters: List[HostFilterRouting], + host_query: SlayerQuery, + public_alias: Optional[str], + public_projection: List[SlotId], + hidden: bool, + subplan_builder: Callable[ + [SlayerQuery, ResolvedSourceBundle], PlannedQuery, + ], + ) -> CrossModelAggregatePlan: + """Build a host-rooted nested sub-plan for a cross-model-FILTERED + local measure (DEV-1503). + + The sub-plan is a ``SlayerQuery`` rooted at the SAME host model with + ``measures=[]`` and the host's dimensions / + time_dimensions. The sub-plan's ``plan_query`` recursion handles + the filter-target join (its ``Column.filter`` will pull in the + joined table at the generator's inline path), the host model's own + ``SlayerModel.filters``, and the per-dimension GROUP BY — producing a + per-grain aggregate that the host base LEFT JOINs back. + + Host query filters are NOT propagated into the sub-plan here — the + host base CTE applies them. The generator's outer-WHERE wrapper + handles aggregate-referencing filters separately (DEV-1503 spec). + """ + # Reconstruct the local measure formula from the AggregateKey. The + # source.path is empty so ``_local_agg_formula`` emits a bare + # ``leaf:agg`` shape (plus any args / kwargs). Carry the user- + # supplied alias through so a host filter referencing the rename + # (``latest_pmt > 500`` for a measure named ``latest_pmt``) binds + # against the same alias in the sub-plan rather than the canonical + # ``latest_payment_last_updated_at`` form. + formula = _local_agg_formula(aggregate_key) + measure_name_for_subplan = public_alias + sub_filters = _classify_subplan_filters(host_filters=host_filters) + rerooted_query = SlayerQuery( + source_model=host_model.name, + measures=[ModelMeasure( + formula=formula, name=measure_name_for_subplan, + )], + dimensions=list(host_query.dimensions or []) or None, + time_dimensions=list(host_query.time_dimensions or []) or None, + filters=sub_filters, + ) + sub_plan = subplan_builder(rerooted_query, bundle) + + grain_pairs = _match_filtered_local_grain_pairs( + host_slots=host_slots, + public_projection=public_projection, + sub_plan=sub_plan, + ) + sub_agg_sid = _find_filtered_local_sub_agg_slot( + sub_plan=sub_plan, formula=formula, host_model=host_model, + ) + cte_schema = _build_filtered_local_cte_schema( + aggregate_key=aggregate_key, host_model=host_model, + ) + + return CrossModelAggregatePlan( + aggregate_slot_id=aggregate_slot_id, + # ``target_model`` is conventionally set to the host name for + # filtered-local; ``cte_root_model`` is the disambiguator the + # renderer reads. + target_model=host_model.name, + cte_root_model=host_model.name, + datasource=host_model.data_source, + join_chain=[], + join_back_pairs=[], + cte_stage_schema=cte_schema, + shared_grain_slots=[host_sid for host_sid, _ in grain_pairs], + applied_filter_ids=[], + where_filter_ids=[], + having_filter_ids=[], + target_model_filters=[], + dropped_filter_warnings=[], + hidden=hidden, + public_alias=public_alias, + rerooted_plan=sub_plan, + rerooted_grain_pairs=grain_pairs, + rerooted_agg_slot_id=sub_agg_sid, + ) + + +# --------------------------------------------------------------------------- +# Cross-model re-rooting (DEV-1450 stage 7b.15e, C1; relocated here in #2) +# --------------------------------------------------------------------------- +# +# When a cross-model aggregate (``policy_amount.total:sum``) is queried with +# host dimensions that are reachable from the TARGET by walking the target's +# own join graph (``policy_amount -> policy -> policy_number``), the +# forward-path CTE ("FROM bare target, GROUP BY forward-path dims only") +# collapses the host dimension to a scalar CROSS JOIN -- every host row gets +# the global aggregate. +# +# The fix mirrors legacy ``_build_rerooted_enriched``: build a full nested +# ``SlayerQuery`` rooted at the target (so all of the target's joins are in +# scope for dimensions AND filters), compile it via ``subplan_builder``, and +# attach the sub-plan to the ``CrossModelAggregatePlan``. The generator +# renders the sub-plan as the ``_cm_*`` CTE and joins it back to the host base +# on the (re-rooted) dimension. Dimensions / filters that don't resolve from +# the target are dropped -- matching legacy's drop-unreachable behaviour. +# +# DEV-1450 #2: this used to be a post-hoc pass in ``stage_planner.plan_query``; +# it now lives behind ``IsolatedCteCrossModelPlanner.plan`` so the +# render-strategy decision (forward vs re-rooted) is owned by the strategy. +# The recursive ``plan_query`` call is injected as ``subplan_builder`` so this +# module does not import ``stage_planner`` (no cycle). + + +def _reroot_ref( + *, model_prefix: Optional[str], name: str, host_model_name: str, + target_model_name: str, +) -> str: + """Re-root one Mode-B ref from the host's perspective to the target's. + + Mirrors legacy ``_build_rerooted_enriched``: + + * host-local (``model_prefix is None``) -> ``.`` (now a + cross-model dim from the target's view), + * on the target itself -> bare ```` (local on target), + * a path through the target -> strip the target prefix, + * any other dotted ref -> kept as-is (resolved via the target's joins). + """ + if model_prefix is None: + return f"{host_model_name}.{name}" + if model_prefix == target_model_name: + return name + if model_prefix.startswith(target_model_name + "."): + return f"{model_prefix[len(target_model_name) + 1:]}.{name}" + return f"{model_prefix}.{name}" + + +def _host_ref_path(model_prefix: Optional[str]) -> Tuple[str, ...]: + """The join path a host ColumnRef / TimeDimension prefix denotes.""" + if not model_prefix: + return () + return tuple(model_prefix.split(".")) + + +def _scalar_formula_literal(value) -> str: + """Render a normalized scalar back into formula text.""" + if isinstance(value, bool): + return "True" if value else "False" + if value is None: + return "None" + if isinstance(value, str): + return repr(value) + return str(value) + + +def _filter_ref_paths(value_key: ValueKey) -> List[Tuple[str, ...]]: + """Join paths of every column-like leaf a (bound) filter references.""" + paths: List[Tuple[str, ...]] = [] + for k in walk_value_keys(value_key): + if isinstance(k, (ColumnKey, ColumnSqlKey, StarKey)): + paths.append(tuple(k.path)) + elif isinstance(k, TimeTruncKey): + paths.append(tuple(column_path(k.column))) + return paths + + +def _local_agg_formula(key: AggregateKey) -> str: + """Reconstruct the LOCAL colon-formula for a cross-model aggregate + (``customers.revenue:sum`` -> ``revenue:sum``) so it can be re-planned + against the target model as a plain local measure. + + Column-valued kwargs (``corr(other=customers.region_id)``) are + re-rooted too: their join path is bound from the HOST, so the leading + agg-source (target) prefix is stripped to express the ref in the + target's local scope (``other=region_id``; a deeper hop keeps its + residual path, ``other=regions.code``). Dropping the path outright + would mis-bind or fail to bind the nested sub-query (CR review).""" + src = key.source + target_path = tuple(getattr(src, "path", ())) + if isinstance(src, StarKey): + base = "*" + elif isinstance(src, ColumnSqlKey): + base = src.column_name + else: # ColumnKey + base = src.leaf + + def _reroot_col_kwarg(v) -> str: + leaf = v.leaf if isinstance(v, ColumnKey) else v.column_name + vpath = tuple(getattr(v, "path", ())) + # Strip the agg-source (target) prefix so the ref is target-local. + residual = ( + vpath[len(target_path):] + if vpath[: len(target_path)] == target_path + else vpath + ) + return ".".join((*residual, leaf)) + + formula = f"{base}:{key.agg}" + parts: List[str] = [] + # Positional args may carry ColumnKey / ColumnSqlKey just like kwargs do + # (rerooting needs path-aware handling on both — CR review). Falling + # through to ``_scalar_formula_literal`` would emit Pydantic-repr noise + # for a column-valued positional arg, mis-binding the nested sub-query. + for a in key.args: + if isinstance(a, (ColumnKey, ColumnSqlKey)): + parts.append(_reroot_col_kwarg(a)) + else: + parts.append(_scalar_formula_literal(a)) + for k, v in key.kwargs: + if isinstance(v, (ColumnKey, ColumnSqlKey)): + parts.append(f"{k}={_reroot_col_kwarg(v)}") + else: + parts.append(f"{k}={_scalar_formula_literal(v)}") + if parts: + formula += "(" + ", ".join(parts) + ")" + return formula + + +_REROOT_BIND_ERRORS = ( + UnknownReferenceError, + AmbiguousReferenceError, + IllegalScopeReferenceError, + ValueError, + NotImplementedError, +) + + +def _maybe_reroot_cross_model_plan( + *, + plan, + query: SlayerQuery, + agg_key: AggregateKey, + bundle: ResolvedSourceBundle, + host_model: SlayerModel, + public_projection: List[str], + subplan_builder: Callable[[SlayerQuery, ResolvedSourceBundle], PlannedQuery], +): + """Attach a re-rooted sub-``PlannedQuery`` to ``plan`` when the host + query carries dimensions reachable from the target by re-rooting through + the target's join graph. Returns ``plan`` unchanged when re-rooting is + unnecessary (only forward-path or genuinely unreachable dims).""" + target_model_name = plan.target_model + target_model = bundle.get_referenced_model(target_model_name) + if target_model is None: + return plan + target_path = tuple(getattr(agg_key.source, "path", ())) + rerooted_bundle = bundle.model_copy(update={"source_model": target_model}) + target_scope = ModelScope(source_model=target_model) + + def _resolvable_ref(ref_str: str) -> Optional[ValueKey]: + try: + return bind_expr( + parse_expr(ref_str), + scope=target_scope, + bundle=rerooted_bundle, + ).value_key + except _REROOT_BIND_ERRORS: + return None + + def _is_forward(path: Tuple[str, ...]) -> bool: + # On the host->target path (handled by the forward-path CTE already). + return bool(path) and path == target_path[: len(path)] + + n_dims = len(query.dimensions or []) + rerooted_dims: List[ColumnRef] = [] + rerooted_tds: List[TimeDimension] = [] + grain_host_sids: List[str] = [] + grain_rerooted_keys: List[ValueKey] = [] + needs_reroot = False + + for i, dim in enumerate(query.dimensions or []): + host_sid = public_projection[i] if i < len(public_projection) else None + host_path = _host_ref_path(dim.model) + rr = _reroot_ref( + model_prefix=dim.model, name=dim.name, + host_model_name=host_model.name, target_model_name=target_model_name, + ) + rr_key = _resolvable_ref(rr) + if rr_key is None: + continue # unreachable from target -> drop + if not _is_forward(host_path): + needs_reroot = True + if host_sid is None: + continue + rerooted_dims.append(ColumnRef(name=rr, label=dim.label)) + grain_host_sids.append(host_sid) + grain_rerooted_keys.append(rr_key) + + for j, td in enumerate(query.time_dimensions or []): + idx = n_dims + j + host_sid = public_projection[idx] if idx < len(public_projection) else None + host_path = _host_ref_path(td.dimension.model) + rr = _reroot_ref( + model_prefix=td.dimension.model, name=td.dimension.name, + host_model_name=host_model.name, target_model_name=target_model_name, + ) + rr_td = TimeDimension( + dimension=ColumnRef(name=rr), + granularity=td.granularity, + date_range=td.date_range, + label=td.label, + ) + try: + rr_key = bind_time_dimension( + rr_td, scope=target_scope, bundle=rerooted_bundle, + ).value_key + except _REROOT_BIND_ERRORS: + continue + if not _is_forward(host_path): + needs_reroot = True + if host_sid is None: + continue + rerooted_tds.append(rr_td) + grain_host_sids.append(host_sid) + grain_rerooted_keys.append(rr_key) + + # Filters. A purely host-local filter (every ref on the host's own + # columns) filters host rows -- it stays at the host base; the join-back + # propagates the cardinality reduction, so adding it to the CTE would risk + # binding a bare name to a same-named TARGET column. A join-traversing + # filter affects the aggregate value and rides into the re-rooted CTE; one + # that reaches OFF the host->target forward path is exactly what the + # forward-path classifier drops, so it also triggers re-rooting (covers a + # cross-model agg filtered through the target's graph with no dimensions). + host_scope = ModelScope(source_model=host_model) + rerooted_filters: List[str] = [] + for f in (query.filters or []): + try: + host_bound = bind_filter( + parse_filter_expr(f), scope=host_scope, bundle=bundle, + ) + except _REROOT_BIND_ERRORS: + continue + host_paths = _filter_ref_paths(host_bound.value_key) + if all(p == () for p in host_paths): + continue # host-local -> applied at the host base only + # The binder strips a same-model self-prefix (C14), so a + # ``.col`` ref binds locally against the target scope without + # any string surgery -- pass the filter through verbatim. + try: + bind_filter( + parse_filter_expr(f), scope=target_scope, bundle=rerooted_bundle, + ) + except _REROOT_BIND_ERRORS: + continue + rerooted_filters.append(f) + if any(p != target_path[: len(p)] for p in host_paths if p): + needs_reroot = True + + if not needs_reroot or not ( + rerooted_dims or rerooted_tds or rerooted_filters + ): + return plan + + rerooted_query = SlayerQuery( + source_model=target_model_name, + measures=[ModelMeasure(formula=_local_agg_formula(agg_key))], + dimensions=rerooted_dims or None, + time_dimensions=rerooted_tds or None, + filters=rerooted_filters or None, + ) + sub_plan = subplan_builder(rerooted_query, rerooted_bundle) + + sub_row_by_key = {s.key: s.id for s in sub_plan.row_slots} + grain_pairs: List[Tuple[str, str]] = [] + for host_sid, rr_key in zip(grain_host_sids, grain_rerooted_keys): + sub_sid = sub_row_by_key.get(rr_key) + if sub_sid is not None: + grain_pairs.append((host_sid, sub_sid)) + + sub_agg_sid = None + for s in sub_plan.aggregate_slots: + if isinstance(s.key, AggregateKey) and not getattr( + s.key.source, "path", (), + ): + sub_agg_sid = s.id + break + if sub_agg_sid is None: + return plan + + return plan.model_copy(update={ + "rerooted_plan": sub_plan, + "rerooted_grain_pairs": grain_pairs, + "rerooted_agg_slot_id": sub_agg_sid, + # The forward-path classifier marked these host filters + # DROP_UNREACHABLE, but the re-rooted CTE re-applies every + # target-reachable filter (and the host base keeps the rest for + # cardinality), so nothing is silently dropped -- clear the now-stale + # warnings and forward-only routing ids. + "dropped_filter_warnings": [], + "where_filter_ids": [], + "having_filter_ids": [], + "applied_filter_ids": [], + }) diff --git a/slayer/engine/measure_expansion.py b/slayer/engine/measure_expansion.py new file mode 100644 index 00000000..fc370780 --- /dev/null +++ b/slayer/engine/measure_expansion.py @@ -0,0 +1,353 @@ +"""Stage 7b.2 (DEV-1450) — pre-bind ModelMeasure expansion. + +Pre-bind AST -> AST rewrite of a ``ParsedExpr`` tree: every ``Ref(name=X)`` +whose ``X`` resolves to a ``ModelMeasure`` (on the model or in +``extra_measures``) is replaced with the recursively-expanded +``ParsedExpr`` produced by ``parse_expr(measure.formula)``. + +Why pre-bind: the binder (``slayer.engine.binding``) raises +``UnknownReferenceError`` for bare measure names because measures are not +columns; running expansion before the binder sees the tree turns those +refs into binder-resolvable column / aggregation nodes. + +Eligibility matrix: + +* Eligible positions: root; ``Arith`` / ``UnaryOp`` / ``Cmp`` / ``BoolOp`` + operands; ``ScalarCall.args``; ``TransformCall.input`` / args / + kwarg values. +* Not eligible: ``DottedRef`` (cross-model dotted paths resolve through + the join graph, not through measure expansion); ``AggCall`` in any + position (``source`` / ``args`` / ``kwargs`` are column-level by + contract); function-name slots on ``TransformCall.op`` / + ``ScalarCall.name`` (those are strings, not ``Ref`` nodes, and would + never match the dispatch even if a measure with the same name + existed); ``Literal`` / ``StarSource``; ``SlayerQuery.order`` entries + (the caller does not pass order entries through this function — order + resolves declared slot names only at the planner layer). + +Recursion controls: + +* Depth limit configurable via ``SLAYER_MEASURE_EXPANSION_DEPTH`` env + var (default ``32``). An explicit ``depth_limit=`` kwarg wins. Exceeded + -> :class:`MeasureRecursionLimitError`. +* Per-chain cycle detection. A measure transitively referencing itself + raises :class:`MeasureCycleError` with the offending chain attached. + +Purity: input ``ParsedExpr`` nodes are frozen Pydantic models; the +function returns a fresh tree. + +Dormant in this commit. Stage 7b.6 (BoundExpr unification) and stage +7b.15 (engine cutover) wire this into the engine pipeline. +""" + +from __future__ import annotations + +import os +from typing import Any, Dict, Optional, Sequence, Tuple, get_args + +from slayer.core.errors import MeasureCycleError, MeasureRecursionLimitError +from slayer.core.models import ModelMeasure, SlayerModel +from slayer.engine.syntax import ( + AggCall, + Arith, + BoolOp, + Cmp, + DottedRef, + Literal, + ParsedExpr, + Ref, + ScalarCall, + StarSource, + TransformCall, + UnaryOp, + parse_expr, +) + +_DEFAULT_DEPTH = 32 +_DEPTH_ENV_VAR = "SLAYER_MEASURE_EXPANSION_DEPTH" + +# Authoritative ParsedExpr node-type tuple, derived from the union in +# slayer.engine.syntax so new node types added there are automatically +# walked by `_maybe_walk` without a silent skip. +_PARSED_EXPR_TYPES: Tuple[type, ...] = get_args(ParsedExpr) + + +def expand_model_measures( + *, + expr: ParsedExpr, + model: SlayerModel, + extra_measures: Sequence[ModelMeasure] = (), + depth_limit: Optional[int] = None, +) -> ParsedExpr: + """Walk ``expr`` and replace bare measure refs with their formula AST. + + See module docstring for the eligibility matrix and recursion rules. + Raises ``ValueError`` if ``depth_limit`` is not a positive integer. + """ + if depth_limit is not None and depth_limit < 1: + raise ValueError( + f"depth_limit must be a positive integer, got {depth_limit!r}." + ) + measures = _collect_named_measures(model=model, extras=extra_measures) + limit = depth_limit if depth_limit is not None else _env_depth_limit() + parse_cache: Dict[str, ParsedExpr] = {} + return _walk( + expr, + measures=measures, + depth_limit=limit, + chain=(), + parse_cache=parse_cache, + ) + + +def _collect_named_measures( + *, + model: SlayerModel, + extras: Sequence[ModelMeasure], +) -> Dict[str, ModelMeasure]: + """Build a ``name -> ModelMeasure`` map. ``extras`` shadow model + measures with the same name. Unnamed measures are not addressable + via bare ref and are skipped. + """ + out: Dict[str, ModelMeasure] = {} + for m in model.measures: + if m.name: + out[m.name] = m + for m in extras: + if m.name: + out[m.name] = m + return out + + +def _env_depth_limit() -> int: + raw = os.environ.get(_DEPTH_ENV_VAR) + if raw is None: + return _DEFAULT_DEPTH + try: + return max(1, int(raw)) + except ValueError: + return _DEFAULT_DEPTH + + +def _walk( + node: ParsedExpr, + *, + measures: Dict[str, ModelMeasure], + depth_limit: int, + chain: Tuple[str, ...], + parse_cache: Dict[str, ParsedExpr], +) -> ParsedExpr: + if isinstance(node, Ref): + return _expand_ref( + node=node, + measures=measures, + depth_limit=depth_limit, + chain=chain, + parse_cache=parse_cache, + ) + if isinstance(node, (DottedRef, StarSource, Literal, AggCall)): + return node + if isinstance(node, TransformCall): + return _walk_transform_call( + node=node, + measures=measures, + depth_limit=depth_limit, + chain=chain, + parse_cache=parse_cache, + ) + if isinstance(node, ScalarCall): + return _walk_scalar_call( + node=node, + measures=measures, + depth_limit=depth_limit, + chain=chain, + parse_cache=parse_cache, + ) + if isinstance(node, Arith): + return node.model_copy( + update={ + "left": _maybe_walk( + node.left, + measures=measures, + depth_limit=depth_limit, + chain=chain, + parse_cache=parse_cache, + ), + "right": _maybe_walk( + node.right, + measures=measures, + depth_limit=depth_limit, + chain=chain, + parse_cache=parse_cache, + ), + } + ) + if isinstance(node, UnaryOp): + return node.model_copy( + update={ + "operand": _maybe_walk( + node.operand, + measures=measures, + depth_limit=depth_limit, + chain=chain, + parse_cache=parse_cache, + ), + } + ) + if isinstance(node, Cmp): + return node.model_copy( + update={ + "left": _maybe_walk( + node.left, + measures=measures, + depth_limit=depth_limit, + chain=chain, + parse_cache=parse_cache, + ), + "right": _maybe_walk( + node.right, + measures=measures, + depth_limit=depth_limit, + chain=chain, + parse_cache=parse_cache, + ), + } + ) + if isinstance(node, BoolOp): + return node.model_copy( + update={ + "operands": tuple( + _maybe_walk( + o, + measures=measures, + depth_limit=depth_limit, + chain=chain, + parse_cache=parse_cache, + ) + for o in node.operands + ) + } + ) + # Unknown node type: leave alone. This branch is unreachable for + # well-formed ParsedExpr but keeps the helper total. + return node + + +def _expand_ref( + *, + node: Ref, + measures: Dict[str, ModelMeasure], + depth_limit: int, + chain: Tuple[str, ...], + parse_cache: Dict[str, ParsedExpr], +) -> ParsedExpr: + if node.name not in measures: + return node + if node.name in chain: + raise MeasureCycleError(chain=list(chain) + [node.name]) + new_chain = chain + (node.name,) + if len(new_chain) > depth_limit: + raise MeasureRecursionLimitError( + chain=list(new_chain), limit=depth_limit + ) + cached = parse_cache.get(node.name) + if cached is None: + cached = parse_expr(measures[node.name].formula) + parse_cache[node.name] = cached + return _walk( + cached, + measures=measures, + depth_limit=depth_limit, + chain=new_chain, + parse_cache=parse_cache, + ) + + +def _walk_transform_call( + *, + node: TransformCall, + measures: Dict[str, ModelMeasure], + depth_limit: int, + chain: Tuple[str, ...], + parse_cache: Dict[str, ParsedExpr], +) -> TransformCall: + new_input = _maybe_walk( + node.input, + measures=measures, + depth_limit=depth_limit, + chain=chain, + parse_cache=parse_cache, + ) + new_args = tuple( + _maybe_walk( + a, + measures=measures, + depth_limit=depth_limit, + chain=chain, + parse_cache=parse_cache, + ) + for a in node.args + ) + new_kwargs = tuple( + ( + k, + _maybe_walk( + v, + measures=measures, + depth_limit=depth_limit, + chain=chain, + parse_cache=parse_cache, + ), + ) + for k, v in node.kwargs + ) + return node.model_copy( + update={"input": new_input, "args": new_args, "kwargs": new_kwargs} + ) + + +def _walk_scalar_call( + *, + node: ScalarCall, + measures: Dict[str, ModelMeasure], + depth_limit: int, + chain: Tuple[str, ...], + parse_cache: Dict[str, ParsedExpr], +) -> ScalarCall: + new_args = tuple( + _maybe_walk( + a, + measures=measures, + depth_limit=depth_limit, + chain=chain, + parse_cache=parse_cache, + ) + for a in node.args + ) + return node.model_copy(update={"args": new_args}) + + +def _maybe_walk( + v: Any, + *, + measures: Dict[str, ModelMeasure], + depth_limit: int, + chain: Tuple[str, ...], + parse_cache: Dict[str, ParsedExpr], +) -> Any: + """Walk ``v`` only if it is a ParsedExpr node; pass scalars through + unchanged. AggCall args/kwargs and the like contain scalars + (Decimal / str / bool) that should not be touched. + """ + if isinstance(v, _PARSED_EXPR_TYPES): + return _walk( + v, + measures=measures, + depth_limit=depth_limit, + chain=chain, + parse_cache=parse_cache, + ) + return v + + +__all__ = ["expand_model_measures"] diff --git a/slayer/engine/normalization.py b/slayer/engine/normalization.py new file mode 100644 index 00000000..1ae48044 --- /dev/null +++ b/slayer/engine/normalization.py @@ -0,0 +1,688 @@ +"""Stage 6 (DEV-1450) — slack normalization layer. + +Rewrites tolerant-but-unambiguous agent input to canonical form before the +typed pipeline sees it, returning every rewrite as a typed +``NormalizationWarning`` (P0). Downstream stages never see the slack form. + +Three seed rules: + +- ``FUNC_STYLE_AGG`` (Mode B only): ``sum(revenue)`` / ``count(*)`` / + ``percentile(amount, p=0.5)`` → colon syntax. Rewrites Mode-B fields + (``ModelMeasure.formula``, ``SlayerQuery.measures[].formula``, + ``SlayerQuery.filters`` entries). + +- ``MISPLACED_MEASURE`` (query shape): bare column-looking entries in + ``SlayerQuery.measures`` that resolve as a column (not a named + ``ModelMeasure``) move to ``SlayerQuery.dimensions``. Mirrors the + existing ``_auto_move_fields_to_dimensions`` heuristic but emits a + structured warning. + +- ``DOT_PATH_IN_SQL`` (Mode A only): sqlglot-AST ``Column`` node in root + scope whose dotted path's leading segment matches a known join target + on the host model → ``__`` alias form (``customers.regions.name`` → + ``customers__regions.name``). Scope-aware via lexical-ancestor walking + so refs inside subqueries / CTE bodies / set-op branches are left + alone. First-segment shadow detection covers CTE names, explicit + ``AS`` aliases, Subquery/CTE FROM sources, and schema/catalog + qualifiers on FROM tables (``FROM customers.regions`` → ``customers`` + shadows). Shadowed cases emit an ambiguity warning without rewriting. + Wired into ``normalize_model`` over ``Column.sql``, ``Column.filter``, + and ``SlayerModel.filters``. + +Each rule emits a ``SlayerNormalizationWarning`` via ``warnings.warn(...)`` +AND appends a ``NormalizationWarning`` payload to the returned result, +so REST / MCP / CLI consumers see the rewrite alongside the response and +``warnings.catch_warnings()`` callers see it via the standard channel. +""" + +from __future__ import annotations + +import re +import warnings as _warnings_module +from typing import List, Optional, Set, Tuple + +import sqlglot +from sqlglot import exp +from sqlglot.optimizer.scope import ScopeType, traverse_scope +from pydantic import BaseModel, ConfigDict, Field + +from slayer.core.enums import BUILTIN_AGGREGATIONS +from slayer.core.models import SlayerModel +from slayer.core.query import SlayerQuery +from slayer.core.refs import IDENT_OR_PATH_RE +from slayer.core.warnings import NormalizationWarning, SlayerNormalizationWarning +from slayer.engine.column_expansion import _root_scope_column_ids + + +# --------------------------------------------------------------------------- +# Result type +# --------------------------------------------------------------------------- + + +class NormalizationResult(BaseModel): + """Output of a normalization pass. + + ``query`` and ``model`` are either the same object the caller passed + in (if no rewrite fired) or a new instance with the slack form + rewritten. ``warnings`` lists one ``NormalizationWarning`` per + rewrite — empty when the input was already canonical. + """ + + model_config = ConfigDict(arbitrary_types_allowed=True) + + query: Optional[SlayerQuery] = None + model: Optional[SlayerModel] = None + warnings: List[NormalizationWarning] = Field(default_factory=list) + + +# --------------------------------------------------------------------------- +# Rule: FUNC_STYLE_AGG +# --------------------------------------------------------------------------- + + +# Aggregation names that are also transform names — the rewrite only fires +# when the inner is a bare identifier, not when it's a colon-form aggregate. +_AMBIGUOUS_AGG_TRANSFORMS = frozenset({"first", "last"}) + +_STRING_LITERAL_RE = re.compile(r"'(?:[^']|'')*'|\"(?:[^\"]|\"\")*\"") + + +def _find_balanced_close(s: str, open_idx: int) -> int: + depth = 0 + in_string = False + string_ch = "" + i = open_idx + while i < len(s): + ch = s[i] + if in_string: + if ch == string_ch: + # Handle '' / "" escapes. + if i + 1 < len(s) and s[i + 1] == string_ch: + i += 2 + continue + in_string = False + elif ch in ("'", '"'): + in_string = True + string_ch = ch + elif ch == "(": + depth += 1 + elif ch == ")": + depth -= 1 + if depth == 0: + return i + i += 1 + return -1 + + +def _split_args(s: str) -> List[str]: + parts: List[str] = [] + depth = 0 + in_string = False + string_ch = "" + current: List[str] = [] + for ch in s: + if in_string: + current.append(ch) + if ch == string_ch: + in_string = False + continue + if ch in ("'", '"'): + in_string = True + string_ch = ch + current.append(ch) + continue + if ch == "(": + depth += 1 + current.append(ch) + continue + if ch == ")": + depth -= 1 + current.append(ch) + continue + if ch == "," and depth == 0: + parts.append("".join(current)) + current = [] + continue + current.append(ch) + if current: + parts.append("".join(current)) + return [p.strip() for p in parts if p.strip()] + + +def _apply_func_style_agg( + formula: str, + *, + location: str, + custom_agg_names: Optional[frozenset[str]] = None, +) -> tuple[str, List[NormalizationWarning]]: + """Rewrite function-style aggregations in ``formula`` to colon syntax. + + Returns ``(rewritten_formula, warnings)`` — ``warnings`` is empty when + nothing changed. + """ + agg_names = BUILTIN_AGGREGATIONS | (custom_agg_names or frozenset()) + sorted_names = sorted(agg_names, key=len, reverse=True) + pattern = re.compile( + r"(? str: + """Rewrite function-style aggregations (``sum(x)`` → ``x:sum``, + ``count(*)`` → ``*:count``) to colon syntax, returning only the rewritten + string. + + Quiet variant of the ``FUNC_STYLE_AGG`` slack rule for read-only, + best-effort consumers (schema-drift cascade attribution, memory entity + tagging) that inspect formulas with the typed Mode-B parser but must NOT + re-surface slack advice to the user — the pipeline path + (``normalize_query`` / ``normalize_model``) is the one that emits + ``SlayerNormalizationWarning``. Returns the formula unchanged when nothing + matches. + """ + with _warnings_module.catch_warnings(): + _warnings_module.simplefilter("ignore", SlayerNormalizationWarning) + rewritten, _ = _apply_func_style_agg( + formula, location="(inspect)", custom_agg_names=custom_agg_names, + ) + return rewritten + + +# --------------------------------------------------------------------------- +# Rule: MISPLACED_MEASURE +# --------------------------------------------------------------------------- + + +def _apply_misplaced_measure( + query: SlayerQuery, + *, + model: Optional[SlayerModel], +) -> tuple[SlayerQuery, List[NormalizationWarning]]: + """Move bare (no-colon, no-function) entries from ``query.measures`` to + ``query.dimensions`` when they name a column on the model that isn't + a ``ModelMeasure`` formula. + + Mirrors the existing ``_auto_move_fields_to_dimensions`` heuristic but + emits a structured warning. When ``model`` is None we can't classify, + so the rule is a no-op. + """ + if not query.measures or model is None: + return query, [] + + measure_formula_names = {m.name for m in model.measures} + column_names = {c.name for c in model.columns} + + new_measures = list(query.measures) + moved_dim_strings: List[str] = [] + emitted: List[NormalizationWarning] = [] + + kept: List = [] + for i, m in enumerate(new_measures): + formula = getattr(m, "formula", None) + if not isinstance(formula, str): + kept.append(m) + continue + if ":" in formula or "(" in formula: + kept.append(m) + continue + # Bare token. If it names a known ModelMeasure formula, keep it as + # a measure. If it names a column on the model, move to dimensions. + bare = formula.strip() + if bare in measure_formula_names: + kept.append(m) + continue + if bare in column_names: + moved_dim_strings.append(bare) + emitted.append(NormalizationWarning( + rule_id="MISPLACED_MEASURE", + original=bare, + normalized=f"dimensions += {bare!r}", + location=f"measures[{i}].formula", + rule_doc_url="docs/agent_input_slack.md#misplaced-measure", + )) + _warnings_module.warn( + SlayerNormalizationWarning(emitted[-1]), stacklevel=2, + ) + continue + # Unknown bare token — leave for downstream resolver to error on. + kept.append(m) + + if not emitted: + return query, [] + + existing_dims = list(query.dimensions or []) + # Append each moved bare column name as a dimension entry. We add as + # plain strings since SlayerQuery.dimensions accepts string entries + # alongside ColumnRefs (the pydantic union validators handle the + # coercion). + new_dimensions = existing_dims + moved_dim_strings + return ( + query.model_copy(update={"measures": kept, "dimensions": new_dimensions}), + emitted, + ) + + +# --------------------------------------------------------------------------- +# Rule: DOT_PATH_IN_SQL (stub for stage 6 — full implementation deferred) +# --------------------------------------------------------------------------- + + +# Node types that have a "natural" root scope sqlglot can analyse directly. +# Any other parsed input (a scalar expression like ``lower(a.b.c)``) gets +# wrapped in a synthetic ``SELECT ... AS _`` for scope traversal — mirrors +# the precedent in ``column_expansion._root_scope_column_ids``. +_STATEMENT_TYPES: Tuple[type, ...] = ( + exp.Select, exp.Union, exp.Intersect, exp.Except, +) + + +def _dot_path_root_scope_analysis( + *, parsed: exp.Expression, +) -> Tuple[Set[int], Set[str]]: + """Return ``(root_col_ids, shadow_names)`` for ``parsed``. + + ``root_col_ids``: ids of ``exp.Column`` nodes whose innermost + scope-defining ancestor is parsed's root scope. Walks lexical + ancestors rather than trusting ``Scope.columns`` (which can include + correlated refs from inner subqueries). + + ``shadow_names``: identifiers defined at the same root scope as: + - CTE definitions, + - FROM/JOIN sources introduced by an explicit ``AS`` alias OR by a + Subquery/CTE source (always alias-like), + - schema/catalog parts of qualified FROM tables (``mydb.foo`` → + ``mydb`` shadows; ``a.b.foo`` → both ``a`` and ``b`` shadow). + Unaliased plain ``FROM customers`` does NOT shadow per spec wording + ("AS alias, CTE name, or schema name"). + """ + if isinstance(parsed, _STATEMENT_TYPES): + scope_id_to_type: dict[int, ScopeType] = {} + for scope in traverse_scope(parsed): + scope_id_to_type[id(scope.expression)] = scope.scope_type + root_scope_node_id = next( + (sid for sid, st in scope_id_to_type.items() if st == ScopeType.ROOT), + None, + ) + if root_scope_node_id is None: + return set(), set() + + root_col_ids: Set[int] = set() + for col in parsed.find_all(exp.Column): + node: Optional[exp.Expression] = col.parent + while node is not None: + if id(node) in scope_id_to_type: + if id(node) == root_scope_node_id: + root_col_ids.add(id(col)) + break + node = node.parent + + shadow_names: Set[str] = set() + for scope in traverse_scope(parsed): + if scope.scope_type != ScopeType.ROOT: + continue + shadow_names |= set(scope.cte_sources) + for src_name, source in scope.sources.items(): + if isinstance(source, exp.Table): + # Explicit AS alias. + if source.alias: + shadow_names.add(src_name) + # Schema / catalog qualifiers on the FROM table. + for part_key in ("db", "catalog"): + part = source.args.get(part_key) + if part is not None: + shadow_names.add(part.name) + else: + # Subquery / CTE-referenced source — always alias-like. + shadow_names.add(src_name) + break + return root_col_ids, shadow_names + return _root_scope_column_ids(parsed=parsed), set() + + +def _apply_dot_path_in_sql( + sql_text: Optional[str], *, location: str, model: Optional[SlayerModel], +) -> Tuple[Optional[str], List[NormalizationWarning]]: + """AST-based, scope-aware DOT_PATH_IN_SQL rewrite. + + Rewrites root-scope dotted refs (``customers.regions.name``) to the + ``__`` alias form (``customers__regions.name``) when the leading + segment is a known join target on ``model``. Refs in subqueries / + CTE-local scopes / set-op branches are left alone (scope-aware). + Refs whose leading segment matches a join target AND is also a + CTE / FROM-alias in the same scope are flagged ambiguous: no + rewrite, one warning carrying ``normalized="(ambiguous: ...)"``. + + Intermediate hops are not validated at normalize-time (no storage + access here); the contract is "first segment matches a join target + on the host model". Downstream join resolution catches an invalid + intermediate the same way it would for the canonical form. + """ + if not sql_text or model is None or not model.joins: + return sql_text, [] + + try: + statements = sqlglot.parse(sql_text) + except Exception: + return sql_text, [] + statements = [s for s in statements if s is not None] + if len(statements) != 1: + # Multi-statement (or empty) — slack input is contractually a single + # scalar expression / predicate. Leave alone. + return sql_text, [] + parsed = statements[0] + + join_target_names = {j.target_model for j in model.joins} + root_col_ids, shadow_names = _dot_path_root_scope_analysis(parsed=parsed) + + emitted: List[NormalizationWarning] = [] + changed = False + + for col in parsed.find_all(exp.Column): + if id(col) not in root_col_ids: + continue + parts = [p.name for p in col.parts] + if len(parts) < 3: + continue + first = parts[0] + if first not in join_target_names: + continue + original = ".".join(parts) + + if first in shadow_names: + payload = NormalizationWarning( + rule_id="DOT_PATH_IN_SQL", + original=original, + normalized="(ambiguous: shadowed by local alias or CTE — not rewritten)", + location=location, + rule_doc_url="docs/agent_input_slack.md#dot-path-in-sql", + ) + emitted.append(payload) + _warnings_module.warn( + SlayerNormalizationWarning(payload), stacklevel=2, + ) + continue + + new_table_name = "__".join(parts[:-1]) + leaf_name = parts[-1] + normalized = f"{new_table_name}.{leaf_name}" + col.set("catalog", None) + col.set("db", None) + col.set("table", exp.to_identifier(new_table_name)) + + payload = NormalizationWarning( + rule_id="DOT_PATH_IN_SQL", + original=original, + normalized=normalized, + location=location, + rule_doc_url="docs/agent_input_slack.md#dot-path-in-sql", + ) + emitted.append(payload) + _warnings_module.warn( + SlayerNormalizationWarning(payload), stacklevel=2, + ) + changed = True + + if not changed: + return sql_text, emitted + return parsed.sql(), emitted + + +# --------------------------------------------------------------------------- +# Top-level entry points +# --------------------------------------------------------------------------- + + +def normalize_query( + query: SlayerQuery, + *, + model: Optional[SlayerModel] = None, + custom_agg_names: Optional[frozenset[str]] = None, +) -> NormalizationResult: + """Apply all enabled slack rules to a ``SlayerQuery``. + + Returns the (possibly rewritten) query and the structured warnings. + Existing in-tree rewriters (notably + ``slayer.core.formula._rewrite_funcstyle_aggregations``) continue to + run during enrichment; in stage 6 they see canonical input and + silently no-op for any input this layer already rewrote. + """ + all_warnings: List[NormalizationWarning] = [] + + # Rule 1: FUNC_STYLE_AGG over Mode-B fields. + new_measures = [] + for i, m in enumerate(query.measures or []): + formula = getattr(m, "formula", None) + if isinstance(formula, str): + rewritten, ws = _apply_func_style_agg( + formula, + location=f"measures[{i}].formula", + custom_agg_names=custom_agg_names, + ) + all_warnings.extend(ws) + if rewritten != formula: + m = m.model_copy(update={"formula": rewritten}) + new_measures.append(m) + + new_filters: List[str] = [] + for i, f in enumerate(query.filters or []): + if isinstance(f, str): + rewritten, ws = _apply_func_style_agg( + f, + location=f"filters[{i}]", + custom_agg_names=custom_agg_names, + ) + all_warnings.extend(ws) + new_filters.append(rewritten) + else: + new_filters.append(f) + + query = query.model_copy(update={ + "measures": new_measures, + "filters": new_filters, + }) + + # Rule 2: MISPLACED_MEASURE. + query, ws = _apply_misplaced_measure(query, model=model) + all_warnings.extend(ws) + + # Rule 3: DOT_PATH_IN_SQL (stub in stage 6). + # Mode-A fields on the query itself are rare — most Mode-A lives on + # the model. Wiring is preserved so future activations need no + # plumbing changes. + + return NormalizationResult(query=query, warnings=all_warnings) + + +def _normalize_model_measures( + model: SlayerModel, *, custom_agg_names: Optional[frozenset[str]], +) -> Tuple[SlayerModel, List[NormalizationWarning]]: + """FUNC_STYLE_AGG over ``model.measures`` (Mode-B). See + :func:`normalize_model` for the ``custom_agg_names`` contract. + """ + if not model.measures: + return model, [] + if custom_agg_names is not None: + custom_names = custom_agg_names + else: + custom_names = frozenset(a.name for a in (model.aggregations or [])) + warnings: List[NormalizationWarning] = [] + new_measures = [] + for i, mm in enumerate(model.measures): + formula = mm.formula + rewritten, ws = _apply_func_style_agg( + formula, + location=f"measures[{i}].formula", + custom_agg_names=custom_names, + ) + warnings.extend(ws) + if rewritten != formula: + mm = mm.model_copy(update={"formula": rewritten}) + new_measures.append(mm) + return model.model_copy(update={"measures": new_measures}), warnings + + +def _normalize_column_dot_paths( + model: SlayerModel, +) -> Tuple[SlayerModel, List[NormalizationWarning]]: + """DOT_PATH_IN_SQL over ``Column.sql`` / ``Column.filter`` (Mode-A).""" + warnings: List[NormalizationWarning] = [] + new_columns = [] + changed = False + for i, c in enumerate(model.columns): + updates: dict = {} + if c.sql is not None: + rewritten_sql, ws = _apply_dot_path_in_sql( + c.sql, location=f"columns[{i}].sql", model=model, + ) + warnings.extend(ws) + if rewritten_sql != c.sql: + updates["sql"] = rewritten_sql + if c.filter is not None: + rewritten_filter, ws = _apply_dot_path_in_sql( + c.filter, location=f"columns[{i}].filter", model=model, + ) + warnings.extend(ws) + if rewritten_filter != c.filter: + updates["filter"] = rewritten_filter + if updates: + c = c.model_copy(update=updates) + changed = True + new_columns.append(c) + if changed: + model = model.model_copy(update={"columns": new_columns}) + return model, warnings + + +def _normalize_model_filter_dot_paths( + model: SlayerModel, +) -> Tuple[SlayerModel, List[NormalizationWarning]]: + """DOT_PATH_IN_SQL over ``SlayerModel.filters`` (Mode-A).""" + if not model.filters: + return model, [] + warnings: List[NormalizationWarning] = [] + new_filters = [] + changed = False + for i, f in enumerate(model.filters): + rewritten, ws = _apply_dot_path_in_sql( + f, location=f"filters[{i}]", model=model, + ) + warnings.extend(ws) + if rewritten != f: + changed = True + new_filters.append(rewritten if rewritten is not None else f) + if changed: + model = model.model_copy(update={"filters": new_filters}) + return model, warnings + + +def normalize_model( + model: SlayerModel, + *, + custom_agg_names: Optional[frozenset[str]] = None, +) -> NormalizationResult: + """Apply slack rules to a ``SlayerModel`` before persistence. + + Mode-A rewrites (``DOT_PATH_IN_SQL``) target ``Column.sql``, + ``Column.filter``, and ``SlayerModel.filters``. Mode-B rewrites + (``FUNC_STYLE_AGG``) target ``ModelMeasure.formula``. The rewrite + semantics match ``normalize_query``. + + ``custom_agg_names`` lets the caller supply the full reachable + aggregation set (model's own aggregations PLUS any defined on joined + models the caller has resolved through storage) so a funcstyle measure + over a joined-model custom aggregation gets rewritten — mirrors + ``normalize_query``'s param. Sharp edges: + + * ``custom_agg_names=None`` (default) → fall back to the model's own + ``aggregations`` (backward-compatible; matches the pre-DEV-1500 + behaviour for direct callers and tests that don't resolve joins). + * ``custom_agg_names=frozenset()`` → empty set is honoured AS-IS: the + model's-own fallback is suppressed. Pass an explicit empty frozenset + only when you want builtins-only recognition. + """ + all_warnings: List[NormalizationWarning] = [] + model, ws = _normalize_model_measures( + model, custom_agg_names=custom_agg_names, + ) + all_warnings.extend(ws) + if model.joins: + model, ws = _normalize_column_dot_paths(model) + all_warnings.extend(ws) + model, ws = _normalize_model_filter_dot_paths(model) + all_warnings.extend(ws) + return NormalizationResult(model=model, warnings=all_warnings) diff --git a/slayer/engine/path_resolution.py b/slayer/engine/path_resolution.py new file mode 100644 index 00000000..4e3e94b3 --- /dev/null +++ b/slayer/engine/path_resolution.py @@ -0,0 +1,105 @@ +"""Stage 3 (DEV-1450) — join-graph walker extracted from query_engine. + +Single source of truth for both dimension and cross-model-measure +resolution (DEV-1369 consolidated the two prior near-duplicates; +DEV-1450 stage 3 lifts the consolidated walker out of the +``SlayerQueryEngine`` class so it's directly importable by the new +binder modules without dragging the engine in). + +Existing call sites in ``slayer/engine/query_engine.py`` keep their +method signature unchanged via a thin shim that supplies +``self._resolve_model`` as the ``resolve_model`` callback. +""" + +from __future__ import annotations + +from typing import Any, Awaitable, Callable, Optional, Tuple + +from slayer.core.models import ModelJoin, SlayerModel + + +class NoJoinError(Exception): + """Sentinel raised by ``walk_join_chain`` when + ``strict_missing_join=False`` and a hop has no matching join. + + Lets lenient callers (dimension resolution) map a missing join to a + ``None`` return without re-walking the path. Strict callers + (cross-model measure resolution) use the ``ValueError`` branch + instead, which carries a more actionable available-joins message. + """ + + def __init__(self, hop_name: str) -> None: + super().__init__(f"no join target named {hop_name!r}") + self.hop_name = hop_name + + +# Type alias for the resolve_model callable. The signature mirrors +# ``SlayerQueryEngine._resolve_model`` so the engine method can be passed +# directly without an adapter. +ResolveModel = Callable[..., Awaitable[SlayerModel]] + + +async def walk_join_chain( + *, + source_model: SlayerModel, + hop_names: list[str], + resolve_model: ResolveModel, + named_queries: Optional[dict] = None, + strict_missing_join: bool = True, +) -> Tuple[SlayerModel, Optional[ModelJoin]]: + """Walk the join graph from ``source_model`` through ``hop_names``, + returning ``(terminal_model, first_join)``. + + Cycle detection: a hop name that already appears on the visited + stack (including ``source_model.name``) raises ``ValueError`` with + the offending path. + + Missing-join behavior: + + * ``strict_missing_join=True`` — raise ``ValueError`` listing the + available joins on the current model. Used by cross-model-measure + resolution where a missing join is a user error worth surfacing. + * ``strict_missing_join=False`` — raise ``NoJoinError`` so the + caller can map to a ``None`` return. Used by dimension resolution + where a missing intermediate join just means the column can't be + reached this way and the caller should keep looking. + + ``resolve_model`` is an async callable + ``(model_name=, named_queries=, prefer_data_source=) -> SlayerModel``. + Typically ``SlayerQueryEngine._resolve_model``. + """ + current_model = source_model + visited = {source_model.name} + first_join: Optional[ModelJoin] = None + + nq: Any = named_queries if named_queries is not None else {} + + for i, hop_name in enumerate(hop_names): + if hop_name in visited: + raise ValueError( + f"Circular join detected while resolving " + f"'{'.'.join(hop_names)}': '{hop_name}' already visited " + f"({' → '.join(visited)} → {hop_name})" + ) + join = next( + (j for j in current_model.joins if j.target_model == hop_name), + None, + ) + if join is None: + if strict_missing_join: + raise ValueError( + f"Model '{current_model.name}' has no join to " + f"'{hop_name}'. Available joins: " + f"{[j.target_model for j in current_model.joins]}" + ) + raise NoJoinError(hop_name) + if i == 0: + first_join = join + current_model = await resolve_model( + model_name=hop_name, + named_queries=nq, + prefer_data_source=current_model.data_source or None, + ) + visited.add(hop_name) + + return current_model, first_join diff --git a/slayer/engine/planned.py b/slayer/engine/planned.py new file mode 100644 index 00000000..5f329396 --- /dev/null +++ b/slayer/engine/planned.py @@ -0,0 +1,371 @@ +"""Stage 7a.1 (DEV-1450) — typed plan shapes consumed by the SQL generator. + +A ``PlannedQuery`` is the final, fully resolved plan that the SQL +generator (stage 7b) compiles to SQL. The plan carries everything a +renderer needs: row slots, aggregate slots, cross-model aggregate +sub-plans, transform layers, filter routing, projection / order / +limit, and an emitted ``StageSchema`` for downstream stages to bind +against. + +Identity-bearing structure is in ``slayer/core/keys.py`` (the +``ValueKey`` family); the planner here associates each key with a +``SlotId`` and the rendering metadata (alias, hidden, label). + +The planning logic that produces a ``PlannedQuery`` lives in other +7a substages — ``planning.py`` (ValueRegistry, TransformLowerer, +ProjectionPlanner), ``cross_model_planner.py`` (I1 strategy), +``stage_planner.py`` (multi-stage DAG). This file is the typed +target. + +These types are dormant in stage 7a — no engine code consumes them +yet. Stage 7b's engine cutover routes through them. +""" + +from __future__ import annotations + +from typing import List, Optional, Tuple + +from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator + +from slayer.core.enums import DataType, JoinType +from slayer.core.errors import UnreachableFilterDroppedWarning +from slayer.core.format import NumberFormat +from slayer.core.keys import Phase, ValueKey +from slayer.core.models import SlayerModel +from slayer.core.scope import StageSchema +from slayer.engine.binding import BoundExpr # re-exported below + + +# Opaque identifier types — kept as plain ``str`` for now. SlotId is +# allocated by the planner's ValueRegistry; BoundFilterId by the +# FilterBinder. The string form keeps tracebacks readable and lets +# tests assert on them without exotic comparisons. +SlotId = str +BoundFilterId = str + + +# --------------------------------------------------------------------------- +# BoundExpr — re-exported from slayer.engine.binding (DEV-1450 stage 7b.6). +# --------------------------------------------------------------------------- +# +# Until stage 7b.6 the planned-side BoundExpr was a separate scaffold +# Pydantic class with an optional ``sql_text`` cache. The binder +# produced its own ``BoundExpr`` shape, so ``ValueSlot.expression`` and +# ``FilterPhase.expression`` could not store binder output directly +# without type unification (Codex HIGH F2 in the earlier round). 7b.6 +# folds the two: the binder's ``BoundExpr(value_key=ValueKey)`` is the +# canonical shape. The render artifact ``sql_text`` is dropped — the +# generator renders from the typed ``value_key`` against the slot +# registry, not a cached string. +__all__ = [ + "BoundExpr", + "BoundFilterId", + "CrossModelAggregatePlan", + "FilterPhase", + "JoinRequirement", + "OrderEntry", + "PlannedQuery", + "SlotId", + "TransformLayer", + "ValueSlot", +] + + +# --------------------------------------------------------------------------- +# ValueSlot +# --------------------------------------------------------------------------- + + +class ValueSlot(BaseModel): + """One materialised slot in a ``PlannedQuery`` (P6). + + Identity comes from ``key`` (a ``ValueKey`` from + ``slayer.core.keys``). Two structurally equal keys share one slot. + Rendering metadata (alias, hidden, label, type) lives here, not on + the key. + + ``declared_name`` is either the user-supplied ``name`` or the + canonical form derived from the formula. ``public_name`` is the + user-facing alias when the slot is part of the public projection + (None for hidden slots). ``public_aliases`` carries multiple + aliases when the same structural key was declared with multiple + explicit names (P4 / C13). + """ + + model_config = ConfigDict(arbitrary_types_allowed=True) + + id: SlotId + key: ValueKey + declared_name: str + public_name: Optional[str] = None + public_aliases: List[str] = Field(default_factory=list) + hidden: bool = False + phase: Phase + label: Optional[str] = None + type: Optional[DataType] = None + expression: Optional[BoundExpr] = None + # DEV-1452 Stage B decision #8 — typed format / description propagated + # by the planner from the source ``ModelMeasure`` / ``Column`` and + # ``_infer_aggregated_format``. Consumed by the migrated + # ``_expand_query_backed_model`` (via the public ``StageSchema``) so + # query-backed virtual-model columns carry the same display metadata + # the legacy enrichment pipeline produced. + format: Optional[NumberFormat] = None + description: Optional[str] = None + + @model_validator(mode="after") + def _hidden_invariant(self) -> "ValueSlot": + # Hidden slots are materialised but never surfaced — they must + # not carry a public_name or public_aliases, otherwise the + # generator would emit them in the public projection. + if self.hidden and (self.public_name is not None or self.public_aliases): + raise ValueError( + f"ValueSlot(id={self.id!r}) is hidden but carries " + f"public_name={self.public_name!r} / " + f"public_aliases={self.public_aliases!r}; hidden slots " + f"must have public_name=None and public_aliases=[]." + ) + return self + + +# --------------------------------------------------------------------------- +# JoinRequirement +# --------------------------------------------------------------------------- + + +class JoinRequirement(BaseModel): + """One hop in a cross-model join chain. + + Mirrors the shape of ``slayer.core.models.ModelJoin`` but is + rooted on the typed-plan side — the planner builds these from + resolved bundle models so the SQL generator never re-walks the + model graph. + """ + + source_model: str + target_model: str + join_pairs: List[List[str]] + join_type: JoinType = JoinType.LEFT + + @field_validator("join_pairs") + @classmethod + def _non_empty(cls, v: List[List[str]]) -> List[List[str]]: + if not v: + raise ValueError("join_pairs must be non-empty") + for i, pair in enumerate(v): + if len(pair) != 2 or not all(isinstance(s, str) and s for s in pair): + raise ValueError( + f"join_pairs[{i}] must be [source_dim, target_dim] " + f"with non-empty strings, got {pair!r}" + ) + return v + + +# --------------------------------------------------------------------------- +# CrossModelAggregatePlan +# --------------------------------------------------------------------------- + + +class CrossModelAggregatePlan(BaseModel): + """Plan for one cross-model aggregate slot (P3 / I1). + + The strategy that populated this plan (today's isolated-CTE form + or a future alternative — see ``cross_model_planner.py``) lives + outside this struct; this is the typed result, not the algorithm. + + Filter routing is route-explicit so the SQL generator (stage 7b) + can render each route without re-classifying: + - ``where_filter_ids`` — host filters propagated to the CTE's WHERE + (decision-table rows: host-local-but-targeted, joined-target-path). + - ``having_filter_ids`` — host filters propagated as HAVING (decision- + table row: cross-model agg-ref on the same target). + - ``target_model_filters`` — the target model's own + ``SlayerModel.filters`` (always-applied WHERE). + ``applied_filter_ids`` is the audit union of where + having for + backward compatibility with the spec's external surface. + + ``hidden=True`` is used for order-only / filter-only refs whose + aggregate value is materialised but not surfaced in the public + projection; ``public_alias`` is ``None`` in that case. + """ + + model_config = ConfigDict(arbitrary_types_allowed=True) + + aggregate_slot_id: SlotId + target_model: str + datasource: str + join_chain: List[JoinRequirement] + join_back_pairs: List[Tuple[ValueKey, ValueKey]] = Field(default_factory=list) + cte_stage_schema: StageSchema + shared_grain_slots: List[SlotId] + applied_filter_ids: List[BoundFilterId] = Field(default_factory=list) + where_filter_ids: List[BoundFilterId] = Field(default_factory=list) + having_filter_ids: List[BoundFilterId] = Field(default_factory=list) + target_model_filters: List[str] = Field(default_factory=list) + dropped_filter_warnings: List[UnreachableFilterDroppedWarning] = Field(default_factory=list) + hidden: bool = False + public_alias: Optional[str] = None + + # DEV-1450 stage 7b.15e (C1) — re-rooted sub-plan. When the host query + # carries dimensions that are reachable from the target by re-rooting + # through the target's join graph (the legacy ``_build_rerooted_enriched`` + # case), the cross-model CTE is rendered as a full nested ``PlannedQuery`` + # rooted at the target (FROM target + joins), preserving the host + # dimension grain instead of collapsing to a scalar CROSS JOIN. ``None`` + # keeps the forward-path "FROM bare target" rendering. + # + # ``rerooted_grain_pairs`` maps (host_dim_slot_id, rerooted_dim_slot_id) + # for the combined LEFT JOIN ON; the generator resolves each side's SQL + # alias independently (host alias vs sub-plan alias need not match). + # ``rerooted_agg_slot_id`` is the sub-plan slot id of the local aggregate; + # the combined SELECT projects it ``AS`` the canonical / public alias. + rerooted_plan: Optional["PlannedQuery"] = None + rerooted_grain_pairs: List[Tuple[SlotId, SlotId]] = Field(default_factory=list) + rerooted_agg_slot_id: Optional[SlotId] = None + + # DEV-1503 — host-rooted CTE for a cross-model-FILTERED local measure. + # The cross-model planner has two distinct cases that produce a nested + # ``rerooted_plan``: + # + # * Cross-model aggregate re-rooting (``target_model`` is the join target, + # the sub-plan is rooted at the target so the target's own join graph + # reaches host dims that the forward-path CTE collapses): ``cte_root_model`` + # stays ``None`` and the renderer uses ``target_model`` for the FROM / + # joins (existing pre-DEV-1503 behaviour). + # * Filtered-local isolation (``AggregateKey.source.path`` is empty but the + # measure's ``Column.filter`` crosses a join, so the aggregate must + # evaluate in its own CTE rooted at the HOST + the filter-target join): + # ``cte_root_model`` is set to the HOST model name and the renderer uses + # the sub-plan's own ``source_relation`` / ``render_source_model`` for + # the FROM. ``target_model`` is conventionally set to the host name in + # this case but the renderer reads ``cte_root_model`` to disambiguate. + cte_root_model: Optional[str] = None + + +# --------------------------------------------------------------------------- +# TransformLayer +# --------------------------------------------------------------------------- + + +class TransformLayer(BaseModel): + """One transform layer in the planned query. + + Window / temporal transforms (``cumsum``, ``time_shift``, + ``rank``, ``lag``, ``lead``, ...) are grouped into layers so the + SQL generator can emit them in the right order (window functions + in an inner SELECT, time_shift as a self-join CTE, etc.). The + layer carries the slot ids that belong to it; rendering details + are decided by the generator per ``op``. + """ + + op: str + slot_ids: List[SlotId] + + +# --------------------------------------------------------------------------- +# FilterPhase +# --------------------------------------------------------------------------- + + +class FilterPhase(BaseModel): + """A bound filter expression routed to its phase (P8). + + ``phase`` is the maximum phase of the slots the filter + references: ROW → WHERE, AGGREGATE → HAVING, POST → post-filter + on the outer SELECT. + + Two carrier modes, mutually exclusive in practice: + + * ``expression`` is a typed ``BoundExpr`` — used for the Mode-B + DSL filters bound by ``bind_filter`` and the planner-emitted + ``BetweenKey`` for ``TimeDimension.date_range``. The renderer + walks the typed value-key tree. + * ``text`` is a Mode-A SQL fragment — used for + ``SlayerModel.filters`` (always-applied WHERE). The renderer + qualifies bare-identifier column refs in ``text_columns`` with + the source-relation alias and emits the result verbatim + (matching legacy ``_build_where_and_having`` qualification). + """ + + id: BoundFilterId + phase: Phase + text: Optional[str] = None + text_columns: Tuple[str, ...] = () + expression: Optional[BoundExpr] = None + + +# --------------------------------------------------------------------------- +# OrderEntry +# --------------------------------------------------------------------------- + + +class OrderEntry(BaseModel): + """One entry in the ORDER BY of a planned query.""" + + slot_id: SlotId + direction: str # "asc" or "desc" + + @field_validator("direction") + @classmethod + def _validate_direction(cls, v: str) -> str: + if v not in ("asc", "desc"): + raise ValueError( + f"OrderEntry.direction must be 'asc' or 'desc', got {v!r}" + ) + return v + + +# --------------------------------------------------------------------------- +# PlannedQuery +# --------------------------------------------------------------------------- + + +class PlannedQuery(BaseModel): + """The fully typed plan for one query stage (P7). + + Consumed by the SQL generator (stage 7b). Carries everything + needed to emit SQL without re-walking the model graph. + + ``stage_schema`` is the projection emitted by this stage — + downstream stages bind against it (P6). Top-level queries that + aren't part of a multi-stage DAG can leave ``stage_schema`` as + ``None``. + """ + + model_config = ConfigDict(arbitrary_types_allowed=True) + + source_relation: str + join_plan: List[JoinRequirement] = Field(default_factory=list) + row_slots: List[ValueSlot] = Field(default_factory=list) + aggregate_slots: List[ValueSlot] = Field(default_factory=list) + cross_model_aggregate_plans: List[CrossModelAggregatePlan] = Field(default_factory=list) + combined_expression_slots: List[ValueSlot] = Field(default_factory=list) + transform_layers: List[TransformLayer] = Field(default_factory=list) + filters_by_phase: List[FilterPhase] = Field(default_factory=list) + projection: List[SlotId] = Field(default_factory=list) + order: List[OrderEntry] = Field(default_factory=list) + limit: Optional[int] = None + offset: Optional[int] = None + stage_schema: Optional[StageSchema] = None + # Stage 7b.10 — the slot id of the active TD (resolved via + # ``_resolve_main_time_dimension``). ``None`` when the stage has no + # time dimension. Time-needing transforms (cumsum / lag / lead / + # first / last / time_shift / consecutive_periods) carry this slot's + # key in ``TransformKey.time_key``; the generator uses it for the + # ``ORDER BY`` clause of the OVER expression. + active_time_dimension_slot_id: Optional[SlotId] = None + # DEV-1450 stage 7b.15d — the concrete ``SlayerModel`` this stage renders + # against, carried from the planner so the generator binds the stage's + # FROM / joins against the SAME model the binder used. For a multi-stage + # DAG this is the stage's OWN source (e.g. ``orders`` for a stage the root + # never reads from), a ModelExtension overlay, or a synthetic model over a + # sibling stage's CTE. ``None`` for a StageSchema-scoped chain stage (the + # generator builds a synthetic model from the upstream schema) and for a + # plain single-model query (the generator uses ``bundle.source_model``). + render_source_model: Optional[SlayerModel] = None + + +# ``CrossModelAggregatePlan.rerooted_plan`` is a forward reference to +# ``PlannedQuery`` (defined above only after the CMA plan). Resolve it now +# that both classes exist (DEV-1450 stage 7b.15e, C1). +CrossModelAggregatePlan.model_rebuild() diff --git a/slayer/engine/planning.py b/slayer/engine/planning.py new file mode 100644 index 00000000..b909c7af --- /dev/null +++ b/slayer/engine/planning.py @@ -0,0 +1,735 @@ +"""Stage 7a.6 (DEV-1450) — ValueRegistry, TransformLowerer, ProjectionPlanner. + +Three composable concerns: + +* ``ValueRegistry`` interns ``ValueKey``s by structural identity. Two + structurally-equal keys share one ``ValueSlot`` (P2). The same key + declared with multiple ``name``s accumulates multiple + ``public_aliases`` on a single slot (P4 / C13). Alias collisions + with source columns / duplicate names are rejected per DEV-1443. + +* ``desugar_change`` / ``desugar_change_pct`` lower sugar transforms + into their underlying form. The inner operand keeps the same + structural identity across all occurrences (DEV-1446) so the + ValueRegistry interns it once. ``partition_by`` threads through to + the underlying ``time_shift`` (C6). + +* ``ProjectionPlanner`` allocates slots for declared measures and + creates hidden slots for refs that appear ONLY in order/filter. + Hidden slots are materialised but trimmed from the public projection. + +Dormant in 7a — no engine wiring. Stage 7a.7's ``stage_planner.py`` +composes these with the cross-model planner to build a +``PlannedQuery``. +""" + +from __future__ import annotations + +from typing import Dict, FrozenSet, List, Optional + +from pydantic import BaseModel, ConfigDict, Field + +from slayer.core.enums import DataType +from slayer.core.format import NumberFormat +from slayer.core.errors import ( + CanonicalAliasShadowsColumnError, + DuplicateMeasureNameError, + MeasureNameCollidesWithColumnError, +) +from slayer.core.keys import ( + AggregateKey, + ArithmeticKey, + BetweenKey, + ColumnKey, + ColumnSqlKey, + InKey, + LiteralKey, + Phase, + ScalarCallKey, + StarKey, + TimeTruncKey, + TransformKey, + ValueKey, + column_leaf, + column_path, + normalize_scalar, +) +from slayer.core.refs import agg_kwarg_canonical_str, canonical_agg_name +from slayer.engine.binding import BoundExpr, BoundFilter +from slayer.engine.planned import SlotId, ValueSlot + +__all__ = [ + "DeclaredMeasure", + "OrderSpec", + "ProjectionPlan", + "ProjectionPlanner", + "ValueRegistry", + "desugar_change", + "desugar_change_pct", + "filter_referenced_slot_ids", + "lower_sugar_transforms", +] + + +# --------------------------------------------------------------------------- +# ValueRegistry +# --------------------------------------------------------------------------- + + +def _fill_missing_metadata( + *, + slot: ValueSlot, + updates: Dict, + label: Optional[str] = None, + type: Optional[DataType] = None, + format: Optional[NumberFormat] = None, + description: Optional[str] = None, +) -> None: + """Populate ``updates`` with each per-field value that is set on the + incoming intern call AND missing on the existing slot (DEV-1452 round-2 + refactor of the previous if-chain to keep + :meth:`ValueRegistry._merge_into_existing` under the S3776 complexity + cap). Mutates ``updates`` in place; never overrides an already-set + field on the slot. + """ + for field_name, new_value in ( + ("label", label), + ("type", type), + ("format", format), + ("description", description), + ): + if getattr(slot, field_name) is None and new_value is not None: + updates[field_name] = new_value + + +class ValueRegistry: + """Interns ``ValueKey``s by structural identity into ``ValueSlot``s. + + Constructor takes ``source_column_names``: the set of column names + on the host model used for the alias-collision validations + (``MeasureNameCollidesWithColumnError``, + ``CanonicalAliasShadowsColumnError``). Pass an empty set when the + host doesn't expose column names (or when the validations should + skip — e.g., in unit tests for the registry in isolation). + """ + + def __init__( + self, + *, + source_column_names: Optional[FrozenSet[str]] = None, + host_model_name: str = "(host)", + ) -> None: + self._source_columns: FrozenSet[str] = ( + source_column_names or frozenset() + ) + self._host_model_name = host_model_name + self._slots: Dict[SlotId, ValueSlot] = {} + self._by_key: Dict[ValueKey, SlotId] = {} + self._declared_names: Dict[str, SlotId] = {} + self._counter = 0 + + def _next_id(self) -> SlotId: + self._counter += 1 + return f"s{self._counter}" + + def intern( + self, + *, + key: ValueKey, + declared_name: str, + phase: Phase, + public_name: Optional[str] = None, + canonical_alias: Optional[str] = None, + hidden: bool = False, + label: Optional[str] = None, + type: Optional[DataType] = None, + expression: Optional["BoundExpr"] = None, + format: Optional[NumberFormat] = None, + description: Optional[str] = None, + ) -> SlotId: + # Alias-collision validations (P4 / DEV-1443). + # Exemption: a dimension whose public name IS its own column + # name (``ColumnKey(path=(), leaf=X)`` declared as ``X``) is the + # column, not a rename of it — collision check skipped. Same + # exemption for a local ``TimeTruncKey`` over that same column + # since a time dimension on ``created_at`` projects the + # (truncated) ``created_at`` column rather than introducing a + # new alias. DEV-1450 stage 7b.13: also exempt + # ``ColumnSqlKey(model=..., column_name=X)`` declared as ``X`` + # -- a derived column (``Column.sql`` set) selected as a + # dimension projects the column unchanged, identical to the + # plain-column case. + is_self_named_dimension = ( + isinstance(key, ColumnKey) + and key.path == () + and public_name == key.leaf + ) or ( + isinstance(key, ColumnSqlKey) + and key.path == () + and public_name == key.column_name + ) or ( + isinstance(key, TimeTruncKey) + and column_path(key.column) == () + and public_name == column_leaf(key.column) + ) + # An UNNAMED ``*:`` (StarKey source) re-aggregation is exempt: its + # canonical alias (``_count``) is a structural marker, not a column + # reference — ``COUNT(*)`` reads no column, so it can't be ambiguous + # with a same-named source column. This is the chain re-count case + # (``*:count`` over a stage that already projects ``_count``). An + # EXPLICIT user name that collides still raises (canonical_alias is set + # only on a rename, so it stays None here for the unnamed form). + is_unnamed_star_agg = ( + isinstance(key, AggregateKey) + and isinstance(getattr(key, "source", None), StarKey) + and canonical_alias is None + ) + if ( + public_name is not None + and public_name in self._source_columns + and not is_self_named_dimension + and not is_unnamed_star_agg + ): + raise MeasureNameCollidesWithColumnError( + name=public_name, model=self._host_model_name, + ) + if ( + canonical_alias is not None + and canonical_alias in self._source_columns + ): + raise CanonicalAliasShadowsColumnError( + formula=declared_name, + canonical=canonical_alias, + model=self._host_model_name, + ) + + existing_sid = self._by_key.get(key) + if existing_sid is not None: + return self._merge_into_existing( + existing_sid=existing_sid, + public_name=public_name, + declared_name=declared_name, + hidden=hidden, + label=label, + type=type, + format=format, + description=description, + ) + + # Fresh slot. Check declared_name collision against a different key. + if public_name is not None: + owner = self._declared_names.get(public_name) + if owner is not None: + raise DuplicateMeasureNameError( + name=public_name, + occurrences=[ + self._slots[owner].declared_name, + declared_name, + ], + ) + + sid = self._next_id() + public_aliases = [public_name] if public_name is not None else [] + slot = ValueSlot( + id=sid, + key=key, + declared_name=declared_name, + public_name=public_name, + public_aliases=public_aliases, + hidden=hidden, + phase=phase, + label=label, + type=type, + expression=expression if expression is not None else BoundExpr(value_key=key), + format=format, + description=description, + ) + self._slots[sid] = slot + self._by_key[key] = sid + if public_name is not None: + self._declared_names[public_name] = sid + return sid + + def _merge_into_existing( + self, + *, + existing_sid: SlotId, + public_name: Optional[str], + declared_name: str, + hidden: bool, + label: Optional[str] = None, + type: Optional[DataType] = None, + format: Optional[NumberFormat] = None, + description: Optional[str] = None, + ) -> SlotId: + slot = self._slots[existing_sid] + updates: Dict = {} + if public_name is not None and public_name not in slot.public_aliases: + owner = self._declared_names.get(public_name) + if owner is not None and owner != existing_sid: + raise DuplicateMeasureNameError( + name=public_name, + occurrences=[ + self._slots[owner].declared_name, + declared_name, + ], + ) + updates["public_aliases"] = list(slot.public_aliases) + [public_name] + if slot.hidden: + updates["hidden"] = False + updates["public_name"] = public_name + self._declared_names[public_name] = existing_sid + elif not hidden and slot.hidden and public_name is None: + # Re-intern as non-hidden — promote to public. + updates["hidden"] = False + # Codex: when a hidden slot is promoted to public, carry the + # display metadata supplied by the public re-intern. Only fill + # missing fields — never overwrite metadata the first intern + # already supplied. + _fill_missing_metadata( + slot=slot, + updates=updates, + label=label, + type=type, + format=format, + description=description, + ) + if updates: + new_slot = slot.model_copy(update=updates) + self._slots[existing_sid] = new_slot + return existing_sid + + def get(self, slot_id: SlotId) -> ValueSlot: + return self._slots[slot_id] + + def find_by_key(self, key: ValueKey) -> Optional[SlotId]: + return self._by_key.get(key) + + @property + def slots(self) -> List[ValueSlot]: + return list(self._slots.values()) + + +# --------------------------------------------------------------------------- +# TransformLowerer +# --------------------------------------------------------------------------- + + +def desugar_change(key: TransformKey) -> ArithmeticKey: + """``change(x)`` → ``x - time_shift(x, periods=-1, [partition_by=…])``. + + The inner ``x`` is identity-preserving — the ``ArithmeticKey`` and + the ``TransformKey`` use the SAME ``ValueKey`` instance, so a + downstream ValueRegistry interns it as one slot (DEV-1446). + + ``partition_by`` (the binder put it on ``key.partition_keys``) + threads through to the underlying ``time_shift`` (C6). ``periods`` + is fixed at ``-1`` (one period back) because ``change`` has no + user-tunable offset. + """ + if key.op != "change": + raise ValueError( + f"desugar_change expected op='change', got {key.op!r}." + ) + inner = key.input + shifted = TransformKey( + op="time_shift", + input=inner, + kwargs=(("periods", normalize_scalar(-1)),), + partition_keys=key.partition_keys, + time_key=key.time_key, + ) + return ArithmeticKey(op="-", operands=(inner, shifted)) + + +def lower_sugar_transforms(key: ValueKey) -> ValueKey: + """Recursively lower ``change`` / ``change_pct`` TransformKeys to + their desugared arithmetic form, preserving the inner aggregate's + structural identity (DEV-1446). Other ValueKey shapes are walked + but otherwise unchanged. + + The desugar functions preserve ``partition_keys`` / ``time_key`` on + the resulting ``time_shift`` TransformKey (DEV-1450 C6), so + ``change(amount:sum, partition_by=region)`` lowers to + ``amount:sum - time_shift(amount:sum, partition_by=region)``. + """ + if isinstance(key, TransformKey): + new_input = lower_sugar_transforms(key.input) + if new_input is not key.input: + key = key.model_copy(update={"input": new_input}) + if key.op == "change": + return desugar_change(key) + if key.op == "change_pct": + return desugar_change_pct(key) + return key + if isinstance(key, ArithmeticKey): + new_ops = tuple(lower_sugar_transforms(op) for op in key.operands) + if all(a is b for a, b in zip(new_ops, key.operands)): + return key + return ArithmeticKey(op=key.op, operands=new_ops) + if isinstance(key, ScalarCallKey): + new_args = tuple( + lower_sugar_transforms(a) + if isinstance( + a, _SLOTTABLE_KIND + (ArithmeticKey, ScalarCallKey, BetweenKey), + ) + else a + for a in key.args + ) + if all(a is b for a, b in zip(new_args, key.args)): + return key + return ScalarCallKey(name=key.name, args=new_args) + if isinstance(key, BetweenKey): + new_col = lower_sugar_transforms(key.column) + new_low = lower_sugar_transforms(key.low) + new_high = lower_sugar_transforms(key.high) + if ( + new_col is key.column + and new_low is key.low + and new_high is key.high + ): + return key + return BetweenKey(column=new_col, low=new_low, high=new_high) + if isinstance(key, InKey): + # DEV-1475: ``InKey.values`` is a literal-only tuple, so it + # carries no sugar to lower; only the LHS column can host a + # rewritable transform. Rebuild only if the column changed. + new_col = lower_sugar_transforms(key.column) + if new_col is key.column: + return key + return InKey(column=new_col, values=key.values, negated=key.negated) + return key + + +def desugar_change_pct(key: TransformKey) -> ArithmeticKey: + """``change_pct(x)`` → ``(x - time_shift(x, periods=-1)) / + NULLIF(time_shift(x, periods=-1), 0)``. + + The divisor is wrapped in ``NULLIF(..., 0)`` so a zero prior-period + value yields NULL instead of a divide-by-zero error / Inf. Same + identity-preservation as ``desugar_change`` — numerator and divisor + share the one ``shifted`` ValueKey instance. + """ + if key.op != "change_pct": + raise ValueError( + f"desugar_change_pct expected op='change_pct', got {key.op!r}." + ) + inner = key.input + shifted = TransformKey( + op="time_shift", + input=inner, + kwargs=(("periods", normalize_scalar(-1)),), + partition_keys=key.partition_keys, + time_key=key.time_key, + ) + numerator = ArithmeticKey(op="-", operands=(inner, shifted)) + guarded_divisor = ScalarCallKey( + name="nullif", args=(shifted, normalize_scalar(0)), + ) + return ArithmeticKey(op="/", operands=(numerator, guarded_divisor)) + + +# --------------------------------------------------------------------------- +# ProjectionPlanner +# --------------------------------------------------------------------------- + + +class DeclaredMeasure(BaseModel): + """One declared measure on a query. + + ``bound`` is the binder's output. ``declared_name`` is the canonical + or user-supplied name. ``public_name`` is the user-facing alias — + set when the user supplied an explicit ``name`` on the measure spec. + + DEV-1452 Stage B decisions #2 + #8: ``type``, ``format``, and + ``description`` carry typed display + slot metadata from the source + ``ModelMeasure`` / ``Column`` so the public slot retains the same + contract the legacy enrichment pipeline produced. ``type`` mirrors + the legacy ``EnrichedMeasure.type`` (count → INT, avg → DOUBLE, + sum/min/max → source column type). + """ + + model_config = ConfigDict(arbitrary_types_allowed=True) + + bound: BoundExpr + declared_name: str + public_name: Optional[str] = None + label: Optional[str] = None + canonical_alias: Optional[str] = None + type: Optional[DataType] = None + format: Optional[NumberFormat] = None + description: Optional[str] = None + + +class OrderSpec(BaseModel): + """One ORDER BY entry on a query.""" + + model_config = ConfigDict(arbitrary_types_allowed=True) + + bound: BoundExpr + direction: str = "asc" + + +class ProjectionPlan(BaseModel): + """ProjectionPlanner output: registry + projection order + filters / order.""" + + model_config = ConfigDict(arbitrary_types_allowed=True) + + registry: "ValueRegistry" + public_projection: List[SlotId] = Field(default_factory=list) + filters: List[BoundFilter] = Field(default_factory=list) + order: List["OrderSpec"] = Field(default_factory=list) + + +_SLOTTABLE_KIND = ( + ColumnKey, ColumnSqlKey, AggregateKey, TransformKey, TimeTruncKey, +) + + +def _iter_slot_deps(key: ValueKey): + """Yield only ``ValueKey``s that need a materialised slot. + + Skips composite-only nodes that the SQL generator inlines: + ``ArithmeticKey`` (operators), ``ScalarCallKey`` (function calls + inlined into SELECT / WHERE), ``LiteralKey``, ``StarKey``. Stops + at ``AggregateKey`` (its inner ``source`` ColumnKey is materialised + inside the aggregate, not as a separate slot). Recurses into + ``TransformKey.input`` so a nested aggregate inside a transform + gets its own hidden slot. + + ``TimeTruncKey`` is itself the materialised slot (the generator + emits the DATE_TRUNC at SELECT time); the inner ColumnKey is not + yielded as a separate dependency — adding a time dimension must + not auto-add the raw column as an output (matches legacy). + """ + if isinstance(key, AggregateKey): + yield key + return + if isinstance(key, TransformKey): + yield key + yield from _iter_slot_deps(key.input) + # Transform aux deps: partition_keys and time_key must be + # materialised as their own slots so the SQL generator (slice + # 7b.10 / 7b.11) can render PARTITION BY / ORDER BY against + # named SELECT projections instead of re-walking the model + # graph. + for pk in key.partition_keys: + yield from _iter_slot_deps(pk) + if key.time_key is not None: + yield from _iter_slot_deps(key.time_key) + return + if isinstance(key, (ColumnKey, ColumnSqlKey, TimeTruncKey)): + yield key + return + if isinstance(key, ArithmeticKey): + for op in key.operands: + yield from _iter_slot_deps(op) + return + if isinstance(key, ScalarCallKey): + for arg in key.args: + if isinstance( + arg, + _SLOTTABLE_KIND + (ArithmeticKey, ScalarCallKey, BetweenKey), + ): + yield from _iter_slot_deps(arg) + return + if isinstance(key, BetweenKey): + # BetweenKey is not itself a slot — the generator inlines it + # into WHERE. Recurse into the column / low / high so the + # underlying ColumnKey shows up as a referenced slot for the + # cross-model routing / hidden-slot pass (Codex F4). + yield from _iter_slot_deps(key.column) + yield from _iter_slot_deps(key.low) + yield from _iter_slot_deps(key.high) + if isinstance(key, InKey): + # DEV-1475: InKey, like BetweenKey, is inlined into WHERE by the + # generator (no public slot of its own). Surface its LHS column + # for cross-model routing / hidden-slot collection; the literal + # RHS values are never slottable. + yield from _iter_slot_deps(key.column) + # StarKey, LiteralKey — never slottable on their own. + + +class ProjectionPlanner: + """Allocate slots for declared measures + hidden slots for refs only + used in order/filter.""" + + def plan( + self, + *, + measures: List[DeclaredMeasure], + filters: List[BoundFilter], + order: List[OrderSpec], + source_column_names: Optional[FrozenSet[str]] = None, + host_model_name: str = "(host)", + ) -> ProjectionPlan: + registry = ValueRegistry( + source_column_names=source_column_names, + host_model_name=host_model_name, + ) + public_projection: List[SlotId] = [] + for m in measures: + sid = registry.intern( + key=m.bound.value_key, + declared_name=m.declared_name, + public_name=m.public_name, + canonical_alias=m.canonical_alias, + phase=m.bound.phase, + label=m.label, + type=m.type, + format=m.format, + description=m.description, + ) + public_projection.append(sid) + # Materialise any auxiliary slot-worthy deps of the measure + # as hidden slots (e.g. the inner AggregateKey of a transform, + # the partition columns, the time_key column). These are + # rendered by the generator into the inner SELECT but not + # surfaced in the public projection. + for dep in _iter_slot_deps(m.bound.value_key): + if dep == m.bound.value_key: + continue + if registry.find_by_key(dep) is None: + registry.intern( + key=dep, + declared_name=_canonical_name(dep), + hidden=True, + phase=dep.phase, + ) + + # Filter and order share the same dependency-selection rule: walk + # the bound expression, intern each slot-worthy key as a hidden + # slot if not already present. + for f in filters: + for dep in _iter_slot_deps(f.value_key): + if registry.find_by_key(dep) is None: + registry.intern( + key=dep, + declared_name=_canonical_name(dep), + hidden=True, + phase=dep.phase, + ) + + for o in order: + for dep in _iter_slot_deps(o.bound.value_key): + if registry.find_by_key(dep) is None: + registry.intern( + key=dep, + declared_name=_canonical_name(dep), + hidden=True, + phase=dep.phase, + ) + + return ProjectionPlan( + registry=registry, + public_projection=public_projection, + filters=filters, + order=order, + ) + + +def _canonical_name(key: ValueKey) -> str: # NOSONAR(S3776) — sequential isinstance dispatch over the closed ValueKey union; each branch is the per-type canonical-name contract. Extracting per-type helpers would scatter the contract. + """Best-effort canonical name for a hidden slot. + + Mirrors the public-alias canonical form used by the engine + elsewhere: ``revenue:sum`` → ``revenue_sum``; ``*:count`` → + ``_count``; ``customers.regions.name`` → flattened ``customers__regions__name``. + """ + if isinstance(key, ColumnKey): + return "__".join(key.path + (key.leaf,)) + if isinstance(key, ColumnSqlKey): + prefix = "__".join(key.path) + "__" if key.path else "" + return f"{prefix}{key.column_name}" + if isinstance(key, TimeTruncKey): + # Legacy alias contract: granularity is encoded in the SQL + # DATE_TRUNC, not in the alias. + return _canonical_name(key.column) + if isinstance(key, AggregateKey): + # DEV-1501: include args/kwargs so parametric aggregates over the + # same value column (``revenue:last(created_at)`` vs + # ``revenue:last(updated_at)``, ``revenue:percentile(p=0.5)`` vs + # ``revenue:percentile(p=0.95)``) get DISTINCT declared names — + # mirrors the cross-model parametric P10 exception. Without this, + # two hidden parametric aggregates collide on a single base-CTE + # alias when materialised. + if isinstance(key.source, StarKey): + measure_name = "*" + else: + leaf = getattr(key.source, "leaf", None) or getattr( + key.source, "column_name", None, + ) + if leaf is None: + return f"_agg_{key.agg}" + measure_name = leaf + agg_args = ( + [agg_kwarg_canonical_str(a) for a in key.args] + if key.args + else None + ) + agg_kwargs = ( + {k: agg_kwarg_canonical_str(v) for k, v in key.kwargs} + if key.kwargs + else None + ) + return canonical_agg_name( + measure_name=measure_name, + aggregation_name=key.agg, + agg_args=agg_args, + agg_kwargs=agg_kwargs, + ) + if isinstance(key, TransformKey): + return f"_{key.op}_inner" + if isinstance(key, ArithmeticKey): + return f"_arith_{key.op}" + if isinstance(key, ScalarCallKey): + return f"_scalar_{key.name}" + if isinstance(key, LiteralKey): + return f"_lit_{key.value}" + if isinstance(key, StarKey): + return "_star" + if isinstance(key, BetweenKey): + # Defensive — BetweenKey shouldn't materialise as a public slot + # in 7b.9; it's always inlined into WHERE by the renderer. + return f"_between_{_canonical_name(key.column)}" + if isinstance(key, InKey): + # Defensive — InKey (DEV-1475) is always inlined into WHERE + # like BetweenKey; never materialises as a public slot. + return f"_in_{_canonical_name(key.column)}" + return "_hidden" + + +ProjectionPlan.model_rebuild() + + +# --------------------------------------------------------------------------- +# Stage 7b.5 — filter → slot id mapping for cross-model planner routing +# --------------------------------------------------------------------------- + + +def filter_referenced_slot_ids( + bound_filter: "BoundFilter", + registry: "ValueRegistry", +) -> "set": + """Return the set of ``SlotId``s that ``bound_filter``'s predicate + references through interned slots. + + Walks the predicate's ``ValueKey`` tree via ``_iter_slot_deps`` — + yielding only slot-worthy keys (``ColumnKey`` / ``ColumnSqlKey`` / + ``AggregateKey`` / ``TransformKey`` / ``TimeTruncKey``) and skipping + composite-only nodes (``ArithmeticKey``, ``ScalarCallKey``, + ``LiteralKey``, ``StarKey``). Each slot-worthy key is looked up in + the registry; keys without an interned slot are silently skipped + (filter literals, hidden registry misses). + + Codex HIGH #3/#4 for DEV-1450: this helper exists so the + cross-model planner gets ``set[SlotId]`` instead of having to + classify ``BoundFilter.referenced_keys`` (which are + pre-interning ``ValueKey``s, not slot ids) or naively walking only + the top-level key (which misses composite-predicate leaves). + """ + result: set = set() + for dep in _iter_slot_deps(bound_filter.value_key): + sid = registry.find_by_key(dep) + if sid is not None: + result.add(sid) + return result diff --git a/slayer/engine/profiling.py b/slayer/engine/profiling.py index b3ef230a..752407d5 100644 --- a/slayer/engine/profiling.py +++ b/slayer/engine/profiling.py @@ -198,9 +198,14 @@ async def _profile_numeric_temporal_columns( """Profile every numeric/temporal column in a single batched min/max query.""" if not columns: return {} + # Deliberately omit ``type`` on the ext columns: DEV-1361's CAST wrap + # on the aggregation expression (``CAST(MIN(ordered_at) AS TIMESTAMP)``) + # is harmful on SQLite, which has no TIMESTAMP type and falls back to + # NUMERIC affinity — coercing ``'2025-01-15'`` to the int ``2025``. The + # profile query only needs the raw min/max value, so we keep the column + # untyped and let the backend return whatever native shape it stores. ext_columns = [ - {"name": f"_slayer_range_{c.name}", "sql": c.sql if c.sql else c.name, - "type": str(c.type)} + {"name": f"_slayer_range_{c.name}", "sql": c.sql if c.sql else c.name} for c in columns ] measures_payload: List[Dict[str, str]] = [] diff --git a/slayer/engine/query_engine.py b/slayer/engine/query_engine.py index c156170d..f0f6725b 100644 --- a/slayer/engine/query_engine.py +++ b/slayer/engine/query_engine.py @@ -21,6 +21,7 @@ SlayerModel, SourceModelOrigin, ) +from slayer.core.warnings import NormalizationWarning from slayer.core.query import ( ColumnRef, SlayerQuery, @@ -31,11 +32,29 @@ CrossModelMeasure, EnrichedMeasure, EnrichedQuery, - public_projection_aliases, ) -from slayer.engine.enrichment import enrich_query +from slayer.engine.enrichment import _collect_reachable_agg_names, enrich_query +from slayer.engine.normalization import normalize_model, normalize_query +from slayer.engine.path_resolution import NoJoinError as _NoJoinError +from slayer.engine.path_resolution import walk_join_chain +from slayer.engine.planned import PlannedQuery +from slayer.engine.response_meta import ( + FieldMetadata as FieldMetadata, # re-export for slayer_client / tests + ResponseAttributes, + _infer_aggregated_format, + build_response_metadata, +) +from slayer.engine.source_bundle import ( + ResolvedSourceBundle, + build_resolved_source_bundle, + expand_query_backed_models_in_bundle, +) +from slayer.engine.stage_ordering import topologically_order_stages +from slayer.engine.stage_planner import plan_stages +from slayer.engine.variables import apply_variables_to_query from slayer.sql.client import SlayerSQLClient -from slayer.sql.generator import SQLGenerator +from slayer.sql.generator import SQLGenerator, generate_planned_stages +from slayer.sql.stage_wrapper import build_flat_rename_wrapper from slayer.storage.base import StorageBackend logger = logging.getLogger(__name__) @@ -63,16 +82,6 @@ ) -class _NoJoinError(Exception): - """Internal sentinel raised by ``_walk_join_chain`` when - ``strict_missing_join=False`` and a hop has no matching join. Lets - callers like ``_resolve_dimension_with_terminal`` map a missing - join to a ``None`` return without re-walking the path.""" - def __init__(self, hop_name: str) -> None: - super().__init__(f"no join target named {hop_name!r}") - self.hop_name = hop_name - - _EXPLAIN_PREFIX = { "postgres": "EXPLAIN ANALYZE", "redshift": "EXPLAIN", @@ -140,24 +149,6 @@ def _build_explain_sql(dialect: str, sql: str) -> str: return f"{prefix} {sql}{suffix}" -class FieldMetadata(BaseModel): - """Metadata for a single field in the query response.""" - - label: Optional[str] = None - format: Optional[NumberFormat] = None - - -class ResponseAttributes(BaseModel): - """Field metadata for a query response, split by type.""" - - dimensions: Dict[str, FieldMetadata] = PydanticField(default_factory=dict) - measures: Dict[str, FieldMetadata] = PydanticField(default_factory=dict) - - def get(self, column: str) -> Optional[FieldMetadata]: - """Look up metadata for a column across both dicts.""" - return self.dimensions.get(column) or self.measures.get(column) - - class SlayerResponse(BaseModel): """Response from a SLayer query.""" @@ -165,6 +156,13 @@ class SlayerResponse(BaseModel): columns: List[str] = PydanticField(default_factory=list) sql: Optional[str] = None attributes: ResponseAttributes = PydanticField(default_factory=ResponseAttributes) + # DEV-1450 stage 6 — slack-normalization warnings. + # Structured payload for each slack rewrite the normalization layer + # performed on the input (function-style aggs, misplaced measures, + # AST-resolvable dotted refs in raw SQL). Empty for queries that + # arrived in canonical form. Surfaced alongside the result so + # REST / MCP / CLI consumers can echo the rewrites back to authors. + warnings: List[NormalizationWarning] = PydanticField(default_factory=list) @model_validator(mode="after") def _populate_columns(self) -> "SlayerResponse": @@ -198,34 +196,24 @@ def to_markdown(self) -> str: return "\n".join([header, separator] + body_lines) -def _infer_aggregated_format( +def _normalize_source_query_stages( model: SlayerModel, - measure_name: str, - aggregation: str, -) -> Optional[NumberFormat]: - """Infer NumberFormat for an aggregated measure based on aggregation type and source measure format. - - Rules: - - count, count_distinct: always INTEGER - - avg, weighted_avg, median: always FLOAT - - sum, min, max, first, last: inherit from source measure - - *:count (measure_name="*"): INTEGER + *, + custom_aggs: Optional[frozenset[str]], +) -> SlayerModel: + """For a query-backed model, run ``normalize_query`` on every stage so + funcstyle aggregations over joined-model custom aggs (DEV-1500) land in + canonical form at save time — ``normalize_model`` itself only walks + ``model.measures``, never ``source_queries``. No-op for table-backed + models. CR PR #153 thread r3330620881. """ - if measure_name == "*": - return NumberFormat(type=NumberFormatType.INTEGER) - - if aggregation in ("count", "count_distinct"): - return NumberFormat(type=NumberFormatType.INTEGER) - - if aggregation in ("avg", "weighted_avg", "median"): - return NumberFormat(type=NumberFormatType.FLOAT) - - # sum, min, max, first, last: inherit from source column's format - source_col = model.get_column(measure_name) - if source_col and source_col.format: - return source_col.format - - return None + if not model.source_queries: + return model + new_stages = [ + normalize_query(stage, model=None, custom_agg_names=custom_aggs).query + for stage in model.source_queries + ] + return model.model_copy(update={"source_queries": new_stages}) class SlayerQueryEngine: @@ -279,191 +267,18 @@ def _scope_named_queries_to_prior( out[k] = v return out - @staticmethod - def _extract_sibling_refs(query: "SlayerQuery", against: set) -> set: - """Names from ``query.source_model`` / inline joins that match ``against``. - - Walks the three shapes ``source_model`` can take — plain string, - dict (``ModelExtension`` or inline ``SlayerModel``), or typed - instance — and collects every name that resolves against the - ``against`` set. Used by both the dependency-graph builder and - the self-reference / root-as-sink validators. - """ - out: set = set() - sm = query.source_model - if isinstance(sm, str): - if sm in against: - out.add(sm) - return out - # Dict shape — disambiguate ModelExtension vs inline SlayerModel - # by presence of ``source_name``. - if isinstance(sm, dict): - src = sm.get("source_name") - if isinstance(src, str) and src in against: - out.add(src) - for j in sm.get("joins") or []: - tgt = j.get("target_model") if isinstance(j, dict) else getattr(j, "target_model", None) - if isinstance(tgt, str) and tgt in against: - out.add(tgt) - return out - # Typed ModelExtension / SlayerModel: source_name lives on - # ModelExtension only; SlayerModel.name is the inline model's - # own identifier, not a reference. - src = getattr(sm, "source_name", None) - if isinstance(src, str) and src in against: - out.add(src) - for j in getattr(sm, "joins", None) or []: - tgt = getattr(j, "target_model", None) - if isinstance(tgt, str) and tgt in against: - out.add(tgt) - return out - - @staticmethod - def _index_query_list_by_name( - rest: List["SlayerQuery"], root: "SlayerQuery", - ) -> Dict[str, "SlayerQuery"]: - """Build ``{name: query}`` for non-final entries, validating that - every non-final entry has a unique name and that the root's - name (if any) doesn't collide. - """ - rest_by_name: Dict[str, "SlayerQuery"] = {} - for q in rest: - if not q.name: - raise ValueError( - "Every non-final entry in a query list must have a " - "'name' (siblings reference each other by name)." - ) - if q.name in rest_by_name: - raise ValueError(f"Duplicate stage name '{q.name}' in query list.") - rest_by_name[q.name] = q - if root.name and root.name in rest_by_name: - raise ValueError( - f"Stage name '{root.name}' is duplicated: the final entry " - f"shares a name with an earlier entry." - ) - return rest_by_name - - @classmethod - def _validate_query_list_invariants( - cls, - queries: List["SlayerQuery"], - rest: List["SlayerQuery"], - root: "SlayerQuery", - sibling_names: set, - ) -> None: - """Reject self-references and any sibling that depends on the root. - - Self-references are caught for every entry (including the root). - Root-as-sink: no non-final stage may reference the root by name. - """ - for q in queries: - if q.name and q.name in cls._extract_sibling_refs(q, {q.name} | sibling_names): - raise ValueError( - f"Stage '{q.name}' references itself — self-references " - f"are not allowed." - ) - if root.name: - referrers = sorted( - q.name for q in rest if root.name in cls._extract_sibling_refs(q, {root.name}) - ) - if referrers: - raise ValueError( - f"The final entry '{root.name}' is the DAG root and must " - f"not be referenced by other stages. Referenced by: " - f"{referrers}." - ) - - @classmethod - def _build_dependency_graph( - cls, - rest_by_name: Dict[str, "SlayerQuery"], - sibling_names: set, - ) -> tuple: - """Build the (in_degree, dependents) adjacency for Kahn's. - - Edge direction: prerequisite → dependent. ``in_degree[X]`` is - the count of siblings ``X`` depends on; ``dependents[X]`` is the - list of siblings that depend on ``X``. - """ - in_degree: Dict[str, int] = dict.fromkeys(rest_by_name, 0) - dependents: Dict[str, List[str]] = {name: [] for name in rest_by_name} - for name, q in rest_by_name.items(): - for prereq in cls._extract_sibling_refs(q, sibling_names): - dependents[prereq].append(name) - in_degree[name] += 1 - return in_degree, dependents - - @staticmethod - def _kahn_sort( - in_degree: Dict[str, int], - dependents: Dict[str, List[str]], - ) -> List[str]: - """Topologically sort by Kahn's algorithm. Cycle → ValueError. - - Mutates ``in_degree`` in place; callers shouldn't reuse it. - The frontier is kept sorted for deterministic output order across - runs. - """ - frontier: List[str] = sorted(n for n, d in in_degree.items() if d == 0) - sorted_names: List[str] = [] - while frontier: - n = frontier.pop(0) - sorted_names.append(n) - unlocked: List[str] = [] - for dep in dependents[n]: - in_degree[dep] -= 1 - if in_degree[dep] == 0: - unlocked.append(dep) - frontier.extend(sorted(unlocked)) - if len(sorted_names) < len(in_degree): - cycle = sorted(set(in_degree) - set(sorted_names)) - raise ValueError( - f"Cycle in query list: stages {cycle} form a cyclic " - f"dependency. The reference graph must be acyclic." - ) - return sorted_names - @classmethod def _topologically_order_queries( cls, queries: List["SlayerQuery"], ) -> List["SlayerQuery"]: - """Re-order a runtime query list so every stage appears after the - siblings it references via ``source_model`` or - ``joins.target_model``. Lets callers submit a DAG in any order — - cycles and self-references are rejected; the input order itself - no longer needs to be a valid topological order. - - The last entry of the input is the entry point / DAG root: its - result is what ``execute`` returns. It stays last; only the - non-final entries are reordered. Stages that aren't reachable - from the root are accepted as utility sub-queries — they flow - through the sort like any other node and remain in the - ``named_queries`` dict (the SQL generator emits them only if - something references them). - - Hand-rolled Kahn's algorithm; no ``graphlib`` dependency. The - actual work is delegated to four single-purpose helpers - (:meth:`_index_query_list_by_name`, - :meth:`_validate_query_list_invariants`, - :meth:`_build_dependency_graph`, :meth:`_kahn_sort`) so this - orchestrator stays under the cognitive-complexity gate. - - Raises ``ValueError`` on: missing ``name`` on any non-final - entry; duplicate stage names; self-references; the root being - depended on by any other stage; or a cycle among non-final - stages. + """Thin classmethod shim — delegates to + :func:`slayer.engine.stage_ordering.topologically_order_stages` + (DEV-1452 Stage B decision #1). Existing callers continue to use + the engine-method surface; the migrated paths import the helper + directly. """ - if len(queries) <= 1: - return list(queries) - rest = list(queries[:-1]) - root = queries[-1] - rest_by_name = cls._index_query_list_by_name(rest, root) - sibling_names: set = set(rest_by_name) - cls._validate_query_list_invariants(queries, rest, root, sibling_names) - in_degree, dependents = cls._build_dependency_graph(rest_by_name, sibling_names) - sorted_names = cls._kahn_sort(in_degree, dependents) - return [rest_by_name[n] for n in sorted_names] + [root] + return topologically_order_stages(queries) async def execute( # NOSONAR S3776 — public dispatch over str/dict/list/SlayerQuery; splitting hides the input-shape contract self, @@ -542,7 +357,13 @@ async def _execute_by_name( f"with source_model='{name}'." ) - stages = list(model.source_queries) + # Codex: stored ``source_queries`` may be in non-topological order + # for ``joins[].target_model`` deps (the save path's + # ``_expand_query_backed_model`` calls ``topologically_order_stages`` + # so it accepts that; ``plan_stages._topo_sort`` only handles + # ``source_model`` deps). Re-use the engine-wide topo-sort here so + # run-by-name matches save-time semantics. + stages = topologically_order_stages(list(model.source_queries)) main_query = stages[-1] named_queries: Dict[str, SlayerQuery] = {} for q in stages[:-1]: @@ -601,97 +422,116 @@ async def _execute_pipeline( # NOSONAR S3776 — linear pipeline (resolve→enr if query.whole_periods_only: query = query.snap_to_whole_periods() - # Resolve model from query.source_model (str, SlayerModel, or ModelExtension). - # Pass query.variables as the outer-vars context for any nested - # query-backed model resolution; runtime_kwarg threads through unchanged. - resolving: set = set() - model = await self._resolve_query_model( - query_model=query.source_model, + # P11 — build the resolved source bundle once. Storage is consulted + # here and only here; the binder then reads from the bundle purely. + bundle = await build_resolved_source_bundle( + query=query, + storage=self.storage, + data_source=prefer_data_source, + runtime_variables=runtime_kwarg, named_queries=named_queries, - _resolving=resolving, + ) + + # Expand every query-backed model in the bundle (source + referenced + # + stage_source) and re-apply root inline_extensions. Shared with + # the migrated ``_expand_query_backed_model`` path (DEV-1452 Stage B + # decision F) so both surfaces consume the identical expansion + # contract — divergence between them silently broke nested + # query-backed targets in the legacy stack. + original_source_model = bundle.source_model + bundle = await expand_query_backed_models_in_bundle( + bundle=bundle, outer_vars=query.variables, runtime_kwarg=runtime_kwarg, - prefer_data_source=prefer_data_source, + dry_run_placeholders=False, + expander=self._expand_query_backed_model, + ) + # ``build_resolved_source_bundle`` raises if the source model can't be + # resolved, so ``source_model`` is always populated here. + model = bundle.source_model + assert model is not None + + # P0 — slack-normalization pass. Rewrites slack-but-unambiguous agent + # input (function-style aggs, misplaced bare measures) to canonical + # form before the typed parser sees it. Each stage normalizes against + # its own resolved model; warnings surface on the response. + sibling_names = set(named_queries) + query, slack_warnings = self._normalize_stage( + query=query, bundle=bundle, sibling_names=sibling_names, + ) + normed_named: Dict[str, SlayerQuery] = {} + for nm, nq in named_queries.items(): + nq2, nq_warnings = self._normalize_stage( + query=nq, bundle=bundle, sibling_names=sibling_names, + ) + normed_named[nm] = nq2 + slack_warnings.extend(nq_warnings) + + # Variable substitution into filters (the only field legacy + # substituted). Root uses the bundle's merged variables; each sibling + # re-merges with its own stage layer (precedence runtime > stage > + # outer > model defaults). + query = apply_variables_to_query( + query=query, variables=bundle.query_variables, ) + root_vars = query.variables + normed_named = { + nm: apply_variables_to_query( + query=nq, + variables={ + # Lowest layer: the stage's OWN source-model defaults (a + # sibling-sourced stage has no resolved model here, so fall + # back to the root model's defaults). + **( + ( + bundle.stage_source_models[nm].query_variables + if nm in bundle.stage_source_models + else (model.query_variables if model else None) + ) + or {} + ), + **(root_vars or {}), + **(nq.variables or {}), + **(runtime_kwarg or {}), + }, + ) + for nm, nq in normed_named.items() + } - # Auto-correct: move bare field names to dimensions if they match - query = await self._auto_move_fields_to_dimensions(query, model, named_queries) + # Plan the DAG (root last) and render the whole chain to one SQL string. + stages = [*normed_named.values(), query] + planned_list = plan_stages(queries=stages, bundle=bundle) + root_planned = planned_list[-1] datasource = await self._resolve_datasource(model=model) - - # Enrich: SlayerQuery + model → EnrichedQuery - enriched = await self._enrich(query=query, model=model, named_queries=named_queries) - - # Generate SQL from EnrichedQuery dialect = self._dialect_for_type(datasource.type) - generator = SQLGenerator(dialect=dialect) - # DEV-1444: this is the final-stage SQL that gets executed and - # shown to the user — pin ``outer`` mode so the projection is - # trimmed to public_projection_aliases(enriched). - sql = generator.generate(enriched=enriched, render_mode="outer") + sql = generate_planned_stages( + planned_list, bundle=bundle, dialect=dialect, + ) logger.debug("Generated SQL:\n%s", sql) - # DEV-1444: the response's attributes + expected_columns must mirror - # the trimmed outer projection — never include hoisted intermediates. - public_aliases = set(public_projection_aliases(enriched)) + # Response metadata (attributes + expected_columns) from the typed + # plan + rendered SQL. expected_columns reads the outer SELECT's + # result keys straight from the SQL; attributes classify each public + # slot dimension-vs-measure with its label / format. + attributes, expected_columns = build_response_metadata( + root_planned=root_planned, bundle=bundle, sql=sql, dialect=dialect, + ) - # Collect field metadata from enriched query, split by type. Each - # entry is included only if its alias is part of the public - # projection (filter-extracted hidden transforms, ORDER-BY - # aggregates, and window-arg hoists are silently dropped). - dim_meta: Dict[str, FieldMetadata] = {} - measure_meta: Dict[str, FieldMetadata] = {} - for d in enriched.dimensions: - if d.alias in public_aliases and (d.label or d.format): - dim_meta[d.alias] = FieldMetadata(label=d.label, format=d.format) - for td in enriched.time_dimensions: - if td.alias in public_aliases and td.label: - dim_meta[td.alias] = FieldMetadata(label=td.label) - for m in enriched.measures: - if m.alias not in public_aliases: - continue - measure_fmt = _infer_aggregated_format( - model=model, - measure_name=m.source_measure_name or m.name, - aggregation=m.aggregation, - ) - if m.label or measure_fmt: - measure_meta[m.alias] = FieldMetadata(label=m.label, format=measure_fmt) - for e in enriched.expressions: - if e.alias not in public_aliases: - continue - measure_meta[e.alias] = FieldMetadata( - label=e.label, - format=NumberFormat(type=NumberFormatType.FLOAT), - ) - for t in enriched.transforms: - if t.alias not in public_aliases: - continue - measure_meta[t.alias] = FieldMetadata( - label=t.label, - format=NumberFormat(type=NumberFormatType.FLOAT), - ) - for cm in enriched.cross_model_measures: - if cm.alias in public_aliases and (cm.label or cm.format): - measure_meta[cm.alias] = FieldMetadata(label=cm.label, format=cm.format) - attributes = ResponseAttributes(dimensions=dim_meta, measures=measure_meta) - - # DEV-1444: expected_columns matches the outer SELECT projection - # exactly (helper-driven). Fall back to the legacy bucket-union if - # ``public_projection`` is empty (e.g. enrichment paths that don't - # yet populate ``user_projection``). - expected_columns = list(public_projection_aliases(enriched)) or ( - [d.alias for d in enriched.dimensions] - + [td.alias for td in enriched.time_dimensions] - + [m.alias for m in enriched.measures if not m.name.startswith(("_inner_", "_ft"))] - + [e.alias for e in enriched.expressions] - + [t.alias for t in enriched.transforms if not t.name.startswith(("_inner_", "_ft"))] - + [cm.alias for cm in enriched.cross_model_measures] + # Models whose live schema a query-time DBAPI error could be attributed + # to (the typed-plan equivalent of the legacy enriched-derived set). + touched = self._touched_models_for_plan( + bundle=bundle, + planned_list=planned_list, + original_source_model=original_source_model, ) # dry_run: return SQL without executing if dry_run: - return SlayerResponse(data=[], columns=expected_columns, sql=sql, attributes=attributes) + return SlayerResponse( + data=[], columns=expected_columns, sql=sql, + attributes=attributes, warnings=slack_warnings, + ) # Execute — reuse SQL client (and its connection pool) per datasource ds_key = datasource.get_connection_string() @@ -706,20 +546,95 @@ async def _execute_pipeline( # NOSONAR S3776 — linear pipeline (resolve→enr rows = await client.execute(sql=explain_sql) except Exception as exc: await self._maybe_raise_schema_drift( - err=exc, model=model, enriched=enriched + err=exc, model=model, touched_models=touched ) raise - return SlayerResponse(data=rows, sql=sql, attributes=attributes) + return SlayerResponse( + data=rows, sql=sql, attributes=attributes, + warnings=slack_warnings, + ) try: rows = await client.execute(sql=sql) except Exception as exc: await self._maybe_raise_schema_drift( - err=exc, model=model, enriched=enriched + err=exc, model=model, touched_models=touched ) raise columns = expected_columns if not rows else [] # fallback for empty results; [] triggers auto-derive - return SlayerResponse(data=rows, columns=columns, sql=sql, attributes=attributes) + return SlayerResponse( + data=rows, columns=columns, sql=sql, attributes=attributes, + warnings=slack_warnings, + ) + + def _normalize_stage( + self, + *, + query: SlayerQuery, + bundle: ResolvedSourceBundle, + sibling_names: "set[str]", + ) -> "tuple[SlayerQuery, list[NormalizationWarning]]": + """Slack-normalize one stage against its resolved model (P0). + + Resolves the stage's source model from the bundle so MISPLACED_MEASURE + and custom-aggregation-aware FUNC_STYLE_AGG see the right column / + aggregation names. A stage sourced from a sibling (a flat StageSchema) + has no ``SlayerModel`` — it normalizes with ``model=None`` (FUNC_STYLE_ + AGG still applies; MISPLACED_MEASURE is a no-op without column names). + """ + sm = query.source_model + model: Optional[SlayerModel] = None + if query.name and query.name in bundle.stage_source_models: + # A named non-root stage resolves to its OWN source model + # (string-concrete or inline) — normalize against that, not the + # root, so MISPLACED_MEASURE / custom-agg rewrites see the right + # columns/aggregations (CR). Sibling-sourced stages are absent + # from stage_source_models, so they fall through to model=None. + model = bundle.stage_source_models[query.name] + elif isinstance(sm, str): + if sm not in sibling_names: + model = bundle.get_referenced_model(sm) + if model is None and ( + bundle.source_model is not None + and bundle.source_model.name == sm + ): + model = bundle.source_model + else: + model = bundle.source_model + custom_aggs: Optional[frozenset[str]] = None + if model is not None: + # DEV-1500: scoped per-stage BFS over the pre-resolved bundle, + # so funcstyle calls over joined-model custom aggregations + # (``rolling_avg(customers.score)`` where ``rolling_avg`` lives + # on ``customers``) are recognised by the slack rewrite. + custom_aggs = bundle.reachable_aggregation_names(start=model) + norm = normalize_query(query, model=model, custom_agg_names=custom_aggs) + out = norm.query if norm.query is not None else query + return out, list(norm.warnings) + + def _touched_models_for_plan( + self, + *, + bundle: ResolvedSourceBundle, + planned_list: "list[PlannedQuery]", + original_source_model: Optional[SlayerModel], + ) -> "set[str]": + """Names of every model this query touched, for schema-drift attribution. + + The bundle's ``referenced_models`` already hold the transitive join + walk plus every sibling-stage base; cross-model aggregate targets come + off each planned stage; query-backed base names are recovered from the + pre-expansion source model. ``_maybe_raise_schema_drift`` widens this + further via ``_expand_join_graph``. + """ + touched: set[str] = {m.name for m in bundle.referenced_models} + for pq in planned_list: + for cmp in pq.cross_model_aggregate_plans: + touched.add(cmp.target_model) + if original_source_model is not None and original_source_model.source_queries: + touched.add(original_source_model.name) + touched |= self._collect_query_backed_base_names(original_source_model) + return touched @staticmethod def _collect_query_backed_base_names(model: SlayerModel) -> "set[str]": @@ -815,20 +730,34 @@ async def _maybe_raise_schema_drift( *, err: BaseException, model: SlayerModel, - enriched: "EnrichedQuery", + enriched: "Optional[EnrichedQuery]" = None, + touched_models: "Optional[set[str]]" = None, ) -> None: """Attribute a query-time exception to schema drift via ``validate_models``. If drift is found in the touched models, raise ``SchemaDriftError`` (with ``err`` as ``__cause__``); otherwise return so the caller re-raises the original exception untouched. + The typed pipeline passes ``touched_models`` directly (computed from + the resolved bundle / plan); the legacy path passes ``enriched`` and + the set is derived from it. Either way the join graph is widened via + ``_expand_join_graph`` before attribution. + Any error from ``validate_models`` itself is swallowed so the original exception is never masked. """ from slayer.core.errors import SchemaDriftError try: - touched = await self._collect_models_touched(model=model, enriched=enriched) + if touched_models is not None: + touched = set(touched_models) + await self._expand_join_graph( + touched=touched, data_source=model.data_source or None + ) + else: + touched = await self._collect_models_touched( + model=model, enriched=enriched + ) # Cross-model measure source models share the parent's DS in # validated queries (cross-DS joins are rejected at resolve # time), so attribution only needs the parent's data_source. @@ -886,7 +815,7 @@ def _build_type_probe_query(self, model: SlayerModel) -> SlayerQuery: measures.append(ModelMeasure(formula=f"{c.name}:{agg}")) return SlayerQuery(source_model=model.name, measures=measures) - async def get_column_types( + async def get_column_types( # NOSONAR(S3776) — linear probe pipeline: query-backed prelude → bundle → expand-nested → plan → render → execute → result-key map-back. Splitting hides the order; each step is its own try/except + early-return so flatness is the easier read. self, model_name: str, data_source: Optional[str] = None, @@ -943,29 +872,68 @@ async def get_column_types( client = self._sql_clients[ds_key] probe_query = self._build_type_probe_query(model=model) + # If the model was expanded by the prelude above (query-backed + # case), it's a virtual sql-mode model — pass it INLINE so the + # bundle resolves against the expanded shape rather than re- + # consulting storage (which still holds the un-expanded source + # plus a potentially stale ``data_source``). + if not model.source_queries: + probe_query = probe_query.model_copy(update={"source_model": model}) try: - enriched = await self._enrich(query=probe_query, model=model) + bundle = await build_resolved_source_bundle( + query=probe_query, + storage=self.storage, + data_source=model.data_source or None, + runtime_variables={}, + named_queries={}, + ) + # Expand any nested query-backed models (join targets, etc.) + # so the planner / generator sees ``sql``-mode shapes. + bundle = await expand_query_backed_models_in_bundle( + bundle=bundle, + outer_vars=None, + runtime_kwarg=None, + dry_run_placeholders=True, + expander=self._expand_query_backed_model, + ) + planned = plan_stages(queries=[probe_query], bundle=bundle) + root = planned[-1] dialect = self._dialect_for_type(datasource.type) - generator = SQLGenerator(dialect=dialect) - # DEV-1444: type probing is a user-visible call site; pin - # ``outer`` mode explicitly so a future default-change cannot - # silently shift type-probe behaviour. - sql = generator.generate(enriched=enriched, render_mode="outer") + sql = generate_planned_stages(planned, bundle=bundle, dialect=dialect) except Exception: - logger.warning("get_column_types enrich/generate failed for model '%s'", model_name) + logger.warning( + "get_column_types plan/generate failed for model '%s'", + model_name, + ) return {} try: raw_types = await client.get_column_types(sql=sql) except Exception: - logger.warning("get_column_types probe failed for model '%s'", model_name) + logger.warning( + "get_column_types probe failed for model '%s'", model_name, + ) return {} - # Map qualified aliases (e.g., "orders.revenue_max") back to bare measure names + # Map qualified aliases (e.g., "orders.revenue_max") back to bare + # measure names. Probe sources can be ColumnKey (.leaf) or + # ColumnSqlKey (.column_name) per DEV-1369 derived columns. result: Dict[str, str] = {} - for em in enriched.measures: - if em.alias in raw_types: - result[em.source_measure_name or em.name] = raw_types[em.alias] + source_relation = root.source_relation + for slot in root.aggregate_slots: + if slot.hidden: + continue + src = getattr(slot.key, "source", None) + bare = ( + getattr(src, "leaf", None) + or getattr(src, "column_name", None) + ) + if bare is None: + continue + public = slot.public_name or slot.declared_name + full = f"{source_relation}.{public}" + if full in raw_types: + result[bare] = raw_types[full] return result def execute_sync( @@ -1215,7 +1183,7 @@ def create_model_from_query_sync( ) ) - async def _expand_query_backed_model( + async def _expand_query_backed_model( # NOSONAR S3776 — linear render pipeline (topo-sort → bundle → expand-nested → normalize → variables → plan → render → wrap); splitting hides the order of operations self, model: SlayerModel, outer_vars: Optional[Dict[str, Any]], @@ -1223,26 +1191,181 @@ async def _expand_query_backed_model( dry_run_placeholders: bool, _resolving: Optional[set], ) -> SlayerModel: - """If ``model`` is query-backed, expand its ``source_queries`` into a - virtual model (with rendered SQL). Otherwise return ``model`` unchanged. - - Read-only — never writes to storage. The persisted cache - (``columns`` / ``backing_query_sql`` / ``data_source``) is populated - only by ``engine.save_model`` / ``create_model_from_query(save=True)``. + """Expand a query-backed ``model`` into a virtual ``sql``-mode model + through the typed pipeline (DEV-1452 Stage B). + + Mirrors ``_execute_pipeline``'s mid-section (bundle → expand-nested + → normalize → variables → ``plan_stages`` → ``generate_planned_stages``) + and wraps the rendered backing SQL in a flat-rename SELECT so the + virtual model exposes downstream-bindable flat columns. Read-only — + never writes to storage; the persisted cache (``columns`` / + ``backing_query_sql`` / ``data_source``) is populated only by + ``engine.save_model`` / ``create_model_from_query(save=True)``. + + ``_resolving`` is preserved for caller signature parity but the + recursion guard moves out of the migrated path. Forward / self / + cycle references in stored ``source_queries`` are caught by + ``topologically_order_stages`` up front; nested query-backed + targets / stage sources are handled by + ``expand_query_backed_models_in_bundle`` (which re-enters this + method recursively via the ``expander`` callback). The two legacy + ContextVars (``_join_target_resolving_var`` / + ``_forbidden_sibling_refs_var``) are not set or read on this path. """ if not model.source_queries: return model - stages = list(model.source_queries) - merged_outer = {**model.query_variables, **(outer_vars or {})} + + # 1. Topo-sort + validate (root-as-sink, joins.target_model + # awareness, inline-nested ``SlayerModel.source_queries`` + # recursion per decision E). + stages = topologically_order_stages(list(model.source_queries)) + final_stage = stages[-1] named_q = {q.name: q for q in stages[:-1] if q.name} - return await self._query_as_model( - inner_query=stages[-1], + + # 2. Build resolved source bundle for the final stage. Stage B + # mirrors the legacy ``_query_as_model`` data_source behaviour: the + # bundle is built WITHOUT a DS hint, so the inner resolution falls + # back to the unique-match / priority-list resolver. This lets + # ``get_column_types`` recover from a stale persisted + # ``model.data_source`` (the cache populator may not have refreshed + # yet); join-target lookups still scope to whichever DS the + # resolved base model actually lives in. + bundle = await build_resolved_source_bundle( + query=final_stage, + storage=self.storage, + data_source=None, + runtime_variables=runtime_kwarg, + outer_variables={**model.query_variables, **(outer_vars or {})}, named_queries=named_q, - override_name=model.name, - _resolving=_resolving, - outer_vars=merged_outer, + ) + + # 3. Expand every query-backed model in the bundle and re-apply + # root inline_extensions. Shared with ``_execute_pipeline``. The + # ``_resolving`` set propagates the in-flight expansion names so + # a query-backed join target that transitively references its + # parent short-circuits via cached ``backing_query_sql`` rather + # than recursing forever. + # + # Codex: pass the bundle's MERGED variables (which already include + # ``{model.query_variables, outer_vars, stage, runtime}`` per + # precedence) rather than the bare stage-level ``final_stage. + # variables``. Otherwise nested expansions / sibling stages lose + # the outer model's ``query_variables`` layer and substitution + # diverges between execute and save-time dry-run. + bundle = await expand_query_backed_models_in_bundle( + bundle=bundle, + outer_vars=bundle.query_variables, runtime_kwarg=runtime_kwarg, dry_run_placeholders=dry_run_placeholders, + expander=self._expand_query_backed_model, + _resolving=(_resolving or set()) | {model.name}, + ) + + # 4. Per-stage normalize + variable substitution. Mirrors + # ``_execute_pipeline:486-535``. + sibling_names = set(named_q) + final_stage, _slack = self._normalize_stage( + query=final_stage, bundle=bundle, sibling_names=sibling_names, + ) + normed_named: Dict[str, SlayerQuery] = {} + for nm, nq in named_q.items(): + nq2, _ = self._normalize_stage( + query=nq, bundle=bundle, sibling_names=sibling_names, + ) + normed_named[nm] = nq2 + final_stage = apply_variables_to_query( + query=final_stage, + variables=bundle.query_variables, + dry_run_placeholders=dry_run_placeholders, + ) + # Codex: ``final_stage.variables`` is the user-supplied stage-level + # dict; ``apply_variables_to_query`` substitutes into filters but + # does not promote merged layers onto ``.variables``. Sibling + # substitution therefore needs ``bundle.query_variables`` (the + # merged ``{runtime > final_stage.variables > model.query_variables + # > outer_vars > source_model_defaults}`` set) — using the bare + # stage dict drops ``model.query_variables`` and dry-run save + # fills sibling filters' ``{var}`` placeholders with ``0``. + normed_named = { + nm: apply_variables_to_query( + query=nq, + variables={ + **( + ( + bundle.stage_source_models[nm].query_variables + if nm in bundle.stage_source_models + else ( + bundle.source_model.query_variables + if bundle.source_model else None + ) + ) + or {} + ), + **(bundle.query_variables or {}), + **(nq.variables or {}), + **(runtime_kwarg or {}), + }, + dry_run_placeholders=dry_run_placeholders, + ) + for nm, nq in normed_named.items() + } + + # 5. Plan + render the DAG. + plan_input = [*normed_named.values(), final_stage] + planned_list = plan_stages(queries=plan_input, bundle=bundle) + root_planned = planned_list[-1] + inner_source_model = bundle.source_model + assert inner_source_model is not None + datasource = await self._resolve_datasource(model=inner_source_model) + dialect = self._dialect_for_type(datasource.type) + rendered = generate_planned_stages( + planned_list, bundle=bundle, dialect=dialect, + ) + + # 6. Wrap with flat-renamed SELECT. Public StageColumn entries + # only — hoisted hidden slots are planner-synthesized + # intermediates the user never declared and must not surface as + # virtual-model columns (P4 closure / decision #3). + public_cols = [ + c for c in ( + root_planned.stage_schema.columns + if root_planned.stage_schema is not None else [] + ) + if c.public_alias is not None + ] + expected = [c.name for c in public_cols] + wrapped_ast = build_flat_rename_wrapper( + source_relation=root_planned.source_relation, + stage_sql=rendered, + expected_columns=expected, + dialect=dialect, + ) + wrapped_sql = wrapped_ast.sql(dialect=dialect, pretty=True) + + # 7. Build virtual model from public StageColumn entries. + # Slot types drive Column.type (decision #2) so ``*:count`` → + # ``INT``, declared ``ModelMeasure.type`` is honored, and source + # column types propagate through ``sum`` / ``min`` / ``max``. + cols = [ + Column( + name=sc.name, + sql=sc.name, + type=sc.type or DataType.DOUBLE, + label=sc.label, + description=sc.description, + format=sc.format, + ) + for sc in public_cols + ] + return SlayerModel( + name=model.name, + sql=wrapped_sql, + data_source=inner_source_model.data_source, + columns=cols, + default_time_dimension=inner_source_model.default_time_dimension, + # source_model_origin intentionally NOT set (decision D): + # the typed pipeline answers DEV-1449 through the flat + # ``StageSchema`` namespace, not via the legacy lineage walk. ) async def _resolve_query_model( # NOSONAR S3776 — type-dispatch on str/SlayerModel/ModelExtension/dict; flat is clearer than per-shape helpers here @@ -1476,13 +1599,120 @@ async def create_model_from_query( # can use the returned model directly. return await self._validate_and_populate_cache(model) + async def _reachable_aggs_for_save( + self, model: SlayerModel, + ) -> Optional[frozenset[str]]: + """Best-effort storage BFS over the join graph collecting custom + aggregation names (DEV-1500). Returns ``None`` when ``model`` has + nothing to normalise (no top-level measures AND no source_queries) + or when the walk yields nothing; absent join targets and storage + errors (incl. ``AmbiguousModelError``) are swallowed so a save + never aborts on a flaky join-target lookup. + + For query-backed models the walk unions the reachable custom-agg + names across every stage's ``source_model`` so the save-time + FUNC_STYLE_AGG rewrite on each stage's measures sees the full + joined-model agg surface (CR PR #153 thread r3330620881). + """ + if not model.measures and not model.source_queries: + return None + + resolver = self._make_join_target_resolver(model.data_source) + collected: set[str] = set(await self._walk_reachable_aggs(model, resolver)) + for stage in model.source_queries or []: + stage_model = await self._resolve_stage_source_model( + stage, data_source=model.data_source, + ) + if stage_model is not None: + collected.update( + await self._walk_reachable_aggs(stage_model, resolver) + ) + return frozenset(collected) if collected else None + + def _make_join_target_resolver( + self, data_source: Optional[str], + ): + """Build the best-effort `resolve_join_target` closure that + ``_collect_reachable_agg_names`` consumes — swallows absent targets + and storage errors (incl. ``AmbiguousModelError``) so a save never + aborts on a flaky join-target lookup.""" + storage = self.storage + + async def _resolver( + target_model_name: str, named_queries, # noqa: ARG001 + ): + if not storage: + return None + try: + if data_source: + target = await storage.get_model( + target_model_name, data_source=data_source, + ) + else: + target = await storage.get_model(target_model_name) + except Exception: # noqa: BLE001 — best-effort; AmbiguousModelError + misc + return None + return (None, target) if target is not None else None + + return _resolver + + @staticmethod + async def _walk_reachable_aggs( + walk_model: SlayerModel, resolver, + ) -> frozenset[str]: + """One swallow-all BFS over ``walk_model``'s join graph collecting + the reachable custom-aggregation names.""" + try: + got = await _collect_reachable_agg_names( + model=walk_model, + resolve_join_target=resolver, + named_queries={}, + ) + except Exception: # noqa: BLE001 — never let the walk abort a save + return frozenset() + return got or frozenset() + + async def _resolve_stage_source_model( + self, stage: SlayerQuery, *, data_source: Optional[str], + ) -> Optional[SlayerModel]: + """For a query-backed model's stage, resolve ``stage.source_model`` + to a concrete ``SlayerModel`` (inline → use as-is; string → look up + in storage; ModelExtension / dict → skip best-effort and fall back + to execute-time normalization).""" + src = stage.source_model + if isinstance(src, SlayerModel): + return src + if not isinstance(src, str): + return None + try: + if data_source: + return await self.storage.get_model(src, data_source=data_source) + return await self.storage.get_model(src) + except Exception: # noqa: BLE001 — best-effort + return None + async def save_model(self, model: SlayerModel) -> SlayerModel: """Persist a SlayerModel through the engine. For query-backed models, rejects user-supplied cache fields and runs save-time dry-run validation before populating the cache. For non- query-backed models, persists as-is. + + DEV-1450 stage 6 — runs the slack-normalization layer over the + incoming model so persisted formulas land in canonical form. The + legacy in-tree rewriters still fire during enrichment for callers + that load and re-execute persisted models. + + DEV-1500 — the FUNC_STYLE_AGG rewrite recognises custom aggregations + defined on joined models via ``_reachable_aggs_for_save`` (best-effort + storage walk; swallowed on missing target / AmbiguousModelError / + storage hiccup). """ + custom_aggs = await self._reachable_aggs_for_save(model) + norm_model = normalize_model(model, custom_agg_names=custom_aggs) + if norm_model.model is not None: + model = norm_model.model + model = _normalize_source_query_stages(model, custom_aggs=custom_aggs) # Capture the *previous* data_source for this name so we can clean # up the old storage entry when a query-backed model's resolved # data_source changes (e.g. its backing query now points at a @@ -1527,18 +1757,20 @@ async def _validate_and_populate_cache(self, model: SlayerModel) -> SlayerModel: """Run save-time dry-run validation on a query-backed model and return a copy with ``columns``, ``backing_query_sql``, and ``data_source`` populated from the virtual model. + + DEV-1452 Stage B — pure delegate to the migrated + ``_expand_query_backed_model`` with ``dry_run_placeholders=True`` + so any required-but-undefaulted ``{var}`` placeholder is filled + with the legacy ``"0"`` sentinel rather than raising at SQL-gen. """ - stages = list(model.source_queries or []) - if not stages: + if not (model.source_queries or []): return model - virtual = await self._query_as_model( - inner_query=stages[-1], - named_queries={q.name: q for q in stages[:-1] if q.name}, - override_name=model.name, - _resolving=set(), + virtual = await self._expand_query_backed_model( + model=model, outer_vars=dict(model.query_variables), runtime_kwarg={}, dry_run_placeholders=True, + _resolving=set(), ) return model.model_copy(update={ "columns": list(virtual.columns), @@ -2009,54 +2241,19 @@ async def _walk_join_chain( named_queries: dict = None, strict_missing_join: bool = True, ) -> "tuple[SlayerModel, ModelJoin | None]": - """Walk the join graph from ``source_model`` through ``hop_names``, - returning ``(terminal_model, first_join)``. Single source of - truth for both dimension and cross-model-measure resolution - (DEV-1369 — consolidates two prior near-duplicate walkers). - - Cycle detection: a hop name that already appears on the visited - stack (including ``source_model.name``) raises ``ValueError`` with - the offending path. - - Missing-join behaviour: - - * ``strict_missing_join=True`` (cross-model-measure callers) — - raise ``ValueError`` listing the available joins. - * ``strict_missing_join=False`` (dimension callers) — raise the - internal :class:`_NoJoinError` sentinel so the caller can map - to a ``None`` return. + """Thin shim — delegates to ``slayer.engine.path_resolution.walk_join_chain``. + + Kept as an instance method so existing call sites + (``self._walk_join_chain(...)``) continue to work unchanged + after the DEV-1450 stage-3 extraction. """ - current_model = source_model - visited = {source_model.name} - first_join: "ModelJoin | None" = None - for i, hop_name in enumerate(hop_names): - if hop_name in visited: - raise ValueError( - f"Circular join detected while resolving " - f"'{'.'.join(hop_names)}': '{hop_name}' already visited " - f"({' → '.join(visited)} → {hop_name})" - ) - join = next( - (j for j in current_model.joins if j.target_model == hop_name), - None, - ) - if join is None: - if strict_missing_join: - raise ValueError( - f"Model '{current_model.name}' has no join to " - f"'{hop_name}'. Available joins: " - f"{[j.target_model for j in current_model.joins]}" - ) - raise _NoJoinError(hop_name) - if i == 0: - first_join = join - current_model = await self._resolve_model( - model_name=hop_name, - named_queries=named_queries or {}, - prefer_data_source=current_model.data_source or None, - ) - visited.add(hop_name) - return current_model, first_join + return await walk_join_chain( + source_model=source_model, + hop_names=hop_names, + resolve_model=self._resolve_model, + named_queries=named_queries, + strict_missing_join=strict_missing_join, + ) async def _auto_move_fields_to_dimensions( self, diff --git a/slayer/engine/response_meta.py b/slayer/engine/response_meta.py new file mode 100644 index 00000000..9e34a14b --- /dev/null +++ b/slayer/engine/response_meta.py @@ -0,0 +1,295 @@ +"""DEV-1450 stage 7b.15d — response metadata from the typed plan. + +The legacy engine derived ``SlayerResponse.attributes`` and +``expected_columns`` from an ``EnrichedQuery``. The typed pipeline has no +``EnrichedQuery``; this module rebuilds the same two artefacts from the root +``PlannedQuery`` plus the final rendered SQL. + +* ``expected_columns`` comes from the final SQL's ``named_selects`` — the + literal result-key columns the rows come back keyed by. Deriving them from + the SQL (rather than re-walking slots) is bulletproof: it is exactly the + outer SELECT projection the generator emitted. +* ``attributes`` (``ResponseAttributes.dimensions`` / ``.measures``) come from + the root ``PlannedQuery``'s public ``ValueSlot``s, mirroring the + ``_full_alias_for_slot`` result-key derivation in ``slayer/sql/generator.py`` + so the keys line up with the rendered projection. + +``FieldMetadata`` / ``ResponseAttributes`` / ``_infer_aggregated_format`` live +here (not in ``query_engine``) so this module imports nothing from the engine — +``query_engine`` re-exports them, keeping the dependency one-directional and +the public import path (``from slayer.engine.query_engine import FieldMetadata``) +unchanged. +""" + +from __future__ import annotations + +from typing import Dict, List, Optional, Tuple + +import sqlglot +from pydantic import BaseModel, Field as PydanticField + +from slayer.core.format import NumberFormat, NumberFormatType +from slayer.core.keys import ( + AggregateKey, + ColumnKey, + ColumnSqlKey, + Phase, + StarKey, + TimeTruncKey, + column_leaf, + column_path, +) +from slayer.core.models import Column, SlayerModel +from slayer.engine.planned import PlannedQuery, ValueSlot +from slayer.engine.source_bundle import ResolvedSourceBundle + + +# --------------------------------------------------------------------------- +# Response metadata types (moved here from query_engine for import hygiene). +# --------------------------------------------------------------------------- + + +class FieldMetadata(BaseModel): + """Metadata for a single field in the query response.""" + + label: Optional[str] = None + format: Optional[NumberFormat] = None + + +class ResponseAttributes(BaseModel): + """Field metadata for a query response, split by type.""" + + dimensions: Dict[str, FieldMetadata] = PydanticField(default_factory=dict) + measures: Dict[str, FieldMetadata] = PydanticField(default_factory=dict) + + def get(self, column: str) -> Optional[FieldMetadata]: + """Look up metadata for a column across both dicts.""" + return self.dimensions.get(column) or self.measures.get(column) + + +def _infer_aggregated_format( + model: SlayerModel, + measure_name: str, + aggregation: str, +) -> Optional[NumberFormat]: + """Infer NumberFormat for an aggregated measure based on aggregation type and source measure format. + + Rules: + - count, count_distinct: always INTEGER + - avg, weighted_avg, median: always FLOAT + - sum, min, max, first, last: inherit from source measure + - *:count (measure_name="*"): INTEGER + """ + if measure_name == "*": + return NumberFormat(type=NumberFormatType.INTEGER) + + if aggregation in ("count", "count_distinct"): + return NumberFormat(type=NumberFormatType.INTEGER) + + if aggregation in ("avg", "weighted_avg", "median"): + return NumberFormat(type=NumberFormatType.FLOAT) + + # sum, min, max, first, last: inherit from source column's format + source_col = model.get_column(measure_name) + if source_col and source_col.format: + return source_col.format + + return None + + +# --------------------------------------------------------------------------- +# expected_columns / attributes from the typed plan +# --------------------------------------------------------------------------- + + +def expected_columns_from_sql(*, sql: str, dialect: str) -> List[str]: + """The outer SELECT's result-key columns, read from the rendered SQL. + + ``named_selects`` returns each projected column's alias (``orders.status``, + ``orders.revenue_sum``, ...) — the exact keys execution returns rows under. + """ + parsed = sqlglot.parse_one(sql, dialect=dialect) + return list(parsed.named_selects) + + +def _model_for_path( + *, bundle: ResolvedSourceBundle, path: Tuple[str, ...] +) -> Optional[SlayerModel]: + """The model a dotted join ``path`` lands on (best-effort). + + Empty path → the host source model. Otherwise the last path segment is + the join target model name; resolve it from the bundle's referenced + models, falling back to the host when absent. + """ + if not path: + return bundle.source_model + return bundle.get_referenced_model(path[-1]) or bundle.source_model + + +def _slot_result_keys(*, slot: ValueSlot, source_relation: str) -> List[str]: + """The public result-key alias(es) for ``slot``. + + Mirrors ``SQLGenerator._full_alias_for_slot``: joined ROW slots emit the + full dotted path (``orders.customers.region``); everything else uses the + slot's public alias(es) — multiple for a C13 multi-name interned slot — + prefixed by the stage's source relation. + """ + key = slot.key + if slot.phase == Phase.ROW: + if isinstance(key, ColumnKey) and key.path: + return [f"{source_relation}." + ".".join(key.path) + f".{key.leaf}"] + if isinstance(key, TimeTruncKey) and column_path(key.column): + return [ + f"{source_relation}." + + ".".join(column_path(key.column)) + + f".{column_leaf(key.column)}" + ] + aliases = slot.public_aliases or [slot.declared_name] + return [f"{source_relation}.{a}" for a in aliases] + + +def _column_for_row_slot( + *, slot: ValueSlot, bundle: ResolvedSourceBundle +) -> Optional[Column]: + """The source ``Column`` backing a ROW slot, for label / format lookup.""" + key = slot.key + if isinstance(key, TimeTruncKey): + key = key.column + if isinstance(key, ColumnKey): + model = _model_for_path(bundle=bundle, path=key.path) + leaf = key.leaf + elif isinstance(key, ColumnSqlKey): + model = bundle.get_referenced_model(key.model) or bundle.source_model + leaf = key.column_name + else: + return None + if model is None: + return None + return model.get_column(leaf) + + +def _owning_model_for_agg_source(*, src, bundle: ResolvedSourceBundle): + """The model that owns an aggregate's source column. + + A ``ColumnSqlKey`` (derived column) carries its owning model name in + ``src.model`` — resolve through that (DEV-1450 #4a/#4b), mirroring + ``_column_for_row_slot``. A ``ColumnKey`` / ``StarKey`` is resolved by + walking ``src.path`` from the host. + """ + if isinstance(src, ColumnSqlKey): + return bundle.get_referenced_model(src.model) or bundle.source_model + return _model_for_path(bundle=bundle, path=getattr(src, "path", ())) + + +def _measure_format( + *, slot: ValueSlot, bundle: ResolvedSourceBundle +) -> Optional[NumberFormat]: + """Number format for a measure slot. + + Aggregate slots inherit via ``_infer_aggregated_format`` (INTEGER for + count(-distinct) / star, FLOAT for avg-family, source-column format for + sum/min/max). Transform / arithmetic / scalar-call slots default to FLOAT, + matching the legacy ``EnrichedQuery`` expression/transform handling. + """ + key = slot.key + if isinstance(key, AggregateKey): + src = key.source + if isinstance(src, StarKey): + measure_name: Optional[str] = "*" + else: + measure_name = getattr(src, "leaf", None) or getattr( + src, "column_name", None + ) + model = _owning_model_for_agg_source(src=src, bundle=bundle) + if measure_name is None or model is None: + return NumberFormat(type=NumberFormatType.FLOAT) + return _infer_aggregated_format( + model=model, measure_name=measure_name, aggregation=key.agg + ) + return NumberFormat(type=NumberFormatType.FLOAT) + + +def _measure_label( + *, slot: ValueSlot, bundle: ResolvedSourceBundle +) -> Optional[str]: + """Label for a measure slot. + + A query measure (``labeled_rev:sum``) inherits its source column's label + when the measure spec carried none — mirroring the legacy enrichment that + propagated ``Column.label`` onto the aggregated field. Star aggregates and + transform / arithmetic slots have no single source column, so they fall + back to the slot's own label (usually ``None``). + """ + if slot.label: + return slot.label + key = slot.key + if isinstance(key, AggregateKey): + src = key.source + if isinstance(src, (ColumnKey, ColumnSqlKey)): + model = _owning_model_for_agg_source(src=src, bundle=bundle) + leaf = getattr(src, "leaf", None) or getattr( + src, "column_name", None, + ) + if model is not None and leaf is not None: + col = model.get_column(leaf) + if col is not None: + return col.label + return None + + +def build_response_metadata( + *, + root_planned: PlannedQuery, + bundle: ResolvedSourceBundle, + sql: str, + dialect: str, +) -> Tuple[ResponseAttributes, List[str]]: + """Build ``(attributes, expected_columns)`` for one executed query. + + ``expected_columns`` is read from the rendered SQL (bulletproof); + ``attributes`` maps each public result key to its ``FieldMetadata``, + classified dimension (ROW-phase slots) vs measure (everything else). + Only keys that actually appear in the rendered projection are surfaced — + a guard against any divergence between this derivation and the generator. + """ + expected_columns = expected_columns_from_sql(sql=sql, dialect=dialect) + public_keys = set(expected_columns) + source_relation = root_planned.source_relation + + dim_meta: Dict[str, FieldMetadata] = {} + measure_meta: Dict[str, FieldMetadata] = {} + + projection_ids = set(root_planned.projection) + candidate_slots = ( + list(root_planned.row_slots) + + list(root_planned.aggregate_slots) + + list(root_planned.combined_expression_slots) + ) + for slot in candidate_slots: + if slot.hidden or slot.id not in projection_ids: + continue + is_dim = slot.phase == Phase.ROW + for rk in _slot_result_keys(slot=slot, source_relation=source_relation): + if rk not in public_keys: + continue + if is_dim: + # Label falls back to the model Column's label when the query + # ColumnRef carried none (legacy ``dim_ref.label or + # dim_def.label``). + col = _column_for_row_slot(slot=slot, bundle=bundle) + label = slot.label or (col.label if col else None) + if isinstance(slot.key, TimeTruncKey): + # Time dimensions carry a label only (legacy parity). + if label: + dim_meta[rk] = FieldMetadata(label=label) + continue + fmt = col.format if col else None + if label or fmt: + dim_meta[rk] = FieldMetadata(label=label, format=fmt) + else: + fmt = _measure_format(slot=slot, bundle=bundle) + label = _measure_label(slot=slot, bundle=bundle) + if label or fmt: + measure_meta[rk] = FieldMetadata(label=label, format=fmt) + + return ResponseAttributes(dimensions=dim_meta, measures=measure_meta), expected_columns diff --git a/slayer/engine/schema_drift.py b/slayer/engine/schema_drift.py index 14ac87da..dffbd2c6 100644 --- a/slayer/engine/schema_drift.py +++ b/slayer/engine/schema_drift.py @@ -17,7 +17,6 @@ import logging from typing import ( Annotated, - Any, Dict, List, Literal, @@ -33,14 +32,7 @@ from sqlglot import exp from slayer.core.enums import DataType -from slayer.core.formula import ( - AggregatedMeasureRef, - ArithmeticField, - MixedArithmeticField, - TransformField, - parse_filter, - parse_formula, -) +from slayer.core.formula import parse_filter from slayer.core.models import ( Column, DatasourceConfig, @@ -54,6 +46,15 @@ _sa_type_is_float, _sa_type_to_data_type, ) +from slayer.engine.normalization import func_style_agg_to_colon +from slayer.engine.syntax import ( + AggCall, + DottedRef, + Ref, + StarSource, + parse_expr, + walk_parsed_refs, +) from slayer.sql.client import SlayerSQLClient logger = logging.getLogger(__name__) @@ -482,70 +483,63 @@ def _extract_column_refs_from_sql(sql: str) -> List[Tuple[Optional[str], str]]: return refs -def _agg_ref_names(agg_refs: Dict[str, AggregatedMeasureRef]) -> Set[str]: - """Names from a ``measure:agg`` placeholder map, excluding ``*``.""" - return {ref.measure_name for ref in agg_refs.values() if ref.measure_name != "*"} - - -def _bare_measure_names( - measure_names: List[str], - agg_refs: Dict[str, AggregatedMeasureRef], - *, - skip_placeholder_prefix: Optional[str] = None, -) -> Set[str]: - """Filter raw ``measure_names`` to the ones that are not colon-syntax - placeholders, optionally also stripping sub-transform placeholders. - """ - out: Set[str] = set() - for n in measure_names: - if n in agg_refs: - continue - if skip_placeholder_prefix and n.startswith(skip_placeholder_prefix): - continue - out.add(n) - return out - +def _parsed_ref_name(node: Union[Ref, DottedRef, AggCall]) -> Optional[str]: + """Textual name of a reference-bearing parse node. -def _walk_field_spec_measure_refs(spec: Any) -> Set[str]: - """Walk a ``FieldSpec`` (parse_formula output) and return the set of - measure_name strings (which may be dotted: ``"customers.revenue"``). + ``AggCall`` collapses to its aggregated source name — the agg itself is + not a column reference, and ``*:count`` (``StarSource`` source) yields + ``None`` because ``*`` is not a real column. Args / kwargs of the + aggregation are opaque (legacy parity). Bare ``Ref`` / ``DottedRef`` + surface as their dotted textual form. """ - if isinstance(spec, AggregatedMeasureRef): - return _agg_ref_names({"_": spec}) - if isinstance(spec, ArithmeticField): - return _agg_ref_names(spec.agg_refs) | _bare_measure_names( - spec.measure_names, spec.agg_refs - ) - if isinstance(spec, MixedArithmeticField): - out = _agg_ref_names(spec.agg_refs) | _bare_measure_names( - spec.measure_names, spec.agg_refs, skip_placeholder_prefix="_t" - ) - for _, t in spec.sub_transforms: - out.update(_walk_field_spec_measure_refs(t)) - return out - if isinstance(spec, TransformField): - return _walk_field_spec_measure_refs(spec.inner) - return set() + if isinstance(node, AggCall): + source = node.source + if isinstance(source, StarSource): + return None + node = source + if isinstance(node, Ref): + return node.name + return ".".join(node.parts) def _measure_formula_refs( - formula: str, - *, - named_measures: Optional[Dict[str, str]] = None, + formula: str, *, custom_agg_names: Optional[Set[str]] = None, ) -> Set[str]: - """Best-effort: parse ``formula`` and return the set of column / measure - names it references. Returns the empty set on any parse failure. - - ``named_measures`` is the map ``{measure_name: formula_text}`` for the - enclosing model — required for bare measure references like - ``aov / *:count`` (where ``aov`` is itself a saved measure on the - model) to parse cleanly. + """Best-effort: parse ``formula`` (Mode-B DSL) and return the set of + column / measure names it references (dotted for cross-model refs, e.g. + ``"customers.revenue"``). Returns the empty set on any parse failure. + + Textual extraction only — no scope binding. The cascade attribution + checks each returned name against the dropped-column / dropped-measure + sets itself, so bare named-measure refs surface by name (``aov``) rather + than being inline-expanded; the cascade reaches the underlying column + through the dropped-measure set in a later fixed-point pass. + + Function-style aggregations on legacy / un-normalized persisted formulas + (``sum(amount)``) are rewritten to colon syntax first via the quiet + ``FUNC_STYLE_AGG`` slack helper — matching the legacy ``parse_formula`` + path. ``custom_agg_names`` lets model-level custom aggregations + (``weighted_avg(amount, weight=qty)``) rewrite too (CR); without them + the call parses as an unknown function and the refs are lost, leaving + the drift cascade incomplete. """ try: - spec = parse_formula(formula, named_measures=named_measures) + parsed = parse_expr( + func_style_agg_to_colon( + formula, + custom_agg_names=( + frozenset(custom_agg_names) if custom_agg_names else None + ), + ) + ) except Exception: return set() - return _walk_field_spec_measure_refs(spec) + out: Set[str] = set() + for node in walk_parsed_refs(parsed): + name = _parsed_ref_name(node) + if name is not None: + out.add(name) + return out def _filter_refs(filter_str: str) -> List[str]: @@ -867,11 +861,23 @@ def _measure_refs_on_base( stage: SlayerQuery, base_name: str, graph: _StageGraph ) -> Set[str]: out: Set[str] = set() + # Custom aggregations on models REACHABLE from the stage source so + # function-style custom aggs in a stage measure rewrite to colon form + # before ref extraction. Scoped to ``graph.reachable`` (not every model + # in the registry) so unrelated custom-agg names can't normalize a + # coincidental function call and produce a false cascade hit (CR). + custom_agg_names: Set[str] = set() + for model_name in graph.reachable: + model = graph.models_by_name.get(model_name) + if model is not None: + custom_agg_names.update(a.name for a in (model.aggregations or [])) for m in stage.measures or []: formula = getattr(m, "formula", None) if not formula: continue - for ref in _measure_formula_refs(formula): + for ref in _measure_formula_refs( + formula, custom_agg_names=custom_agg_names, + ): attributed = _attribute_ref_to_base( ref=ref, base_name=base_name, graph=graph ) @@ -1303,17 +1309,46 @@ def _first_dropped_cause( return None +def _reachable_agg_names_from_state( + *, start: SlayerModel, state: "_CascadeState", +) -> Set[str]: + """Sync BFS over ``state.models_by_name`` collecting custom aggregation + names reachable from ``start`` via the join graph. DEV-1500 — lets the + measure-cascade rule recognise function-style references to custom + aggregations defined on joined models (``rolling_avg(customers.score)`` + where ``rolling_avg`` lives on the joined ``customers``). Visited-guarded, + unbounded depth; absent targets are skipped (best-effort). + """ + names: Set[str] = set() + visited: Set[str] = set() + queue: List[SlayerModel] = [start] + while queue: + current = queue.pop(0) + if current.name in visited: + continue + visited.add(current.name) + if current.aggregations: + names.update(a.name for a in current.aggregations) + for join in current.joins: + if join.target_model in visited: + continue + nxt = state.models_by_name.get(join.target_model) + if nxt is not None: + queue.append(nxt) + return names + + def _cascade_measures(*, model: SlayerModel, state: _CascadeState) -> bool: """Rule 2: ``ModelMeasure.formula`` referencing a dropped column or dropped measure.""" changed = False - named_measures = {m.name: m.formula for m in model.measures if m.name} dropped_set = state.dropped_measures.get(model.name, set()) + custom_agg_names = _reachable_agg_names_from_state(start=model, state=state) for measure in model.measures: if measure.name is None or measure.name in dropped_set: continue refs = _measure_formula_refs( - measure.formula, named_measures=named_measures + measure.formula, custom_agg_names=custom_agg_names, ) cause = _first_dropped_cause(refs=refs, model=model, state=state) if cause is None: diff --git a/slayer/engine/source_bundle.py b/slayer/engine/source_bundle.py new file mode 100644 index 00000000..3c6e7acc --- /dev/null +++ b/slayer/engine/source_bundle.py @@ -0,0 +1,558 @@ +"""Stage 2 (DEV-1450) — ResolvedSourceBundle: eagerly resolved query inputs (P11). + +The orchestrator builds this once at the top of execute; the binder reads +from it purely. No ContextVar machinery, no callback re-resolution — the +binder is provably scope-only because everything it needs is in the bundle. + +Contents (per DEV-1450 spec): +- Source model (the host of the query). +- All other referenced models (joined targets, sibling stage hosts). +- Inline ``ModelExtension`` overlays (extra columns / measures / joins). +- Named query siblings (raw ``SlayerQuery``s; the stage planner compiles + each to its own ``StageSchema`` as siblings are traversed in + topological order). +- ``query_variables`` (merged precedence: runtime > stage > outer > model). +- Datasource hint (the ``data_source=`` kwarg that wins over the priority + list). + +Per I2 of the DEV-1450 execution plan, ``source_model`` is ``Optional`` +from day one. DEV-1450's binder asserts ``source_model is not None``; +the type-level optionality is the extension point for a future +anchor-less mode. +""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union + +from pydantic import BaseModel, ConfigDict, Field + +from slayer.core.enums import DataType +from slayer.core.models import Column, ModelJoin, ModelMeasure, SlayerModel +from slayer.core.query import ModelExtension, SlayerQuery +from slayer.engine.variables import merge_query_variables + +if TYPE_CHECKING: + from slayer.core.scope import StageSchema + from slayer.storage.base import StorageBackend + +logger = logging.getLogger(__name__) + + +class ResolvedSourceBundle(BaseModel): + """Eagerly resolved inputs to one query execution (P11).""" + + model_config = ConfigDict(arbitrary_types_allowed=True) + + source_model: Optional[SlayerModel] = None + referenced_models: List[SlayerModel] = Field(default_factory=list) + inline_extensions: List[ModelExtension] = Field(default_factory=list) + named_queries: Dict[str, SlayerQuery] = Field(default_factory=dict) + # DEV-1450 stage 7b.15d — per-named-stage resolved source model, keyed by + # stage name. Populated for siblings whose source resolves to a concrete + # model (a stored model, an inline ``SlayerModel``, or a ``ModelExtension`` + # over a stored base). Siblings sourced FROM another sibling (chain or a + # ``ModelExtension`` over a sibling) are omitted — the planner resolves + # those against the upstream ``StageSchema`` at plan time. Lets each stage + # in a heterogeneous DAG bind against its OWN source rather than the root's. + stage_source_models: Dict[str, SlayerModel] = Field(default_factory=dict) + query_variables: Dict[str, Any] = Field(default_factory=dict) + datasource_hint: Optional[str] = None + + def get_referenced_model(self, name: str) -> Optional[SlayerModel]: + """Linear lookup by name. The list is small (handful of joined + models per query), so the O(n) scan is fine. + """ + for m in self.referenced_models: + if m.name == name: + return m + return None + + def reachable_aggregation_names( + self, *, start: SlayerModel, + ) -> Optional[frozenset[str]]: + """Custom aggregation names reachable from ``start`` via the join + graph, resolved against this bundle's pre-loaded + ``referenced_models``. Sync mirror of + ``enrichment._collect_reachable_agg_names`` — used by the slack + FUNC_STYLE_AGG normalizer so a custom aggregation defined on a + joined model (e.g. ``rolling_avg(customers.score)``) is recognised + and rewritten to colon form. + + BFS, visited-guarded by model name, unbounded depth. Join targets + absent from ``referenced_models`` are skipped (best-effort, matches + the rest of the bundle). Returns ``None`` when nothing reachable + carries any custom aggregation; the contract is "None when empty", + not an empty frozenset. + + Scoping is per call: a stage normalises against its own source + model, so a stage only sees aggregations reachable from the model + it actually queries — not the union across every sibling stage. + """ + names: set[str] = set() + visited: set[str] = set() + queue: list[SlayerModel] = [start] + while queue: + current = queue.pop(0) + if current.name in visited: + continue + visited.add(current.name) + if current.aggregations: + names.update(a.name for a in current.aggregations) + for join in current.joins: + if join.target_model in visited: + continue + nxt = self.get_referenced_model(join.target_model) + if nxt is not None: + queue.append(nxt) + return frozenset(names) if names else None + + +# Anything accepted as ``SlayerQuery.source_model``. +SourceSpec = Union[str, SlayerModel, ModelExtension, Dict[str, Any]] + + +def _apply_extension_overlay( + base: SlayerModel, ext: ModelExtension +) -> SlayerModel: + """Extend ``base`` with the extra columns / measures / joins of ``ext``. + + Mirrors the ``ModelExtension`` branch of + ``SlayerQueryEngine._resolve_query_model`` so the typed pipeline sees the + same overlaid model the legacy path produced. + """ + extra_cols = [ + Column.model_validate(c) if isinstance(c, dict) else c + for c in (ext.columns or []) + ] + extra_measures = [ + ModelMeasure.model_validate(m) if isinstance(m, dict) else m + for m in (ext.measures or []) + ] + extra_joins = [ + ModelJoin.model_validate(j) if isinstance(j, dict) else j + for j in (ext.joins or []) + ] + return base.model_copy( + update={ + "columns": list(base.columns) + extra_cols, + "measures": list(base.measures) + extra_measures, + "joins": list(base.joins) + extra_joins, + } + ) + + +def _source_name_if_sibling( + spec: SourceSpec, sibling_names: "set[str] | Dict[str, Any]" +) -> Optional[str]: + """Return the sibling stage name a ``source_model`` spec reads from, if any. + + Covers the bare-string form (``source_model="kpis"``) AND the + ``ModelExtension`` / dict-with-``source_name`` form + (``source_model={"source_name": "kpis", ...}``) — both reference a sibling + when the name is in ``sibling_names``. Returns ``None`` otherwise. + """ + if isinstance(spec, str): + return spec if spec in sibling_names else None + if isinstance(spec, ModelExtension): + return spec.source_name if spec.source_name in sibling_names else None + if isinstance(spec, dict) and isinstance(spec.get("source_name"), str): + nm = spec["source_name"] + return nm if nm in sibling_names else None + return None + + +def _follow_sibling_chain( + spec: SourceSpec, named_queries: Dict[str, SlayerQuery] +) -> SourceSpec: + """Resolve a ``source_model`` that points at a named sibling stage down + to the real base spec it ultimately reads from. + + The bundle's ``source_model`` must be the real base the root chain bottoms + out at — not a sibling name — so a query-backed datasource lookup and the + single-model binding path resolve correctly. Follows both the bare-string + sibling form and the ``ModelExtension`` / dict-over-sibling form down to + the first spec that does NOT read from a sibling. A cycle raises + ``ValueError`` (mirrors the legacy ``_resolve_model`` circular-reference + guard). + """ + seen: List[str] = [] + while True: + sib = _source_name_if_sibling(spec, named_queries) + if sib is None: + return spec + if sib in seen: + chain = " -> ".join([*seen, sib]) + raise ValueError( + f"Circular reference detected in source_queries DAG: {chain}" + ) + seen.append(sib) + spec = named_queries[sib].source_model + + +async def _resolve_source_spec( + spec: SourceSpec, + *, + storage: "StorageBackend", + data_source: Optional[str], +) -> SlayerModel: + """Resolve any ``source_model`` spec to a concrete ``SlayerModel``. + + Storage-only and read-only (P11). Handles the four input shapes the + public API accepts: stored-model name, inline ``SlayerModel``, + ``ModelExtension`` overlay, and the dict forms of both. + """ + if isinstance(spec, SlayerModel): + return spec + if isinstance(spec, ModelExtension): + base = await storage.get_model(spec.source_name, data_source=data_source) + if base is None: + raise ValueError(f"Model '{spec.source_name}' not found") + return _apply_extension_overlay(base, spec) + if isinstance(spec, str): + model = await storage.get_model(spec, data_source=data_source) + if model is None: + raise ValueError(f"Model '{spec}' not found") + return model + if isinstance(spec, dict): + if "source_name" in spec: + ext = ModelExtension.model_validate(spec) + return await _resolve_source_spec( + ext, storage=storage, data_source=data_source + ) + return SlayerModel.model_validate(spec) + raise ValueError(f"Invalid source_model type: {type(spec)!r}") + + +async def build_resolved_source_bundle( + *, + query: SlayerQuery, + storage: "StorageBackend", + data_source: Optional[str] = None, + runtime_variables: Optional[Dict[str, Any]] = None, + outer_variables: Optional[Dict[str, Any]] = None, + named_queries: Optional[Dict[str, SlayerQuery]] = None, +) -> ResolvedSourceBundle: + """Eagerly assemble the :class:`ResolvedSourceBundle` for one execution (P11). + + Resolves the query's source model (every input shape), walks the join + graph transitively to collect every model the binder may hop through, + threads the named-query sibling map, merges the variable layers, and + records the datasource hint. Storage is consulted here and only here; + the binder then reads from the bundle purely. + + Variable precedence (highest first): runtime > query (stage) > outer > + source-model defaults. ``outer_variables`` is the enclosing-query layer + for a query-backed model resolved as a nested source — left ``None`` for + plain top-level execution. + """ + named_queries = named_queries or {} + sibling_names = set(named_queries) + + # The bundle's source_model is the real base the root chain bottoms out + # at — follow the sibling chain past any named-stage indirection. When the + # ROOT source is a ``ModelExtension`` over a NON-sibling base, the overlay + # is recorded in ``inline_extensions`` so the engine can re-apply it AFTER + # a query-backed base expands (expansion derives columns from the backing + # query and would otherwise drop the overlay's extra columns). + root_spec = _follow_sibling_chain(query.source_model, named_queries) + inline_extensions: List[ModelExtension] = [] + if _source_name_if_sibling(root_spec, sibling_names) is None: + ext = _as_extension_over_nonsibling(root_spec, sibling_names) + if ext is not None: + inline_extensions.append(ext) + source_model = await _resolve_source_spec( + root_spec, storage=storage, data_source=data_source + ) + + # Joins never cross datasource boundaries: scope the graph walk by the + # source model's own data_source (matches engine._expand_join_graph), + # falling back to the execution hint only when the model carries none. + walk_ds = source_model.data_source or data_source or None + + referenced_models = await _collect_referenced_models( + source_model=source_model, + named_queries=named_queries, + storage=storage, + data_source=walk_ds, + ) + + # Per-named-stage source models — each non-sibling-sourced sibling resolves + # to its OWN concrete model so heterogeneous DAGs (stage A over ``orders``, + # stage B over ``customers``) bind each stage against the right host. + stage_source_models: Dict[str, SlayerModel] = {} + for nm, nq in named_queries.items(): + if _source_name_if_sibling(nq.source_model, sibling_names) is not None: + continue # sibling-sourced: planner resolves via upstream StageSchema + # A non-sibling-sourced stage's source MUST resolve to a concrete model; + # a failure here (typoed / missing model) is a genuine error, not a + # best-effort skip — swallowing it would silently fall back to the root + # source and emit wrong SQL when column names overlap. + stage_source_models[nm] = await _resolve_source_spec( + nq.source_model, storage=storage, data_source=walk_ds or data_source + ) + + query_variables = merge_query_variables( + runtime=runtime_variables, + stage=query.variables, + outer=outer_variables, + model_defaults=source_model.query_variables, + ) + + return ResolvedSourceBundle( + source_model=source_model, + referenced_models=referenced_models, + inline_extensions=inline_extensions, + named_queries=dict(named_queries), + stage_source_models=stage_source_models, + query_variables=query_variables, + datasource_hint=data_source, + ) + + +async def _collect_referenced_models( + *, + source_model: SlayerModel, + named_queries: Dict[str, SlayerQuery], + storage: "StorageBackend", + data_source: Optional[str], +) -> List[SlayerModel]: + """Transitive join-graph walk (Kahn-free BFS), best-effort. + + Seeds: the source model, plus the real base model of every named sibling + stage (so a sibling's own ``plan_query`` can hop through targets the root + stage never touches). Follows each model's ``joins[].target_model`` within + ``data_source``; absent targets are skipped silently (mirrors + ``SlayerQueryEngine._expand_join_graph``). The source model is returned + first so ``get_referenced_model`` finds the host before any same-named + join target. + """ + # Models we already hold concretely (host + each sibling's real base, + # overlay-resolved), keyed by name. Resolving siblings through + # ``_resolve_source_spec`` means an extension-added join on a sibling is + # walked too. Best-effort: a sibling whose base is absent is skipped. + preseeded: Dict[str, SlayerModel] = {source_model.name: source_model} + for sib in named_queries.values(): + spec = _follow_sibling_chain(sib.source_model, named_queries) + try: + sib_model = await _resolve_source_spec( + spec, storage=storage, data_source=data_source + ) + except ValueError as exc: + logger.debug("sibling source resolution failed for %r: %s", spec, exc) + continue + preseeded.setdefault(sib_model.name, sib_model) + + collected: Dict[str, SlayerModel] = {} + visited: set[str] = set() + frontier: List[str] = list(preseeded) + while frontier: + name = frontier.pop() + if name in visited: + continue + visited.add(name) + model = preseeded.get(name) + if model is None: + try: + model = await storage.get_model(name, data_source=data_source) + except Exception as exc: # best-effort; absent target is fine + logger.debug("join-target lookup failed for %r: %s", name, exc) + model = None + if model is None: + continue + collected.setdefault(name, model) + for join in model.joins: + if join.target_model not in visited: + frontier.append(join.target_model) + + ordered = [source_model] + ordered.extend(m for n, m in collected.items() if n != source_model.name) + return ordered + + +def _as_extension_over_nonsibling( + spec: SourceSpec, sibling_names: "set[str]" +) -> Optional[ModelExtension]: + """Return the ``ModelExtension`` if ``spec`` overlays a NON-sibling base. + + Used to record the root overlay so the engine can re-apply it after a + query-backed base expands. Returns ``None`` for plain strings, inline + models, and overlays over a sibling (those are handled by the planner). + """ + if isinstance(spec, ModelExtension): + ext = spec + elif isinstance(spec, dict) and isinstance(spec.get("source_name"), str): + ext = ModelExtension.model_validate(spec) + else: + return None + if ext.source_name in sibling_names: + return None + return ext + + +def synthetic_model_from_stage_schema( + *, name: str, schema: "StageSchema", data_source: str +) -> SlayerModel: + """A stand-in ``SlayerModel`` whose ``sql_table`` is a stage's CTE name and + whose columns are that stage's flat output columns. + + Lets the binder / cross-model planner resolve a join (or cross-model ref) + targeting a sibling stage, and the generator emit ``FROM AS `` / + ``LEFT JOIN ...`` — the stage is materialised as a CTE elsewhere + (``generate_planned_stages``); this is the rendering vehicle for that CTE + relation. ``StageColumn.name`` is already the ``__``-flattened downstream + bind name, so the synthetic column names match how downstream refs bind. + """ + return SlayerModel( + name=name, + data_source=data_source or "_stage", + sql_table=name, + columns=[ + Column(name=c.name, type=c.type or DataType.DOUBLE) + for c in schema.columns + ], + ) + + +def stage_bundle_with_siblings( + *, + bundle: ResolvedSourceBundle, + source_model: SlayerModel, + sibling_schemas: Dict[str, "StageSchema"], + data_source: str, +) -> ResolvedSourceBundle: + """Per-stage bundle: ``source_model`` is the stage's own host; synthetic + sibling models (one per already-emitted ``StageSchema``) are threaded into + ``referenced_models`` so a join / cross-model ref to a sibling resolves. + + The host comes first (``get_referenced_model`` finds it before any same- + named join target), then the synthetic siblings, then the original bundle's + referenced models (minus any shadowed by the host or a synthetic sibling). + """ + synths = [ + synthetic_model_from_stage_schema( + name=n, schema=s, data_source=data_source + ) + for n, s in sibling_schemas.items() + ] + shadow = {source_model.name} | {s.name for s in synths} + referenced = ( + [source_model] + + synths + + [m for m in bundle.referenced_models if m.name not in shadow] + ) + return bundle.model_copy( + update={"source_model": source_model, "referenced_models": referenced} + ) + + +async def expand_query_backed_models_in_bundle( # NOSONAR(S3776) — three sequential expansion blocks (source + referenced + stage_source) that mutate the bundle in series. Each block guards on its own condition; splitting forces ``bundle`` to ping-pong through helpers without simplifying anything. The inner recursion guard is the only shared state. + *, + bundle: ResolvedSourceBundle, + outer_vars: Optional[Dict[str, Any]], + runtime_kwarg: Optional[Dict[str, Any]], + dry_run_placeholders: bool, + expander, + _resolving: Optional["set[str]"] = None, +) -> ResolvedSourceBundle: + """Expand every query-backed model in the bundle and re-apply any root + ``ModelExtension`` overlay (DEV-1452 Stage B decision F). + + Three expansion blocks, mirroring ``_execute_pipeline`` exactly: + + 1. Source model — if ``bundle.source_model.source_queries`` is set, + expand to ``sql``-mode and re-apply every ``bundle.inline_extensions`` + overlay (expansion derives columns from the backing query and would + otherwise drop the overlay's extra columns). + 2. Referenced models — every join / cross-model target with + ``source_queries`` set expands so the generator renders it as a + backing-SQL subquery rather than a bare table. + 3. Stage source models — a non-root stage whose own source is a stored + query-backed model likewise expands to ``sql``-mode before the + planner binds it. + + ``expander`` is the callback that performs the actual per-model + expansion. ``_execute_pipeline`` and the migrated + ``_expand_query_backed_model`` both pass + ``self._expand_query_backed_model``; the callback signature is + ``async (model, *, outer_vars, runtime_kwarg, dry_run_placeholders, + _resolving) -> SlayerModel``. + + Returns a fresh ``ResolvedSourceBundle`` with the expanded models; + ``inline_extensions`` is preserved as-is for traceability but the + overlay has already been folded into ``source_model``. + + ``_resolving`` is the recursion guard: a set of model names already + being expanded in this asyncio task. A query-backed join target that + transitively references its parent (or itself) is short-circuited + with the cached ``backing_query_sql`` if available, otherwise left + unchanged — mirrors the legacy ``_render_query_backed_join_target`` + contract. + """ + resolving: "set[str]" = set(_resolving) if _resolving is not None else set() + + async def _expand_or_short_circuit(model: SlayerModel) -> SlayerModel: + if model.name in resolving: + # Re-entry: use cached backing SQL if available, otherwise + # return unchanged (the binder / generator will surface a + # clear error when no sql_table / sql is set). + if model.backing_query_sql: + return model.model_copy( + update={"sql": model.backing_query_sql}, + ) + return model + resolving.add(model.name) + try: + return await expander( + model=model, + outer_vars=outer_vars, + runtime_kwarg=runtime_kwarg, + dry_run_placeholders=dry_run_placeholders, + _resolving=resolving, + ) + finally: + resolving.discard(model.name) + + # 1. Source model + inline_extensions re-apply. + if bundle.source_model is not None and bundle.source_model.source_queries: + expanded = await _expand_or_short_circuit(bundle.source_model) + for ext in bundle.inline_extensions: + expanded = _apply_extension_overlay(expanded, ext) + bundle = bundle.model_copy( + update={ + "source_model": expanded, + "referenced_models": [expanded] + + [ + m + for m in bundle.referenced_models + if m.name != expanded.name + ], + } + ) + source_model = bundle.source_model + + # 2. Referenced models (skip the source model itself — handled above). + if source_model is not None and any( + rm.name != source_model.name and rm.source_queries + for rm in bundle.referenced_models + ): + expanded_refs: List[SlayerModel] = [] + for rm in bundle.referenced_models: + if rm.name != source_model.name and rm.source_queries: + rm = await _expand_or_short_circuit(rm) + expanded_refs.append(rm) + bundle = bundle.model_copy(update={"referenced_models": expanded_refs}) + + # 3. Stage source models. + if any(m.source_queries for m in bundle.stage_source_models.values()): + expanded_stage_sources: Dict[str, SlayerModel] = {} + for nm, sm in bundle.stage_source_models.items(): + if sm.source_queries: + sm = await _expand_or_short_circuit(sm) + expanded_stage_sources[nm] = sm + bundle = bundle.model_copy( + update={"stage_source_models": expanded_stage_sources} + ) + + return bundle diff --git a/slayer/engine/stage_ordering.py b/slayer/engine/stage_ordering.py new file mode 100644 index 00000000..4aacbc96 --- /dev/null +++ b/slayer/engine/stage_ordering.py @@ -0,0 +1,238 @@ +"""DEV-1452 Stage B — Kahn topo-sort for stored / runtime ``source_queries`` +stage lists. + +Extracted from ``SlayerQueryEngine._topologically_order_queries`` so the +migrated ``_expand_query_backed_model`` / ``_validate_and_populate_cache`` +can validate stored ``source_queries`` with the same fault-tolerance +contract the runtime ``execute(query=list[...])`` path uses. Decisions +#1 + E of the Stage B plan: + +* Last stage stays root / sink. Cycles, self-references, duplicate + names, and root referenced by another stage all raise ``ValueError`` + with messages that name the offending stage. Forward references — + a stage that names a sibling appearing later in the input list — are + reordered, not rejected: that is the whole point of the topo-sort, + and the user-facing contract is that stored / runtime stage lists + may be supplied in any topologically valid order. +* Sibling refs are walked recursively through inline ``SlayerModel`` + (typed or dict), ``ModelExtension`` (typed or dict), and ``ModelJoin`` + shapes inside ``joins[].target_model``. An inline-nested + ``source_queries`` list contributes edges from the enclosing stage to + any sibling referenced inside. + +The classmethod shim on ``SlayerQueryEngine`` delegates here so existing +call sites (``execute(query=list[...])`` at query_engine.py:469) remain +unchanged. +""" +from __future__ import annotations + +from typing import Any, Dict, List, Set + + +def _extract_sibling_refs(query: Any, against: Set[str]) -> Set[str]: + """Collect every sibling name referenced by ``query`` that appears in + ``against``. Walks ``source_model`` (string, typed ``SlayerModel`` / + ``ModelExtension``, or dict) plus any ``joins[].target_model``. + + When ``source_model`` is an inline ``SlayerModel`` carrying its own + ``source_queries``, recurses into each inner stage so a sibling name + hidden inside a nested stage still surfaces as an edge from the + enclosing stage. Same recursion through ``ModelExtension`` shapes. + + The traversal is purely structural — it never touches storage and + never raises on a missing-sibling reference (the caller validates + edges against ``against``). + """ + out: Set[str] = set() + _walk_spec(query.source_model, against, out) + return out + + +def _walk_spec(spec: Any, against: Set[str], out: Set[str]) -> None: # NOSONAR(S3776) — recursive walker over five ``source_model`` shapes (str / dict-ModelExtension / dict-inline-SlayerModel / typed ModelExtension / typed SlayerModel) plus nested ``source_queries`` recursion. Splitting fragments the per-shape contract; the recursion + isinstance dispatch IS the function. + """Recursively collect sibling refs from a ``source_model`` spec. + + Handled shapes: + * ``str`` — bare sibling name. + * ``dict`` — disambiguates by key shape: + - ``source_name`` present → ``ModelExtension`` form. + - ``source_queries`` present → inline ``SlayerModel`` form. + - otherwise treated as inline ``SlayerModel``. + * Typed ``SlayerModel`` — walk ``joins[].target_model`` AND every + inner stage's ``source_model`` in ``source_queries``. + * Typed ``ModelExtension`` — walk ``source_name`` and ``joins``. + """ + if spec is None: + return + if isinstance(spec, str): + if spec in against: + out.add(spec) + return + if isinstance(spec, dict): + if "source_name" in spec: + src = spec.get("source_name") + if isinstance(src, str) and src in against: + out.add(src) + for j in spec.get("joins") or []: + tgt = ( + j.get("target_model") if isinstance(j, dict) + else getattr(j, "target_model", None) + ) + if isinstance(tgt, str) and tgt in against: + out.add(tgt) + return + # Inline SlayerModel-as-dict (presence of ``source_queries`` is + # the discriminator from a typed-shape dict; also handles the + # legitimate no-source_queries inline-model dict by inspecting + # ``joins`` directly). + for j in spec.get("joins") or []: + tgt = ( + j.get("target_model") if isinstance(j, dict) + else getattr(j, "target_model", None) + ) + if isinstance(tgt, str) and tgt in against: + out.add(tgt) + for inner_q in spec.get("source_queries") or []: + inner_spec = ( + inner_q.get("source_model") + if isinstance(inner_q, dict) + else getattr(inner_q, "source_model", None) + ) + _walk_spec(inner_spec, against, out) + return + # Typed shapes — ModelExtension vs SlayerModel discriminated by the + # presence of ``source_name`` (only ModelExtension has it). + src = getattr(spec, "source_name", None) + if isinstance(src, str) and src in against: + out.add(src) + for j in getattr(spec, "joins", None) or []: + tgt = getattr(j, "target_model", None) + if isinstance(tgt, str) and tgt in against: + out.add(tgt) + # Inline SlayerModel may itself carry ``source_queries``; recurse so + # references hidden inside any nested stage's ``source_model`` + # surface as edges from the enclosing stage. + for inner_q in getattr(spec, "source_queries", None) or []: + inner_spec = getattr(inner_q, "source_model", None) + _walk_spec(inner_spec, against, out) + + +def _index_query_list_by_name(rest: List[Any], root: Any) -> Dict[str, Any]: + """Build ``{name: query}`` for the non-final entries. Validates that + every non-final stage has a unique non-empty name and that the + root's name (if any) doesn't collide with a sibling. + """ + rest_by_name: Dict[str, Any] = {} + for q in rest: + if not q.name: + raise ValueError( + "Every non-final entry in a query list must have a " + "'name' (siblings reference each other by name)." + ) + if q.name in rest_by_name: + raise ValueError(f"Duplicate stage name '{q.name}' in query list.") + rest_by_name[q.name] = q + if root.name and root.name in rest_by_name: + raise ValueError( + f"Stage name '{root.name}' is duplicated: the final entry " + f"shares a name with an earlier entry." + ) + return rest_by_name + + +def _validate_query_list_invariants( + queries: List[Any], + rest: List[Any], + root: Any, + sibling_names: Set[str], +) -> None: + """Reject self-references and any sibling that depends on the root. + + Self-references are caught for every entry (including the root). + Root-as-sink: no non-final stage may reference the root by name. + """ + for q in queries: + if q.name and q.name in _extract_sibling_refs(q, {q.name} | sibling_names): + raise ValueError( + f"Stage '{q.name}' references itself — self-references " + f"are not allowed." + ) + if root.name: + referrers = sorted( + q.name for q in rest if root.name in _extract_sibling_refs(q, {root.name}) + ) + if referrers: + raise ValueError( + f"The final entry '{root.name}' is the DAG root and must " + f"not be referenced by other stages. Referenced by: " + f"{referrers}." + ) + + +def _build_dependency_graph( + rest_by_name: Dict[str, Any], + sibling_names: Set[str], +) -> "tuple[Dict[str, int], Dict[str, List[str]]]": + """Build the (in_degree, dependents) adjacency for Kahn's algorithm.""" + in_degree: Dict[str, int] = dict.fromkeys(rest_by_name, 0) + dependents: Dict[str, List[str]] = {name: [] for name in rest_by_name} + for name, q in rest_by_name.items(): + for prereq in _extract_sibling_refs(q, sibling_names): + dependents[prereq].append(name) + in_degree[name] += 1 + return in_degree, dependents + + +def _kahn_sort( + in_degree: Dict[str, int], + dependents: Dict[str, List[str]], +) -> List[str]: + """Topologically sort by Kahn's algorithm. Cycle → ``ValueError``. + + The frontier is kept sorted for deterministic output order across runs. + """ + frontier: List[str] = sorted(n for n, d in in_degree.items() if d == 0) + sorted_names: List[str] = [] + while frontier: + n = frontier.pop(0) + sorted_names.append(n) + unlocked: List[str] = [] + for dep in dependents[n]: + in_degree[dep] -= 1 + if in_degree[dep] == 0: + unlocked.append(dep) + frontier.extend(sorted(unlocked)) + if len(sorted_names) < len(in_degree): + cycle = sorted(set(in_degree) - set(sorted_names)) + raise ValueError( + f"Cycle in query list: stages {cycle} form a cyclic " + f"dependency. The reference graph must be acyclic." + ) + return sorted_names + + +def topologically_order_stages(queries: List[Any]) -> List[Any]: + """Re-order a query list so every stage appears after the siblings it + references via ``source_model`` or ``joins[].target_model``. + + The final entry is the entry point / DAG root: it stays last. Only + the non-final entries are reordered. Stages that aren't reachable + from the root are accepted as utility sub-queries — they flow + through the sort like any other node. + + Raises ``ValueError`` on: missing ``name`` on a non-final entry; + duplicate stage names; self-references; the root being depended on + by any other stage; or a cycle among non-final stages. + """ + if len(queries) <= 1: + return list(queries) + rest = list(queries[:-1]) + root = queries[-1] + rest_by_name = _index_query_list_by_name(rest, root) + sibling_names: Set[str] = set(rest_by_name) + _validate_query_list_invariants(queries, rest, root, sibling_names) + in_degree, dependents = _build_dependency_graph(rest_by_name, sibling_names) + sorted_names = _kahn_sort(in_degree, dependents) + return [rest_by_name[n] for n in sorted_names] + [root] + + +__all__ = ["topologically_order_stages"] diff --git a/slayer/engine/stage_planner.py b/slayer/engine/stage_planner.py new file mode 100644 index 00000000..309cd2d2 --- /dev/null +++ b/slayer/engine/stage_planner.py @@ -0,0 +1,1613 @@ +"""Stage 7a.7 (DEV-1450) — multi-stage source_queries planner. + +Orchestrates a list of ``SlayerQuery`` stages into a list of +``PlannedQuery``s, the typed input the SQL generator (stage 7b) will +consume. + +Per-stage pipeline: + + raw SlayerQuery → parse (per measure / filter / order) → bind → + ProjectionPlanner → PlannedQuery (+ emitted StageSchema) + +Multi-stage: + +* Stages are topologically sorted so each stage appears after the + siblings it references via ``source_model``. +* Downstream stages bind against the upstream ``StageSchema`` (P6) — + flat namespace, no dotted-join walking. ``IllegalScopeReferenceError`` + on dotted refs (DEV-1449). +* Each stage's ``StageSchema`` columns use the user-supplied ``name`` + (or canonical alias) as the column ``name`` (DEV-1448). + +Dormant in 7a — no engine wiring. Stage 7b's engine cutover flips +``engine.execute`` / ``engine.save_model`` over to ``plan_stages``. +""" + +from __future__ import annotations + +from typing import Dict, FrozenSet, List, Optional, Tuple, Union + +from slayer.core.enums import DataType +from slayer.core.format import NumberFormat +from slayer.core.errors import ( + AmbiguousReferenceError, + UnknownReferenceError, +) +from slayer.core.keys import ( + AggregateKey, + ArithmeticKey, + BetweenKey, + ColumnKey, + ColumnSqlKey, + InKey, + LiteralKey, + Phase, + ScalarCallKey, + StarKey, + TimeTruncKey, + TransformKey, + ValueKey, + normalize_scalar, +) +from slayer.core.models import SlayerModel +from slayer.core.query import ModelExtension, SlayerQuery, TimeDimension +from slayer.core.refs import agg_kwarg_canonical_str, canonical_agg_name +from slayer.core.scope import ModelScope, StageColumn, StageSchema +from slayer.engine.binding import ( + BoundExpr as BinderBoundExpr, + BoundFilter, + bind_expr, + bind_filter, + bind_time_dimension, + walk_value_keys, +) +from slayer.engine.cross_model_planner import ( + CrossModelPlanner, + HostFilterRouting, + IsolatedCteCrossModelPlanner, +) +from slayer.engine.measure_expansion import expand_model_measures +from slayer.engine.response_meta import _infer_aggregated_format +from slayer.engine.planned import ( + BoundExpr as PlannedBoundExpr, + FilterPhase, + OrderEntry, + PlannedQuery, + TransformLayer, + ValueSlot, +) +from slayer.engine.planning import ( + DeclaredMeasure, + OrderSpec, + ProjectionPlanner, + _iter_slot_deps, + filter_referenced_slot_ids, + lower_sugar_transforms, +) +from slayer.engine.source_bundle import ( + ResolvedSourceBundle, + _apply_extension_overlay, + _source_name_if_sibling, + stage_bundle_with_siblings, + synthetic_model_from_stage_schema, +) +from slayer.engine.syntax import parse_expr, parse_filter_expr +from slayer.sql.sql_expr import has_window_function +from slayer.sql.sql_predicate import parse_sql_predicate + + +__all__ = ["plan_query", "plan_stages"] + + +# Stage 7b.10 — TIME_NEEDING transform ops that require a resolvable +# time dimension to render their OVER ``ORDER BY``. Mirrors the legacy +# ``TIME_TRANSFORMS`` set at ``slayer/core/formula.py:33``. +_TIME_NEEDING_TRANSFORM_OPS = frozenset({ + "cumsum", + "change", + "change_pct", + "time_shift", + "first", + "last", + "lag", + "lead", + "consecutive_periods", +}) + + +def _attach_time_keys( + key: ValueKey, *, td_key: TimeTruncKey, +) -> ValueKey: + """Walk ``key``; for every ``TransformKey`` whose op needs a time + dimension and whose ``time_key`` is ``None``, return a copy with + ``time_key=td_key``. Identity-preserving when nothing changes. + + Mirrors ``lower_sugar_transforms``' walker shape so identity + semantics line up: nested TransformKey/ArithmeticKey/ScalarCallKey/ + BetweenKey trees are rebuilt only on the path containing a patch. + """ + if isinstance(key, TransformKey): + new_input = _attach_time_keys(key.input, td_key=td_key) + out = key + if new_input is not key.input: + out = out.model_copy(update={"input": new_input}) + if out.op in _TIME_NEEDING_TRANSFORM_OPS and out.time_key is None: + out = out.model_copy(update={"time_key": td_key}) + return out + if isinstance(key, ArithmeticKey): + new_ops = tuple( + _attach_time_keys(o, td_key=td_key) for o in key.operands + ) + if all(a is b for a, b in zip(new_ops, key.operands)): + return key + return ArithmeticKey(op=key.op, operands=new_ops) + if isinstance(key, ScalarCallKey): + new_args = tuple( + _attach_time_keys(a, td_key=td_key) + if isinstance( + a, (TransformKey, ArithmeticKey, ScalarCallKey, BetweenKey), + ) + else a + for a in key.args + ) + if all(a is b for a, b in zip(new_args, key.args)): + return key + return ScalarCallKey(name=key.name, args=new_args) + if isinstance(key, BetweenKey): + nc = _attach_time_keys(key.column, td_key=td_key) + nl = _attach_time_keys(key.low, td_key=td_key) + nh = _attach_time_keys(key.high, td_key=td_key) + if nc is key.column and nl is key.low and nh is key.high: + return key + return BetweenKey(column=nc, low=nl, high=nh) + if isinstance(key, InKey): + # DEV-1475: ``InKey.values`` is a literal-only tuple — no + # transforms to attach a time key to. Only the LHS column path + # can carry a transform; rebuild only if it changed. + nc = _attach_time_keys(key.column, td_key=td_key) + if nc is key.column: + return key + return InKey(column=nc, values=key.values, negated=key.negated) + return key + + +def _find_unresolved_time_needing_op(key: ValueKey) -> Optional[str]: + """Return the op name of the first time-needing TransformKey reached + that has ``time_key is None``, or ``None`` if every time-needing + transform in the tree is resolved. + """ + if isinstance(key, TransformKey): + if key.op in _TIME_NEEDING_TRANSFORM_OPS and key.time_key is None: + return key.op + return _find_unresolved_time_needing_op(key.input) + if isinstance(key, ArithmeticKey): + for o in key.operands: + found = _find_unresolved_time_needing_op(o) + if found: + return found + return None + if isinstance(key, ScalarCallKey): + for a in key.args: + if isinstance( + a, (TransformKey, ArithmeticKey, ScalarCallKey, BetweenKey), + ): + found = _find_unresolved_time_needing_op(a) + if found: + return found + return None + if isinstance(key, BetweenKey): + for k in (key.column, key.low, key.high): + found = _find_unresolved_time_needing_op(k) + if found: + return found + return None + if isinstance(key, InKey): + # DEV-1475: only the LHS column can host a time-needing + # transform; the RHS values are literals. + return _find_unresolved_time_needing_op(key.column) + return None + + +def plan_query( # NOSONAR(S3776) — planner entry-point dispatcher. The DEV-1503 addition is a small trigger-predicate branch + a kwarg pass-through; the function's pre-existing complexity is owned by the multi-stage scope / bundle / projection / filter-routing wiring it orchestrates and is tracked as a separate refactor. + *, + query: SlayerQuery, + bundle: ResolvedSourceBundle, + scope: Optional[Union[ModelScope, StageSchema]] = None, + cross_model_planner: Optional[CrossModelPlanner] = None, + stage_schemas: Optional[Dict[str, StageSchema]] = None, + disable_dev1503_isolation: bool = False, +) -> PlannedQuery: + """Compile one ``SlayerQuery`` into a typed ``PlannedQuery``. + + ``scope`` defaults to a ``ModelScope`` over ``bundle.source_model``; + pass an explicit ``StageSchema`` to bind against an upstream stage. + ``stage_schemas`` is a name → StageSchema map used by + ``plan_stages`` to wire multi-stage references. + + ``disable_dev1503_isolation`` (DEV-1503) suppresses the trigger that + isolates cross-model-FILTERED local measures into a per-measure CTE. + Threaded through ``subplan_builder`` whenever a cross-model strategy + recurses for a host-rooted (or target-rooted) nested sub-plan, so the + sub-plan's same filtered measure is rendered inline (not infinitely + re-isolated) and so re-rooting of a genuine cross-model aggregate + doesn't redundantly DEV-1503-isolate the target's own filter joins. + """ + stage_schemas = stage_schemas or {} + cross_model_planner = ( + cross_model_planner or IsolatedCteCrossModelPlanner() + ) + + if scope is None: + source = query.source_model + if isinstance(source, str) and source in stage_schemas: + scope = stage_schemas[source] + else: + scope = ModelScope(source_model=bundle.source_model) + + # The generator must render this stage's FROM / joins against the SAME + # model the binder used. For a ModelScope that's the (possibly overlaid / + # synthetic) host; for a StageSchema chain stage it's None (the generator + # builds a synthetic model from the upstream schema). + render_source_model = ( + scope.source_model if isinstance(scope, ModelScope) else None + ) + + # Downstream stages bind against a flat StageSchema — ``__`` is legal + # in their refs (the upstream's flattened multi-hop aliases); model- + # scoped stages keep the P1 rejection. + flat_scope = isinstance(scope, StageSchema) + + declared_measures = _declared_measures_from_query( + query=query, scope=scope, bundle=bundle, + ) + + # DEV-1450 stage 7b.8 — alias lookup for ORDER BY resolution. + # A user-supplied order column may reference the declared measure + # by its public name (user-supplied ``name``), declared name + # (canonical OR user), or canonical alias. The order pass below + # checks this map BEFORE falling back to ``bind_expr`` so refs to + # aggregate aliases like ``amount_sum`` resolve through the + # projection registry rather than against model scope (where they + # don't exist as columns). + declared_alias_to_bound: Dict[str, BinderBoundExpr] = {} + for dm in declared_measures: + for alias in (dm.public_name, dm.declared_name, dm.canonical_alias): + if alias is not None: + declared_alias_to_bound.setdefault(alias, dm.bound) + + # DEV-1450 stage 7b.15 (DEV-1445, C5): declared-MEASURE aliases a + # filter may reference by name. A filter ``rev >= 100`` for a measure + # declared ``{"formula": "customers.revenue:sum", "name": "rev"}`` + # interns ``rev`` onto the cross-model aggregate slot rather than + # failing to resolve against the model columns; the dotted/colon form + # already interns structurally, so both forms share one slot (P2/P4). + # + # Only MEASURE aliases enter this map — never dimension / time- + # dimension names. A time dimension's declared name IS its raw column + # (e.g. ``created_at``), so a WHERE filter ``created_at <= '...'`` + # (such as the one ``snap_to_whole_periods`` injects) must resolve to + # the raw column, not to the truncated dimension slot. ``declared_ + # measures`` is built in dim → time-dim → measure order, so the + # measure entries are the tail past the dim/time-dim prefix. + n_dims = len(query.dimensions or []) + n_tds = len(query.time_dimensions or []) + filter_alias_map: Dict[str, ValueKey] = {} + for dm in declared_measures[n_dims + n_tds:]: + for alias in (dm.public_name, dm.declared_name, dm.canonical_alias): + if alias is not None: + filter_alias_map.setdefault(alias, dm.bound.value_key) + + # DEV-1450 stage 7b.9 — filter list construction in legacy WHERE + # order: date_range filters first, then SlayerModel.filters + # (Mode-A SQL), then user query filters (Mode-B DSL). The legacy + # generator emits date_range BEFORE iterating ``enriched.filters`` + # (slayer/sql/generator.py:2527 vs :2540), and ``enriched.filters`` + # itself is model filters then query filters (enrichment.py:1192). + # + # ``bound_filters`` carries the typed-BoundFilter entries (date_range + # + query filters) for the cross-model routing and projection + # planner passes. Model filters bypass ``bound_filters`` since + # they're Mode-A SQL text without a typed value-key — they're + # appended directly to ``filters_by_phase`` between the two + # bound-filter buckets. + bound_filters: List[BoundFilter] = [] + # Parallel to bound_filters — original query-filter text for user + # filters (None for date_range bounds, which are synthesized from + # TimeDimension.date_range and have no caller-visible source string). + # Wired into HostFilterRouting.text below so cross_model_planner can + # recover ROW-phase user-filter text WITHOUT slicing host_query.filters + # (which has not been deduped, unlike bound_filters) — CR PR #153 + # thread r3350000254. + bound_filter_texts: List[Optional[str]] = [] + text_filter_entries: List[FilterPhase] = [] + + # 1. date_range filters (one per TD with a 2-element date_range) + for td in (query.time_dimensions or []): + if not td.date_range or len(td.date_range) != 2: + continue + if not isinstance(scope, ModelScope): + continue + bf = _build_date_range_filter(td=td, scope=scope, bundle=bundle) + bound_filters.append(bf) + bound_filter_texts.append(None) + n_date_range = len(bound_filters) + + # 2. SlayerModel.filters — Mode-A SQL, always-applied WHERE. + if isinstance(scope, ModelScope) and scope.source_model is not None: + for j, mf in enumerate(scope.source_model.filters or []): + text_filter_entries.append(_validate_model_filter( + mf=mf, idx=j, model=scope.source_model, + )) + + # 3. user query filters (Mode-B DSL). + # + # DEV-1450 stage 7b.15 (DEV-1445): two filter strings that bind to the + # same structural ``ValueKey`` are one predicate (P2). The alias and + # dotted/colon forms of a renamed cross-model aggregate ref + # (``rev >= 100`` and ``customers.revenue:sum >= 100``) intern onto the + # same slot, so emitting both would duplicate the HAVING clause — + # dedupe by bound key, keeping first occurrence. + for f in (query.filters or []): + if not isinstance(f, str): + continue + bf = bind_filter( + parsed=parse_filter_expr(f, allow_dunder=flat_scope), + scope=scope, + bundle=bundle, + alias_map=filter_alias_map, + ) + if any(existing.value_key == bf.value_key for existing in bound_filters): + continue + bound_filters.append(bf) + bound_filter_texts.append(f) + + order_specs = [] + for o in (query.order or []): + col_name = o.column.name + full_name = o.column.full_name + # Prefer declared-measure alias resolution over model-scope + # binding (DEV-1450 stage 7b.8 — gap fix): aggregate canonical + # aliases like ``amount_sum`` are not columns on the model, so + # ``bind_expr`` would raise. The alias map covers user-supplied + # ``name``, canonical alias, and the declared name itself. + # + # DEV-1450 stage 7b.15 (DEV-1443/1445): a cross-model order key + # written ``customers.revenue:sum`` is coerced by ``OrderItem`` + # to ColumnRef(model="customers", name="revenue_sum"), so the + # leaf alone (``col_name``) never matches the declared canonical + # ``customers.revenue_sum``. Try the full dotted form too, then + # fall back to binding the preserved colon/path ``raw_formula`` + # so the order key interns onto the same cross-model aggregate + # slot (P2/P4) rather than raising. + if col_name in declared_alias_to_bound: + bo = declared_alias_to_bound[col_name] + elif full_name in declared_alias_to_bound: + bo = declared_alias_to_bound[full_name] + elif _flatten_dotted(full_name) in declared_alias_to_bound: + # A joined dimension / time dimension is declared under its + # flattened ``__`` form (``stores.opened_at`` → + # ``stores__opened_at``; DEV-1449 / C4). An ORDER BY entry + # written in dotted form must intern onto that same declared + # slot rather than binding the raw column as a fresh slot. + bo = declared_alias_to_bound[_flatten_dotted(full_name)] + elif f"_{col_name}" in declared_alias_to_bound: + # ``*:count`` surfaces as the alias ``_count`` (the ``*`` is + # dropped, the leading ``_`` kept as a marker); users naturally + # order by the bare ``count``. Mirror the legacy + # ``_resolve_order_column`` ``_name`` fallback. + bo = declared_alias_to_bound[f"_{col_name}"] + elif o.raw_formula: + bo = bind_expr( + parsed=parse_expr(o.raw_formula, allow_dunder=flat_scope), + scope=scope, + bundle=bundle, + ) + else: + # Bind the FULL reference (``customers.region``), not just the + # leaf — otherwise a structured dotted ORDER ColumnRef without a + # raw_formula rebinds as ``region`` and hits the wrong host + # column or fails as ambiguous (CR). + bo = bind_expr( + parsed=parse_expr(full_name, allow_dunder=flat_scope), + scope=scope, + bundle=bundle, + ) + order_specs.append(OrderSpec(bound=bo, direction=o.direction)) + + # Stage 7b.10 — attach the active TD as ``time_key`` on every + # time-needing TransformKey (cumsum / lag / lead / first / last / + # time_shift / consecutive_periods / change / change_pct) whose + # binder-output left ``time_key`` as ``None``. Closes the 7b.4 + # carry-over gap: ``_bind_transform`` does not have query / scope + # context to resolve the TD, so the planner does it here after all + # binding completes. Validation mirrors legacy + # ``enrichment.py:564-569`` -- any time-needing transform with no + # resolvable TD raises with the legacy phrase. + active_td_key: Optional[TimeTruncKey] = None + if isinstance(scope, ModelScope) and scope.source_model is not None: + active_td = _resolve_main_time_dimension( + query=query, model=scope.source_model, + ) + if active_td is not None: + active_td_bound = bind_time_dimension( + td=active_td, scope=scope, bundle=bundle, + ) + atd_key = active_td_bound.value_key + assert isinstance(atd_key, TimeTruncKey) + active_td_key = atd_key + + if active_td_key is not None: + declared_measures = [ + DeclaredMeasure( + bound=BinderBoundExpr( + value_key=_attach_time_keys( + dm.bound.value_key, td_key=active_td_key, + ), + ), + declared_name=dm.declared_name, + public_name=dm.public_name, + label=dm.label, + canonical_alias=dm.canonical_alias, + type=dm.type, + format=dm.format, + description=dm.description, + ) + for dm in declared_measures + ] + bound_filters = [ + BoundFilter( + value_key=_attach_time_keys( + bf.value_key, td_key=active_td_key, + ), + phase=bf.phase, + referenced_keys=tuple( + walk_value_keys( + _attach_time_keys( + bf.value_key, td_key=active_td_key, + ), + ), + ), + ) + for bf in bound_filters + ] + order_specs = [ + OrderSpec( + bound=BinderBoundExpr( + value_key=_attach_time_keys( + spec.bound.value_key, td_key=active_td_key, + ), + ), + direction=spec.direction, + ) + for spec in order_specs + ] + + # Validation: any time-needing transform that still has + # ``time_key=None`` after patching means there was no resolvable TD. + for bucket in ( + [dm.bound.value_key for dm in declared_measures], + [bf.value_key for bf in bound_filters], + [spec.bound.value_key for spec in order_specs], + ): + for vk in bucket: + op = _find_unresolved_time_needing_op(vk) + if op is not None: + raise ValueError( + f"Transform '{op}' requires an unambiguous time " + f"dimension. Add a single time_dimensions entry, or " + f"set main_time_dimension to select among multiple " + f"time dimensions." + ) + + # Sugar lowering for ``change`` / ``change_pct`` runs AFTER the + # patching pass so the desugared ``time_shift`` inherits the patched + # ``time_key`` (DEV-1446 identity preservation still holds — the + # inner AggregateKey instance is not rebuilt by lowering). + declared_measures = [ + DeclaredMeasure( + bound=BinderBoundExpr( + value_key=lower_sugar_transforms(dm.bound.value_key), + ), + declared_name=dm.declared_name, + public_name=dm.public_name, + label=dm.label, + canonical_alias=dm.canonical_alias, + type=dm.type, + format=dm.format, + description=dm.description, + ) + for dm in declared_measures + ] + bound_filters = [ + BoundFilter( + value_key=lower_sugar_transforms(bf.value_key), + phase=bf.phase, + referenced_keys=tuple( + walk_value_keys(lower_sugar_transforms(bf.value_key)), + ), + ) + for bf in bound_filters + ] + order_specs = [ + OrderSpec( + bound=BinderBoundExpr( + value_key=lower_sugar_transforms(spec.bound.value_key), + ), + direction=spec.direction, + ) + for spec in order_specs + ] + + source_col_names = _source_column_names(scope) + host_model_name = _host_model_name(scope) + + projection = ProjectionPlanner().plan( + measures=declared_measures, + filters=bound_filters, + order=order_specs, + source_column_names=source_col_names, + host_model_name=host_model_name, + ) + + row_slots, agg_slots, combined_slots = _bucket_slots( + projection.registry.slots, + ) + + # Build filters_by_phase in legacy WHERE order: + # 1. date_range bound filters (bound_filters[:n_date_range]) + # 2. model.filters (text_filter_entries) + # 3. user query bound filters (bound_filters[n_date_range:]) + # bound_filter_ids preserves the mapping back to bound_filters for + # the cross-model routing pass that follows (text_filter_entries + # are excluded — model filters never feed cross-model routing). + filters_by_phase: List[FilterPhase] = [] + bound_filter_ids: List[str] = [] + for i, bf in enumerate(bound_filters[:n_date_range]): + fid = f"f{i}" + filters_by_phase.append( + FilterPhase( + id=fid, phase=bf.phase, text=None, + expression=PlannedBoundExpr(value_key=bf.value_key), + ), + ) + bound_filter_ids.append(fid) + filters_by_phase.extend(text_filter_entries) + for i, bf in enumerate(bound_filters[n_date_range:], start=n_date_range): + fid = f"f{i}" + filters_by_phase.append( + FilterPhase( + id=fid, phase=bf.phase, text=None, + expression=PlannedBoundExpr(value_key=bf.value_key), + ), + ) + bound_filter_ids.append(fid) + # Stage 7b.5 — cross-model planner wiring. For every aggregate slot + # whose source carries a non-empty join path (cross-model agg-ref + # like ``customers.revenue:sum``), invoke the cross_model_planner + # to produce a CrossModelAggregatePlan with explicit WHERE/HAVING/ + # target_model_filters routes. HostFilterRouting records carry the + # post-projection slot ids each filter references (via + # filter_referenced_slot_ids — Codex HIGH #3/#4 fold-in). + # host_filter_routings only carries entries that have a typed + # BoundFilter (date_range + user filters). Model.filters (text-only) + # are always row-phase host-local WHERE and never need to be routed + # to a cross-model CTE — they're skipped here. + host_filter_routings: List[HostFilterRouting] = [] + for fid, bf, ftext in zip( + bound_filter_ids, bound_filters, bound_filter_texts, + ): + host_filter_routings.append(HostFilterRouting( + filter_id=fid, + phase=bf.phase, + referenced_slot_ids=sorted(filter_referenced_slot_ids( + bf, projection.registry, + )), + text=ftext, + )) + + cross_model_plans = [] + host_slots_for_classifier = projection.registry.slots + for slot in agg_slots: + key = slot.key + if not isinstance(key, AggregateKey): + continue + agg_path = getattr(key.source, "path", ()) + # DEV-1503 — extended trigger predicate. Invoke the cross-model planner + # when the aggregate's source carries a non-empty join path (cross- + # model aggregate, existing behaviour) OR when the aggregate's + # ``column_filter_key`` references a non-anchor join path (cross- + # model-FILTERED local measure — the new filtered-local isolation + # case). The typed ``referenced_join_paths`` field is computed at + # binder time by ``compute_column_filter_join_paths``. + has_cross_model_filter = ( + not disable_dev1503_isolation + and key.column_filter_key is not None + and bool(key.column_filter_key.referenced_join_paths) + ) + if not agg_path and not has_cross_model_filter: + continue + # DEV-1450 #2: re-rooting (C1) is owned by the strategy. We hand it + # the host query, the public projection, and a sub-plan builder so it + # can compile a nested re-rooted PlannedQuery when the host carries + # dimensions / filters reachable only through the TARGET's join graph; + # otherwise it returns the forward plan unchanged. The builder is the + # same ``plan_query`` recursion the post-hoc pass used, injected here + # so cross_model_planner.py needn't import stage_planner. + # + # DEV-1503 — the subplan_builder ALWAYS suppresses DEV-1503 isolation: + # for filtered-local isolation, the host-rooted sub-plan contains the + # same filtered measure and would otherwise recurse infinitely; for + # the existing cross-model re-rooting case, the sub-plan's target- + # rooted local aggregate would redundantly DEV-1503-isolate its own + # filter joins (already handled by the surrounding cross-model CTE). + reroot_enabled = ( + isinstance(scope, ModelScope) and scope.source_model is not None + ) + plan = cross_model_planner.plan( + aggregate_slot_id=slot.id, + aggregate_key=key, + bundle=bundle, + host_slots=host_slots_for_classifier, + host_filters=host_filter_routings, + public_alias=slot.public_name, + hidden=slot.hidden, + host_query=query if reroot_enabled else None, + public_projection=( + projection.public_projection if reroot_enabled else None + ), + subplan_builder=( + (lambda q, b: plan_query( + query=q, bundle=b, cross_model_planner=cross_model_planner, + disable_dev1503_isolation=True, + )) + if reroot_enabled else None + ), + ) + cross_model_plans.append(plan) + + order_entries = [] + for spec in order_specs: + sid = projection.registry.find_by_key(spec.bound.value_key) + if sid is not None: + order_entries.append( + OrderEntry(slot_id=sid, direction=spec.direction), + ) + + transform_layers = _emit_transform_layers(slots=projection.registry.slots) + stage_schema = _emit_stage_schema( + query=query, projection=projection, + ) + source_relation = ( + query.source_model + if isinstance(query.source_model, str) + else host_model_name + ) + + # Stage 7b.10 — surface the active TD's slot id so the generator can + # render ``ORDER BY `` in OVER clauses without re-walking + # the model graph. ``None`` when there is no TD (validation already + # ran above; we only reach here if no time-needing transform exists). + active_td_slot_id = ( + projection.registry.find_by_key(active_td_key) + if active_td_key is not None + else None + ) + + return PlannedQuery( + source_relation=source_relation, + row_slots=row_slots, + aggregate_slots=agg_slots, + cross_model_aggregate_plans=cross_model_plans, + combined_expression_slots=combined_slots, + transform_layers=transform_layers, + filters_by_phase=filters_by_phase, + projection=projection.public_projection, + order=order_entries, + limit=query.limit, + offset=query.offset, + stage_schema=stage_schema, + active_time_dimension_slot_id=active_td_slot_id, + render_source_model=render_source_model, + ) + + +def _coerce_extension(spec) -> ModelExtension: + """Coerce a ``ModelExtension`` / dict-with-``source_name`` to a typed + ``ModelExtension`` (for overlaying onto a synthetic sibling model).""" + if isinstance(spec, ModelExtension): + return spec + return ModelExtension.model_validate(spec) + + +def _stage_scope_and_bundle( + *, + query: SlayerQuery, + bundle: ResolvedSourceBundle, + stage_schemas: Dict[str, StageSchema], + data_source: str, + is_root: bool, +) -> "Tuple[Union[ModelScope, StageSchema], ResolvedSourceBundle]": + """Resolve one DAG stage's ``(scope, per-stage bundle)``. + + Each stage binds against its OWN source — not the root's — so a + heterogeneous DAG (stage A over ``orders``, stage B over ``customers``) + resolves each host correctly. Synthetic models for already-planned sibling + stages are threaded into the per-stage bundle so a join / cross-model ref + that targets a sibling resolves against the sibling's flat output columns. + """ + src = query.source_model + sibling_names = set(stage_schemas) + sib = _source_name_if_sibling(src, sibling_names) + + # 1. ``ModelExtension`` / dict OVER a sibling stage: overlay the extra + # columns / measures / joins onto a synthetic model of the sibling CTE + # and bind ModelScope-style (so derived overlay columns resolve). + if sib is not None and not isinstance(src, str): + base = synthetic_model_from_stage_schema( + name=sib, schema=stage_schemas[sib], data_source=data_source, + ) + overlaid = _apply_extension_overlay(base, _coerce_extension(src)) + others = {n: s for n, s in stage_schemas.items() if n != sib} + sb = stage_bundle_with_siblings( + bundle=bundle, source_model=overlaid, + sibling_schemas=others, data_source=data_source, + ) + return ModelScope(source_model=overlaid), sb + + # 2. Bare-string sibling source (chain): bind against the upstream flat + # StageSchema (P6 / DEV-1449). The synthetic upstream model is the + # per-stage host for any cross-model planning / generation consistency. + if isinstance(src, str) and src in stage_schemas: + synth = synthetic_model_from_stage_schema( + name=src, schema=stage_schemas[src], data_source=data_source, + ) + others = {n: s for n, s in stage_schemas.items() if n != src} + sb = stage_bundle_with_siblings( + bundle=bundle, source_model=synth, + sibling_schemas=others, data_source=data_source, + ) + return stage_schemas[src], sb + + # 3. Model-scoped: the stage's own resolved source model. The root uses the + # bundle's source_model (the chain bottoms out at the root's source); + # a named sibling uses its pre-resolved per-stage model. + if is_root: + stage_model = bundle.source_model + else: + stage_model = bundle.stage_source_models.get(query.name) or bundle.source_model + sb = stage_bundle_with_siblings( + bundle=bundle, source_model=stage_model, + sibling_schemas=stage_schemas, data_source=data_source, + ) + return ModelScope(source_model=stage_model), sb + + +def plan_stages( + *, + queries: List[SlayerQuery], + bundle: ResolvedSourceBundle, + cross_model_planner: Optional[CrossModelPlanner] = None, +) -> List[PlannedQuery]: + """Plan a multi-stage DAG. Topo sort, then plan each stage against its own + resolved source + the synthetic models of its already-planned siblings.""" + if len(queries) == 1: + return [plan_query( + query=queries[0], + bundle=bundle, + cross_model_planner=cross_model_planner, + )] + ordered = _topo_sort(queries) + root = ordered[-1] + data_source = ( + (bundle.source_model.data_source if bundle.source_model else None) + or "_stage" + ) + stage_schemas: Dict[str, StageSchema] = {} + results: List[PlannedQuery] = [] + for q in ordered: + scope, stage_bundle = _stage_scope_and_bundle( + query=q, + bundle=bundle, + stage_schemas=stage_schemas, + data_source=data_source, + is_root=q is root, + ) + planned = plan_query( + query=q, + bundle=stage_bundle, + scope=scope, + cross_model_planner=cross_model_planner, + stage_schemas=stage_schemas, + ) + results.append(planned) + if q.name and planned.stage_schema is not None: + stage_schemas[q.name] = planned.stage_schema + return results + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _format_description_for_dimension( + *, scope: Union[ModelScope, StageSchema], full_name: str, +) -> Tuple[Optional[NumberFormat], Optional[str]]: + """Lift ``format`` / ``description`` for a plain dimension off the + source ``Column``. Returns ``(None, None)`` when the ref can't be + resolved (joined / time-truncated / stage-scoped refs) — those + paths surface their metadata through ``response_meta`` instead. + + DEV-1452 Stage B decision #8 — the planner threads these into the + public slot so the migrated query-backed virtual model carries the + same display contract the legacy enrichment pipeline did. + """ + if not isinstance(scope, ModelScope) or scope.source_model is None: + return None, None + if "." in full_name: + return None, None + col = scope.source_model.get_column(full_name) + if col is None: + return None, None + return col.format, col.description + + +_COUNT_AGGREGATIONS: FrozenSet[str] = frozenset({"count", "count_distinct"}) +_FLOAT_AGGREGATIONS: FrozenSet[str] = frozenset({ + "avg", "weighted_avg", "median", + "stddev_samp", "stddev_pop", "var_samp", "var_pop", + "corr", "covar_samp", "covar_pop", "percentile", +}) + + +def _infer_aggregated_type( + *, + model: SlayerModel, + measure_name: Optional[str], + aggregation: str, +) -> Optional[DataType]: + """Type for an aggregated measure slot. Mirrors + ``_infer_aggregated_format`` (decision #2 of the Stage B plan): + + * ``*:count`` (measure_name=``"*"``) → ``INT`` + * ``count`` / ``count_distinct`` → ``INT`` + * ``avg`` / ``weighted_avg`` / ``median`` / parametric / stat aggs → + ``DOUBLE`` + * ``sum`` / ``min`` / ``max`` / ``first`` / ``last`` → inherit from + source column type (DOUBLE if absent). + """ + if measure_name == "*": + return DataType.INT + if aggregation in _COUNT_AGGREGATIONS: + return DataType.INT + if aggregation in _FLOAT_AGGREGATIONS: + return DataType.DOUBLE + # sum / min / max / first / last — preserve source column type. + if measure_name is None: + return None + col = model.get_column(measure_name) + if col is not None and col.type is not None: + return col.type + return None + + +def _format_description_for_measure_formula( + *, scope: Union[ModelScope, StageSchema], bound, +) -> Tuple[Optional[NumberFormat], Optional[str]]: + """Lift ``format`` / ``description`` for a measure formula. The + aggregation-aware format comes from ``_infer_aggregated_format`` when + the bound expression is a bare local aggregate; description follows + the source ``Column`` (sum / min / max preserve documentation). + """ + if not isinstance(scope, ModelScope) or scope.source_model is None: + return None, None + if not isinstance(bound.value_key, AggregateKey): + return None, None + src = bound.value_key.source + if isinstance(src, StarKey): + # ``*:count`` — INTEGER format inferred by helper; no description. + return ( + _infer_aggregated_format( + model=scope.source_model, + measure_name="*", + aggregation=bound.value_key.agg, + ), + None, + ) + if not isinstance(src, (ColumnKey, ColumnSqlKey)): + return None, None + if getattr(src, "path", ()): # cross-model — handled by response_meta + return None, None + bare = getattr(src, "leaf", None) or getattr(src, "column_name", None) + if bare is None: + return None, None + fmt = _infer_aggregated_format( + model=scope.source_model, + measure_name=bare, + aggregation=bound.value_key.agg, + ) + col = scope.source_model.get_column(bare) + desc = col.description if col is not None else None + return fmt, desc + + +def _type_for_measure_formula( + *, scope: Union[ModelScope, StageSchema], bound, +) -> Optional[DataType]: + """Lift ``type`` for a measure-formula slot. + + Mirrors ``_format_description_for_measure_formula`` — sources the + type from ``_infer_aggregated_type`` for local aggregates so the + migrated query-backed virtual model carries ``*:count → INT``, + ``avg → DOUBLE``, ``sum → source column type``. Declared + ``ModelMeasure.type`` overrides — that flow runs through + ``expand_model_measures`` so the bound key already carries the + declared type via the source column lookup. + """ + if not isinstance(scope, ModelScope) or scope.source_model is None: + return None + if not isinstance(bound.value_key, AggregateKey): + return None + src = bound.value_key.source + if isinstance(src, StarKey): + return _infer_aggregated_type( + model=scope.source_model, + measure_name="*", + aggregation=bound.value_key.agg, + ) + if not isinstance(src, (ColumnKey, ColumnSqlKey)): + return None + if getattr(src, "path", ()): # cross-model — handled elsewhere + return None + bare = getattr(src, "leaf", None) or getattr(src, "column_name", None) + if bare is None: + return None + return _infer_aggregated_type( + model=scope.source_model, + measure_name=bare, + aggregation=bound.value_key.agg, + ) + + +def _joined_column_type( + *, source_model: SlayerModel, full_name: str, bundle: ResolvedSourceBundle, +) -> Optional[DataType]: + """Best-effort type of a dotted (joined) dimension by walking the join + chain — mirrors ``binding._resolve_dotted`` (``parts[:-1]`` are join + hops matched on ``target_model``, ``parts[-1]`` is the leaf column), + but returns ``None`` on any miss instead of raising. The binder has + already validated the ref, so this is a guard rather than a primary + check. + """ + parts = full_name.split(".") + if parts and parts[0] == source_model.name: # C14 self-prefix strip + parts = parts[1:] + if not parts: + return None + *hops, leaf = parts + current = source_model + visited = {current.name} + for hop in hops: + if not any(j.target_model == hop for j in current.joins): + return None + nxt = bundle.get_referenced_model(hop) + if nxt is None or nxt.name in visited: + return None + visited.add(nxt.name) + current = nxt + col = current.get_column(leaf) + return col.type if col is not None else None + + +def _type_for_dimension( + *, + scope: Union[ModelScope, StageSchema], + full_name: str, + bundle: ResolvedSourceBundle, +) -> Optional[DataType]: + """Lift ``type`` for a dimension. Local refs read the source column; + joined (dotted) refs walk the join chain to the terminal column's + type. Returning ``None`` for joined refs (the old behaviour) made + ``_query_as_model`` coerce them to ``DOUBLE``, mistyping joined + string / temporal dimensions on the persisted virtual model. + """ + if not isinstance(scope, ModelScope) or scope.source_model is None: + return None + if "." in full_name: + return _joined_column_type( + source_model=scope.source_model, full_name=full_name, bundle=bundle, + ) + col = scope.source_model.get_column(full_name) + return col.type if col is not None else None + + +def _saved_model_measure_type( + *, scope: Union[ModelScope, StageSchema], formula: str, +) -> Optional[DataType]: + """Lift the explicit ``type`` from a saved ``ModelMeasure`` when the + query formula is a bare reference to one. + + ``expand_model_measures`` rewrites ``adjusted_total`` to the saved + measure's underlying formula AST but doesn't surface the measure's + explicit ``type=`` to downstream consumers — so an explicit + ``ModelMeasure(formula="amount:sum * 1.0", type=DataType.DOUBLE)`` + on a reusable named measure would otherwise be lost unless + ``_type_for_measure_formula`` happens to infer the same value. This + helper rescues it by re-looking-up the saved measure here. + + Only fires when the formula text is itself a bare identifier + matching a ``ModelMeasure.name`` on the source model; arithmetic / + function-call / colon-suffix formulas always fall through to + inference (the saved-measure type only applies when the user + references the saved measure by name directly). + """ + if not isinstance(scope, ModelScope) or scope.source_model is None: + return None + bare = formula.strip() + if not bare.isidentifier(): + return None + saved = scope.source_model.get_measure(bare) + return saved.type if saved is not None else None + + +def _declared_measures_from_query( + *, + query: SlayerQuery, + scope: Union[ModelScope, StageSchema], + bundle: ResolvedSourceBundle, +) -> List[DeclaredMeasure]: + # Downstream stages bind against a flat StageSchema whose columns ARE + # the ``__``-flattened multi-hop aliases of the upstream stage, so + # ``__`` is legal in their refs (P5 / DEV-1449). Model-scoped stages + # keep the P1 rejection. + flat_scope = isinstance(scope, StageSchema) + declared: List[DeclaredMeasure] = [] + for d in (query.dimensions or []): + full = d.full_name + bound = bind_expr( + parsed=parse_expr(full, allow_dunder=flat_scope), + scope=scope, + bundle=bundle, + ) + flat_name = _flatten_dotted(full) + fmt, desc = _format_description_for_dimension( + scope=scope, full_name=full, + ) + dim_type = _type_for_dimension( + scope=scope, full_name=full, bundle=bundle, + ) + declared.append(DeclaredMeasure( + bound=bound, + declared_name=flat_name, + public_name=flat_name, + label=d.label, + type=dim_type, + format=fmt, + description=desc, + )) + # Time dimensions follow dimensions in the public projection — matches + # the legacy ``user_projection`` order (dims, then time dims, then + # measures). + for td in (query.time_dimensions or []): + full = td.dimension.full_name + bound = bind_time_dimension(td=td, scope=scope, bundle=bundle) + flat_name = _flatten_dotted(full) + declared.append(DeclaredMeasure( + bound=bound, + declared_name=flat_name, + public_name=flat_name, + label=td.label, + type=DataType.TIMESTAMP, + )) + for m in (query.measures or []): + formula = m.formula + explicit_name = m.name + parsed = parse_expr(formula, allow_dunder=flat_scope) + # DEV-1450 stage 7b.8 — pre-bind ModelMeasure expansion. A bare + # ``Ref`` whose name matches a saved ``ModelMeasure`` on the + # host model is rewritten to the measure's formula AST so the + # binder resolves the underlying columns. Only applies against + # ModelScope (downstream stages bind against StageSchema and + # don't expose saved measures). + if isinstance(scope, ModelScope) and scope.source_model is not None: + parsed = expand_model_measures( + expr=parsed, + model=scope.source_model, + ) + bound = bind_expr(parsed=parsed, scope=scope, bundle=bundle) + # Stage 7b.10: sugar-lowering of ``change`` / ``change_pct`` now + # runs in ``plan_query`` AFTER time-key patching, so the inner + # ``time_shift`` inherits a patched ``time_key`` instead of + # ``None``. Identity-preservation for the inner aggregate slot + # (DEV-1446) still holds — ``lower_sugar_transforms`` keeps the + # inner ``AggregateKey`` instance unchanged. + canonical = _canonical_alias_for_formula(formula, bound=bound) + declared_name = explicit_name or canonical + public_name = explicit_name or canonical + fmt, desc = _format_description_for_measure_formula( + scope=scope, bound=bound, + ) + # Codex: type-priority chain (highest wins): + # 1. ``m.type`` — user-supplied override on the query measure. + # 2. Saved ``ModelMeasure.type`` — when the query formula is a + # bare reference to a reusable saved measure on the source + # model, that measure's explicit type wins over inference. + # ``expand_model_measures`` rewrites the AST but drops the + # source measure's type metadata; re-look-up here. + # 3. ``_type_for_measure_formula`` — aggregation-aware inference. + # Mirrors how the legacy ``EnrichedMeasure.type`` honored an + # explicit type before falling back to inference. + m_type = ( + m.type + or _saved_model_measure_type(scope=scope, formula=formula) + or _type_for_measure_formula(scope=scope, bound=bound) + ) + declared.append(DeclaredMeasure( + bound=bound, + declared_name=declared_name, + public_name=public_name, + label=m.label, + canonical_alias=canonical if explicit_name else None, + type=m_type, + format=fmt, + description=desc, + )) + return declared + + +def _topo_sort(queries: List[SlayerQuery]) -> List[SlayerQuery]: + """Kahn's algorithm: order stages so each appears after its + siblings it references via ``source_model``. + + Raises ``ValueError`` on: + * duplicate stage names, + * a cycle in the dependency graph. + + Stages without a ``name`` (typically the final / root) are appended + last in input order. + """ + if len(queries) <= 1: + return list(queries) + named = [q for q in queries if q.name] + names = [q.name for q in named] + duplicates = sorted({n for n in names if names.count(n) > 1}) + if duplicates: + raise ValueError( + f"Duplicate stage names in source_queries DAG: {duplicates}" + ) + by_name = {q.name: q for q in named} + in_degree = {q.name: 0 for q in named} + edges: Dict[str, List[str]] = {q.name: [] for q in named} + for q in named: + # A stage depends on a sibling when its ``source_model`` reads from it — + # either the bare-string form OR a ``ModelExtension`` / dict over the + # sibling. Capturing both keeps the topo order + cycle detection correct + # for extension-over-sibling stages (not just join-target deps, which + # the engine's runtime list sorter handles upstream). + dep = _source_name_if_sibling(q.source_model, by_name) + if dep is not None and dep != q.name: + in_degree[q.name] += 1 + edges[dep].append(q.name) + sorted_names: List[str] = [] + queue = [n for n, d in in_degree.items() if d == 0] + while queue: + n = queue.pop(0) + sorted_names.append(n) + for dep in edges[n]: + in_degree[dep] -= 1 + if in_degree[dep] == 0: + queue.append(dep) + if len(sorted_names) != len(in_degree): + remaining = sorted(set(in_degree) - set(sorted_names)) + raise ValueError( + f"Cycle detected in source_queries DAG involving stages: " + f"{remaining}" + ) + sorted_named = [by_name[n] for n in sorted_names] + unnamed = [q for q in queries if q.name is None] + return sorted_named + unnamed + + +def _flatten_dotted(name: str) -> str: + return name.replace(".", "__") + + +def _canonical_alias_for_formula( + formula: str, + *, + bound: Optional[BinderBoundExpr] = None, +) -> str: + """Compute the canonical public alias for a measure formula. + + Mirrors ``canonical_agg_name`` for any formula whose bound root is + an ``AggregateKey`` (covers bare ``revenue:sum`` AND parametric + forms like ``revenue:percentile(p=0.5)`` / ``corr(other=quantity)``). + Pre-binding text-shape recognition is used only as a fallback when + no bound expression is supplied. For arbitrary formulas + (transforms, arithmetic), sanitise the formula text so the alias + remains a valid identifier. + + DEV-1450 stage 7b.13: parametric aggregations route through + ``canonical_agg_name`` so kwargs are sanitised consistently with the + legacy enrichment path (``p=0.5`` -> ``_p_0_5``). Without this, the + naive text-replace fallback below leaks the ``=`` literally into the + alias (``amount_percentile_p=0_5_``), breaking parity. + """ + if bound is not None and isinstance(bound.value_key, AggregateKey): + key = bound.value_key + if isinstance(key.source, StarKey): + measure_name: Optional[str] = "*" + # Cross-model star (``customers.*:count``) carries its join + # path so the canonical alias keeps the ``customers.`` prefix + # (result key ``orders.customers._count``). + path: Tuple[str, ...] = key.source.path + else: + # ColumnKey exposes ``.leaf``; ColumnSqlKey exposes + # ``.column_name``. Both shapes can appear as aggregate + # sources (the synth adapter rejects ``ColumnSqlKey`` with + # a typed deferral; the planner still needs to derive an + # alias before the generator runs). Mirror ``_canonical_name`` + # at ``planning.py:540-545``. + measure_name = ( + getattr(key.source, "leaf", None) + or getattr(key.source, "column_name", None) + ) + path = getattr(key.source, "path", ()) + if measure_name is not None: + prefix = ".".join(path) + "." if path else "" + # Local aggregates retain the kwarg suffix to match legacy + # ``enrichment.py:349`` (``percentile(p=0.5)`` -> + # ``_p_0_5``). Cross-model aggregates ALSO retain it -- the + # legacy ``query_engine.py:2160`` drops it, causing CTE alias + # collision on two parametric variants, which the 7b.5 fix + # corrected at the planner layer. Result-key shape for + # cross-model parametric aggs therefore diverges from + # legacy in this slice (no parity tests for that combination; + # structural correctness over bit-identical legacy output). + return prefix + canonical_agg_name( + measure_name=measure_name, + aggregation_name=key.agg, + agg_args=[agg_kwarg_canonical_str(a) for a in key.args] or None, + agg_kwargs={ + k: agg_kwarg_canonical_str(v) for k, v in key.kwargs + } or None, + ) + # Fall through to text-based path -- AggregateKey source is + # neither StarKey, ColumnKey, nor ColumnSqlKey (shouldn't be + # reachable in practice; the binder restricts sources to + # those three shapes). + text = formula.strip() + if ":" in text and "(" not in text: + base, agg = text.rsplit(":", 1) + return canonical_agg_name( + measure_name=base, aggregation_name=agg, + ) + return ( + text.replace(".", "_").replace(":", "_").replace(" ", "_") + .replace("(", "_").replace(")", "_").replace(",", "_") + ) + + +def _source_column_names( + scope: Union[ModelScope, StageSchema], +) -> FrozenSet[str]: + if isinstance(scope, ModelScope) and scope.source_model is not None: + return frozenset(c.name for c in scope.source_model.columns) + if isinstance(scope, StageSchema): + return frozenset(c.name for c in scope.columns) + return frozenset() + + +def _host_model_name( + scope: Union[ModelScope, StageSchema], +) -> str: + if isinstance(scope, ModelScope) and scope.source_model is not None: + return scope.source_model.name + if isinstance(scope, StageSchema): + return scope.relation_name + return "(stage)" + + +def _bucket_slots(slots: List[ValueSlot]): + row: List[ValueSlot] = [] + agg: List[ValueSlot] = [] + combined: List[ValueSlot] = [] + for s in slots: + if s.phase == Phase.ROW: + row.append(s) + elif s.phase == Phase.AGGREGATE: + agg.append(s) + else: + combined.append(s) + return row, agg, combined + + +def _emit_stage_schema( + *, + query: SlayerQuery, + projection, +) -> StageSchema: + """Build the StageSchema from the projection plan. + + Only public slots appear (hidden slots are trimmed). One column per + occurrence in ``public_projection`` so multi-alias declarations + (same key with two ``name``s) emit one column per alias rather + than two copies of ``public_aliases[0]``. + """ + columns: List[StageColumn] = [] + alias_idx: Dict[str, int] = {} + for sid in projection.public_projection: + slot = projection.registry.get(sid) + if slot.hidden: + continue + idx = alias_idx.setdefault(sid, 0) + if idx < len(slot.public_aliases): + alias = slot.public_aliases[idx] + else: + alias = slot.declared_name + alias_idx[sid] = idx + 1 + # The downstream bind name + CTE column name are the ``__``-flattened + # form so a later stage can reference a cross-model aggregate + # (``customers.revenue_sum`` → ``customers__revenue_sum``), matching + # how dimensions already flatten and how the legacy virtual-model + # rename exposed these columns (P5/DEV-1449). ``public_alias`` keeps + # the dotted result-key form. Dimensions / local / user-named + # measures have no dot, so flattening is a no-op for them. + flat = _flatten_dotted(alias) + # Two distinct public columns that flatten to the same downstream + # name (e.g. a joined ``customers.region`` and a literal model column + # ``customers__region`` via the C11 carve-out) would make the stage's + # CTE column ambiguous. Surface it instead of silently binding the + # first match downstream. + if any(c.name == flat for c in columns): + raise ValueError( + f"Stage column name collision on {flat!r}: two projected " + f"columns flatten to the same downstream name. Give one an " + f"explicit measure `name` to disambiguate." + ) + columns.append(StageColumn( + name=flat, + sql_alias=flat, + public_alias=alias, + type=slot.type, + label=slot.label, + hidden=False, + format=slot.format, + description=slot.description, + )) + relation_name = query.name or "(unnamed_stage)" + return StageSchema(relation_name=relation_name, columns=columns) + + +def _emit_transform_layers(*, slots: List[ValueSlot]) -> List[TransformLayer]: + """One TransformLayer per ``TransformKey`` slot, emitted in + dependency order (innermost transform first). + + Nested transforms (``cumsum(change(amount:sum))``) require + per-slot layers so the generator can render the inner window / + self-join before the outer one consumes it. Repeated ops at + different nesting levels stay in separate layers; collapsing by + op would lose the ordering invariant. + + Per-slot transform metadata (partition_keys, time_key, args, + kwargs) lives on the slot's ``key`` (TransformKey); the generator + slices read it from there. + """ + transform_slots = [ + s for s in slots if isinstance(s.key, TransformKey) + ] + # Topological order: a slot whose TransformKey.input references + # another slot's key must come AFTER that other slot. Walk + # `_iter_slot_deps` to discover dependencies among transform slots. + slot_by_key = {s.key: s for s in transform_slots} + in_degree = {s.id: 0 for s in transform_slots} + deps_of: Dict[str, List[str]] = {s.id: [] for s in transform_slots} + for s in transform_slots: + # The slot's transform depends on whatever transform slots + # appear inside its ValueKey tree (e.g. cumsum(change(...))'s + # cumsum slot depends on the change/time_shift slot). + for dep in _iter_slot_deps(s.key): + if dep is s.key or not isinstance(dep, TransformKey): + continue + dep_slot = slot_by_key.get(dep) + if dep_slot is None: + continue + deps_of[dep_slot.id].append(s.id) + in_degree[s.id] += 1 + # Kahn's algorithm: start from independent layers. + ready = [s.id for s in transform_slots if in_degree[s.id] == 0] + ordered_ids: List[str] = [] + while ready: + nxt = ready.pop(0) + ordered_ids.append(nxt) + for child in deps_of[nxt]: + in_degree[child] -= 1 + if in_degree[child] == 0: + ready.append(child) + # Fallback: any remaining slots (shouldn't happen with the typed + # pipeline's identity-via-key, but guard) get appended in input order. + seen = set(ordered_ids) + for s in transform_slots: + if s.id not in seen: + ordered_ids.append(s.id) + by_id = {s.id: s for s in transform_slots} + return [ + TransformLayer(op=by_id[sid].key.op, slot_ids=[sid]) + for sid in ordered_ids + ] + + +# --------------------------------------------------------------------------- +# Stage 7b.3c — date_range → filter + main-TD disambiguation +# --------------------------------------------------------------------------- + + +def _validate_model_filter( + *, + mf: str, + idx: int, + model: SlayerModel, +) -> FilterPhase: + """Validate a ``SlayerModel.filters`` entry and emit a text-only + ``FilterPhase`` for it. + + Replicates legacy validation (``slayer/engine/enrichment.py:1138-1219``): + + * ``parse_sql_predicate`` rejects DSL constructs (colon aggregation, + transform calls) and raw ``OVER(...)`` window functions. + * Reject references to a ``ModelMeasure`` declared on the same + model — model filters are WHERE-clause SQL, can't reference + aggregates (legacy ``enrichment.py:1147-1153``). + * Reject references to a column whose ``Column.sql`` contains a + window function (legacy ``enrichment.py:1205-1219``). + * DEV-1450 follow-up #4b: references to a NON-windowed derived + ``Column.sql`` column are now accepted — the generator inlines the + column's expanded SQL at render time + (``SQLGenerator._render_model_filter_sql``) and pulls any joins the + expansion crosses into the FROM, matching legacy + ``resolve_filter_columns``. + """ + parsed = parse_sql_predicate(mf) + measure_names = {m.name for m in (model.measures or [])} + windowed_columns = { + c.name for c in model.columns + if c.sql and has_window_function(c.sql) + } + for col in parsed.columns: + if col in measure_names: + raise ValueError( + f"Model filter {mf!r} references measure {col!r}. " + f"Model filters can only reference table columns (WHERE). " + f"Use query-level filters for measure conditions." + ) + if col in windowed_columns: + raise ValueError( + f"Model filter {mf!r} references column {col!r} whose " + f"SQL contains a window function. Factor it into a " + f"multi-stage source_queries model or use a rank-family " + f"transform at query time." + ) + return FilterPhase( + id=f"mf{idx}", + phase=Phase.ROW, + text=mf, + text_columns=tuple(parsed.columns), + expression=None, + ) + + +def _build_date_range_filter( + *, + td: TimeDimension, + scope: ModelScope, + bundle: ResolvedSourceBundle, +) -> BoundFilter: + """Build a row-phase ``BoundFilter`` from a ``TimeDimension``'s + ``date_range``. + + The predicate binds against the bare underlying ``ColumnKey`` + (not the ``TimeTruncKey``) so generator slice 7b.11 can apply the + filter to the outer projection while the shifted self-join CTE + reads raw data. Shape: + + BetweenKey(column=col, low=start, high=end) + + Inclusive on both sides — matches legacy ``column BETWEEN start + AND end``. The typed BetweenKey lets the SQL generator emit + ``exp.Between`` rather than ``col >= start AND col <= end``, + closing the syntactic parity gap with the legacy generator + (DEV-1450 stage 7b.9). + + Bound literals are normalised via ``normalize_scalar``; strings + pass through unchanged. + """ + full = td.dimension.full_name + parsed = parse_expr(full) + bound_col_expr = bind_expr(parsed=parsed, scope=scope, bundle=bundle) + col_key = bound_col_expr.value_key + # DEV-1450 #4a: a derived (Column.sql) temporal column binds to a + # ColumnSqlKey; the BetweenKey accepts both kinds and the generator + # renders a ColumnSqlKey by expanding (`` BETWEEN ...``). + if not isinstance(col_key, (ColumnKey, ColumnSqlKey)): + raise ValueError( + f"date_range filter for TimeDimension {full!r} expected a " + f"column reference; got {type(col_key).__name__}." + ) + + start, end = td.date_range[0], td.date_range[1] + predicate = BetweenKey( + column=col_key, + low=LiteralKey(value=normalize_scalar(start)), + high=LiteralKey(value=normalize_scalar(end)), + ) + refs = tuple(walk_value_keys(predicate)) + phase = max((k.phase for k in refs), default=predicate.phase) + return BoundFilter( + value_key=predicate, phase=phase, referenced_keys=refs, + ) + + +def _resolve_main_time_dimension( + *, + query: SlayerQuery, + model: SlayerModel, +) -> Optional[TimeDimension]: + """Resolve the active time dimension for transform / windowing. + + * 0 TDs → ``None``. + * 1 TD → that TD (``query.main_time_dimension`` is ignored — + matches legacy semantics). + * 2+ TDs: + * ``query.main_time_dimension`` set → match by ``full_name`` + first, then by ``leaf``; raise ``UnknownReferenceError`` if + neither matches. + * Else ``model.default_time_dimension`` set → match by leaf; + return ``None`` if it doesn't match a TD in this query + (legacy graceful no-op — the default points at a column the + user didn't include in this query's time_dimensions). + * Else → ``None``. + """ + tds = list(query.time_dimensions or []) + if not tds: + return None + if len(tds) == 1: + return tds[0] + + if query.main_time_dimension: + target = query.main_time_dimension + # Prefer full-name (more specific) over leaf match. + for td in tds: + if td.dimension.full_name == target: + return td + leaf_matches = [td for td in tds if td.dimension.name == target] + if len(leaf_matches) == 1: + return leaf_matches[0] + if len(leaf_matches) > 1: + # Ambiguous: multiple TDs share the same leaf (e.g. + # ``customers.created_at`` and ``payments.created_at``). + # Force the user to disambiguate via full_name. + raise AmbiguousReferenceError( + name=target, + candidates=[td.dimension.full_name for td in leaf_matches], + ) + raise UnknownReferenceError( + name=target, + scope_kind="TimeDimension", + scope_summary=( + f"time_dimensions: " + f"{[td.dimension.full_name for td in tds]}" + ), + suggestion=None, + ) + + default = model.default_time_dimension + if default: + # Legacy ``_resolve_time_alias`` returns + # ``f"{model.name}.{default_time_dimension}"``, which only points + # at the host model — never at a joined TD. Preserve that: prefer + # a host-local TD (``td.dimension.model is None``) over any + # joined TD that happens to share the leaf name. + for td in tds: + if td.dimension.model is None and td.dimension.name == default: + return td + return None diff --git a/slayer/engine/syntax.py b/slayer/engine/syntax.py new file mode 100644 index 00000000..5c34b019 --- /dev/null +++ b/slayer/engine/syntax.py @@ -0,0 +1,949 @@ +"""Stage 7a.3 (DEV-1450) — Mode-B Python-AST parser. + +Public entry point: ``parse_expr(text: str) -> ParsedExpr``. + +The parser consumes a Mode-B expression string (the SLayer DSL used in +``ModelMeasure.formula``, ``SlayerQuery.measures``, +``SlayerQuery.filters``, …) and emits a typed ``ParsedExpr`` tree. It +is PURE syntax — no scope resolution, no named-measure expansion, no +function-style aggregation rewriting (those are upstream concerns: the +slack normalization layer does function-style → colon; the binder +handles scope and named-measure expansion). + +Pipeline order (per query / model save): + + raw → slack normalize → parse_expr → bind → plan → SQL + +Mode-B grammar: + +* bare identifier (``revenue``) +* dotted path (``customers.regions.name``) +* colon aggregation (``revenue:sum``, ``*:count``, + ``price:weighted_avg(weight=qty)``, ``revenue:last(ordered_at)``) +* transform call (``cumsum``, ``lag``, ``rank``, ``time_shift``, …) +* scalar function (closed allowlist from ``SCALAR_FUNCTIONS``) +* arithmetic / comparison / boolean / unary +* parenthesised grouping + +Rejections (per DEV-1450 spec): + +* Function calls not in SCALAR_FUNCTIONS / transforms / aggregations → + ``UnknownFunctionError``. +* Raw ``OVER(...)`` clauses → ``IllegalWindowInFilterError``. +* ``__`` in any user-supplied identifier → ``ValueError`` (reserved for + internal join-path aliases on the SQL side). +* Chained comparisons (``1 < x < 10``) → ``ValueError``; the user + splits as ``1 < x and x < 10``. + +ParsedExpr family: ``Ref`` / ``DottedRef`` / ``StarSource`` / +``Literal`` / ``AggCall`` / ``TransformCall`` / ``ScalarCall`` / +``Arith`` / ``UnaryOp`` / ``Cmp`` / ``BoolOp``. All are frozen +Pydantic models with value-based equality so tests assert via ``==``. + +Dormant in stage 7a — no engine code calls ``parse_expr`` yet. The +binder (stage 7a.5) is the first consumer. +""" + +from __future__ import annotations + +import ast +import re +from decimal import Decimal +from typing import Any, Dict, Iterator, List, Optional, Tuple, Union + +from pydantic import BaseModel, ConfigDict + +from slayer.core.errors import IllegalWindowInFilterError, UnknownFunctionError +from slayer.core.formula import ALL_TRANSFORMS +from slayer.core.keys import SCALAR_FUNCTIONS + + +# --------------------------------------------------------------------------- +# ParsedExpr family +# --------------------------------------------------------------------------- + + +class _BaseNode(BaseModel): + model_config = ConfigDict(frozen=True, arbitrary_types_allowed=True) + + +class Ref(_BaseNode): + name: str + + +class DottedRef(_BaseNode): + parts: Tuple[str, ...] + + +class StarSource(_BaseNode): + pass + + +class Literal(_BaseNode): + value: Union[Decimal, str, bool, None] = None + + +class TupleLit(_BaseNode): + """A literal-only tuple/list RHS for ``IN`` / ``NOT IN`` filters (DEV-1475). + + Only emitted on the right-hand side of a ``Cmp`` whose op is ``in`` or + ``not in``. Every ``elements`` entry is a ``Literal`` — references and + expressions on the RHS are rejected at parse time so the binder can + fold the predicate into a single ``InKey`` with a tuple of + ``LiteralKey``. Empty tuples are rejected too (an empty IN is a SQL + quirk that varies by dialect; reject early with a clear message). + """ + + elements: Tuple[Literal, ...] + + +class AggCall(_BaseNode): + source: Union[Ref, DottedRef, StarSource] + agg: str + args: Tuple[Any, ...] = () + kwargs: Tuple[Tuple[str, Any], ...] = () + + +class TransformCall(_BaseNode): + op: str + input: Any + args: Tuple[Any, ...] = () + kwargs: Tuple[Tuple[str, Any], ...] = () + + +class ScalarCall(_BaseNode): + name: str + args: Tuple[Any, ...] = () + + +class Arith(_BaseNode): + op: str + left: Any + right: Any + + +class UnaryOp(_BaseNode): + op: str + operand: Any + + +class Cmp(_BaseNode): + op: str + left: Any + right: Any + + +class BoolOp(_BaseNode): + op: str + operands: Tuple[Any, ...] + + +ParsedExpr = Union[ + Ref, DottedRef, StarSource, Literal, TupleLit, + AggCall, TransformCall, ScalarCall, + Arith, UnaryOp, Cmp, BoolOp, +] + + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- + + +_PLACEHOLDER_PREFIX = "__slayer_agg_" +_PLACEHOLDER_RE = re.compile(rf"^{_PLACEHOLDER_PREFIX}(\d+)__$") +_OVER_RE = re.compile(r"\bOVER\s*\(", re.IGNORECASE) +# Python-string-literal matcher (handles backslash escapes). Mode-B expressions +# use Python string syntax, so a SQL-style ('' / "") doubling matcher would +# both miss ``"x \" OVER("`` and false-positive on the ``OVER(`` inside it +# (Codex). Used by ``_normalize_sql_filter_operators``, ``_preprocess_colons``, +# and the raw-``OVER(`` pre-scan. +_PY_STRING_LITERAL_RE = re.compile(r"'(?:\\.|[^'\\])*'|\"(?:\\.|[^\"\\])*\"") +_COLON_AGG_RE = re.compile( + r"(\*|[a-zA-Z_]\w*(?:\.[a-zA-Z_]\w*)*(?:\.\*)?)" # source: * / ident / dotted + r":" + r"([a-zA-Z_]\w*)" + # No (args) consumption — Python's AST handles that. +) + +_FILTER_KEYWORDS = frozenset({"and", "or", "not", "in", "is"}) +_SCAN_TOKEN_RE = re.compile( + r"(?P\s+)" + r"|(?P'(?:\\.|[^'\\])*'|\"(?:\\.|[^\"\\])*\")" + r"|(?P[A-Za-z_]\w*)" + r"|(?P\d+(?:\.\d*)?|\.\d+)" + r"|(?P==|<=|>=|!=)" + r"|(?P=)" + r"|(?P\()" + r"|(?P\))" + r"|(?P\[)" + r"|(?P\])" + r"|(?P,)" + r"|(?P.)", + re.DOTALL, +) + + +_BIN_OP_MAP: Dict[type, str] = { + ast.Add: "+", ast.Sub: "-", ast.Mult: "*", ast.Div: "/", + ast.Mod: "%", ast.Pow: "**", ast.FloorDiv: "//", +} +_CMP_OP_MAP: Dict[type, str] = { + ast.Eq: "==", ast.NotEq: "!=", + ast.Lt: "<", ast.LtE: "<=", + ast.Gt: ">", ast.GtE: ">=", + # ``IS`` / ``IS NOT`` (Codex review): the filter normalizer lowers SQL + # ``IS NULL`` / ``IS NOT NULL`` to Python ``is None`` / ``is not None``; + # without these entries the AST converter raised on ``ast.Is`` / + # ``ast.IsNot`` and any DSL filter using the SQL-style spelling failed + # to plan. The downstream SQL generator renders ``is`` / ``is not`` + # against a ``None`` literal as ``IS NULL`` / ``IS NOT NULL``. + ast.Is: "is", ast.IsNot: "is not", + # DEV-1475: SQL-style ``IN`` / ``NOT IN`` with a literal-tuple RHS. + # ``_normalize_sql_filter_operators`` already lowercases ``IN`` / + # ``NOT IN`` to the Python keywords; the AST then carries ``ast.In`` + # / ``ast.NotIn`` here. The ``ast.Compare`` branch enforces the + # tuple/list-only RHS shape and validates that every element is a + # literal. + ast.In: "in", ast.NotIn: "not in", +} + + +# --------------------------------------------------------------------------- +# Public entry point +# --------------------------------------------------------------------------- + + +def parse_expr(text: str, *, allow_dunder: bool = False) -> ParsedExpr: + """Parse a Mode-B expression string into a ``ParsedExpr``. + + ``allow_dunder`` permits ``__`` in identifiers. It defaults to + ``False`` (P1: Mode-B user input rejects ``__``; use single-dot DSL + paths). The stage planner sets it ``True`` only when binding a + downstream stage against a flat ``StageSchema`` (P5/DEV-1449), whose + columns ARE the ``__``-flattened multi-hop aliases of the upstream + stage (``customers__region``). Legality there is the binder's + concern (the column must exist in the upstream schema). + + Raises: + ValueError: empty input, syntax error, unsupported AST node, + chained comparison, or ``__`` in a user identifier (unless + ``allow_dunder``). + UnknownFunctionError: function call not in + ``SCALAR_FUNCTIONS`` / ``ALL_TRANSFORMS``. + IllegalWindowInFilterError: raw ``OVER(...)`` clause anywhere + in ``text``. + """ + if not text or not text.strip(): + raise ValueError("Empty Mode-B expression.") + + # Scan for a raw window clause AFTER blanking string literals (Python + # syntax, so escapes count), so a value like ``status == 'OVER('`` or + # ``status == "x \" OVER("`` isn't mistaken for window usage (CR / Codex). + if _OVER_RE.search(_PY_STRING_LITERAL_RE.sub("", text)): + raise IllegalWindowInFilterError( + filter_expr=text, + source="raw OVER(...) is not allowed in Mode-B DSL", + suggestion=( + "use a transform instead (rank, percent_rank, dense_rank, " + "ntile, cumsum, lag, lead, time_shift, …)." + ), + ) + + preprocessed, agg_map = _preprocess_colons(text) + + try: + py_ast = ast.parse(preprocessed, mode="eval").body + except SyntaxError as e: + raise ValueError( + f"Invalid Mode-B expression {text!r}: {e}" + ) + + if not allow_dunder: + _reject_dunder_in_ast(py_ast, original=text) + + return _convert(py_ast, agg_map=agg_map, original=text) + + +def parse_filter_expr(text: str, *, allow_dunder: bool = False) -> ParsedExpr: + """Parse a Mode-B *filter* string, accepting SQL operator spellings. + + Filters historically accepted SQL-style operators (``=``, ``<>``, ``NULL``, + and the keyword forms ``AND`` / ``OR`` / ``NOT`` / ``IS`` / ``IN``) + alongside the Python spellings. This wrapper normalizes those to their + Python equivalents (string-literal-aware, so quoted contents are + untouched) and then delegates to :func:`parse_expr`. Measures / order use + ``parse_expr`` directly — only filters get the SQL-operator leniency, + matching the legacy ``parse_filter`` contract. + """ + return parse_expr(_normalize_sql_filter_operators(text), allow_dunder=allow_dunder) + + +def _normalize_sql_filter_operators(text: str) -> str: + """Rewrite SQL operator spellings to Python ones outside string literals. + + ``NULL`` → ``None``; ``IS`` / ``NOT`` / ``AND`` / ``OR`` / ``IN`` → + lowercase; standalone ``=`` → ``==``; ``<>`` → ``!=``. Replicated from the + legacy ``slayer.core.formula._preprocess_sql_operators`` so the typed + pipeline doesn't depend on the module DEV-1452 deletes. + """ + # CR review: use the escape-aware Python-string matcher so backslash- + # escaped quotes don't leak ``IS`` / ``IN`` / ``AND`` rewrites into + # the string body (``"x \" IN ("``). + parts = _PY_STRING_LITERAL_RE.split(text) + literals = _PY_STRING_LITERAL_RE.findall(text) + result: List[str] = [] + for i, part in enumerate(parts): + part = re.sub(r"\bNULL\b", "None", part, flags=re.IGNORECASE) + for kw in ("IS", "NOT", "AND", "OR", "IN"): + part = re.sub(rf"\b{kw}\b", kw.lower(), part, flags=re.IGNORECASE) + part = part.replace("<>", "!=") + # SQL ``||`` concat → Python ``|`` (BitOr), reinterpreted as a + # ``concat(...)`` ScalarCall in ``_convert``. ``|`` binds tighter + # than comparisons in Python just as ``||`` does in SQL, so + # ``a || b = 'x'`` and ``a | b == 'x'`` group identically. + part = part.replace("||", "|") + result.append(part) + if i < len(literals): + result.append(literals[i]) + # DEV-1492: the `=` → `==` rewrite runs on the rejoined string with a + # call-paren-aware scanner that leaves keyword-argument `=` alone + # inside non-scalar calls (transforms / parametric aggregations). Run + # last so the scanner sees lowercased keywords and the post-`<>` / + # post-`||` text — only its own pass can touch literal-spanning + # paren context correctly. + return _rewrite_comparison_equals("".join(result)) + + +def _classify_paren( + hist: List[Tuple[str, str]], +) -> Tuple[bool, Optional[str]]: + """Classify an open ``(`` as a CALL or GROUPING paren. + + A ``(`` is a call paren when the previous significant token is a + bare identifier (callable name) or a callable-suffix token (``)`` + / ``]``). Lowercase keywords (``and`` / ``or`` / ``not`` / ``in`` + / ``is``) do NOT make the next ``(`` a call paren — that's why + ``not(...)`` and ``x in (...)`` carry grouping parens. + + DEV-1492 iteration 3: a colon-aggregation context + (``revenue:first(...)``) makes the call a parametric aggregation + regardless of the callee name. ``first`` and ``last`` sit in both + :data:`ALL_TRANSFORMS` and the built-in aggregation set + (``_AMBIGUOUS_AGG_TRANSFORMS`` in ``slayer/core/formula.py``); + after a ``:`` they are always aggregations, never transforms. + Drop the callee to ``None`` so :func:`_is_kwarg_equals` takes the + aggregation/unknown branch (kwargs preserved after ``(`` or + ``,``). + """ + prev_kind = hist[-1][0] if hist else None + prev_text = hist[-1][1] if hist else "" + is_call = prev_kind == "NAME" or prev_text in (")", "]") + callee = hist[-1][1] if (is_call and prev_kind == "NAME") else None + if callee is not None and len(hist) >= 2 and hist[-2] == ("OTHER", ":"): + callee = None + return is_call, callee + + +def _is_kwarg_equals( + stack: List[Tuple[bool, Optional[str]]], + hist: List[Tuple[str, str]], +) -> bool: + """Whether a lone ``=`` is a Python keyword-argument separator. + + Three callee classes get different treatment (DEV-1492 iteration 2): + + * **Scalar** (``callee in SCALAR_FUNCTIONS``) — never a kwarg; + scalars reject keyword args by design. + * **Transform** (``callee in ALL_TRANSFORMS``) — the first + positional is always the value to transform, so a kwarg can + only appear AFTER a ``,``. This preserves the documented + predicate-input form (``consecutive_periods(status = 'paid')`` + where the SQL ``=`` is part of the predicate, not a kwarg). + * **Aggregation or unknown** — kwarg can be the first arg + (``weighted_avg(weight=qty)``, ``percentile(p=0.5)``), so the + ``=`` may follow either ``(`` or ``,``. + """ + top = stack[-1] if stack else None + if top is None or not top[0]: + return False + callee = top[1] + if callee is not None and callee in SCALAR_FUNCTIONS: + return False + prev_kind = hist[-1][0] if hist else None + if prev_kind != "NAME": + return False + prev_prev_kind = hist[-2][0] if len(hist) >= 2 else None + if callee is not None and callee in ALL_TRANSFORMS: + return prev_prev_kind == "COMMA" + return prev_prev_kind in ("LPAREN", "COMMA") + + +def _push_hist(hist: List[Tuple[str, str]], kind: str, text: str) -> None: + """Append ``(kind, text)`` and trim ``hist`` to the last 2 entries.""" + hist.append((kind, text)) + if len(hist) > 2: + del hist[0] + + +def _handle_pass_through( + m: "re.Match[str]", + out: List[str], + stack: List[Tuple[bool, Optional[str]]], + hist: List[Tuple[str, str]], +) -> None: + """Whitespace and string literals: emit verbatim, don't touch hist.""" + out.append(m.group(0)) + + +def _handle_ident( + m: "re.Match[str]", + out: List[str], + stack: List[Tuple[bool, Optional[str]]], + hist: List[Tuple[str, str]], +) -> None: + ident = m.group(0) + out.append(ident) + _push_hist(hist, "KW" if ident in _FILTER_KEYWORDS else "NAME", ident) + + +def _handle_other( + m: "re.Match[str]", + out: List[str], + stack: List[Tuple[bool, Optional[str]]], + hist: List[Tuple[str, str]], +) -> None: + """Numbers, compound ops (``==``/``<=``/``>=``/``!=``), brackets, and + any single char not otherwise classified (``:``, ``.``, ``+``, ``-``, + ``*``, ``/``, ``%``, ``<``, ``>``, ``!``, ``|``, ``{``, ``}``, ...).""" + text = m.group(0) + out.append(text) + _push_hist(hist, "OTHER", text) + + +def _handle_comma( + m: "re.Match[str]", + out: List[str], + stack: List[Tuple[bool, Optional[str]]], + hist: List[Tuple[str, str]], +) -> None: + out.append(",") + _push_hist(hist, "COMMA", ",") + + +def _handle_lparen( + m: "re.Match[str]", + out: List[str], + stack: List[Tuple[bool, Optional[str]]], + hist: List[Tuple[str, str]], +) -> None: + stack.append(_classify_paren(hist)) + out.append("(") + _push_hist(hist, "LPAREN", "(") + + +def _handle_rparen( + m: "re.Match[str]", + out: List[str], + stack: List[Tuple[bool, Optional[str]]], + hist: List[Tuple[str, str]], +) -> None: + if stack: + stack.pop() + out.append(")") + _push_hist(hist, "OTHER", ")") + + +def _handle_op_eq( + m: "re.Match[str]", + out: List[str], + stack: List[Tuple[bool, Optional[str]]], + hist: List[Tuple[str, str]], +) -> None: + out.append("=" if _is_kwarg_equals(stack, hist) else "==") + _push_hist(hist, "OTHER", "=") + + +_HANDLERS: Dict[str, Any] = { + "ws": _handle_pass_through, + "string": _handle_pass_through, + "ident": _handle_ident, + "number": _handle_other, + "op_eq2": _handle_other, + "op_eq": _handle_op_eq, + "lparen": _handle_lparen, + "rparen": _handle_rparen, + "lbrack": _handle_other, + "rbrack": _handle_other, + "comma": _handle_comma, + "other": _handle_other, +} + + +def _rewrite_comparison_equals(text: str) -> str: + """Rewrite SQL-style ``=`` to Python ``==`` except where Python would + treat the ``=`` as a keyword argument inside a non-scalar call. + + Per the architecture (``docs/architecture/parsing.md`` — scalars + reject kwargs, only ``AggCall`` / ``TransformCall`` carry kwargs), + a lone ``=`` is preserved (kwarg) iff: + + 1. the innermost open paren is a CALL paren (see + :func:`_classify_paren`), + 2. the callee of that innermost call is NOT in + :data:`SCALAR_FUNCTIONS` — scalars never accept kwargs, so a + ``=`` inside them is the user's SQL comparison + (``coalesce(status = 'paid', False)``), + 3. the ``=`` is immediately preceded (skipping whitespace) by an + identifier preceded (skipping whitespace) by the call's ``(`` + or by a ``,`` at that paren depth — Python's keyword-argument + grammar (see :func:`_is_kwarg_equals`). + + Compound operators (``==``, ``<=``, ``>=``, ``!=``) are emitted + verbatim by the tokenizer so their ``=`` is never touched. String + literals are tokenized as a unit (Python single/double-quoted with + backslash escapes) and pass through without perturbing the paren + stack or the previous-significant-token history. + + Token-class dispatch goes through :data:`_HANDLERS` keyed by the + regex's ``lastgroup`` (each token-class group is named and the + arms are mutually exclusive, so ``lastgroup`` is exactly the + matched arm). + """ + out: List[str] = [] + # Each frame: (is_call, callee). ``callee`` is the identifier text + # preceding the call's ``(``, or ``None`` (e.g. a callable expression + # like ``f()(x=1)`` where the prior token is ``)``). + stack: List[Tuple[bool, Optional[str]]] = [] + # Trailing window of the last 2 significant tokens (oldest-first). + # Categories: NAME (bare identifier), KW (lowercase Python keyword), + # LPAREN, COMMA, OTHER (everything else; string literals don't enter + # the history). + hist: List[Tuple[str, str]] = [] + for m in _SCAN_TOKEN_RE.finditer(text): + _HANDLERS[m.lastgroup](m, out, stack, hist) + return "".join(out) + + +# --------------------------------------------------------------------------- +# Reference walk (best-effort textual extraction) +# --------------------------------------------------------------------------- + + +def walk_parsed_refs( + parsed: ParsedExpr, +) -> Iterator[Union[Ref, DottedRef, AggCall]]: + """Yield every reference-bearing leaf node in a ``ParsedExpr`` tree. + + Yields ``Ref`` (bare identifier), ``DottedRef`` (dotted join path), and + ``AggCall`` (colon-syntax aggregation) nodes — the leaves a formula + actually references. This is the scope-free counterpart to the binder's + ``walk_value_keys``: callers that only need the *names* a formula touches + (schema-drift cascade attribution, memory entity tagging) walk the parse + tree directly instead of binding it against a scope. + + Descent rules (chosen to match the legacy ``parse_formula`` / + ``FieldSpec`` walk exactly): + + * ``AggCall`` is yielded as a unit — the aggregation's source / args / + kwargs are NOT descended (``weighted_avg(weight=quantity)`` surfaces + ``price``, never ``quantity``). + * ``TransformCall`` descends ``input`` only; positional args, kwargs, and + ``partition_by`` columns are opaque. + * ``ScalarCall`` descends every positional arg (``coalesce`` / ``nullif`` + wrapping aggregated or bare refs). + * ``Arith`` / ``UnaryOp`` / ``Cmp`` / ``BoolOp`` descend their operands. + * ``Literal`` and ``StarSource`` yield nothing. + """ + if isinstance(parsed, (Ref, DottedRef, AggCall)): + yield parsed + return + if isinstance(parsed, TransformCall): + yield from walk_parsed_refs(parsed.input) + return + if isinstance(parsed, ScalarCall): + for a in parsed.args: + yield from walk_parsed_refs(a) + return + if isinstance(parsed, Arith): + yield from walk_parsed_refs(parsed.left) + yield from walk_parsed_refs(parsed.right) + return + if isinstance(parsed, Cmp): + yield from walk_parsed_refs(parsed.left) + yield from walk_parsed_refs(parsed.right) + return + if isinstance(parsed, UnaryOp): + yield from walk_parsed_refs(parsed.operand) + return + if isinstance(parsed, BoolOp): + for op in parsed.operands: + yield from walk_parsed_refs(op) + return + # Literal / StarSource / TupleLit → no references. + # ``TupleLit`` carries only ``Literal`` elements by construction + # (see the ``ast.Compare`` branch of ``_convert``) so the walk stops + # here without descending — same shape as ``Literal``. + + +# --------------------------------------------------------------------------- +# Internals +# --------------------------------------------------------------------------- + + +def _reject_dunder_in_ast(node: ast.AST, *, original: str) -> None: + """Walk the parsed AST and reject any user identifier containing ``__``. + + Robust to string literals (they're ``ast.Constant`` nodes, not + identifier nodes) and to placeholder names generated by the colon + preprocessor (filtered by ``_PLACEHOLDER_PREFIX``). Walks ``Name`` + (``foo__bar``), ``Attribute.attr`` (``customers.foo__bar``), and + ``keyword.arg`` (``f(weight__bad=…)``). + """ + def _check(token: str) -> None: + if "__" in token and not token.startswith(_PLACEHOLDER_PREFIX): + raise ValueError( + f"Mode-B expression {original!r} contains double-" + f"underscore in identifier {token!r}: `__` is reserved " + f"for internal join-path aliases on the SQL side. Use " + f"single-dot DSL paths (e.g. `customers.region`) in " + f"queries and ModelMeasure formulas." + ) + + for child in ast.walk(node): + if isinstance(child, ast.Name): + _check(child.id) + elif isinstance(child, ast.Attribute): + _check(child.attr) + elif isinstance(child, ast.keyword) and child.arg is not None: + _check(child.arg) + + +def _preprocess_colons( + text: str, +) -> Tuple[str, Dict[int, Tuple[Union[Ref, DottedRef, StarSource], str]]]: + """Replace ``:`` with placeholder identifiers. + + Captures source kind + agg name. Any trailing ``(args)`` is left in + place so Python's AST parses it naturally as a Call. String literal + spans are skipped — the literal text is user data, not DSL syntax. + """ + agg_map: Dict[int, Tuple[Union[Ref, DottedRef, StarSource], str]] = {} + counter = [0] + literal_spans = [ + # CR review: use the escape-aware matcher so backslash-escaped + # quotes don't leak ``:sum`` colon rewrites into the string body. + (m.start(), m.end()) for m in _PY_STRING_LITERAL_RE.finditer(text) + ] + + def _in_literal(pos: int) -> bool: + return any(s <= pos < e for s, e in literal_spans) + + def _replace(match: re.Match) -> str: + if _in_literal(match.start()): + return match.group(0) + source_str = match.group(1) + agg_name = match.group(2) + source: Union[Ref, DottedRef, StarSource] + if source_str == "*": + source = StarSource() + elif "." in source_str: + source = DottedRef(parts=tuple(source_str.split("."))) + else: + source = Ref(name=source_str) + idx = counter[0] + counter[0] += 1 + agg_map[idx] = (source, agg_name) + return f"{_PLACEHOLDER_PREFIX}{idx}__" + + return _COLON_AGG_RE.sub(_replace, text), agg_map + + +def _convert(node: ast.AST, *, agg_map: Dict, original: str) -> ParsedExpr: # NOSONAR(S3776) — one-pass dispatch over ast node kinds (Constant/Name/Compare/BinOp/UnaryOp/BoolOp/Call/Attribute…) producing typed ParsedExpr; the branches are flat and short, and splitting hides the exhaustive ast-kind coverage one read scans for. Surfaces ParsedExpr.kind contract directly. + if isinstance(node, ast.Constant): + return _convert_constant(node, original=original) + + if isinstance(node, ast.Name): + m = _PLACEHOLDER_RE.match(node.id) + if m: + idx = int(m.group(1)) + source, agg = agg_map[idx] + return AggCall(source=source, agg=agg) + return Ref(name=node.id) + + if isinstance(node, ast.Attribute): + parts = _flatten_attribute(node, original=original) + return DottedRef(parts=tuple(parts)) + + if isinstance(node, ast.Call): + return _convert_call(node, agg_map=agg_map, original=original) + + if isinstance(node, ast.BinOp): + op_type = type(node.op) + if op_type is ast.BitOr: + # SQL ``||`` concat operator (``parse_filter_expr`` normalizes + # ``||`` → ``|``). Desugar to the existing ``concat`` scalar + # call so binding + per-dialect SQL emission are fully reused. + return ScalarCall( + name="concat", + args=( + _convert(node.left, agg_map=agg_map, original=original), + _convert(node.right, agg_map=agg_map, original=original), + ), + ) + if op_type not in _BIN_OP_MAP: + raise ValueError( + f"Invalid Mode-B expression {original!r}: unsupported " + f"binary operator {op_type.__name__}." + ) + return Arith( + op=_BIN_OP_MAP[op_type], + left=_convert(node.left, agg_map=agg_map, original=original), + right=_convert(node.right, agg_map=agg_map, original=original), + ) + + if isinstance(node, ast.UnaryOp): + op_type = type(node.op) + if op_type is ast.USub: + return UnaryOp( + op="-", + operand=_convert(node.operand, agg_map=agg_map, original=original), + ) + if op_type is ast.UAdd: + # `+x` is a no-op; collapse to the operand directly. + return _convert(node.operand, agg_map=agg_map, original=original) + if op_type is ast.Not: + return UnaryOp( + op="not", + operand=_convert(node.operand, agg_map=agg_map, original=original), + ) + raise ValueError( + f"Invalid Mode-B expression {original!r}: unsupported unary " + f"operator {op_type.__name__}." + ) + + if isinstance(node, ast.Compare): + if len(node.ops) != 1 or len(node.comparators) != 1: + raise ValueError( + f"Invalid Mode-B expression {original!r}: chained " + f"comparisons are not supported. Each Cmp must be a " + f"single comparison; split (e.g.) `1 < x < 10` into " + f"`1 < x and x < 10`." + ) + op_type = type(node.ops[0]) + if op_type not in _CMP_OP_MAP: + raise ValueError( + f"Invalid Mode-B expression {original!r}: unsupported " + f"comparison operator {op_type.__name__}." + ) + # DEV-1475: ``IN`` / ``NOT IN`` carry a literal-only tuple RHS + # (``status in ('completed', 'pending')``). Reject scalar RHS, + # empty RHS, and any non-literal element so the binder can fold + # the predicate into a single ``InKey`` with confidence. Signed + # numerics (``amount in (-1, -2)``) are admitted by collapsing + # ``UnaryOp(USub/UAdd, Constant(int|float))`` to a signed + # ``Literal`` at validation time (Codex review). + if op_type in (ast.In, ast.NotIn): + rhs_node = node.comparators[0] + if not isinstance(rhs_node, (ast.Tuple, ast.List)): + raise ValueError( + f"Invalid Mode-B expression {original!r}: the right-" + f"hand side of ``in`` / ``not in`` must be a tuple/" + f"list literal (e.g. ``status in ('a', 'b')``); got " + f"{type(rhs_node).__name__}." + ) + if not rhs_node.elts: + raise ValueError( + f"Invalid Mode-B expression {original!r}: empty " + f"tuple is not allowed on the right-hand side of " + f"``in`` / ``not in`` (dialect-dependent SQL); use " + f"a non-empty literal tuple." + ) + elements: List[Literal] = [] + for elt in rhs_node.elts: + converted = _convert_in_rhs_element( + elt, agg_map=agg_map, original=original, + ) + elements.append(converted) + return Cmp( + op=_CMP_OP_MAP[op_type], + left=_convert(node.left, agg_map=agg_map, original=original), + right=TupleLit(elements=tuple(elements)), + ) + return Cmp( + op=_CMP_OP_MAP[op_type], + left=_convert(node.left, agg_map=agg_map, original=original), + right=_convert(node.comparators[0], agg_map=agg_map, original=original), + ) + + if isinstance(node, ast.BoolOp): + op_str = "and" if isinstance(node.op, ast.And) else "or" + operands = tuple( + _convert(v, agg_map=agg_map, original=original) for v in node.values + ) + return BoolOp(op=op_str, operands=operands) + + raise ValueError( + f"Invalid Mode-B expression {original!r}: unsupported AST node " + f"{type(node).__name__}." + ) + + +def _convert_in_rhs_element( + node: ast.AST, *, agg_map: Dict, original: str, +) -> Literal: + """Convert one element of an ``IN`` / ``NOT IN`` literal-tuple RHS. + + The Python parser emits a negative numeric literal as + ``UnaryOp(USub, Constant(int|float))`` rather than a bare + ``Constant`` with a negative value, so a strict ``isinstance(_, + Literal)`` check against ``_convert``'s output would reject + ``amount in (-1, -2)``. This helper collapses the sign onto the + inner numeric before the literal check (Codex review). Boolean and + string literals are unaffected. + """ + if isinstance(node, ast.UnaryOp) and isinstance(node.op, (ast.USub, ast.UAdd)): + inner = node.operand + if ( + isinstance(inner, ast.Constant) + and isinstance(inner.value, (int, float)) + and not isinstance(inner.value, bool) + ): + signed = -inner.value if isinstance(node.op, ast.USub) else inner.value + if isinstance(signed, int): + return Literal(value=Decimal(signed)) + return Literal(value=Decimal(str(signed))) + converted = _convert(node, agg_map=agg_map, original=original) + if not isinstance(converted, Literal): + raise ValueError( + f"Invalid Mode-B expression {original!r}: every element on " + f"the right-hand side of ``in`` / ``not in`` must be a " + f"literal (string, number, or boolean); got " + f"{type(converted).__name__}." + ) + return converted + + +def _convert_constant(node: ast.Constant, *, original: str) -> Literal: + val = node.value + if isinstance(val, bool): + return Literal(value=val) + if val is None: + return Literal(value=None) + if isinstance(val, int): + return Literal(value=Decimal(val)) + if isinstance(val, float): + return Literal(value=Decimal(str(val))) + if isinstance(val, str): + return Literal(value=val) + raise ValueError( + f"Invalid Mode-B expression {original!r}: unsupported literal " + f"type {type(val).__name__}." + ) + + +def _flatten_attribute( + node: ast.Attribute, *, original: str, +) -> List[str]: + parts: List[str] = [node.attr] + cur: ast.AST = node.value + while isinstance(cur, ast.Attribute): + parts.append(cur.attr) + cur = cur.value + if isinstance(cur, ast.Name): + parts.append(cur.id) + else: + raise ValueError( + f"Invalid Mode-B expression {original!r}: unsupported " + f"attribute base {type(cur).__name__}." + ) + return list(reversed(parts)) + + +def _convert_kwarg_value(node: ast.AST, *, agg_map: Dict, original: str): + """Convert a call keyword-argument value. + + List / tuple values (e.g. ``partition_by=[region, channel]`` for the + rank family) convert to a tuple of converted elements so the parser + accepts the documented multi-column transform-kwarg grammar instead of + raising on the bare ``ast.List`` node; scalar values convert normally. + """ + if isinstance(node, (ast.List, ast.Tuple)): + return tuple( + _convert(e, agg_map=agg_map, original=original) for e in node.elts + ) + return _convert(node, agg_map=agg_map, original=original) + + +def _convert_call( + node: ast.Call, *, agg_map: Dict, original: str, +) -> ParsedExpr: + if not isinstance(node.func, ast.Name): + raise ValueError( + f"Invalid Mode-B expression {original!r}: function calls " + f"with non-name callee are not supported." + ) + func_name = node.func.id + + args = tuple( + _convert(a, agg_map=agg_map, original=original) for a in node.args + ) + # Reject ``**kwargs`` dictionary unpacking (``kw.arg is None``) rather + # than silently dropping it (CR) — a dropped ``**`` would change call + # semantics without warning. + if any(kw.arg is None for kw in node.keywords): + raise ValueError( + f"Invalid Mode-B expression {original!r}: dictionary unpacking " + f"(**kwargs) is not supported in calls." + ) + kwargs = tuple( + (kw.arg, _convert_kwarg_value(kw.value, agg_map=agg_map, original=original)) + for kw in node.keywords + if kw.arg is not None # guarded above; narrows kw.arg to str + ) + + # Aggregation placeholder? + m = _PLACEHOLDER_RE.match(func_name) + if m: + idx = int(m.group(1)) + source, agg = agg_map[idx] + return AggCall(source=source, agg=agg, args=args, kwargs=kwargs) + + # Transform? + if func_name in ALL_TRANSFORMS: + if not args: + raise ValueError( + f"Invalid Mode-B expression {original!r}: transform " + f"{func_name!r} requires at least one positional argument " + f"(the value to transform)." + ) + return TransformCall( + op=func_name, + input=args[0], + args=args[1:], + kwargs=kwargs, + ) + + # Scalar function? + if func_name in SCALAR_FUNCTIONS: + if kwargs: + raise ValueError( + f"Invalid Mode-B expression {original!r}: scalar function " + f"{func_name!r} does not accept keyword arguments. Pass " + f"values positionally." + ) + return ScalarCall(name=func_name, args=args) + + # Otherwise — unknown. + raise UnknownFunctionError( + name=func_name, + location=original, + suggestion=( + f"Mode-B accepts only the closed scalar allowlist " + f"({sorted(SCALAR_FUNCTIONS)}), transforms " + f"({sorted(ALL_TRANSFORMS)}), and colon-syntax aggregations " + f"(e.g. `revenue:sum`). Function-style aggregations like " + f"`sum(revenue)` are normalised by the slack layer; if you " + f"see this error for one, slack normalization was bypassed." + ), + ) diff --git a/slayer/engine/variables.py b/slayer/engine/variables.py new file mode 100644 index 00000000..7213e15b --- /dev/null +++ b/slayer/engine/variables.py @@ -0,0 +1,104 @@ +"""Stage 7b.1 (DEV-1450) — variable substitution in the new pipeline. + +Moves the ``{var}`` placeholder substitution that previously lived inside +the legacy enrichment path (``slayer.engine.enrichment.py:1162``) into a +small, pipeline-friendly module. + +Public surface: + +- :func:`merge_query_variables` collapses the four configured variable + layers (model defaults < outer query < stage query < runtime kwarg) + into the effective dict that populates + ``ResolvedSourceBundle.query_variables``. Precedence: runtime > stage > + outer > model_defaults. +- :func:`apply_variables_to_query` returns a copy of the input + ``SlayerQuery`` with ``{var}`` substituted in its ``filters`` list. The + helper always returns a fresh ``SlayerQuery`` instance for predictable + pipeline semantics. ``dry_run_placeholders=True`` fills any unresolved + valid placeholder with the legacy ``"0"`` sentinel instead of raising + — used by save-time dry-run SQL generation. Invalid placeholder names + still raise regardless of ``dry_run_placeholders``. + +Scope deliberately matches the legacy enrichment scope — +``SlayerQuery.filters`` is the only field this helper substitutes into. +Formula text, ``Column.sql``, ``Column.filter``, and +``SlayerModel.filters`` are NOT variable-substituted today, and this +module preserves that contract. + +The dormant module is unwired from the engine in this commit; stage +7b.15 (engine cutover) makes it the substitution path used by +``engine.execute`` and ``engine.save_model``. +""" + +from __future__ import annotations + +from typing import Any, Dict, Optional + +from slayer.core.query import ( + SlayerQuery, + extract_placeholder_names, + substitute_variables, +) + +_PLACEHOLDER_FILL_VALUE = "0" + + +def merge_query_variables( + *, + runtime: Optional[Dict[str, Any]], + stage: Optional[Dict[str, Any]], + outer: Optional[Dict[str, Any]], + model_defaults: Optional[Dict[str, Any]], +) -> Dict[str, Any]: + """Collapse the four variable layers into the effective dict. + + Precedence (highest wins): runtime > stage > outer > model_defaults. + ``None`` and empty-dict layers are identities. + """ + return { + **(model_defaults or {}), + **(outer or {}), + **(stage or {}), + **(runtime or {}), + } + + +def apply_variables_to_query( + *, + query: SlayerQuery, + variables: Optional[Dict[str, Any]] = None, + dry_run_placeholders: bool = False, +) -> SlayerQuery: + """Return a copy of ``query`` with ``{var}`` substituted in ``filters``. + + The returned ``SlayerQuery`` is always a fresh instance, including in + the no-op cases (``query.filters`` is ``None`` / empty / contains no + placeholders). ``variables=None`` is normalized to an empty dict. + When ``dry_run_placeholders=True``, unresolved valid placeholders are + filled with ``"0"`` instead of raising — the legacy save-time + dry-run behaviour. Invalid placeholder names still raise + ``ValueError`` regardless of ``dry_run_placeholders``, because the + dry-run shortcut is for missing *values*, not for bypassing name + validation. + """ + if query.filters is None: + return query.model_copy() + + effective: Dict[str, Any] = dict(variables or {}) + if dry_run_placeholders: + for placeholder in extract_placeholder_names(query): + effective.setdefault(placeholder, _PLACEHOLDER_FILL_VALUE) + + substituted = [ + substitute_variables(filter_str=f, variables=effective) + for f in query.filters + ] + return query.model_copy(update={"filters": substituted}) + + +__all__ = [ + "apply_variables_to_query", + "extract_placeholder_names", + "merge_query_variables", + "substitute_variables", +] diff --git a/slayer/memories/resolver.py b/slayer/memories/resolver.py index 3d676516..832d0183 100644 --- a/slayer/memories/resolver.py +++ b/slayer/memories/resolver.py @@ -33,19 +33,23 @@ from pydantic import BaseModel -from slayer.core.enums import BUILTIN_AGGREGATIONS -from slayer.core.errors import EntityResolutionError -from slayer.core.formula import ( - AggregatedMeasureRef, - ArithmeticField, - FieldSpec, - MixedArithmeticField, - TransformField, - parse_formula, +from slayer.core.errors import ( + EntityResolutionError, + UnknownFunctionError, ) from slayer.core.models import SlayerModel from slayer.core.query import ColumnRef, SlayerQuery, TimeDimension from slayer.core.refs import strip_agg_suffix as _strip_agg_suffix +from slayer.engine.enrichment import _collect_reachable_agg_names +from slayer.engine.normalization import func_style_agg_to_colon +from slayer.engine.syntax import ( + AggCall, + ParsedExpr, + Ref, + StarSource, + parse_expr, + walk_parsed_refs, +) from slayer.memories.models import ( MEMORY_CANONICAL_PREFIX as _MEMORY_PREFIX, _validate_memory_id_charset, @@ -412,17 +416,34 @@ async def resolve_entity( # NOSONAR(S3776) — single linear dispatch matching # --------------------------------------------------------------------------- -def _formula_aggregated_refs(field: FieldSpec) -> Iterable[AggregatedMeasureRef]: - """Yield every ``AggregatedMeasureRef`` reachable inside ``field``.""" - if isinstance(field, AggregatedMeasureRef): - yield field - elif isinstance(field, TransformField): - yield from _formula_aggregated_refs(field.inner) - elif isinstance(field, (ArithmeticField, MixedArithmeticField)): - yield from field.agg_refs.values() - if isinstance(field, MixedArithmeticField): - for _ph, sub in field.sub_transforms: - yield from _formula_aggregated_refs(sub) +def _formula_entity_tokens(parsed: ParsedExpr) -> Iterable[str]: + """Yield every entity *token* referenced by a parsed Mode-B formula. + + Colon-syntax aggregations surface as ``":"`` (``*:count`` + included verbatim), bare and dotted refs surface as their textual form. + Each token is fed one-by-one into ``resolve_entity`` downstream, which + canonicalises it (stripping the agg suffix, collapsing ``*:count`` to the + model, walking dotted join paths). No binding — the resolver does its own + resolution. + + Aggregation args / kwargs (e.g. ``weighted_avg(weight=quantity)``) and + transform partition columns are opaque (legacy parity) — only the + aggregated source / inner value surfaces. + """ + for node in walk_parsed_refs(parsed): + if isinstance(node, AggCall): + source = node.source + if isinstance(source, StarSource): + src_name = "*" + elif isinstance(source, Ref): + src_name = source.name + else: # DottedRef + src_name = ".".join(source.parts) + yield f"{src_name}:{node.agg}" + elif isinstance(node, Ref): + yield node.name + else: # DottedRef + yield ".".join(node.parts) _FILTER_AGG_SUFFIX_RE = re.compile(r":\w+(?:\([^)]*\))?") @@ -528,28 +549,49 @@ def _add(forms: Iterable[str]) -> None: _add(result.canonical_forms) warnings.extend(result.warnings) - # 4. measures — parse each formula and walk for AggregatedMeasureRef. - named_measures = { - m.name: m.formula - for m in source_model.measures - if m.name is not None - } - extra_agg_names = frozenset( - a.name for a in source_model.aggregations - ) | BUILTIN_AGGREGATIONS + # 4. measures — parse each formula (Mode-B DSL) and resolve each + # referenced entity token. Function-style aggregations are rewritten to + # colon syntax first (quiet FUNC_STYLE_AGG slack helper). DEV-1500 — the + # custom-agg name set includes aggregations defined on *joined* models + # too, so entity extraction over a formula like + # ``rolling_avg(customers.score)`` (where ``rolling_avg`` lives on + # ``customers``) still produces the ``customers.score`` token instead of + # falling through to the unknown-function branch. + async def _resolve_join_target_for_resolver( + target_model_name: str, named_queries, # noqa: ARG001 + ): + try: + target = await storage.get_model( + target_model_name, data_source=source_model.data_source, + ) + except Exception: # noqa: BLE001 — best-effort; resolver must not raise + return None + if target is None: + return None + return (None, target) + + try: + reachable_agg_names = await _collect_reachable_agg_names( + model=source_model, + resolve_join_target=_resolve_join_target_for_resolver, + named_queries={}, + ) + except Exception: # noqa: BLE001 — never let the walk break extraction + reachable_agg_names = None + custom_agg_names = reachable_agg_names or frozenset() for m in query.measures or []: if m.formula is None: continue try: - parsed = parse_formula( - m.formula, - extra_agg_names=extra_agg_names, - named_measures=named_measures or None, + parsed = parse_expr( + func_style_agg_to_colon( + m.formula, custom_agg_names=custom_agg_names + ) ) - except ValueError: - # Formula didn't parse as colon syntax — fall back to - # treating the bare formula text as an entity reference - # (handles ``formula="aov"``-style refs to named measures). + # IllegalWindowInFilterError is-a ValueError, so ValueError covers it. + except (ValueError, UnknownFunctionError): + # Not parseable as Mode-B DSL — fall back to treating the whole + # formula text as a single entity reference. result = await resolve_entity( m.formula, storage=storage, @@ -558,14 +600,9 @@ def _add(forms: Iterable[str]) -> None: _add(result.canonical_forms) warnings.extend(result.warnings) continue - for ref in _formula_aggregated_refs(parsed): - agg_token = ( - f"{ref.measure_name}:{ref.aggregation_name}" - if ref.aggregation_name - else ref.measure_name - ) + for token in _formula_entity_tokens(parsed): result = await resolve_entity( - agg_token, + token, storage=storage, source_model=source_model, ) diff --git a/slayer/pg_facade/connection.py b/slayer/pg_facade/connection.py index 322d2537..32320b32 100644 --- a/slayer/pg_facade/connection.py +++ b/slayer/pg_facade/connection.py @@ -18,7 +18,7 @@ import logging import re import struct -from typing import Dict, List, Optional +from typing import Dict, Iterator, List, Optional, Tuple import sqlglot import sqlglot.errors @@ -58,6 +58,153 @@ _BACKEND_PID = 1 _BACKEND_SECRET = 0 _PARAM_PLACEHOLDER = re.compile(r"\$(\d+)") + + +def _skip_line_comment(sql: str, i: int) -> int: + """``-- … newline`` (or EOF).""" + nl = sql.find("\n", i + 2) + return len(sql) if nl < 0 else nl + 1 + + +def _skip_block_comment(sql: str, i: int) -> int: + """``/* … */``, Postgres-style nested.""" + n = len(sql) + depth = 1 + i += 2 + while i < n and depth > 0: + if sql[i] == "/" and i + 1 < n and sql[i + 1] == "*": + depth += 1 + i += 2 + elif sql[i] == "*" and i + 1 < n and sql[i + 1] == "/": + depth -= 1 + i += 2 + else: + i += 1 + return i + + +def _skip_unquoted_identifier(sql: str, i: int) -> int: + """Postgres: letter|_ then alnum|_|$. Treats ``metric$1`` as one token.""" + n = len(sql) + i += 1 + while i < n and (sql[i].isalnum() or sql[i] in ("_", "$")): + i += 1 + return i + + +def _skip_e_string(sql: str, i: int) -> int: + """``E'…'`` / ``e'…'`` — backslash escapes AND ``''`` escape.""" + n = len(sql) + i += 2 # skip prefix + opening quote + while i < n: + if sql[i] == "\\" and i + 1 < n: + i += 2 + continue + if sql[i] == "'": + if i + 1 < n and sql[i + 1] == "'": + i += 2 + continue + return i + 1 + i += 1 + return i + + +def _skip_single_quoted_string(sql: str, i: int) -> int: + """``'…'`` — only the ``''`` doubled-quote escape (no backslash).""" + n = len(sql) + i += 1 + while i < n: + if sql[i] == "'": + if i + 1 < n and sql[i + 1] == "'": + i += 2 + continue + return i + 1 + i += 1 + return i + + +def _skip_double_quoted_identifier(sql: str, i: int) -> int: + """``"…"`` — only the ``""`` escape.""" + n = len(sql) + i += 1 + while i < n: + if sql[i] == '"': + if i + 1 < n and sql[i + 1] == '"': + i += 2 + continue + return i + 1 + i += 1 + return i + + +def _try_skip_dollar_quoted(sql: str, i: int) -> Optional[int]: + """If sql[i:] opens a ``$tag$ … $tag$`` literal, return the index just + past the closing tag; otherwise return None (so the caller can fall + through to ``$N`` placeholder matching). Requires a word boundary + before the opening ``$`` — otherwise ``ident$1`` would be misread.""" + n = len(sql) + prev = sql[i - 1] if i > 0 else "" + if prev.isalnum() or prev == "_": + return None + j = i + 1 + while j < n and (sql[j].isalnum() or sql[j] == "_"): + j += 1 + if j >= n or sql[j] != "$": + return None + tag = sql[i : j + 1] + end = sql.find(tag, j + 1) + return n if end < 0 else end + len(tag) + + +def _handle_dollar( + sql: str, i: int, +) -> Tuple[int, Optional[Tuple[int, int, int]]]: + """Resolve a ``$`` at ``i`` to either a dollar-quoted-string skip or a + ``$N`` placeholder match (or a single-char advance if neither). Returns + ``(new_i, placeholder_or_None)`` so the main walker stays a flat + dispatch.""" + dq_end = _try_skip_dollar_quoted(sql, i) + if dq_end is not None: + return dq_end, None + m = _PARAM_PLACEHOLDER.match(sql, i) + if m: + return m.end(), (int(m.group(1)), m.start(), m.end()) + return i + 1, None + + +def _iter_param_placeholders(sql: str) -> Iterator[Tuple[int, int, int]]: + """Yield ``(param_index, start, end)`` for every ``$N`` placeholder in + ``sql`` that is **not** inside a string literal, quoted identifier, + dollar-quoted string, or line/block comment. + + Mirrors Postgres' tokenizer well enough for standard-shape SQL coming + from libpq / asyncpg / psql / JDBC clients. Closes the regex-only path + that rewrote ``$N`` tokens inside literals and comments (Codex review on + PR #153). Per-lexical-context skip helpers (``_skip_*`` / ``_handle_*``) + own the state-machine; the main loop is a flat ``c → helper`` dispatch. + """ + i = 0 + n = len(sql) + while i < n: + c = sql[i] + if c == "-" and i + 1 < n and sql[i + 1] == "-": + i = _skip_line_comment(sql, i) + elif c == "/" and i + 1 < n and sql[i + 1] == "*": + i = _skip_block_comment(sql, i) + elif c in ("E", "e") and i + 1 < n and sql[i + 1] == "'": + i = _skip_e_string(sql, i) + elif c.isalpha() or c == "_": + i = _skip_unquoted_identifier(sql, i) + elif c == "'": + i = _skip_single_quoted_string(sql, i) + elif c == '"': + i = _skip_double_quoted_identifier(sql, i) + elif c == "$": + i, ph = _handle_dollar(sql, i) + if ph is not None: + yield ph + else: + i += 1 # The single schema the facade advertises (matches pg_namespace / current_schema). PUBLIC_SCHEMA = "public" @@ -413,13 +560,19 @@ def _substitute_params(self, stmt: _PreparedStatement, bind: proto.BindMessage) ) literals.append(literal_for_substitution(value)) - def repl(match: "re.Match[str]") -> str: - idx = int(match.group(1)) + # Walk placeholders in lexical order, skipping ones inside string + # literals / quoted identifiers / dollar-quoted strings / comments. + parts: List[str] = [] + last = 0 + for idx, start, end in _iter_param_placeholders(stmt.sql): + parts.append(stmt.sql[last:start]) if 1 <= idx <= len(literals): - return literals[idx - 1] - return match.group(0) - - return _PARAM_PLACEHOLDER.sub(repl, stmt.sql) + parts.append(literals[idx - 1]) + else: + parts.append(stmt.sql[start:end]) + last = end + parts.append(stmt.sql[last:]) + return "".join(parts) async def _handle_describe(self, msg: proto.DescribeMessage) -> None: if msg.kind == "S": @@ -675,7 +828,7 @@ def _resolve_param_oids(stmt: _PreparedStatement) -> List[int]: """ declared = stmt.parameter_oids max_idx = max( - (int(m.group(1)) for m in _PARAM_PLACEHOLDER.finditer(stmt.sql)), + (idx for idx, _, _ in _iter_param_placeholders(stmt.sql)), default=0, ) count = max(len(declared), max_idx) diff --git a/slayer/sql/generator.py b/slayer/sql/generator.py index f5eb4ba6..d192e97e 100644 --- a/slayer/sql/generator.py +++ b/slayer/sql/generator.py @@ -8,7 +8,7 @@ import copy import logging import re -from typing import List, Optional +from typing import Any, Dict, List, Optional, Set, Tuple import sqlglot from sqlglot import exp @@ -19,8 +19,229 @@ DataType, TimeGranularity, ) +from pydantic import BaseModel, ConfigDict + +from slayer.core.errors import AggregationNotAllowedError +from slayer.core.models import Aggregation +from slayer.core.refs import agg_kwarg_canonical_str +from slayer.engine.column_expansion import ( + _is_trivial_base, + _walk_path_to_target_sync, + collect_root_scope_joined_paths, + expand_derived_refs_sync, +) from slayer.engine.enriched import EnrichedMeasure, EnrichedQuery, public_projection_aliases +from slayer.engine.source_bundle import ( + stage_bundle_with_siblings, + synthetic_model_from_stage_schema, +) from slayer.sql.sqlite_dialect import rewrite_sqlite_json_extract +from slayer.sql.stage_wrapper import build_flat_rename_wrapper + + +class AggRenderSpec(BaseModel): + """DEV-1452 — typed input record for the dialect-aware aggregation + helpers (``_build_agg``, ``_build_percentile``, ``_build_stat_agg``, + ``_build_formula_agg``, ``_resolve_value_sql``, ``_resolve_agg_param``, + ``_build_ranked_subquery_from_planned``). + + Decouples the helpers from ``EnrichedMeasure`` so the legacy enrichment + pipeline can be deleted without forking dialect SQL emission. Carries + exactly the 11 fields the helpers empirically read; ``EnrichedMeasure`` + fields outside this set (``agg_args``, ``source_measure_name``, + ``distinct``, ``window``, ``user_declared``, ``label``, + ``filter_columns``) are deliberately NOT carried — ``count_distinct`` + dispatches on the agg name, and the positional time arg for + ``first`` / ``last`` is pre-resolved into ``time_column`` at spec-build + time. + """ + + model_config = ConfigDict(frozen=True, arbitrary_types_allowed=True) + + sql: str | None + """Column SQL expression (``Column.sql`` or its bare name); ``None`` for + ``*:count`` (renders as ``COUNT(*)``). + + Typed as ``str | None`` (not ``Optional[str]``) deliberately — the + field is **required** at construction; the explicit nullable form + documents that and dodges Sonar's S8396 false-positive on the + Pydantic-v2 ``Optional[X]``-implies-default-None misconception.""" + + name: str + """Source column name — qualified under ``model_name`` when ``sql`` is + None or a bare identifier. Empty for star-source aggregates.""" + + model_name: str + """Qualifier for unqualified column refs in ``sql`` / ``filter_sql`` / + aggregation params — the source relation.""" + + aggregation: str + """Aggregation name (``sum`` / ``count`` / ``percentile`` / …). Empty + string for the non-aggregation bare-column branch.""" + + alias: str + """Result-column alias used by the filtered first/last ranked-subquery + bookkeeping (``filtered_rn_map``, ``filtered_match_map`` lookups).""" + + aggregation_def: Optional[Aggregation] = None + """Custom-aggregation definition (formula + params) for aggregations + outside the built-in set. ``None`` for built-ins.""" + + agg_kwargs: Dict[str, str] = {} + """Query-time aggregation parameter overrides (already stringified via + ``agg_kwarg_canonical_str`` at spec-build time).""" + + filter_sql: Optional[str] = None + """Column-filter predicate (``Column.filter``) wired in at aggregation + time; the helpers wrap the aggregate as ``SUM(CASE WHEN THEN + END)``.""" + + time_column: Optional[str] = None + """Explicit time column for first/last ranking (overrides the query's + default). Pre-resolved from ``AggregateKey.args`` for the planner path.""" + + type: Optional[DataType] = None + """Declared outer-result type — when set, callers wrap the final + aggregate expression in ``CAST AS `` via ``_wrap_cast_for_type``.""" + + column_type: Optional[DataType] = None + """Source column's declared type — wraps the inner (pre-aggregation) + expression in CAST when the column.sql is a non-bare expression (e.g. + ``json_extract(...)``). Distinct from ``type`` which wraps the outer + aggregate.""" + + +class FirstLastRenderState(BaseModel): + """DEV-1501 — bundle of maps produced by + ``_build_first_last_base_select`` (host base) or + ``_render_cross_model_cte`` (cross-model CTE) that the HAVING render + path needs to thread into ``_build_agg`` so a HAVING aggregate + references the same ``_first_rn`` / ``_last_rn{suffix}`` column the + SELECT projects (instead of bare ``_last_rn``, which collapses + distinct time-column specs). + """ + + model_config = ConfigDict(frozen=True, arbitrary_types_allowed=True) + + rn_suffix_map: Dict[str, str] = {} + """Effective-time-column → rn suffix (``""`` / ``"_2"`` / …). Empty + when no first/last aggregates are in scope.""" + + default_time_col_sql: Optional[str] = None + """Fallback time column when a spec has no explicit ``time_column``. + ``None`` when every first/last spec carries an explicit arg.""" + + filtered_rn_map: Dict[str, str] = {} + """Per-spec-alias → dedicated rn column for filtered first/last + aggregates (Column.filter wired in at aggregation time).""" + + filtered_match_map: Dict[str, str] = {} + """Per-spec-alias → match-flag column for filtered first/last.""" + + agg_synth_alias: Optional[str] = None + """DEV-1501 Group A.3 — only set for the cross-model CTE single-agg + case. The cross-model CTE projects exactly one aggregate; if HAVING + references the same key, ``_render_filter_value_key_in_target_scope`` + must rebuild the synth with THIS alias so the ``filtered_rn_map`` / + ``filtered_match_map`` lookups (keyed by the synth alias) hit. Host- + base callers leave this ``None`` and rely on ``aliases_by_slot_id`` + threaded through ``_build_where_having_from_planned`` instead.""" + + +def _iter_first_last_leaves(key) -> "list": # NOSONAR(S3776) — sequential isinstance dispatch over the closed ValueKey union; each branch is the per-type recursion contract for surfacing first/last AggregateKey leaves. Extracting per-type helpers would scatter the contract. + """DEV-1501 (Codex round 3): walk a composite ValueKey for first / + last ``AggregateKey`` leaves. + + Composite aggregate slots (``ArithmeticKey`` / ``ScalarCallKey``) + aren't separately materialised — their operand AggregateKeys are + inlined at the composite render path. Without surfacing the leaves + here, the ranked-subquery builder wouldn't see their distinct time + columns and the composite render would resolve every operand to + bare ``_last_rn``. + + Returns local first/last leaves only (cross-model operands raise in + the composite render path; row / literal / scalar-call / + transform / between / in branches recurse into operands without + surfacing themselves). + """ + from slayer.core.keys import ( + AggregateKey, + ArithmeticKey, + BetweenKey, + InKey, + ScalarCallKey, + ) + + out: list = [] + + def _walk(k) -> None: + if isinstance(k, AggregateKey): + if k.agg in ("first", "last") and not getattr( + k.source, "path", (), + ): + out.append(k) + return + if isinstance(k, ArithmeticKey): + for o in k.operands: + _walk(o) + return + if isinstance(k, ScalarCallKey): + for a in k.args: + _walk(a) + return + if isinstance(k, BetweenKey): + _walk(k.column) + _walk(k.low) + _walk(k.high) + return + if isinstance(k, InKey): + _walk(k.column) + # LiteralKey / ColumnKey / TimeTruncKey / TransformKey / etc.: + # not a first/last operand carrier; stop recursing. + + _walk(key) + return out + + +def _agg_render_spec_from_enriched(em: "EnrichedMeasure") -> AggRenderSpec: + """Adapt a legacy ``EnrichedMeasure`` to the typed ``AggRenderSpec`` for + the refactored dialect helpers (DEV-1452 Stage A). + + Pure field-mapping shim — drops the fields the helpers don't consume + (``agg_args``, ``source_measure_name``, ``distinct``, ``window``, + ``user_declared``, ``label``, ``filter_columns``). The legacy + ``SQLGenerator.generate(enriched=...)`` path uses this to keep emitting + byte-identical SQL through the refactored helpers; the shim is deleted + in Stage D along with the rest of the legacy pipeline. + """ + return AggRenderSpec( + sql=em.sql, + name=em.name, + model_name=em.model_name, + aggregation=em.aggregation, + alias=em.alias, + aggregation_def=em.aggregation_def, + agg_kwargs=dict(em.agg_kwargs), + filter_sql=em.filter_sql, + time_column=em.time_column, + type=em.type, + column_type=em.column_type, + ) + + +def _render_scalar_literal(v: Any) -> exp.Expression: + """Render a Python scalar (None / bool / int / float / Decimal / str) + as a bare sqlglot literal node. Used by the POST-phase filter renderer + for ``LiteralKey.value`` AND any non-key arg inside ``ScalarCallKey``. + """ + from decimal import Decimal + if v is None: + return exp.Null() + if isinstance(v, bool): + return exp.true() if v else exp.false() + if isinstance(v, (int, float, Decimal)): + return exp.Literal.number(str(v)) + return exp.Literal.string(str(v)) def _wrap_cast_for_type(expr: exp.Expression, dt: Optional[DataType]) -> exp.Expression: @@ -49,6 +270,24 @@ def _wrap_cast_for_type(expr: exp.Expression, dt: Optional[DataType]) -> exp.Exp return expr return exp.Cast(this=expr, to=exp.DataType(this=target)) + +def _filter_cast_type(dt: Optional[DataType]) -> Optional[DataType]: + """The CAST target to use when rendering a derived column inside a + WHERE / HAVING predicate (DEV-1450 #4a). + + Temporal types (``DATE`` / ``TIMESTAMP``) are suppressed: in a filter + the derived expression is COMPARED, not type-enforced, and + ``CAST(text AS TIMESTAMP)`` on SQLite gives the expression NUMERIC + affinity — truncating a string timestamp to its leading year and + breaking ``BETWEEN`` / comparison. A base temporal column in the same + position is never cast (it renders as a bare ``exp.Column``), so this + keeps the derived form on par. Non-temporal types pass through so a + derived numeric / boolean column still gets its enforcing CAST. + """ + if dt in (DataType.DATE, DataType.TIMESTAMP): + return None + return dt + logger = logging.getLogger(__name__) # Maps aggregation name (string) → SQL function name. @@ -83,6 +322,26 @@ def _wrap_cast_for_type(expr: exp.Expression, dt: Optional[DataType]) -> exp.Exp # Subset of _STAT_AGG_NAMES that take two columns (LHS + `other=` kwarg). _TWO_ARG_STAT_AGGS: frozenset[str] = frozenset({"corr", "covar_samp", "covar_pop"}) +# DEV-1450 stage 7b.13: aggregations dispatched through the built-in +# path (``_build_agg`` -> ``_build_*`` family). A name in this set always +# resolves to a built-in renderer; a name NOT in the set MUST resolve to +# a model-level ``Aggregation`` definition (``SlayerModel.aggregations``) +# or it's a hard error. Model-level overrides for built-in names ARE +# permitted and get threaded into ``AggRenderSpec.aggregation_def`` so +# ``_resolve_agg_param`` honours their default params (CodeRabbit +# fold-in on DEV-1452 PR #144 — the prior "synth adapter doesn't +# propagate aggregation_def for built-ins" TODO is now done). +# +# Name kept as ``_LOCAL_SLICE`` for grep continuity with 7b.8-7b.12 +# call sites and tests; the set is no longer local-only. +_BUILTIN_BAREARG_AGGS_LOCAL_SLICE: frozenset[str] = frozenset({ + "sum", "avg", "min", "max", "count", "count_distinct", "median", + "percentile", "weighted_avg", + "corr", "covar_samp", "covar_pop", + "stddev_samp", "stddev_pop", "var_samp", "var_pop", + "first", "last", +}) + # DEV-1337: dialects with native single-arg `log10(x)` / `log2(x)`. sqlglot # normalises both into a generic ``Log(this=Literal(base), expression=arg)`` # AST and re-emits as ``LOG(base, x)`` for almost every dialect, which @@ -116,6 +375,11 @@ def _wrap_cast_for_type(expr: exp.Expression, dt: Optional[DataType]) -> exp.Exp # join site that follows the same pattern. _SQL_COL_SEP = ",\n " +# Repeated SQL keyword fragments — extracted so the same literal isn't +# duplicated across CTE / window emission sites (Sonar S1192). +_SQL_WITH = "WITH " +_SQL_PARTITION_BY = "PARTITION BY " + # Matches safe aggregation parameter values: identifiers, qualified names, numeric literals. _SAFE_AGG_PARAM_RE = re.compile( r'^(?:' @@ -656,7 +920,7 @@ def _build_combined(self, enriched: EnrichedQuery, select = select.select(td_expr.as_(td.alias)) group_exprs.append(td_expr) - agg_expr, _ = self._build_agg(measure=cm.measure) + agg_expr, _ = self._build_agg(_agg_render_spec_from_enriched(cm.measure)) # DEV-1361: cast the cross-model agg result if a result type # was declared on the source ModelMeasure. agg_expr = _wrap_cast_for_type(agg_expr, cm.measure.type) @@ -766,7 +1030,7 @@ def _build_combined(self, enriched: EnrichedQuery, group_exprs.append(col_expr) agg_expr, _ = self._build_agg( - measure=unfiltered, + _agg_render_spec_from_enriched(unfiltered), rn_suffix_map=rn_suffix_map, default_time_col=enriched.last_agg_time_column, ) @@ -788,7 +1052,7 @@ def _build_combined(self, enriched: EnrichedQuery, select = select.select(td_expr.as_(td.alias)) group_exprs.append(td_expr) - agg_expr, _ = self._build_agg(measure=unfiltered) + agg_expr, _ = self._build_agg(_agg_render_spec_from_enriched(unfiltered)) agg_expr = _wrap_cast_for_type(agg_expr, measure.type) select = select.select(agg_expr.as_(measure.alias)) @@ -1333,7 +1597,7 @@ def _generate_base(self, enriched: EnrichedQuery, if skip_isolated and (_has_cross_model_filter(measure) or _is_windowed_measure(measure)): continue # Will be handled in its own CTE agg_expr, is_agg = self._build_agg( - measure=measure, + _agg_render_spec_from_enriched(measure), rn_suffix_map=rn_suffix_map, default_time_col=enriched.last_agg_time_column, filtered_rn_map=filtered_rn_map, @@ -1588,7 +1852,7 @@ def _generate_with_computed(self, enriched: EnrichedQuery, # Build final CTE clause cte_strs = [f"{name} AS (\n{sql}\n)" for name, sql in ctes] - cte_clause = "WITH " + ",\n".join(cte_strs) + cte_clause = _SQL_WITH + ",\n".join(cte_strs) final_cte = ctes[-1][0] @@ -1801,7 +2065,7 @@ def _build_transform_sql(t) -> str: # NOSONAR S3776 — flat dispatch over tran time_col = f'"{t.time_alias}"' if t.time_alias else None partition_cols = getattr(t, "partition_aliases", []) or [] partition_clause = ( - "PARTITION BY " + ", ".join(f'"{a}"' for a in partition_cols) + _SQL_PARTITION_BY + ", ".join(f'"{a}"' for a in partition_cols) if partition_cols else "" ) @@ -2136,30 +2400,30 @@ def _resolve_sql( return exp.Column(this=exp.to_identifier(sql), table=exp.to_identifier(model_name)) return _wrap_cast_for_type(self._parse(sql), type) - def _resolve_value_sql(self, measure: "EnrichedMeasure") -> str: - """Resolve ``measure.sql`` (or ``measure.name``) into a fully-qualified + def _resolve_value_sql(self, spec: AggRenderSpec) -> str: + """Resolve ``spec.sql`` (or ``spec.name``) into a fully-qualified SQL string for the value column. Mirrors what ``_build_agg`` does for the standard sum/avg/min/max path so the dialect-aware builders (median/percentile/stat-aggs/formula) emit the same qualified identifiers. """ return self._resolve_sql( - sql=measure.sql, - name=measure.name, - model_name=measure.model_name, - type=measure.column_type, + sql=spec.sql, + name=spec.name, + model_name=spec.model_name, + type=spec.column_type, ).sql(dialect=self.dialect) def _resolve_agg_param( self, - measure: "EnrichedMeasure", + spec: AggRenderSpec, *, name: str, agg_name: str, ) -> str: """Pull a named aggregation parameter, with query-time SQL-injection validation and model-level-default fallback. Returns the SQL string - with bare identifiers qualified under ``measure.model_name`` (via + with bare identifiers qualified under ``spec.model_name`` (via ``_resolve_sql``); qualified names and numeric literals pass through unchanged. Raises ``ValueError`` if neither source supplies the parameter — reused by ``_build_percentile`` (``p=``) and @@ -2167,11 +2431,11 @@ def _resolve_agg_param( ``weight=`` flow. """ raw: Optional[str] = None - if name in measure.agg_kwargs: - raw = measure.agg_kwargs[name] + if name in spec.agg_kwargs: + raw = spec.agg_kwargs[name] _validate_agg_param_value(raw, name, agg_name) - elif measure.aggregation_def: - for param in measure.aggregation_def.params: + elif spec.aggregation_def: + for param in spec.aggregation_def.params: if param.name == name: raw = param.sql break @@ -2182,65 +2446,73 @@ def _resolve_agg_param( f"(e.g., 'measure:{agg_name}({name}=column)')." ) return self._resolve_sql( - sql=raw, name=raw, model_name=measure.model_name, + sql=raw, name=raw, model_name=spec.model_name, ).sql(dialect=self.dialect) def _build_agg( self, - measure: EnrichedMeasure, + spec: AggRenderSpec, rn_suffix_map: Optional[dict[str, str]] = None, default_time_col: Optional[str] = None, filtered_rn_map: Optional[dict[str, str]] = None, filtered_match_map: Optional[dict[str, str]] = None, ) -> tuple[exp.Expression, bool]: - """Build an aggregation expression from an enriched measure.""" - agg_name = measure.aggregation + """Build an aggregation expression from an AggRenderSpec.""" + agg_name = spec.aggregation if not agg_name: # Not an aggregation — raw expression - if measure.sql: + if spec.sql: return self._resolve_sql( - sql=measure.sql, - name=measure.name, - model_name=measure.model_name, - type=measure.column_type, + sql=spec.sql, + name=spec.name, + model_name=spec.model_name, + type=spec.column_type, ), False return exp.Column( - this=exp.to_identifier(measure.name), - table=exp.to_identifier(measure.model_name), + this=exp.to_identifier(spec.name), + table=exp.to_identifier(spec.model_name), ), False # --- first/last: MAX(CASE WHEN _rn = 1 THEN col END) --- if agg_name in ("first", "last"): col_expr = self._resolve_sql( - sql=measure.sql, - name=measure.name, - model_name=measure.model_name, - type=measure.column_type, + sql=spec.sql, + name=spec.name, + model_name=spec.model_name, + type=spec.column_type, ) col = col_expr.sql(dialect=self.dialect) suffix = "" - if rn_suffix_map and default_time_col: - effective_tc = measure.time_column or default_time_col - suffix = rn_suffix_map.get(effective_tc, "") + if rn_suffix_map is not None: + # DEV-1501: when no default ranking time column is in scope, + # every first/last spec is guaranteed to carry an explicit + # ``time_column`` (validated in + # ``_build_first_last_base_select``); so the suffix lookup + # must not gate on ``default_time_col`` being truthy, else + # distinct-time-column specs all collapse to ``_last_rn``. + effective_tc = spec.time_column or default_time_col + if effective_tc is not None: + suffix = rn_suffix_map.get(effective_tc, "") rn_col = f"_first_rn{suffix}" if agg_name == "first" else f"_last_rn{suffix}" # For filtered first/last, use the dedicated ROW_NUMBER column # that pushes non-matching rows to the bottom of the ranking. - # Look up by alias (unique per enriched measure) so two filtered - # measures sharing source/agg but with different filters map to - # their own respective rank columns. Use the per-measure match - # flag (also projected by the ranked subquery) instead of - # re-emitting measure.filter_sql here — the filter can reference - # joined-table columns that are not in scope outside the subquery. - if measure.filter_sql and filtered_rn_map: - filtered_rn = filtered_rn_map.get(measure.alias, rn_col) + # Look up by alias (unique per spec) so two filtered specs + # sharing source/agg but with different filters map to their + # own respective rank columns. Use the per-spec match flag + # (also projected by the ranked subquery) instead of + # re-emitting spec.filter_sql here — the filter can reference + # joined-table columns that are not in scope outside the + # subquery. + if spec.filter_sql and filtered_rn_map: + filtered_rn = filtered_rn_map.get(spec.alias, rn_col) match_col = ( - filtered_match_map.get(measure.alias) + filtered_match_map.get(spec.alias) if filtered_match_map else None ) # Fall back to the raw filter expression only if no match flag # was projected (legacy callers); accepts the leak risk. - filter_clause = f"{match_col} = 1" if match_col else measure.filter_sql + filter_clause = f"{match_col} = 1" if match_col else spec.filter_sql case_sql = ( f"MAX(CASE WHEN {filtered_rn} = 1 AND {filter_clause} " f"THEN {col} END)" @@ -2248,7 +2520,7 @@ def _build_agg( else: # ``col`` is already a fully-qualified SQL expression resolved # via ``_resolve_sql`` earlier in this branch, so we don't need - # to re-prefix ``measure.model_name``. (DEV-1333.) + # to re-prefix ``spec.model_name``. (DEV-1333.) case_sql = f"MAX(CASE WHEN {rn_col} = 1 THEN {col} END)" return self._parse(case_sql), True @@ -2258,39 +2530,39 @@ def _build_agg( # SQLite/ClickHouse/MySQL) so it gets its own builder rather than # going through the BUILTIN_AGGREGATION_FORMULAS path. if agg_name == "percentile": - return self._build_percentile(measure), True + return self._build_percentile(spec), True # Statistical aggregates also dispatch to a dedicated builder so # the SQLite-UDF / native-function / NotImplementedError split # mirrors _build_median. if agg_name in _STAT_AGG_NAMES: - return self._build_stat_agg(measure), True - return self._build_formula_agg(measure, agg_name), True + return self._build_stat_agg(spec), True + return self._build_formula_agg(spec, agg_name), True # --- Resolve inner expression --- - if agg_name == "count" and measure.sql is None: + if agg_name == "count" and spec.sql is None: # COUNT(*) — if filtered, use COUNT(CASE WHEN filter THEN 1 END) - if measure.filter_sql: - case_sql = f"CASE WHEN {measure.filter_sql} THEN 1 END" + if spec.filter_sql: + case_sql = f"CASE WHEN {spec.filter_sql} THEN 1 END" inner = self._parse(case_sql) else: inner = exp.Star() - elif measure.sql: + elif spec.sql: inner = self._resolve_sql( - sql=measure.sql, - name=measure.name, - model_name=measure.model_name, - type=measure.column_type, + sql=spec.sql, + name=spec.name, + model_name=spec.model_name, + type=spec.column_type, ) else: inner = exp.Column( - this=exp.to_identifier(measure.name), - table=exp.to_identifier(measure.model_name), + this=exp.to_identifier(spec.name), + table=exp.to_identifier(spec.model_name), ) - # --- Apply measure-level filter as CASE WHEN wrapper --- - if measure.filter_sql and not (agg_name == "count" and measure.sql is None): + # --- Apply spec-level filter as CASE WHEN wrapper --- + if spec.filter_sql and not (agg_name == "count" and spec.sql is None): inner_sql = inner.sql(dialect=self.dialect) - case_sql = f"CASE WHEN {measure.filter_sql} THEN {inner_sql} END" + case_sql = f"CASE WHEN {spec.filter_sql} THEN {inner_sql} END" inner = self._parse(case_sql) # --- count_distinct --- @@ -2313,12 +2585,12 @@ def _build_agg( agg_class = agg_class_map[agg_func] return agg_class(this=inner), True - def _build_formula_agg(self, measure: EnrichedMeasure, agg_name: str) -> exp.Expression: + def _build_formula_agg(self, spec: AggRenderSpec, agg_name: str) -> exp.Expression: """Build SQL for formula-based aggregations (weighted_avg, custom).""" # Get formula: from aggregation_def or built-in formula = None - if measure.aggregation_def and measure.aggregation_def.formula: - formula = measure.aggregation_def.formula + if spec.aggregation_def and spec.aggregation_def.formula: + formula = spec.aggregation_def.formula elif agg_name in BUILTIN_AGGREGATION_FORMULAS: formula = BUILTIN_AGGREGATION_FORMULAS[agg_name] @@ -2330,12 +2602,12 @@ def _build_formula_agg(self, measure: EnrichedMeasure, agg_name: str) -> exp.Exp # Collect param values: query-time overrides > aggregation_def defaults param_defaults = {} - if measure.aggregation_def: - param_defaults = {p.name: p.sql for p in measure.aggregation_def.params} - params = {**param_defaults, **measure.agg_kwargs} + if spec.aggregation_def: + param_defaults = {p.name: p.sql for p in spec.aggregation_def.params} + params = {**param_defaults, **spec.agg_kwargs} # Validate query-time parameter values to prevent SQL injection - for pname, pval in measure.agg_kwargs.items(): + for pname, pval in spec.agg_kwargs.items(): _validate_agg_param_value(pval, pname, agg_name) # Validate required params @@ -2349,22 +2621,22 @@ def _build_formula_agg(self, measure: EnrichedMeasure, agg_name: str) -> exp.Exp ) # Resolve {value} and {param_name} via _resolve_sql so bare identifiers - # are qualified under measure.model_name (matching the standard - # sum/avg/min/max path). When the measure carries a row-level filter, + # are qualified under spec.model_name (matching the standard + # sum/avg/min/max path). When the spec carries a row-level filter, # wrap row-level references (the value AND any column-ref params) in # CASE WHEN so non-matching rows contribute NULL to all terms — but # leave literal-default params unwrapped, since `(CASE WHEN ... THEN # 100 END)` for a constant `scale=100` would turn it into a row # expression and break grouped SQL semantics. - col_expr = _wrap_filter(self._resolve_value_sql(measure), measure.filter_sql) + col_expr = _wrap_filter(self._resolve_value_sql(spec), spec.filter_sql) substituted = formula.replace("{value}", col_expr) for param_name, param_val in params.items(): param_ast = self._resolve_sql( - sql=param_val, name=param_val, model_name=measure.model_name, + sql=param_val, name=param_val, model_name=spec.model_name, ) param_expr = param_ast.sql(dialect=self.dialect) - if measure.filter_sql and not isinstance(param_ast, exp.Literal): - param_expr = _wrap_filter(param_expr, measure.filter_sql) + if spec.filter_sql and not isinstance(param_ast, exp.Literal): + param_expr = _wrap_filter(param_expr, spec.filter_sql) substituted = substituted.replace(f"{{{param_name}}}", param_expr) return self._parse(substituted) @@ -2385,20 +2657,20 @@ def _build_median(self, inner: exp.Expression) -> exp.Expression: # Postgres, DuckDB, and most others: PERCENTILE_CONT return self._parse(f"PERCENTILE_CONT(0.5) WITHIN GROUP (ORDER BY {inner_sql})") - def _build_percentile(self, measure: "EnrichedMeasure") -> exp.Expression: + def _build_percentile(self, spec: AggRenderSpec) -> exp.Expression: """Build a PERCENTILE_CONT(p) aggregation expression (dialect-dependent). - ``p`` comes from ``measure.agg_kwargs['p']`` (validated against + ``p`` comes from ``spec.agg_kwargs['p']`` (validated against SQL injection) or from a model-level ``Aggregation`` default. - Filter handling mirrors ``_build_formula_agg``: when the measure + Filter handling mirrors ``_build_formula_agg``: when the spec carries a row-level filter, the value column is wrapped in ``CASE WHEN ... END`` so non-matching rows contribute NULL and are ignored by the aggregate. Both the value column and ``p`` flow through ``_resolve_sql`` so bare identifiers are qualified - under ``measure.model_name`` and numeric literals pass through + under ``spec.model_name`` and numeric literals pass through unchanged. """ - p = self._resolve_agg_param(measure, name="p", agg_name="percentile") + p = self._resolve_agg_param(spec, name="p", agg_name="percentile") # `p` must be a numeric literal in [0, 1]. Without this guard a # caller could pass `measure:percentile(p=quantity)` (or a model- # level default like `p=pg_sleep(10)` that bypasses @@ -2426,7 +2698,7 @@ def _build_percentile(self, measure: "EnrichedMeasure") -> exp.Expression: "Use MariaDB or compute the value client-side." ) - col_expr = _wrap_filter(self._resolve_value_sql(measure), measure.filter_sql) + col_expr = _wrap_filter(self._resolve_value_sql(spec), spec.filter_sql) if self.dialect == "sqlite": # Provided by the percentile_cont(value, p) UDF registered on connect. @@ -2439,7 +2711,7 @@ def _build_percentile(self, measure: "EnrichedMeasure") -> exp.Expression: return self._parse(sql_str) - def _build_stat_agg(self, measure: "EnrichedMeasure") -> exp.Expression: + def _build_stat_agg(self, spec: AggRenderSpec) -> exp.Expression: """Build SQL for the statistical aggregations added in DEV-1317. Handles ``stddev_samp``, ``stddev_pop``, ``var_samp``, ``var_pop`` @@ -2452,14 +2724,14 @@ def _build_stat_agg(self, measure: "EnrichedMeasure") -> exp.Expression: SQLite) so generator output resolves at runtime. Both legs flow through ``_resolve_sql`` so bare identifiers are - qualified under ``measure.model_name`` (matches the standard + qualified under ``spec.model_name`` (matches the standard sum/avg/min/max path). Filter handling mirrors ``_build_percentile`` / ``_build_formula_agg``: a row-level filter wraps the value AND the ``other`` column in ``CASE WHEN filter THEN col END`` so non-matching rows contribute NULL — which the aggregates skip. """ - agg_name = measure.aggregation + agg_name = spec.aggregation # Resolve the `other=` kwarg before the MySQL guard so that a # missing-required-param error takes priority over the @@ -2469,8 +2741,8 @@ def _build_stat_agg(self, measure: "EnrichedMeasure") -> exp.Expression: other_expr: Optional[str] = None if agg_name in _TWO_ARG_STAT_AGGS: other_expr = _wrap_filter( - self._resolve_agg_param(measure, name="other", agg_name=agg_name), - measure.filter_sql, + self._resolve_agg_param(spec, name="other", agg_name=agg_name), + spec.filter_sql, ) if agg_name in _TWO_ARG_STAT_AGGS and self.dialect == "mysql": @@ -2480,7 +2752,7 @@ def _build_stat_agg(self, measure: "EnrichedMeasure") -> exp.Expression: f"Use MariaDB or compute the value client-side." ) - col_expr = _wrap_filter(self._resolve_value_sql(measure), measure.filter_sql) + col_expr = _wrap_filter(self._resolve_value_sql(spec), spec.filter_sql) if agg_name in _TWO_ARG_STAT_AGGS: sql_str = f"{agg_name.upper()}({col_expr}, {other_expr})" @@ -2554,7 +2826,7 @@ def _build_where_and_having( for m in enriched.measures: if m.name == col_name: agg_expr, _ = self._build_agg( - measure=m, + _agg_render_spec_from_enriched(m), rn_suffix_map=rn_suffix_map, default_time_col=enriched.last_agg_time_column, filtered_rn_map=filtered_rn_map, @@ -2598,3 +2870,7262 @@ def _build_where_and_having( having_clause = self._parse_predicate(having_sql) return where_clause, having_clause + + # ====================================================================== + # DEV-1450 stage 7b.8 — PlannedQuery → SQL. + # + # The legacy generator (everything above) consumes EnrichedQuery. This + # new entry point consumes the typed PlannedQuery from + # slayer/engine/stage_planner.py. The two paths coexist until the + # engine cutover (stage 7b.15) flips the default path. + # + # 7b.8 scope: local-only single-model queries — row-phase dims, local + # aggregates, Mode-B row filters, ORDER BY / LIMIT / OFFSET, dim-only + # dedup. Cross-model, time dimensions, transforms, and aggregate + # filtering raise NotImplementedError with an explicit stage marker + # so silent parity drift is impossible. + # ====================================================================== + + def generate_from_planned( # NOSONAR(S3776) — top-level dispatch over cross-model / transform-chain / plain branches plus the conditional outer-trim wrap. Each branch is a coherent compilation strategy; extracting would scatter the shared planned_query / slots_by_id / aliases_by_slot_id state across helpers without simplifying anything. + self, + planned_query, + *, + bundle, + ) -> str: + """Render a typed ``PlannedQuery`` to SQL. + + Mirrors the local-only branch of ``_generate_base`` but reads + from typed PlannedQuery fields (``row_slots`` / ``aggregate_slots`` + / ``filters_by_phase`` / ``order`` / ``transform_layers``) + instead of ``EnrichedQuery``. Reuses legacy dialect helpers + (``_resolve_sql`` / ``_build_agg`` / ``_wrap_cast_for_type`` / + ``_parse_predicate`` / ``_build_date_trunc``) so dialect-specific + behavior is rendered identically to the legacy ``generate()`` + path — the parity oracle in ``tests/parity_oracle.py`` pins + this contract. + + Stage 7b.10 adds window-transform rendering: when + ``planned_query.transform_layers`` is non-empty, the base SELECT + is emitted as ``WITH base AS (...)``, Kahn-batched step CTEs + carry the window functions, and an outer wrap projects in + user-spec order. POST-phase filters that reference transform + slots wrap as ``SELECT * FROM (...) AS _filtered WHERE ...``. + ``time_shift`` / ``consecutive_periods`` layers raise + ``NotImplementedError`` with a ``7b.11`` marker. + """ + + source_model = bundle.source_model + if source_model is None: + raise ValueError( + "generate_from_planned requires bundle.source_model to be set", + ) + source_relation = planned_query.source_relation + + if planned_query.cross_model_aggregate_plans: + return self._render_with_cross_model_plans( + planned_query=planned_query, bundle=bundle, + ) + + # 7b.10 — fail fast on transform ops this slice does not render + # (time_shift / consecutive_periods belong to 7b.11). Walks + # ``transform_layers`` for an explicit op match AND walks every + # ``TransformKey.input`` reachable from public slots so a + # ``change`` desugared into ``time_shift`` raises with the same + # marker. + self._validate_window_transform_ops_for_7b10( + planned_query=planned_query, + ) + + slots_by_id = { + s.id: s + for s in ( + list(planned_query.row_slots) + + list(planned_query.aggregate_slots) + + list(planned_query.combined_expression_slots) + ) + } + + # 7b.10 — slot key -> id lookup. ``PlannedQuery`` does not carry + # the ``ValueRegistry``, so the generator builds its own map. + # Used for resolving ``TransformKey.input`` / ``partition_keys`` / + # ``time_key`` references to step-CTE aliases. + slot_id_by_key: Dict[Any, str] = { + s.key: s.id for s in slots_by_id.values() + } + + public_proj_set: Set[str] = set(planned_query.projection) + # 7b.10 / DEV-1501 — base CTE projects hidden slots referenced as + # transform inputs / partition_keys / time_key / filter operands + # (AGGREGATE + POST phase) / order targets so step CTEs, HAVING, + # and the outer ORDER BY can name them. In the NO-transform path + # we additionally pass ``aggregates_only=True`` so only + # AggregateKey leaves get pulled in from order/filter walks — a + # hidden ROW order target (e.g. ``ORDER BY customer_id`` with + # ``customer_id`` not projected) would otherwise materialise into + # GROUP BY and silently change query grain. Hidden ROW order + # targets in the no-transform path keep raising NotImplementedError + # at the inline ORDER BY render path. + no_transform = not bool(planned_query.transform_layers) + extra_materialize_ids = self._collect_base_aux_slot_ids( + planned_query=planned_query, + slot_id_by_key=slot_id_by_key, + slots_by_id=slots_by_id, + include_order=True, + aggregates_only=no_transform, + ) + base_render_order = list(planned_query.projection) + [ + sid for sid in extra_materialize_ids if sid not in public_proj_set + ] + + # Build the base SELECT body. ``aliases_by_slot_id`` is a list + # of full aliases per slot, in projection visit order — needed + # so duplicate public_aliases on a single interned slot (DEV-1450 + # C13: two declared measures with the same key + different names) + # survive the CTE chain. ``available_alias_by_slot_id`` is the + # canonical "pick one" map used by transform-input / time-key / + # partition-key / order-entry lookups (any alias of the slot + # refers to the same column value, so any will do). + ( + base_select, + aliases_by_slot_id, + has_aggregation, + group_by_keys, + where_consumed, + first_last_state, + ) = self._build_base_select_for_planned( + planned_query=planned_query, + bundle=bundle, + source_model=source_model, + source_relation=source_relation, + base_render_order=base_render_order, + slots_by_id=slots_by_id, + ) + + where_clause, having_clause = self._build_where_having_from_planned( + planned_query=planned_query, + source_relation=source_relation, + source_model=source_model, + bundle=bundle, + first_last_state=first_last_state, + aliases_by_slot_id=aliases_by_slot_id, + ) + + # ``where_consumed`` is True for the first/last ranked-subquery path: + # the WHERE is applied INSIDE the ranked subquery (it must filter raw + # rows before ranking), so re-applying it on the outer SELECT would be + # both redundant and — for filters that should narrow the ranked set — + # semantically wrong. + if where_clause is not None and not where_consumed: + base_select = base_select.where(where_clause) + + # Match legacy _generate_base:1375 — dim-only-dedup OR + # has_aggregation triggers GROUP BY (dim-only emits GROUP BY + # before LIMIT so unique dim tuples can't silently drop past + # row N). + dim_only_dedup = bool(group_by_keys) and not has_aggregation + needs_group_by = has_aggregation or dim_only_dedup + if needs_group_by and group_by_keys: + for gb in group_by_keys.values(): + base_select = base_select.group_by(gb) + + if having_clause is not None: + base_select = base_select.having(having_clause) + + # No transforms → existing pre-7b.10 path: apply ORDER/LIMIT + # directly on the base select. DEV-1501: when the base + # materialised hidden order/filter aggregate slots (slot ids in + # ``base_render_order`` not in ``planned_query.projection``), + # wrap the base in an outer SELECT that trims to the public + # projection and moves ORDER BY / LIMIT / OFFSET to the outer + # level — mirrors the transform path's outer wrap shape, minus + # the step CTE chain. + if not planned_query.transform_layers: + public_slot_ids = set(planned_query.projection) + has_hidden_materialised = any( + sid not in public_slot_ids for sid in base_render_order + ) + if has_hidden_materialised: + return self._build_outer_trim_wrap_sql( + base_select=base_select, + planned_query=planned_query, + source_relation=source_relation, + aliases_by_slot_id=aliases_by_slot_id, + slots_by_id=slots_by_id, + bundle=bundle, + ) + base_select = self._apply_order_limit_from_planned( + select=base_select, + planned_query=planned_query, + source_relation=source_relation, + slots_by_id=slots_by_id, + source_model=source_model, + bundle=bundle, + aliases_by_slot_id=aliases_by_slot_id, + ) + return base_select.sql(dialect=self.dialect, pretty=True) + + # 7b.10 — transform layers present. Build the CTE chain. + base_cte_sql = base_select.sql(dialect=self.dialect, pretty=True) + ctes: list[tuple[str, str]] = [("base", base_cte_sql)] + # "Pick one" map for transform-input / time-key / partition-key / + # order-entry / POST-filter lookups. Initialised from the first + # alias of every materialised slot. + available_alias_by_slot_id: Dict[str, str] = { + sid: aliases[0] + for sid, aliases in aliases_by_slot_id.items() + if aliases + } + + pending_layers = list(planned_query.transform_layers) + step_num = 0 + # 7b.11 — gather a global view of WHERE-able row-phase filters + # for the shifted CTE (which re-aggregates the source and needs + # the same WHERE minus BetweenKey date_range filters). Built + # once outside the loop since the source filters don't change + # across layers. + shifted_where_parts = self._build_shifted_cte_where_parts( + planned_query=planned_query, + source_relation=source_relation, + source_model=source_model, + bundle=bundle, + ) + while pending_layers: + ready_window: list = [] + ready_time_shift: list = [] + ready_cp: list = [] + not_ready: list = [] + for layer in pending_layers: + if not self._transform_layer_deps_ready( + layer=layer, + slots_by_id=slots_by_id, + slot_id_by_key=slot_id_by_key, + available_alias_by_slot_id=available_alias_by_slot_id, + ): + not_ready.append(layer) + elif layer.op == "time_shift": + ready_time_shift.append(layer) + elif layer.op == "consecutive_periods": + ready_cp.append(layer) + else: + ready_window.append(layer) + if not (ready_window or ready_time_shift or ready_cp): + pending_ops = [layer.op for layer in pending_layers] + raise RuntimeError( + f"DEV-1450 stage 7b.11: transform layer dependencies " + f"could not be resolved; pending ops: {pending_ops!r}.", + ) + # --- Window batch (one step CTE per Kahn batch) ---------- + if ready_window: + step_num += 1 + step_name = f"step{step_num}" + prev_cte = ctes[-1][0] + carry_aliases_sorted = sorted( + a for aliases in aliases_by_slot_id.values() for a in aliases + ) + step_parts = [f'"{a}"' for a in carry_aliases_sorted] + for layer in ready_window: + for slot_id in layer.slot_ids: + slot = slots_by_id[slot_id] + alias = ( + slot.public_aliases[0] + if slot.public_aliases + else slot.declared_name + ) + full_alias = f"{source_relation}.{alias}" + window_sql = self._render_window_transform_sql( + slot=slot, + slots_by_id=slots_by_id, + slot_id_by_key=slot_id_by_key, + available_alias_by_slot_id=available_alias_by_slot_id, + planned_query=planned_query, + ) + if slot.type is not None: + wrapped = _wrap_cast_for_type( + self._parse(window_sql), slot.type, + ) + window_sql = wrapped.sql(dialect=self.dialect) + step_parts.append(f'{window_sql} AS "{full_alias}"') + aliases_by_slot_id.setdefault(slot_id, []).append( + full_alias, + ) + available_alias_by_slot_id.setdefault( + slot_id, full_alias, + ) + step_sql = ( + "SELECT\n " + + _SQL_COL_SEP.join(step_parts) + + f"\nFROM {prev_cte}" + ) + ctes.append((step_name, step_sql)) + # --- time_shift layers (each gets shifted_ + sjoin_ pair) - + for layer in ready_time_shift: + for slot_id in layer.slot_ids: + slot = slots_by_id[slot_id] + self._emit_time_shift_ctes_for_planned( + slot=slot, + ctes=ctes, + slots_by_id=slots_by_id, + slot_id_by_key=slot_id_by_key, + available_alias_by_slot_id=available_alias_by_slot_id, + aliases_by_slot_id=aliases_by_slot_id, + source_model=source_model, + source_relation=source_relation, + shifted_where_parts=shifted_where_parts, + planned_query=planned_query, + bundle=bundle, + ) + # --- consecutive_periods layers (cp_reset_ + cp_value_ pair) + for layer in ready_cp: + for slot_id in layer.slot_ids: + slot = slots_by_id[slot_id] + self._emit_consecutive_periods_ctes_for_planned( + slot=slot, + ctes=ctes, + slots_by_id=slots_by_id, + slot_id_by_key=slot_id_by_key, + available_alias_by_slot_id=available_alias_by_slot_id, + aliases_by_slot_id=aliases_by_slot_id, + planned_query=planned_query, + source_relation=source_relation, + ) + pending_layers = not_ready + + # 7b.11 — materialise POST-phase ArithmeticKey / ScalarCallKey + # slots that the user projected but no transform layer rendered. + # ``change(amount:sum)`` lowers to ``amount:sum - time_shift(...)``; + # the time_shift slot is rendered as a self-join CTE pair, but + # the outer ArithmeticKey slot that subtracts them needs its + # own step CTE. Same shape covers ``change_pct`` (division of + # arithmetic operands) and any future POST-phase non-transform + # slot the planner emits. + from slayer.core.keys import ( + ArithmeticKey as _ArithKey, + ScalarCallKey as _ScalarKey, + TransformKey as _TKey, + ) + unmaterialised: list = [] + for cslot in planned_query.combined_expression_slots: + if isinstance(cslot.key, _TKey): + # Transform-key slots are materialised by transform_layers. + continue + if cslot.id in aliases_by_slot_id: + continue + if isinstance(cslot.key, (_ArithKey, _ScalarKey)): + unmaterialised.append(cslot) + if unmaterialised: + step_num += 1 + step_name = f"step{step_num}" + prev_cte = ctes[-1][0] + carry_aliases_sorted = sorted( + a for aliases in aliases_by_slot_id.values() for a in aliases + ) + step_parts = [f'"{a}"' for a in carry_aliases_sorted] + for cslot in unmaterialised: + alias = ( + cslot.public_aliases[0] + if cslot.public_aliases + else cslot.declared_name + ) + full_alias = f"{source_relation}.{alias}" + rendered = self._render_value_key_against_aliases( + key=cslot.key, + slot_id_by_key=slot_id_by_key, + available_alias_by_slot_id=available_alias_by_slot_id, + ) + expr_sql = rendered.sql(dialect=self.dialect) + if cslot.type is not None: + wrapped = _wrap_cast_for_type( + self._parse(expr_sql), cslot.type, + ) + expr_sql = wrapped.sql(dialect=self.dialect) + step_parts.append(f'{expr_sql} AS "{full_alias}"') + aliases_by_slot_id.setdefault(cslot.id, []).append( + full_alias, + ) + available_alias_by_slot_id.setdefault( + cslot.id, full_alias, + ) + step_sql = ( + "SELECT\n " + + _SQL_COL_SEP.join(step_parts) + + f"\nFROM {prev_cte}" + ) + ctes.append((step_name, step_sql)) + + # Inner SELECT inside _outer wrap: ALL carried aliases sorted + # (matches legacy _generate_with_computed:1607). + final_cte = ctes[-1][0] + inner_sorted = sorted( + a for aliases in aliases_by_slot_id.values() for a in aliases + ) + inner_sql = ( + "SELECT\n " + + _SQL_COL_SEP.join(f'"{a}"' for a in inner_sorted) + + f"\nFROM {final_cte}" + ) + + cte_clause = ( + _SQL_WITH + + ",\n".join(f"{name} AS (\n{sql}\n)" for name, sql in ctes) + ) + chain_sql = f"{cte_clause}\n{inner_sql}" + + # POST-phase filter wrap (filters referencing transform / arith + # slots). Mirrors legacy _generate_with_computed:1627-1648 — + # ``SELECT * FROM () AS _filtered WHERE ``. + post_filter_conditions = self._render_post_phase_filter_conditions( + planned_query=planned_query, + slot_id_by_key=slot_id_by_key, + available_alias_by_slot_id=available_alias_by_slot_id, + ) + if post_filter_conditions: + chain_sql = ( + f"SELECT *\nFROM (\n{chain_sql}\n) AS _filtered" + f"\nWHERE {_SQL_AND_JOINER.join(post_filter_conditions)}" + ) + + # Outer SELECT in user-projection order (public slots only). + # Per-slot index walks each slot's public_aliases so duplicate + # interned names (DEV-1450 C13) both surface in the result. + public_aliases_user_order: list[str] = [] + outer_alias_index: Dict[str, int] = {} + for sid in planned_query.projection: + slot = slots_by_id[sid] + if slot.hidden: + continue + all_aliases = aliases_by_slot_id.get(sid, []) + if not all_aliases: + continue + idx = outer_alias_index.setdefault(sid, 0) + alias = ( + all_aliases[idx] if idx < len(all_aliases) else all_aliases[-1] + ) + outer_alias_index[sid] = idx + 1 + public_aliases_user_order.append(alias) + outer_sql = ( + "SELECT\n " + + _SQL_COL_SEP.join(f'"{a}"' for a in public_aliases_user_order) + + f"\nFROM (\n{chain_sql}\n) AS _outer" + ) + + # ORDER BY / LIMIT / OFFSET on the outermost wrap. + return self._apply_order_limit_to_planned_sql_string( + sql=outer_sql, + planned_query=planned_query, + slots_by_id=slots_by_id, + available_alias_by_slot_id=available_alias_by_slot_id, + ) + + # ----------------------------------------------------------------- + # Stage 7b.10 helpers + # ----------------------------------------------------------------- + + @staticmethod + def _validate_window_transform_ops_for_7b10(*, planned_query) -> None: + """Validate transform-layer op scope. + + 7b.11 lifted ``time_shift`` and ``consecutive_periods`` from + the deferred set — both render through dedicated self-join / + staged-window CTE pairs. The deferred set is now empty; the + function stays in place as a safety net for follow-up ops + added by later slices. + + It also enforces the **composite-input** rule that survives + from 7b.10: + + * ``time_shift`` requires a slottable leaf input (the legacy + self-join CTE re-aggregates the source — composite expressions + would need an inner expression layer). + * ``consecutive_periods`` accepts a slottable leaf OR a top-level + comparison ``ArithmeticKey`` (the boolean predicate shape + ``amount:sum > 0`` is the canonical user form). Other + composite shapes (numeric subtraction, scalar calls) are + rejected with a ``composite-input transforms`` marker so the + test suite's per-op composite assertions pin a unified message. + """ + from slayer.core.keys import ( + AggregateKey, + ArithmeticKey, + BetweenKey, + ColumnKey, + ColumnSqlKey, + InKey, + ScalarCallKey, + TimeTruncKey, + TransformKey, + ) + + # 7b.11 lifted these — placeholder set for future slices. + deferred: set = set() + + leaf_kinds = (ColumnKey, ColumnSqlKey, AggregateKey, TimeTruncKey) + # Keep aligned with _emit_consecutive_periods_ctes_for_planned — + # the renderer dispatches arithmetic ops via _compose_arithmetic_op + # which supports these binary comparisons only. + _COMPARISON_OPS = {"==", "!=", "<", "<=", ">", ">="} + + def _walk(key) -> Optional[str]: + if isinstance(key, TransformKey): + if key.op in deferred: + return key.op + return _walk(key.input) + if isinstance(key, ArithmeticKey): + for o in key.operands: + found = _walk(o) + if found: + return found + return None + if isinstance(key, ScalarCallKey): + for a in key.args: + if isinstance( + a, + (TransformKey, ArithmeticKey, ScalarCallKey, BetweenKey, InKey), + ): + found = _walk(a) + if found: + return found + return None + if isinstance(key, BetweenKey): + for k in (key.column, key.low, key.high): + found = _walk(k) + if found: + return found + return None + if isinstance(key, InKey): + # DEV-1475: only LHS column can host a deferred transform. + return _walk(key.column) + return None + + # Explicit layer ops + composite-input enforcement. + for layer in planned_query.transform_layers: + if layer.op in deferred: + raise NotImplementedError( + f"DEV-1450 stage 7b.11: transform op {layer.op!r} " + f"(self-join CTE) deferred to a follow-up slice.", + ) + if layer.op in ("time_shift", "consecutive_periods"): + # Walk the layer's slot ids and assert their TransformKey + # inputs satisfy the per-op composite-input rule. + slots_map = { + s.id: s + for s in ( + list(planned_query.row_slots) + + list(planned_query.aggregate_slots) + + list(planned_query.combined_expression_slots) + ) + } + for sid in layer.slot_ids: + slot = slots_map.get(sid) + if slot is None or not isinstance(slot.key, TransformKey): + continue + inner = slot.key.input + if isinstance(inner, leaf_kinds): + continue + if ( + layer.op == "consecutive_periods" + and isinstance(inner, ArithmeticKey) + and inner.op in _COMPARISON_OPS + ): + # Boolean predicate shape — accepted. + continue + raise ValueError( + f"Nesting a transform inside {layer.op!r} " + f"(input={type(inner).__name__}) is not supported. " + f"Compute the inner transform in an earlier stage of " + f"a multi-stage `source_queries` model and reference " + f"its output in this stage." + ) + + # Reachable trees of every slot we'll need to render. + slots = ( + list(planned_query.row_slots) + + list(planned_query.aggregate_slots) + + list(planned_query.combined_expression_slots) + ) + for slot in slots: + found_op = _walk(slot.key) + if found_op is not None: + raise NotImplementedError( + f"DEV-1450 stage 7b.11: transform op {found_op!r} " + f"(reached via slot id={slot.id!r}, key=" + f"{type(slot.key).__name__}) deferred to a follow-up " + f"slice.", + ) + + @staticmethod + def _collect_base_aux_slot_ids( # NOSONAR(S3776) — recursive ValueKey walker (nested ``_collect_from``) over the closed key union plus three top-level passes (transform layers / phase-gated filter deps / order deps). Each pass is one decision; extracting them would scatter the slot-dep contract. + *, + planned_query, + slot_id_by_key: Dict[Any, str], + slots_by_id: Dict[str, Any], + include_order: bool = True, + aggregates_only: bool = False, + ) -> Set[str]: + """Return slot ids the base CTE must project beyond the public + projection. + + Walks every ``TransformKey`` in ``transform_layers`` for its + ``input`` / ``partition_keys`` / ``time_key`` deps; walks every + AGGREGATE- and POST-phase ``FilterPhase.expression`` for + slot-worthy deps; walks ``OrderEntry.slot_id`` keys when + ``include_order`` is True. Only ``ColumnKey`` / ``ColumnSqlKey`` + / ``TimeTruncKey`` / ``AggregateKey`` slot ids are returned + (those that the base CTE renders); transform slot ids are + excluded since they're materialised in step CTEs. + + DEV-1501: ``aggregates_only=True`` narrows leaf collection to + ``AggregateKey`` slots ONLY (row leaves on order/filter paths are + skipped). Used by the no-transform path so that materialising a + hidden order/filter aggregate does NOT accidentally pull a hidden + ROW dep into ``base_render_order`` (which would add it to GROUP + BY and silently change query grain). Composites still recurse so + their AggregateKey operands surface. + """ + from slayer.core.keys import ( + AggregateKey, + ArithmeticKey, + BetweenKey, + ColumnKey, + ColumnSqlKey, + InKey, + Phase, + ScalarCallKey, + TimeTruncKey, + TransformKey, + ) + + if aggregates_only: + base_kinds: Tuple[type, ...] = (AggregateKey,) + else: + base_kinds = (ColumnKey, ColumnSqlKey, TimeTruncKey, AggregateKey) + out: Set[str] = set() + + def _collect_from(key) -> None: + if isinstance(key, base_kinds): + sid = slot_id_by_key.get(key) + if sid is not None: + out.add(sid) + return + # ``aggregates_only`` mode: still SKIP non-aggregate row leaves + # at the leaf level — but the composite/walker branches below + # continue to recurse so their nested AggregateKey operands + # surface. + if aggregates_only and isinstance( + key, (ColumnKey, ColumnSqlKey, TimeTruncKey), + ): + return + if isinstance(key, TransformKey): + _collect_from(key.input) + for p in key.partition_keys: + _collect_from(p) + if key.time_key is not None: + _collect_from(key.time_key) + return + if isinstance(key, ArithmeticKey): + for o in key.operands: + _collect_from(o) + return + if isinstance(key, ScalarCallKey): + for a in key.args: + if isinstance( + a, + ( + TransformKey, ArithmeticKey, ScalarCallKey, + BetweenKey, InKey, ColumnKey, ColumnSqlKey, + TimeTruncKey, AggregateKey, + ), + ): + _collect_from(a) + return + if isinstance(key, BetweenKey): + _collect_from(key.column) + _collect_from(key.low) + _collect_from(key.high) + return + if isinstance(key, InKey): + # DEV-1475: only the LHS column references a slot; the + # RHS values are bare literals with no slot identity. + _collect_from(key.column) + return + # LiteralKey / StarKey / unknown: nothing to materialise. + + # Transform layer deps. + for layer in planned_query.transform_layers: + for slot_id in layer.slot_ids: + slot = slots_by_id.get(slot_id) + if slot is None: + continue + key = slot.key + if isinstance(key, TransformKey): + _collect_from(key.input) + for p in key.partition_keys: + _collect_from(p) + if key.time_key is not None: + _collect_from(key.time_key) + + # Filter deps for AGGREGATE-phase (HAVING) and POST-phase filters + # (the latter only in the transform path, where + # ``_render_post_phase_filter_conditions`` actually applies them). + # A hidden ``revenue:last(...) > 100`` HAVING aggregate needs the + # same ranked-subquery materialisation as the ORDER BY path, so + # its AggregateKey must reach ``base_render_order`` alongside the + # projected and order-only ones (DEV-1501). POST-phase walk is + # gated on the presence of transforms — POST filters reference + # ``TransformKey`` and so are planner-unreachable in no-transform + # queries; walking them anyway would silently materialise their + # operands without applying the filter (CodeRabbit DEV-1501 PR + # #159 Group B). + has_transforms = bool(planned_query.transform_layers) + for fp in planned_query.filters_by_phase: + if fp.phase == Phase.AGGREGATE: + pass # walk + elif fp.phase == Phase.POST and has_transforms: + pass # walk + else: + continue + if fp.expression is not None: + _collect_from(fp.expression.value_key) + + # 7b.10 / DEV-1501 — order hidden refs reach the base CTE so + # ORDER BY can resolve via materialised aliases. Walk + # ``OrderEntry.slot_id`` → that slot's key (so any transform / + # arithmetic inside also surfaces its base deps). In the + # no-transform path the caller passes ``aggregates_only=True`` + # so hidden ROW order targets are NOT pulled in (they stay + # inline-rendered or raise NotImplementedError downstream). + if include_order: + for oe in planned_query.order: + slot = slots_by_id.get(oe.slot_id) + if slot is None: + continue + _collect_from(slot.key) + + return out + + @staticmethod + def _transform_layer_deps_ready( + *, + layer, + slots_by_id: Dict[str, Any], + slot_id_by_key: Dict[Any, str], + available_alias_by_slot_id: Dict[str, str], + ) -> bool: + """A layer is ready when every slot-worthy dep its TransformKeys + reference (``input`` + ``partition_keys`` + ``time_key``) has + an alias materialised in a prior CTE. + """ + from slayer.core.keys import ( + AggregateKey, + ArithmeticKey, + BetweenKey, + ColumnKey, + ColumnSqlKey, + InKey, + ScalarCallKey, + TimeTruncKey, + TransformKey, + ) + + slotted_kinds = ( + ColumnKey, ColumnSqlKey, TimeTruncKey, AggregateKey, TransformKey, + ) + + def _ready(key) -> bool: + if isinstance(key, slotted_kinds): + sid = slot_id_by_key.get(key) + if sid is None: + # Not interned as a slot — can be inlined. + return True + return sid in available_alias_by_slot_id + if isinstance(key, ArithmeticKey): + return all(_ready(o) for o in key.operands) + if isinstance(key, ScalarCallKey): + for a in key.args: + if isinstance( + a, + ( + TransformKey, ArithmeticKey, ScalarCallKey, + BetweenKey, InKey, ColumnKey, ColumnSqlKey, + TimeTruncKey, AggregateKey, + ), + ) and not _ready(a): + return False + return True + if isinstance(key, BetweenKey): + return all( + _ready(k) for k in (key.column, key.low, key.high) + ) + if isinstance(key, InKey): + # DEV-1475: only LHS column needs slot readiness; RHS + # values are literals (always ready). + return _ready(key.column) + return True + + for slot_id in layer.slot_ids: + slot = slots_by_id.get(slot_id) + if slot is None or not isinstance(slot.key, TransformKey): + continue + tk = slot.key + if not _ready(tk.input): + return False + for p in tk.partition_keys: + if not _ready(p): + return False + if tk.time_key is not None and not _ready(tk.time_key): + return False + return True + + def _build_base_select_for_planned( # NOSONAR(S3776) — join-path collection and derived-dim expansion are extracted to helpers; the residual is the one cohesive per-slot ROW/AGGREGATE projection + GROUP-BY assembly pass. + self, + *, + planned_query, + bundle, + source_model, + source_relation: str, + base_render_order: List[str], + slots_by_id: Dict[str, Any], + skip_cross_model_aggs: bool = False, + skip_filter_ids: Optional[Set[str]] = None, + ): + """Build the base SELECT (sqlglot ``Select``) for ``generate_from_planned``. + + Iterates ``base_render_order`` (public projection followed by + aux materialisation slot ids), rendering each ROW / AGGREGATE + slot. POST-phase slots are skipped — step CTEs render them. + + Returns ``(base_select, aliases_by_slot_id, has_aggregation, + group_by_keys)``. ``aliases_by_slot_id`` is a list per slot to + preserve duplicate public aliases (DEV-1450 C13). + + DEV-1450 stage 7b.12: joined ROW slots (ColumnKey.path != () + and TimeTruncKey.column.path != ()) are rendered by walking + the bundle's join graph and emitting ``LEFT JOIN`` clauses in + the FROM. ``skip_cross_model_aggs=True`` is passed by the + cross-model orchestrator so the ``_base`` CTE omits AGGREGATE + slots that live in a per-plan ``_cm_*`` CTE. + """ + from slayer.core.enums import TimeGranularity + from slayer.core.keys import ( + AggregateKey, + ColumnKey, + ColumnSqlKey, + Phase, + TimeTruncKey, + ) + + # Walk row slots to collect every joined path so the FROM + # clause carries the needed LEFT JOINs in one pass. + needed_join_paths = self._collect_joined_paths_for_base( + base_render_order=base_render_order, + slots_by_id=slots_by_id, + source_model=source_model, + source_relation=source_relation, + bundle=bundle, + ) + # Pre-expand derived (ColumnSqlKey) ROW + TIME dimensions: inline + # sibling/joined derived refs (DEV-1333 / DEV-1410) and pull any joins + # their SQL crosses into the FROM (appended to ``needed_join_paths``). + derived_expr_by_sid = self._expand_derived_row_dims( + base_render_order=base_render_order, slots_by_id=slots_by_id, + source_relation=source_relation, source_model=source_model, + bundle=bundle, needed_join_paths=needed_join_paths, + ) + # WHERE-phase filters referencing joined columns (direct, derived, or + # Mode-A ``__`` paths) pull their joins into the FROM too. Filters + # routed to a cross-model ``_cm_*`` CTE (``skip_filter_ids``) are + # applied there, not on ``_base`` — pulling their join into ``_base`` + # would add an unused (and, for one-to-many joins, cardinality- + # changing) LEFT JOIN. + for p in self._collect_filter_join_paths( + planned_query=planned_query, source_model=source_model, + source_relation=source_relation, bundle=bundle, + skip_filter_ids=skip_filter_ids, + ): + if p not in needed_join_paths: + needed_join_paths.append(p) + # DEV-1494: a ``Column.filter`` on an aggregated measure becomes a + # CASE-WHEN wrapper; pull any join its predicate crosses (directly, or via + # a derived ref) into the FROM — including filtered aggregates nested in + # composite (arithmetic / scalar-call) AGGREGATE-phase keys. + for p in self._collect_column_filter_join_paths( + base_render_order=base_render_order, slots_by_id=slots_by_id, + source_relation=source_relation, source_model=source_model, + bundle=bundle, + ): + if p not in needed_join_paths: + needed_join_paths.append(p) + # DEV-1502: an AGGREGATE slot whose SOURCE is a derived + # (``ColumnSqlKey``) column whose ``Column.sql`` crosses a join + # (``customers__regions.population``) needs the same discovery the + # dimension path performs (DEV-1484). Symmetric with the filter case + # above; render-time expansion in ``_build_agg_render_spec_from_planned`` + # already qualifies the body, so this pass only closes the join-discovery + # gap. Cross-model aggregate sources are skipped — they're owned by the + # per-plan ``_cm_*`` CTE (the symmetric in-CTE discovery gap is tracked + # separately). + for p in self._collect_aggregate_source_join_paths( + base_render_order=base_render_order, slots_by_id=slots_by_id, + source_relation=source_relation, source_model=source_model, + bundle=bundle, + ): + if p not in needed_join_paths: + needed_join_paths.append(p) + from_clause, base_joins = self._build_from_and_joins( + source_model=source_model, + source_relation=source_relation, + joined_paths=needed_join_paths, + bundle=bundle, + ) + + # DEV-1450: first/last AGGREGATIONS rank rows via a ROW_NUMBER + # subquery (mirrors legacy ``_generate_base`` + ``_build_last_ + # ranked_from``). + if self._has_first_last_aggregate( + base_render_order=base_render_order, slots_by_id=slots_by_id, + ): + # DEV-1503 — local first/last in ``_base`` alongside cross-model + # CTEs is supported: the ranked subquery wraps ``_base`` for the + # local measures; cross-model / filtered-local aggregates are + # deferred to their per-plan ``_cm_*`` CTEs (their slot ids are + # excluded from ``base_render_order`` by the caller, so + # ``_build_first_last_base_select`` never sees them and emits no + # dangling references). + return self._build_first_last_base_select( + planned_query=planned_query, + bundle=bundle, + source_model=source_model, + source_relation=source_relation, + base_render_order=base_render_order, + slots_by_id=slots_by_id, + from_clause=from_clause, + base_joins=base_joins, + skip_filter_ids=skip_filter_ids, + ) + + select_columns: list[exp.Expression] = [] + group_by_keys: Dict[str, exp.Expression] = {} + has_aggregation = False + alias_index: Dict[str, int] = {} + aliases_by_slot_id: Dict[str, List[str]] = {} + + def _record_alias(sid: str, full_alias: str) -> None: + aliases_by_slot_id.setdefault(sid, []).append(full_alias) + + for sid in base_render_order: + slot = slots_by_id[sid] + # DEV-1450 stage 7b.12: joined ROW slots emit the FULL + # dotted result-key form (``orders.customers.region_id``). + # The planner emits a flat ``customers__region_id`` + # declared_name for downstream stage binding (DEV-1449 / C4 + # contract), but the public projection alias must preserve + # the dotted path for the result-key contract (P10). Local + # slots keep the existing ``.`` + # form. + full_alias = self._full_alias_for_slot( + slot=slot, + source_relation=source_relation, + alias_index=alias_index, + ) + + if slot.phase == Phase.ROW: + key = slot.key + if isinstance(key, ColumnKey): + col_expr = self._joined_or_local_dim_expr( + path=key.path, + leaf=key.leaf, + source_model=source_model, + source_relation=source_relation, + bundle=bundle, + ) + select_columns.append(col_expr.copy().as_(full_alias)) + group_by_keys.setdefault(sid, col_expr) + _record_alias(sid, full_alias) + elif isinstance(key, TimeTruncKey): + col_expr = self._raw_time_col_expr_for_planned( + time_column=key.column, + source_model=source_model, + source_relation=source_relation, + bundle=bundle, + ) + trunc_expr = self._build_date_trunc( + col_expr=col_expr, + granularity=TimeGranularity(key.granularity), + ) + select_columns.append(trunc_expr.copy().as_(full_alias)) + group_by_keys.setdefault(sid, trunc_expr) + _record_alias(sid, full_alias) + elif isinstance(key, ColumnSqlKey): + # A derived column (``Column.sql`` set) used as a dimension, + # e.g. ``ratio = A.bar / B.foo_normalized`` (cross-table) or + # ``c2 = c1 * 2`` (sibling-derived chain). Local + # (``path == ()``) derived columns are pre-expanded above + # (sibling/joined refs inlined, joins pulled in); fall back + # to the non-expanded resolution for any other shape. + col_expr = derived_expr_by_sid.get(sid) + if col_expr is None: + col_expr = self._dim_column_expr_from_planned( + source_model=source_model, + source_relation=source_relation, + leaf=key.column_name, + ) + select_columns.append(col_expr.copy().as_(full_alias)) + group_by_keys.setdefault(sid, col_expr) + _record_alias(sid, full_alias) + else: + raise NotImplementedError( + f"DEV-1450 stage 7b.10+: row-phase key type " + f"{type(key).__name__} not supported in the " + f"local-only / time-dim slice." + ) + + elif slot.phase == Phase.AGGREGATE: + key = slot.key + if not isinstance(key, AggregateKey): + # AGGREGATE-phase composite (arithmetic / scalar-call of + # aggregates, e.g. ``expensenet:avg + benchmarkexp:avg``). + # Render inline; cast the whole composite once. + composite, any_agg = self._render_aggregate_composite_expr( + key=key, + slot=slot, + source_model=source_model, + source_relation=source_relation, + bundle=bundle, + ) + if any_agg: + composite = _wrap_cast_for_type(composite, slot.type) + has_aggregation = True + select_columns.append(composite.copy().as_(full_alias)) + _record_alias(sid, full_alias) + continue + agg_path = getattr(key.source, "path", ()) + if agg_path: + if skip_cross_model_aggs: + # Cross-model aggregate; rendered by the per-plan + # ``_cm_*`` CTE. Skip in the host base. + continue + raise NotImplementedError( + f"DEV-1450 stage 7b.12: cross-model aggregate " + f"(source.path={agg_path!r}) reached the local " + f"base SELECT path. The cross-model orchestrator " + f"should have routed this through `_render_with_" + f"cross_model_plans`." + ) + # DEV-1450 stage 7b.12: ``column_filter_key`` is now + # propagated into the synthetic EnrichedMeasure's + # ``filter_sql`` field so ``_build_agg`` wraps the + # aggregate as ``SUM(CASE WHEN THEN col END)``. + synth = self._build_agg_render_spec_from_planned( + slot=slot, + key=key, + source_model=source_model, + source_relation=source_relation, + full_alias=full_alias, + bundle=bundle, + ) + agg_expr, is_agg = self._build_agg(synth) + if is_agg: + agg_expr = _wrap_cast_for_type(agg_expr, slot.type) + has_aggregation = True + select_columns.append(agg_expr.copy().as_(full_alias)) + _record_alias(sid, full_alias) + else: + # POST-phase slot in projection — handled by step CTEs. + # Don't add to base select; step CTE will materialise. + continue + + base_select = exp.Select() + for col in select_columns: + base_select = base_select.select(col) + base_select = base_select.from_(from_clause) + for join_expr, on_expr, join_type in base_joins: + base_select = base_select.join( + join_expr, on=on_expr, join_type=join_type, + ) + return ( + base_select, aliases_by_slot_id, has_aggregation, group_by_keys, + False, None, + ) + + def _has_first_last_aggregate( + self, *, base_render_order: List[str], slots_by_id: Dict[str, Any], + ) -> bool: + """True if any LOCAL ``first`` / ``last`` AGGREGATE slot appears in + the base render order — directly as an ``AggregateKey`` slot OR + as an operand inside a composite (``ArithmeticKey`` / + ``ScalarCallKey``) aggregate slot. + + Cross-model first/last (non-empty ``source.path``) is excluded — + it is not rendered by the ranked-subquery path (each cross-model + aggregate has its own CTE). DEV-1501 (Codex round 4): composite- + only first/last (e.g. ``last(created_at) + last(updated_at)`` + with no direct sibling) must still trigger the ranked-subquery + path; without composite-aware detection the composite render + would emit ``MAX(CASE WHEN _last_rn = 1 …)`` referencing a + column the bare-FROM never projects. + """ + from slayer.core.keys import AggregateKey, Phase + + for sid in base_render_order: + slot = slots_by_id.get(sid) + if slot is None or slot.phase != Phase.AGGREGATE: + continue + key = slot.key + if ( + isinstance(key, AggregateKey) + and key.agg in ("first", "last") + and not getattr(key.source, "path", ()) + ): + return True + # Composite slot (no direct AggregateKey): walk for first/ + # last AggregateKey leaves. The composite render needs the + # ranked subquery so each operand's ``_first_rn`` / + # ``_last_rn{suffix}`` column exists. + if not isinstance(key, AggregateKey) and _iter_first_last_leaves(key): + return True + return False + + def _resolve_ranking_time_column_from_planned( + self, + *, + base_render_order: List[str], + slots_by_id: Dict[str, Any], + source_model, + source_relation: str, + bundle, + ) -> Optional[str]: + """Resolve the default ORDER-BY time column for first/last + ROW_NUMBER ranking (mirrors legacy ``_resolve_last_agg_time``). + + Precedence (matching legacy): the first ``DATE`` / ``TIMESTAMP`` + regular dimension, then the first time-dimension slot's raw column, + then the model's ``default_time_dimension``. Returns the qualified + SQL string (e.g. ``"orders.created_at"`` / ``"stores.opened_at"``), + or ``None`` when nothing temporal is in scope. + + (The legacy ``main_time_dimension`` short-circuit and the + filter-referenced-date fallback are corner cases the spec permits + diverging on; they are not reproduced here.) + """ + from slayer.core.keys import ColumnKey, Phase, TimeTruncKey + + for sid in base_render_order: + slot = slots_by_id[sid] + if slot.phase == Phase.ROW and isinstance(slot.key, ColumnKey): + model = source_model + for hop in slot.key.path: + nxt = bundle.get_referenced_model(hop) + if nxt is None: + model = None + break + model = nxt + if model is None: + continue + col_def = next( + (c for c in model.columns if c.name == slot.key.leaf), None, + ) + if col_def is not None and col_def.type in ( + DataType.DATE, DataType.TIMESTAMP, + ): + return self._joined_or_local_dim_expr( + path=slot.key.path, leaf=slot.key.leaf, + source_model=source_model, + source_relation=source_relation, bundle=bundle, + ).sql(dialect=self.dialect) + for sid in base_render_order: + slot = slots_by_id[sid] + if slot.phase == Phase.ROW and isinstance(slot.key, TimeTruncKey): + return self._raw_time_col_expr_for_planned( + time_column=slot.key.column, source_model=source_model, + source_relation=source_relation, bundle=bundle, + ).sql(dialect=self.dialect) + if source_model.default_time_dimension: + return f"{source_relation}.{source_model.default_time_dimension}" + return None + + def _resolve_explicit_time_col( # NOSONAR(S3776) — sequential isinstance dispatch over ColumnKey (bare ref → ``__``-joined path alias) and ColumnSqlKey (derived column → bare-ident-qualify vs complex-emit-verbatim). Extracting the per-shape branches would scatter the time-arg resolution contract; each branch is one decision. + self, + *, + key, + source_model, + source_relation: str, + bundle=None, + ) -> Optional[str]: + """Resolve the explicit positional time arg on a ``first`` / ``last`` + aggregate into a SQL string suitable for ``ORDER BY`` inside the + ranked subquery. + + Handles both bare-column refs (``ColumnKey`` — + ``amount:last(created_at)``) and derived-column refs + (``ColumnSqlKey`` — ``amount:last(net_amount_date)`` where + ``net_amount_date`` has a non-trivial ``Column.sql``). For derived + columns the column's ``Column.sql`` is materialised through + ``_expand_derived_column_sql`` (when ``bundle`` is available) so + inner bare refs qualify to ``source_relation`` and joined refs to + their ``__``-path alias — a complex expression like + ``date(created_at)`` can't go ambiguous against a same-named column + on a joined table inside the ranked subquery. Without a ``bundle`` + it falls back to bare-ident qualification / verbatim emit. + + Returns ``None`` for non-first/last aggs and when ``key.args`` is + empty or its first element is neither a ``ColumnKey`` nor a + ``ColumnSqlKey``. Cross-model paths on derived time args + (``ColumnSqlKey`` with non-empty ``path``) raise + ``NotImplementedError`` rather than silently emitting against the + wrong relation alias — that case is tracked alongside bug (c) of + the four-bug Stage B package in DEV-1476. + """ + from slayer.core.keys import ColumnKey, ColumnSqlKey + + if key.agg not in ("first", "last"): + return None + for a in key.args: + if isinstance(a, ColumnKey): + relation = "__".join(a.path) if a.path else source_relation + return f"{relation}.{a.leaf}" + if isinstance(a, ColumnSqlKey): + if a.path: + raise NotImplementedError( + f"Cross-model derived time column " + f"(path={a.path!r}, column={a.column_name!r}) on " + f"first/last positional arg is not yet supported " + f"by the ranked-subquery builder; tracked as " + f"DEV-1476." + ) + col = next( + (c for c in source_model.columns if c.name == a.column_name), + None, + ) + if col is None: + raise ValueError( + f"Derived time column {a.column_name!r} (positional " + f"arg of {key.agg!r}) not found on model " + f"{source_model.name!r}." + ) + if bundle is not None: + # Qualify inner bare refs against ``source_relation`` (and + # joined refs to their ``__``-path alias) so a complex + # derived time expression can't bind to the wrong table + # inside the ranked subquery's joins — same expansion the + # aggregate-source path uses. + return self._expand_derived_column_sql( + source_model=source_model, + source_relation=source_relation, + column_name=a.column_name, + bundle=bundle, + ) + # No bundle (defensive): bare-ident qualify, else emit verbatim. + col_sql = col.sql if col.sql else col.name + if col_sql.isidentifier(): + return f"{source_relation}.{col_sql}" + return self._parse(col_sql).sql(dialect=self.dialect) + # Unrecognised positional arg type — leave time_column unset and + # let _build_ranked_subquery_from_planned fall back to the + # query's default ranking column. + break + return None + + def _build_ranked_subquery_from_planned( # NOSONAR(S3776) — Group 2 already factored the per-spec ROW_NUMBER passes into _build_unfiltered_rn_columns / _build_filtered_rn_columns; what's left is exp.Select / from / joins / where assembly that has to live in one place. + self, + *, + source_relation: str, + default_time_col_sql: str, + partition_exprs: List[exp.Expression], + extra_projections: List[Tuple[str, exp.Expression]], + synth_specs: List[AggRenderSpec], + from_clause: exp.Expression, + base_joins: List, + where_clause: Optional[exp.Expression], + ) -> Tuple[exp.Expression, dict, dict, dict]: + """Build the ROW_NUMBER-ranked subquery that wraps the source for + first/last aggregation (planned-native port of + ``_build_last_ranked_from``). + + Projects ``source_relation.*`` plus the supplied ``extra_projections`` + (truncated time dimensions / joined dimensions referenced by the + outer SELECT) plus one ``ROW_NUMBER`` column per distinct + (effective-time-column, agg) pair. Filtered first/last measures get + a dedicated ranking column (non-matching rows pushed to the bottom) + and a boolean match flag. WHERE is applied INSIDE so it filters raw + rows before ranking. Returns ``(subquery, rn_suffix_map, + filtered_rn_map, filtered_match_map)``. + """ + partition_clause = "" + if partition_exprs: + partition_clause = _SQL_PARTITION_BY + ", ".join( + p.sql(dialect=self.dialect) for p in partition_exprs + ) + + select_exprs: List[exp.Expression] = [ + exp.Column(this=exp.Star(), table=exp.to_identifier(source_relation)), + ] + for alias, e in extra_projections: + select_exprs.append(e.copy().as_(alias)) + + unfiltered_exprs, rn_suffix_map = self._build_unfiltered_rn_columns( + synth_specs=synth_specs, + default_time_col_sql=default_time_col_sql, + partition_clause=partition_clause, + ) + select_exprs.extend(unfiltered_exprs) + + filtered_exprs, filtered_rn_map, filtered_match_map = ( + self._build_filtered_rn_columns( + synth_specs=synth_specs, + default_time_col_sql=default_time_col_sql, + partition_clause=partition_clause, + ) + ) + select_exprs.extend(filtered_exprs) + + inner = exp.Select() + for e in select_exprs: + inner = inner.select(e) + inner = inner.from_(from_clause) + for join_expr, on_expr, join_type in base_joins: + inner = inner.join(join_expr, on=on_expr, join_type=join_type) + if where_clause is not None: + inner = inner.where(where_clause) + subquery = exp.Subquery( + this=inner, alias=exp.to_identifier(source_relation), + ) + return subquery, rn_suffix_map, filtered_rn_map, filtered_match_map + + def _build_unfiltered_rn_columns( + self, + *, + synth_specs: List[AggRenderSpec], + default_time_col_sql: str, + partition_clause: str, + ) -> Tuple[List[exp.Expression], Dict[str, str]]: + """One ``ROW_NUMBER`` projection per distinct effective time column + for the unfiltered ``first`` / ``last`` specs. + + Each unique effective time column gets a stable suffix in render + order (first sorted gets ``""``, then ``"_2"``, ...); the same + time column shared by both ``first`` and ``last`` produces two + projections (`_first_rn{suffix}` ASC, `_last_rn{suffix}` DESC). + Returns ``(rn_select_exprs, rn_suffix_map)``. + """ + time_col_agg_types: Dict[str, set] = {} + for m in synth_specs: + if m.aggregation in ("first", "last") and not m.filter_sql: + eff = m.time_column or default_time_col_sql + time_col_agg_types.setdefault(eff, set()).add(m.aggregation) + sorted_tcs = sorted(time_col_agg_types) + rn_suffix_map: Dict[str, str] = { + tc: ("" if i == 0 else f"_{i + 1}") + for i, tc in enumerate(sorted_tcs) + } + rn_exprs: List[exp.Expression] = [] + for tc in sorted_tcs: + suffix = rn_suffix_map[tc] + if "last" in time_col_agg_types[tc]: + rn_exprs.append( + self._parse( + f"ROW_NUMBER() OVER ({partition_clause} " + f"ORDER BY {tc} DESC)" + ).as_(f"_last_rn{suffix}") + ) + if "first" in time_col_agg_types[tc]: + rn_exprs.append( + self._parse( + f"ROW_NUMBER() OVER ({partition_clause} " + f"ORDER BY {tc} ASC)" + ).as_(f"_first_rn{suffix}") + ) + return rn_exprs, rn_suffix_map + + def _build_filtered_rn_columns( + self, + *, + synth_specs: List[AggRenderSpec], + default_time_col_sql: str, + partition_clause: str, + ) -> Tuple[List[exp.Expression], Dict[str, str], Dict[str, str]]: + """One dedicated ``ROW_NUMBER`` + match-flag projection per distinct + ``(filter, time, agg)`` triple for the filtered ``first`` / ``last`` + specs. + + Filtered first/last needs to push non-matching rows past the + winners; emits ``ROW_NUMBER() OVER (... ORDER BY CASE WHEN + THEN 0 ELSE 1 END,