From 3e3a6acc17f7c0305600cbaecb240edb86da9398 Mon Sep 17 00:00:00 2001 From: tchivs Date: Sat, 5 Sep 2026 17:14:02 +0800 Subject: [PATCH 1/2] fix: validate metric dependency references Resolve metric dependencies iteratively and report cycles or unknown references with their dependency chain and defining file. Preserve forward references, shared dependencies, and case-insensitive metric names with execution-based regression coverage. Signed-off-by: tchivs --- docs/concepts/metrics/definition.md | 4 +- sqlmesh/core/metric/definition.py | 102 ++++++---- tests/core/metric/test_metric.py | 287 ++++++++++++++-------------- 3 files changed, 211 insertions(+), 182 deletions(-) diff --git a/docs/concepts/metrics/definition.md b/docs/concepts/metrics/definition.md index 81d3f2a729..a05596bc2b 100644 --- a/docs/concepts/metrics/definition.md +++ b/docs/concepts/metrics/definition.md @@ -59,7 +59,7 @@ Because the `prod.users.country` and `prod.searches.num_searches` models have sp Metrics can perform additional operations/calculations with other metrics. -In this example, the third metric `clicks_per_search` is calculated by dividing the first metric `total_searches` by the second metric `total_clicks`: +In this example, the third metric `clicks_per_search` is calculated by dividing `total_clicks` by `total_searches`: ```sql linenums="1" METRIC ( @@ -78,6 +78,8 @@ METRIC ( ); ``` +Metric references are case insensitive, including quoted references, and may refer to metrics defined later in the project. Shared dependencies are resolved once during expansion. An unknown dependency or a dependency cycle raises a configuration error with the dependency chain and the path of the definition containing the invalid reference. + ## Properties The `METRIC` definition supports the following keys. The `name` and `expression` keys are required. diff --git a/sqlmesh/core/metric/definition.py b/sqlmesh/core/metric/definition.py index 6119a883ed..7924dc3fe1 100644 --- a/sqlmesh/core/metric/definition.py +++ b/sqlmesh/core/metric/definition.py @@ -99,43 +99,71 @@ def to_metric( self, metas: t.Dict[str, MetricMeta], metrics: UniqueKeyDict[str, Metric] ) -> Metric: """Converts a metric meta into a fully expanded and standalone metric.""" - metric_refs = {} - agg_or_ref = False - - for node in self.expression.walk(): - if isinstance(node, exp.Alias): - _raise_metric_config_error( - f"Alias found for metric '{self.name}' which is not allowed", self._path - ) - elif isinstance(node, exp.AggFunc): - agg_or_ref = True - elif isinstance(node, exp.Column) and not node.table: - agg_or_ref = True - ref = node.sql(dialect=self.dialect) - - if ref not in metrics: - metrics[ref] = metas[ref].to_metric(metas, metrics) - - metric_refs[node] = metrics[ref] - - if not agg_or_ref: - _raise_metric_config_error( - f"Metric '{self.name}' missing an aggregation or metric ref", self._path - ) - - if metric_refs: - expanded = self.expression.copy() - for column in expanded.find_all(exp.Column): - metric = metric_refs.get(column) - - if metric: - column.replace(metric.expanded.copy()) - else: - expanded = exp.alias_(self.expression, self.name) - - metric = Metric(**self.dict(), expanded=expanded) - metric._path = self._path - return metric + # Suspend each expression walk while resolving a dependency so cycles of + # any depth can be diagnosed without using the Python call stack. + stack: t.List[t.Tuple[MetricMeta, t.Iterator[exp.Expr], t.Dict[exp.Column, str], bool]] = [ + (self, self.expression.walk(), {}, False) + ] + visiting = {self.name} + + while True: + meta, nodes, metric_refs, agg_or_ref = stack.pop() + + for node in nodes: + if isinstance(node, exp.Alias): + _raise_metric_config_error( + f"Alias found for metric '{meta.name}' which is not allowed", meta._path + ) + elif isinstance(node, exp.AggFunc): + agg_or_ref = True + elif isinstance(node, exp.Column) and not node.table: + agg_or_ref = True + ref = node.name.lower() + metric_refs[node] = ref + + if ref not in metrics: + is_cycle = ref in visiting + if is_cycle or ref not in metas: + dependency_path = " -> ".join( + [frame[0].name for frame in stack] + [meta.name, ref] + ) + if is_cycle: + message = f"Metric dependency cycle detected: {dependency_path}" + else: + message = ( + f"Unknown metric '{ref}' referenced by metric '{meta.name}' " + f"(dependency path: {dependency_path})" + ) + _raise_metric_config_error(message, meta._path) + + dependency = metas[ref] + stack.append((meta, nodes, metric_refs, agg_or_ref)) + stack.append((dependency, dependency.expression.walk(), {}, False)) + visiting.add(ref) + break + else: + if not agg_or_ref: + _raise_metric_config_error( + f"Metric '{meta.name}' missing an aggregation or metric ref", meta._path + ) + + if metric_refs: + expanded = meta.expression.copy() + for column in expanded.find_all(exp.Column): + reference = metric_refs.get(column) + if reference is not None: + column.replace(metrics[reference].expanded.copy()) + else: + expanded = exp.alias_(meta.expression, meta.name) + + metric = Metric(**meta.dict(), expanded=expanded) + metric._path = meta._path + visiting.remove(meta.name) + + if not stack: + return metric + + metrics[meta.name] = metric class Metric(MetricMeta, frozen=True): diff --git a/tests/core/metric/test_metric.py b/tests/core/metric/test_metric.py index 51c97fbe3d..d6aca441ea 100644 --- a/tests/core/metric/test_metric.py +++ b/tests/core/metric/test_metric.py @@ -1,156 +1,155 @@ +from pathlib import Path + import pytest from sqlmesh.core import dialect as d -from sqlmesh.core.metric import expand_metrics, load_metric_ddl -from sqlmesh.core.metric.definition import _get_measure_and_dim_tables +from sqlmesh.core.metric import expand_metrics, load_metric_ddl, rewrite +from sqlmesh.core.reference import ReferenceGraph +from sqlmesh.utils import UniqueKeyDict from sqlmesh.utils.errors import ConfigError -def test_load_metric_ddl(): - a = d.parse_one( - """ - -- description a - METRIC ( - name A, - expression SUM(x), - owner b - ); - """ - ) - - meta = load_metric_ddl(a, dialect="") - assert meta.name == "a" - assert meta.expression.sql() == "SUM(x)" - assert meta.owner == "b" - assert meta.description == "description a" - - -def test_load_invalid(): - with pytest.raises( - ConfigError, match=r"Only METRIC\(...\) statements are allowed. Found SELECT" - ): - load_metric_ddl( - d.parse_one( - """ - SELECT 1; - """ - ), - dialect="", +@pytest.mark.parametrize( + "statement, tokens", + [ + ("SELECT 1", ("METRIC", "SELECT")), + ("METRIC(name invalid, expression 1)", ("invalid", "aggregation", "ref")), + ], +) +def test_load_invalid(statement, tokens): + path = Path("metrics/invalid.sql") + with pytest.raises(ConfigError) as exc_info: + load_metric_ddl(d.parse_one(statement), dialect="duckdb", path=path).to_metric({}, {}) + + message = str(exc_info.value) + assert str(path) in message + assert all(token in message for token in tokens) + + +def _load_metas(definitions, dialect="duckdb"): + metas = UniqueKeyDict("metrics") + for name, expression in definitions: + meta = load_metric_ddl( + d.parse_one(f"METRIC(name {name}, expression {expression})", dialect=dialect), + dialect=dialect, + path=Path("metrics") / f"{name.lower()}.sql", ) - - with pytest.raises(ConfigError, match=r"Metric 'a' missing an aggregation or metric ref."): - load_metric_ddl( - d.parse_one( - """ - METRIC ( - name a, - expression 1 - ) - """ - ), - dialect="", - ).to_metric({}, {}) - - -def test_expand_metrics(): - expressions = d.parse( - """ - -- description a - METRIC ( - name a, - expression SUM(model.x), - owner b - ); - - -- description b - METRIC ( - name b, - expression COUNT(DISTINCT model.y), - owner b - ); - - -- description c - METRIC ( - name c, - expression a / b, - owner b - ); - - -- description d - METRIC ( - name d, - expression c + 1, - owner b - ); - """ - ) - - metas = {} - for expr in expressions: - meta = load_metric_ddl(expr, dialect="") metas[meta.name] = meta + return metas - metrics = expand_metrics(metas) - - metric_a = metrics["a"] - assert metric_a.name == "a" - assert metric_a.expression.sql() == "SUM(model.x)" - assert metric_a.expanded.sql() == "SUM(model.x) AS a" - assert metric_a.formula.sql() == "a AS a" - assert metric_a.owner == "b" - assert metric_a.description == "description a" - - metric_b = metrics["b"] - assert metric_b.name == "b" - assert metric_b.expression.sql() == "COUNT(DISTINCT model.y)" - assert metric_b.expanded.sql() == "COUNT(DISTINCT model.y) AS b" - assert metric_b.formula.sql() == "b AS b" - - metric_c = metrics["c"] - assert metric_c.name == "c" - assert metric_c.expression.sql() == "a / b" - assert metric_c.expanded.sql() == "SUM(model.x) AS a / COUNT(DISTINCT model.y) AS b" - assert metric_c.formula.sql() == "a / b AS c" - - metric_d = metrics["d"] - assert metric_d.expression.sql() == "c + 1" - assert metric_d.expanded.sql() == "SUM(model.x) AS a / COUNT(DISTINCT model.y) AS b + 1" - assert metric_d.formula.sql() == "a / b + 1 AS d" - - assert metric_d.aggs == { - d.parse_one("SUM(model.x) AS a"): ("model", ()), - d.parse_one("COUNT(DISTINCT model.y) AS b"): ("model", ()), - } - - metas = {} - for expr in expressions: - meta = load_metric_ddl(expr, dialect="snowflake") - metas[meta.name] = meta - # Checks that metric names are not normalized according to the target dialect - snowflake_metrics = expand_metrics(metas) - assert all(metric_name.islower() for metric_name in snowflake_metrics) - - metric_c = snowflake_metrics["c"] - assert metric_c.name == "c" - assert metric_c.expression.sql() == "a / b" - assert metric_c.expanded.sql() == "SUM(model.x) AS a / COUNT(DISTINCT model.y) AS b" - assert metric_c.formula.sql() == "a / b AS c" - - -def test_get_measure_and_dim_tables(): - assert _get_measure_and_dim_tables(d.parse_one("SUM(a.x)")) == ("a", ()) - assert _get_measure_and_dim_tables(d.parse_one("SUM(a.x + a.y)")) == ("a", ()) - assert _get_measure_and_dim_tables(d.parse_one("SUM(a.x + b.y)")) == ("a", ("b",)) - assert _get_measure_and_dim_tables(d.parse_one("c.z + SUM(a.x)")) == ("a", ("c",)) - assert _get_measure_and_dim_tables(d.parse_one("SUM(IF(c.z = 'dim', a.x, 0))")) == ( - "a", - ("c",), +def _execute_metrics(metrics, names): + import duckdb + + query = rewrite( + "SELECT " + ", ".join(f"METRIC({name})" for name in names) + " FROM __semantic.__table", + graph=ReferenceGraph([]), + metrics=metrics, + dialect="duckdb", ) - assert _get_measure_and_dim_tables( - d.parse_one("SUM(IF(c.z = 'dim' AND b.y > 0, (a.x + a.x) + 3, 0))") - ) == ("a", ("c", "b")) - assert _get_measure_and_dim_tables(d.parse_one("SUM(CASE b.y WHEN 1 THEN a.x ELSE 0 END)")) == ( - "a", - ("b",), + with duckdb.connect(":memory:") as connection: + connection.execute("CREATE TABLE facts(amount INT, category VARCHAR)") + connection.execute("INSERT INTO facts VALUES (10, 'a'), (15, 'a'), (5, 'b')") + return connection.execute(query.sql(dialect="duckdb")).fetchall() + + +@pytest.mark.parametrize("direct", [False, True]) +def test_forward_and_diamond_dependencies(direct): + metas = _load_metas( + [ + ("root", "ratio + adjusted + total / count"), + ("ratio", "total / count"), + ("adjusted", "ratio + 1"), + ("total", "SUM(facts.amount)"), + ("count", "COUNT(DISTINCT facts.category)"), + ] ) + if direct: + root = metas.pop("root") + metrics = UniqueKeyDict("metrics") + metrics[root.name] = root.to_metric(metas, metrics) + else: + metrics = expand_metrics(metas) + + assert _execute_metrics(metrics, ["root", "ratio", "adjusted", "total", "count"]) == [ + (46.0, 15.0, 16.0, 30, 2) + ] + + +@pytest.mark.parametrize("dialect", ["duckdb", "snowflake"]) +def test_case_insensitive_metric_dependencies(dialect): + metas = _load_metas( + [ + ("RaTiO", 'ToTaL / "COUNT"'), + ("TOTAL", "SUM(facts.amount)"), + ("Count", "COUNT(DISTINCT facts.category)"), + ], + dialect=dialect, + ) + + assert _execute_metrics(expand_metrics(metas), ["ratio", "total", "count"]) == [(15.0, 30, 2)] + + +def test_direct_expansion_reuses_resolved_dependencies(): + total = _load_metas([("total", "SUM(facts.amount)")])["total"] + metrics = UniqueKeyDict("metrics") + metrics[total.name] = total.to_metric({}, metrics) + doubled = _load_metas([("doubled", "total + total")])["doubled"] + metrics[doubled.name] = doubled.to_metric({}, metrics) + + assert _execute_metrics(metrics, ["total", "doubled"]) == [(30, 60)] + + +@pytest.mark.parametrize("direct", [False, True]) +@pytest.mark.parametrize( + "definitions, dependency_path, source_path, diagnostic", + [ + ( + [("root", "root")], + "root -> root", + "metrics/root.sql", + "cycle", + ), + ( + [("root", "x"), ("x", "y"), ("y", "x")], + "root -> x -> y -> x", + "metrics/y.sql", + "cycle", + ), + ( + [("root", "intermediate"), ("intermediate", "missing_metric + 1")], + "root -> intermediate -> missing_metric", + "metrics/intermediate.sql", + "unknown", + ), + ], + ids=["self_cycle", "nested_cycle", "unknown_dependency"], +) +def test_invalid_metric_dependencies(direct, definitions, dependency_path, source_path, diagnostic): + metas = _load_metas(definitions) + + with pytest.raises(ConfigError) as exc_info: + if direct: + root = metas.pop("root") + root.to_metric(metas, UniqueKeyDict("metrics")) + else: + expand_metrics(metas) + + message = str(exc_info.value) + assert diagnostic in message.lower() + assert dependency_path in message + assert source_path in message + + +def test_long_dependency_cycle(): + depth = 1100 + metas = _load_metas([(f"m{i}", f"m{(i + 1) % depth}") for i in range(depth)]) + + with pytest.raises(ConfigError) as exc_info: + expand_metrics(metas) + + message = str(exc_info.value) + assert "cycle" in message.lower() + assert "m0 -> m1 -> m2" in message + assert "m1099 -> m0" in message + assert "metrics/m1099.sql" in message From d8e1b009da79d79ccface1bd7bfd451f57db01b0 Mon Sep 17 00:00:00 2001 From: tchivs Date: Sat, 5 Sep 2026 17:17:20 +0800 Subject: [PATCH 2/2] fix: preserve filters and grouping across metric sources Resolve complete predicates for every contributing fact before aggregation. Preserve full-join grouping keys, including NULL and composite keys, and retain outer SQL scope references. Use PostgreSQL-compatible composite equality for nullable grouping keys and add native model-to-query acceptance coverage. Reject unsupported query shapes explicitly while retaining the Metrics prototype status. Signed-off-by: tchivs --- docs/concepts/metrics/overview.md | 14 +- docs/development.md | 36 + sqlmesh/core/metric/rewriter.py | 354 ++++++--- sqlmesh/core/reference.py | 8 +- .../integration/test_integration_metrics.py | 165 +++++ tests/core/metric/test_rewriter.py | 679 ++++++++++++------ tests/core/test_reference.py | 12 + 7 files changed, 952 insertions(+), 316 deletions(-) create mode 100644 tests/core/engine_adapter/integration/test_integration_metrics.py diff --git a/docs/concepts/metrics/overview.md b/docs/concepts/metrics/overview.md index 8364bdd74a..68322b23d0 100644 --- a/docs/concepts/metrics/overview.md +++ b/docs/concepts/metrics/overview.md @@ -41,7 +41,7 @@ FROM __semantic.__table -- special table for simple metric queries GROUP BY ds ``` -When that model query is run, SQLMesh uses its semantic understanding of the query and metrics definitions to generate the code that is actually executed by the SQL engine: +SQLMesh expands the metric into SQL equivalent to the following query (generated aliases may differ): ``` sql linenums="1" SELECT @@ -60,3 +60,15 @@ FROM ( ``` SQLMesh automatically generates the correct join to use values from both the `sushi.orders` and `sushi.customers` tables. + +## Filters and grouping across metrics + +Queries against `__semantic.__table` treat its columns as logical dimensions. A `WHERE` predicate is resolved independently for every contributing fact source and applied before aggregation, including when the filtered dimension is not in `GROUP BY`. Derived metrics are calculated from those filtered aggregates. + +Dimension resolution prefers a column on the fact itself. Otherwise, the column must resolve to one reachable model through the configured grains and references. An explicit dimension-table alias selects that model instead of a same-named fact column. Unknown, unreachable, or ambiguous dimensions are rejected rather than ignored. + +Fact aggregates are combined with a full join by default. Group keys present only in a later fact are retained, and matching `NULL` group keys are combined, including composite keys. Missing metric values remain `NULL`; they are not automatically converted to zero. + +On PostgreSQL, full joins use composite-key equality to retain `NULL` groups without the planner restriction on `IS NOT DISTINCT FROM` join conditions. Corresponding dimensions must have matching PostgreSQL types; use an explicit cast in the grouping expression when models expose different types. Metric arithmetic follows the definition's SQL dialect: use a numeric or floating-point cast for fractional ratios of integer counts, and `NULLIF(denominator, 0)` when a zero denominator should produce `NULL`. + +The prototype rejects subqueries in metric `WHERE` filters, grouping sets, and reference paths that cannot be compiled into safe matching-key joins. This includes multi-hop paths that change reference keys. These checks do not restrict ordinary SQL scopes without `METRIC` expressions. They are not a substitute for application authorization. diff --git a/docs/development.md b/docs/development.md index d80ef60689..b2a81f3c92 100644 --- a/docs/development.md +++ b/docs/development.md @@ -68,6 +68,42 @@ Run more comprehensive tests that run on each commit: make slow-test ``` +### PostgreSQL Metrics acceptance + +The native Metrics suite loads a temporary SQLMesh project, applies its models to PostgreSQL, compiles queries with `Context.rewrite`, and checks results returned by PostgreSQL. It uses the existing `inttest_postgres` gateway and `postgres` / `docker` test markers. Only scheduling metadata uses an isolated in-memory DuckDB connection; model and metric SQL run on PostgreSQL. Temporary model schemas are cleaned up by the integration fixtures. + +With the project's PostgreSQL test service available (see `make engine-postgres-up`), run: + +```bash +pytest tests/core/engine_adapter/integration/test_integration_metrics.py -q +``` + +On a Linux Docker host, a separate test container can instead be run without publishing a database port. After activating the development virtual environment: + +```bash +( + set -eu + name="sqlmesh-metrics-pg-$$" + docker network create --internal "$name" + trap 'docker rm -f "$name" >/dev/null 2>&1 || true; docker network rm "$name" >/dev/null 2>&1 || true' EXIT + docker run -d --rm --name "$name" --network "$name" \ + -e POSTGRES_HOST_AUTH_METHOD=trust postgres:16-alpine + ready=false + for attempt in $(seq 1 30); do + if docker exec "$name" pg_isready -h 127.0.0.1 -U postgres; then + ready=true + break + fi + sleep 1 + done + "$ready" + export DOCKER_HOSTNAME="$(docker inspect -f '{{range .NetworkSettings.Networks}}{{.IPAddress}}{{end}}' "$name")" + pytest tests/core/engine_adapter/integration/test_integration_metrics.py -q +) +``` + +The trust-authenticated container is for disposable tests only: its internal network has no published host port or persistent data volume. Do not point this suite at a business database. Personal SQLMesh gateway overrides must not redirect `inttest_postgres` away from the intended test instance. + ## Documentation In order to run the documentation server, you will need to install the dependencies by running the following command. diff --git a/sqlmesh/core/metric/rewriter.py b/sqlmesh/core/metric/rewriter.py index 6c9ec429a8..59369ac5e7 100644 --- a/sqlmesh/core/metric/rewriter.py +++ b/sqlmesh/core/metric/rewriter.py @@ -2,14 +2,18 @@ import typing as t -from sqlglot import exp +from sqlglot import Dialect, exp from sqlglot.dialects.dialect import DialectType -from sqlglot.optimizer import Scope, find_all_in_scope, optimize +from sqlglot.dialects.postgres import Postgres +from sqlglot.errors import OptimizeError +from sqlglot.optimizer import find_all_in_scope, optimize +from sqlglot.optimizer.normalize_identifiers import normalize_identifiers from sqlglot.optimizer.optimize_joins import optimize_joins from sqlglot.optimizer.qualify import qualify from sqlmesh.core import dialect as d from sqlmesh.core.metric.definition import Metric, remove_namespace +from sqlmesh.utils.errors import ConfigError if t.TYPE_CHECKING: from sqlmesh.core.reference import ReferenceGraph @@ -32,11 +36,37 @@ def __init__( self.metrics = metrics self.dialect = dialect self.join_type = join_type - self.semantic_name = f"{semantic_schema}.{semantic_table}" + self.semantic_name = exp.table_name( + normalize_identifiers( + exp.to_table(f"{semantic_schema}.{semantic_table}"), dialect=dialect + ) + ) + + def _group_join_condition(self, left: t.List[exp.Expr], right: t.List[exp.Expr]) -> exp.Expr: + if not left: + return exp.true() + if type(Dialect.get_or_raise(self.dialect)) is Postgres: + # PostgreSQL cannot plan a FULL JOIN on IS NOT DISTINCT FROM. + # Composite equality is merge-joinable and treats NULL fields as equal, + # unlike SQL row-constructor comparison without the RECORD casts. + return exp.EQ( + this=exp.Cast( + this=exp.Anonymous(this="ROW", expressions=[key.copy() for key in left]), + to=exp.DataType.build("record", dialect="postgres", udt=True), + ), + expression=exp.Cast( + this=exp.Anonymous(this="ROW", expressions=[key.copy() for key in right]), + to=exp.DataType.build("record", dialect="postgres", udt=True), + ), + ) + return exp.and_( + *(exp.NullSafeEQ(this=a.copy(), expression=b.copy()) for a, b in zip(left, right)) + ) def rewrite(self, expression: exp.Expr) -> exp.Expr: - for select in list(expression.find_all(exp.Select)): - self._expand(select) + for select in reversed(list(expression.find_all(exp.Select))): + if next(find_all_in_scope(select, d.MetricAgg), None) is not None: + self._expand(select) return expression @@ -45,7 +75,9 @@ def _build_sources(self, projections: t.List[exp.Expr]) -> SourceAggsAndJoins: for projection in projections: for ref in find_all_in_scope(projection, d.MetricAgg): - metric = self.metrics[ref.this.name] + metric = self.metrics.get(ref.this.name.lower()) + if metric is None: + raise ConfigError(f"Unknown metric '{ref.this.name}'") ref.replace(metric.formula.this) for agg, (measure, dims) in metric.aggs.items(): @@ -57,131 +89,247 @@ def _build_sources(self, projections: t.List[exp.Expr]) -> SourceAggsAndJoins: return sources def _expand(self, select: exp.Select) -> None: - base = select.args["from_"].this.find(exp.Table) + from_ = select.args.get("from_") + if from_ is None or not isinstance(from_.this, exp.Table): + raise ConfigError("Metric queries require a table source") + base = from_.this base_alias = base.alias_or_name base_name = exp.table_name(base) + logical_alias = base_alias if base_name == self.semantic_name else "" - sources: SourceAggsAndJoins = ( - {} if base_name == self.semantic_name else {base_name: (set(), {})} - ) + where = select.args.get("where") + if where is not None and where.find(exp.Query) is not None: + raise ConfigError("Subqueries in metric WHERE filters are not supported") + + sources: SourceAggsAndJoins = {} if logical_alias else {base_name: (set(), {})} sources.update(self._build_sources(select.selects)) + if next(find_all_in_scope(select, d.MetricAgg), None) is not None: + raise ConfigError("METRIC references must appear in the SELECT projections") group = select.args.pop("group", None) group_by = group.expressions if group else [] - - mapping = { - remove_namespace(exp.table_name(source.assert_is(exp.Table))): name - for name, source in Scope(select).references - if name != base_alias - } - - explicit_joins = {exp.table_name(join.this): join for join in select.args.pop("joins", [])} - + if group and any(value for key, value in group.args.items() if key != "expressions"): + raise ConfigError("Grouping sets are not supported in metric queries") + + explicit_joins = {} + aliases = {} if logical_alias else {base_alias: base_name} + for join in select.args.pop("joins", []): + if not isinstance(join.this, exp.Table): + raise ConfigError("Metric dimension joins require table sources") + target = exp.table_name(join.this) + if target in explicit_joins or target == base_name: + raise ConfigError(f"Ambiguous metric dimension join to '{target}'") + explicit_joins[target] = join + aliases[join.this.alias_or_name] = target + + select.set("where", None) + # Stable private names let computed grains be joined as values, not re-evaluated + # against aggregate rows, and avoid collisions with named metric aggregates. + used_names = {agg.alias_or_name for aggs, _ in sources.values() for agg in aggs} + group_names = [] + for i in range(len(group_by)): + key = f"__metric_group_{i}" + while key in used_names: + key += "_" + used_names.add(key) + group_names.append(key) + + merged_keys: t.List[exp.Expr] = [] for i, (name, (aggs, joins)) in enumerate(sources.items()): - source: exp.Expr = exp.to_table(name) table_name = remove_namespace(name) - - if not isinstance(source, exp.Select): - source = exp.Select().from_( - exp.alias_(source, table_name, table=True, copy=False), copy=False + source_aliases = {model: alias for alias, model in aliases.items()} + source_aliases[name] = table_name + joins.update({target: join.copy() for target, join in explicit_joins.items()}) + grain = [ + self._resolve_columns(e.copy(), name, logical_alias, aliases, source_aliases, joins) + for e in group_by + ] + predicate = ( + self._resolve_columns( + where.this.copy(), name, logical_alias, aliases, source_aliases, joins ) - - joins.update(explicit_joins) - query = self._add_joins(source, name, joins, group_by, mapping).select( - *sorted(aggs, key=str), copy=False + if where is not None + else None ) + for join in list(joins.values()): + if join is not None and join.args.get("on") is not None: + join.set( + "on", + self._resolve_columns( + join.args["on"], name, logical_alias, aliases, source_aliases, joins + ), + ) + query = exp.select().from_( + exp.alias_(exp.to_table(name), table_name, table=True, copy=False), copy=False + ) + self._add_joins(query, name, joins, source_aliases) + query.select( + *(exp.alias_(e, key, copy=False) for e, key in zip(grain, group_names)), + *sorted(aggs, key=str), + copy=False, + ) + if grain: + query.group_by(*(e.copy() for e in grain), copy=False) + if predicate is not None: + query.where(predicate, copy=False) if not query.selects: - query.select("*", copy=False) - + query.select(exp.Literal.number(1), copy=False).distinct(copy=False) + + # Metric aggregates use canonical model aliases; explicit dimension aliases + # must also be honored inside conditional aggregates. + aggregate_aliases = { + remove_namespace(model): alias for model, alias in source_aliases.items() + } + for agg in aggs: + for column in find_all_in_scope(agg, exp.Column): + if column.table in aggregate_aliases: + column.set("table", exp.to_identifier(aggregate_aliases[column.table])) + + outer_alias = f"__metric_source_{i}" + keys: t.List[exp.Expr] = [exp.column(key, table=outer_alias) for key in group_names] if i == 0: - where = select.args.pop("where", None) - - if where: - query.where(_replace_table(where.this, table_name, base_alias), copy=False) - - select.from_(query.subquery(base_alias, copy=False), copy=False) + select.from_(query.subquery(outer_alias, copy=False), copy=False) + merged_keys = keys else: select.join( query, - on=[e.eq(_replace_table(e.copy(), table_name, base_alias)) for e in group_by], # type: ignore + on=self._group_join_condition(merged_keys, keys), join_type=self.join_type, - join_alias=table_name, + join_alias=outer_alias, copy=False, ) + merged_keys = [ + exp.Coalesce(this=left, expressions=[right]) + for left, right in zip(merged_keys, keys) + ] + + replacements = dict(zip(group_by, merged_keys)) + + def replace_grain(node: exp.Expr) -> exp.Expr: + if node in replacements: + return replacements[node].copy() + # Scalar subqueries belong to a separate SQL scope. + return node.copy() if isinstance(node, exp.Query) else node + + for projection in select.selects: + output_name = projection.output_name + rewritten = projection.transform(replace_grain) + if output_name and rewritten.output_name != output_name: + rewritten = exp.alias_(rewritten, output_name, copy=False) + projection.replace(rewritten) + for clause in ("order", "having", "qualify", "distinct"): + if select.args.get(clause) is not None: + select.set(clause, select.args[clause].transform(replace_grain)) + if select.args.get("windows"): + select.set( + "windows", [window.transform(replace_grain) for window in select.args["windows"]] + ) - for node in find_all_in_scope(query, exp.Column, exp.TableAlias): # type: ignore[arg-type,var-annotated] - if isinstance(node, exp.Column): - if node.table in mapping: - node.set("table", exp.to_identifier(mapping[node.table])) - else: - if node.name in mapping: - node.set("this", exp.to_identifier(mapping[node.name])) + def _resolve_columns( + self, + expression: exp.Expr, + source: str, + logical_alias: str, + aliases: t.Dict[str, str], + source_aliases: t.Dict[str, str], + joins: t.Dict[str, t.Optional[exp.Join]], + ) -> exp.Expr: + for column in find_all_in_scope(expression, exp.Column): + try: + models = self.graph.models_for_column(source, column.name) + except KeyError: + models = [] + + target = aliases.get(column.table) + if target is not None: + if target not in models: + raise ConfigError( + f"Cannot resolve metric dimension '{column}' from '{source}' via '{target}'" + ) + elif column.table and column.table != logical_alias: + raise ConfigError(f"Unknown metric dimension alias '{column.table}' in '{column}'") + elif source in models: + target = source + elif len(models) == 1: + target = models[0] + elif len(models) > 1: + raise ConfigError( + f"Ambiguous metric dimension '{column}' from '{source}': {', '.join(models)}" + ) + else: + raise ConfigError(f"Cannot resolve metric dimension '{column}' from '{source}'") + + if target != source: + joins.setdefault(target, None) + column.set( + "table", exp.to_identifier(source_aliases.get(target, remove_namespace(target))) + ) + return expression def _add_joins( self, source: exp.Select, name: str, joins: t.Dict[str, t.Optional[exp.Join]], - group_by: t.List[exp.Expr], - mapping: t.Dict[str, str], - ) -> exp.Select: - grain = [e.copy() for e in group_by] - table_name = remove_namespace(name) - mapping = {v: k for k, v in mapping.items()} - - for expr in grain: - for node in expr.walk(): - if isinstance(node, exp.Column): - models = self.graph.models_for_column(name, node.name) - - if name in models: - node.args["table"] = exp.to_identifier(table_name) - elif models: - t = mapping.get(node.table) - model = next( - (model for model in models if remove_namespace(model) == t), - models[0], - ) - node.args["table"] = exp.to_identifier(t or remove_namespace(model)) - if model not in joins: - joins[model] = None - - for target, join in joins.items(): + aliases: t.Dict[str, str], + ) -> None: + joined = {name} + for target in joins: + if target in joined: + continue path = self.graph.find_path(name, target) - for i in range(len(path) - 1): - a_ref = path[i] - b_ref = path[i + 1] - a_model_alias = remove_namespace(a_ref.model_name) - b_model_alias = remove_namespace(b_ref.model_name) - + if ( + not path + or path[0].model_name != name + or path[-1].model_name != target + or any(a.name != b.name for a, b in zip(path, path[1:])) + ): + raise ConfigError(f"Cannot safely join metric dimension '{target}' from '{name}'") + for a_ref, b_ref in zip(path, path[1:]): + if b_ref.model_name in joined: + continue a = a_ref.expression.copy() - a.set("table", exp.to_identifier(a_model_alias)) b = b_ref.expression.copy() - b.set("table", exp.to_identifier(b_model_alias)) - on = a.eq(b) - - if join: + if isinstance(a, exp.Alias): + a = a.this + if isinstance(b, exp.Alias): + b = b.this + for expression, model in ((a, a_ref.model_name), (b, b_ref.model_name)): + for column in expression.find_all(exp.Column): + column.set( + "table", exp.to_identifier(aliases.get(model, remove_namespace(model))) + ) + on: exp.Condition = a.eq(b) + explicit = joins.get(b_ref.model_name) + if explicit is not None: + join = explicit.copy() + if join.args.get("on") is not None and join.args["on"] != on: + on = exp.and_(on, join.args["on"]) join.set("on", on) + join.set("using", None) source.append("joins", join) else: source.join( b_ref.model_name, on=on, join_type="LEFT", - join_alias=b_model_alias, + join_alias=aliases.get( + b_ref.model_name, remove_namespace(b_ref.model_name) + ), dialect=self.dialect, copy=False, ) - - return source.select(*grain, copy=False).group_by(*grain, copy=False) + joined.add(b_ref.model_name) -def _replace_table(node: exp.Expr, table: str, base_alias: str) -> exp.Expr: - for column in find_all_in_scope(node, exp.Column): - if column.table == base_alias: - column.args["table"] = exp.to_identifier(table) - return node +def _prepare_metric_references(expression: exp.Expr) -> exp.Expr: + # A METRIC argument names a metric, not a column in one of the query's tables. + for ref in expression.find_all(d.MetricAgg): + if not isinstance(ref.this, (exp.Column, exp.Identifier)): + raise ConfigError(f"Invalid metric reference '{ref.this}'") + ref.set("this", exp.to_identifier(ref.this.name)) + return expression def rewrite( @@ -191,14 +339,20 @@ def rewrite( dialect: t.Optional[str] = "", ) -> exp.Expr: rewriter = Rewriter(graph=graph, metrics=metrics, dialect=dialect) - - return optimize( - d.parse_one(sql, dialect=dialect) if isinstance(sql, str) else sql, - dialect=dialect, - quote_identifiers=False, - rules=( - qualify, - rewriter.rewrite, - optimize_joins, - ), - ) + expression = d.parse_one(sql, dialect=dialect) if isinstance(sql, str) else sql + try: + return optimize( + expression, + dialect=dialect, + quote_identifiers=False, + rules=( + _prepare_metric_references, + qualify, + rewriter.rewrite, + optimize_joins, + ), + ) + except OptimizeError as ex: + if expression.find(d.MetricAgg) is None: + raise + raise ConfigError(f"Cannot resolve metric query: {ex}") from ex diff --git a/sqlmesh/core/reference.py b/sqlmesh/core/reference.py index 9e93ce7b38..48f6890a4e 100644 --- a/sqlmesh/core/reference.py +++ b/sqlmesh/core/reference.py @@ -127,6 +127,11 @@ def find_path(self, source: str, target: str, max_depth: int = 3) -> t.List[Refe ref_name = path[-1].name for model_name in sorted(self._ref_models[ref_name]): + if model_name == target: + ref = self._model_refs[model_name][ref_name] + if model_name not in visited and not (many and not ref.unique): + return path + [ref] + continue for ref in self._model_refs[model_name].values(): # paths cannot have loops or contain many to many refs if model_name in visited or (many and not ref.unique): @@ -134,9 +139,6 @@ def find_path(self, source: str, target: str, max_depth: int = 3) -> t.List[Refe new_path = path + [ref] - if model_name == target: - return new_path - if len(new_path) < max_depth: queue.append(new_path) diff --git a/tests/core/engine_adapter/integration/test_integration_metrics.py b/tests/core/engine_adapter/integration/test_integration_metrics.py new file mode 100644 index 0000000000..e63c200076 --- /dev/null +++ b/tests/core/engine_adapter/integration/test_integration_metrics.py @@ -0,0 +1,165 @@ +# SPDX-License-Identifier: Apache-2.0 +from datetime import datetime +from decimal import Decimal +from pathlib import Path + +import pytest +from sqlglot import exp + +from sqlmesh.core.config import Config, ModelDefaultsConfig +from tests.core.engine_adapter.integration import ( + ENGINES_BY_NAME, + TestContext, + generate_pytest_params, +) + + +@pytest.fixture(params=list(generate_pytest_params(ENGINES_BY_NAME["postgres"]))) +def ctx(request, create_test_context): + yield from create_test_context(*request.param) + + +@pytest.fixture +def metric_project(ctx: TestContext, tmp_path: Path): + schema = exp.to_table(ctx.schema()).sql("postgres") + models = tmp_path / "models" + metrics = tmp_path / "metrics" + models.mkdir() + metrics.mkdir() + rows = { + "a": "(1,'A','web',10.50,'2024-02-28 23:30:00'), (2,'A','app',20.50,'2024-02-29 00:30:00'), (3,'B','web',100.00,'2024-03-01 00:00:00'), (4,NULL,NULL,11.00,'2024-03-01 00:00:00')", + "b": "(11,'A','web',2.00,'2024-02-28 23:30:00'), (12,'A','app',3.00,'2024-02-29 00:30:00'), (13,'A','app',5.00,'2024-02-29 01:30:00'), (14,'B','web',50.00,'2024-03-01 00:00:00'), (15,'C','app',7.00,'2024-03-02 00:00:00'), (16,NULL,NULL,2.00,'2024-03-01 00:00:00')", + "c": "(21,'B','web',4.00,'2024-03-01 00:00:00'), (22,'C','app',7.00,'2024-03-02 00:00:00'), (23,'D','web',8.00,'2024-03-03 00:00:00'), (24,NULL,NULL,3.00,'2024-03-01 00:00:00')", + } + for name, values in rows.items(): + (models / f"{name}.sql").write_text( + f""" + MODEL (name {schema}.{name}, kind FULL, grain {name}_id, references org); + SELECT id::INT AS {name}_id, org::TEXT AS org, channel::TEXT AS channel, + amount::DECIMAL(12,2) AS amount, occurred_at::TIMESTAMP AS occurred_at + FROM (VALUES {values}) AS data(id, org, channel, amount, occurred_at) + """ + ) + (models / "organizations.sql").write_text( + f""" + MODEL (name {schema}.organizations, kind FULL, grain org); + SELECT org::TEXT AS org, region::TEXT AS region, status::TEXT AS status + FROM (VALUES ('A','north','ACTIVE'),('B','south','INACTIVE'), + ('C','north','ACTIVE'),('D','south','INACTIVE')) AS data(org,region,status) + """ + ) + (models / "quoted.sql").write_text( + f""" + MODEL (name {schema}.\"QuotedFacts\", kind FULL, grain \"Id\"); + SELECT 1 AS "Id", 'A'::TEXT AS "Org", 1.25::DECIMAL(12,2) AS "Amount" + """ + ) + (metrics / "metrics.sql").write_text( + "\n".join( + f"METRIC(name {name}_sum, expression SUM({schema}.{name}.amount));" for name in rows + ) + + f""" + METRIC(name ratio, expression a_sum / b_sum); + METRIC(name a_count, expression COUNT({schema}.a.a_id)); + METRIC(name b_count, expression COUNT({schema}.b.b_id)); + METRIC(name count_ratio, expression CAST(a_count AS DOUBLE) / NULLIF(b_count, 0)); + METRIC(name active_amount, expression SUM(IF({schema}.organizations.status = 'ACTIVE', {schema}.a.amount, 0))); + METRIC(name quoted_amount, expression SUM({schema}."QuotedFacts"."Amount")); + """ + ) + + def configure(gateway: str, config: Config) -> None: + config.model_defaults = ModelDefaultsConfig(dialect="postgres") + + context = ctx.create_context(path=tmp_path, config_mutator=configure) + try: + context.plan(auto_apply=True, no_prompts=True) + yield context, schema + finally: + context.close() + + +def _query(project, sql): + context, _ = project + return context.engine_adapter.fetchall(context.rewrite(sql).sql("postgres")) + + +def test_metrics_postgres_filtered_ratio(metric_project): + assert _query( + metric_project, + "SELECT METRIC(a_sum), METRIC(b_sum), METRIC(ratio), METRIC(count_ratio) " + "FROM __semantic.__table s WHERE s.org = 'A'", + ) == [(Decimal("31.00"), Decimal("10.00"), Decimal("3.1"), pytest.approx(2 / 3))] + assert _query( + metric_project, + "SELECT s.org, METRIC(a_sum), METRIC(b_sum), METRIC(ratio) " + "FROM __semantic.__table s WHERE s.channel='web' GROUP BY s.org ORDER BY s.org", + ) == [ + ("A", Decimal("10.50"), Decimal("2.00"), Decimal("5.25")), + ("B", Decimal("100.00"), Decimal("50.00"), Decimal("2.0")), + ] + + +def test_metrics_postgres_empty_scope(metric_project): + assert _query( + metric_project, + "SELECT METRIC(a_count), METRIC(b_count), METRIC(ratio), METRIC(count_ratio) " + "FROM __semantic.__table s WHERE s.org = 'missing'", + ) == [(0, 0, None, None)] + + +def test_metrics_postgres_three_facts_null_keys(metric_project): + assert _query( + metric_project, + "SELECT s.org, s.channel, METRIC(a_sum), METRIC(b_sum), METRIC(c_sum) " + "FROM __semantic.__table s GROUP BY s.org,s.channel " + "ORDER BY s.org NULLS FIRST,s.channel NULLS FIRST", + ) == [ + (None, None, Decimal("11"), Decimal("2"), Decimal("3")), + ("A", "app", Decimal("20.5"), Decimal("8"), None), + ("A", "web", Decimal("10.5"), Decimal("2"), None), + ("B", "web", Decimal("100"), Decimal("50"), Decimal("4")), + ("C", "app", None, Decimal("7"), Decimal("7")), + ("D", "web", None, None, Decimal("8")), + ] + + +def test_metrics_postgres_dimension_filter(metric_project): + _, schema = metric_project + assert _query( + metric_project, + f"SELECT s.org,METRIC(a_sum),METRIC(b_sum),METRIC(active_amount) " + f"FROM __semantic.__table s LEFT JOIN {schema}.organizations o ON s.org=o.org " + "WHERE o.region='north' GROUP BY s.org ORDER BY s.org", + ) == [("A", Decimal("31"), Decimal("10"), Decimal("31")), ("C", None, Decimal("7"), None)] + + +def test_metrics_postgres_time_buckets(metric_project): + assert _query( + metric_project, + "SELECT DATE_TRUNC('day',s.occurred_at) AS day,METRIC(a_sum),METRIC(b_sum) " + "FROM __semantic.__table s " + "WHERE s.occurred_at >= CAST('2024-02-29' AS TIMESTAMP) " + "AND s.occurred_at < CAST('2024-03-02' AS TIMESTAMP) " + "GROUP BY DATE_TRUNC('day',s.occurred_at) ORDER BY day", + ) == [ + (datetime(2024, 2, 29), Decimal("20.5"), Decimal("8")), + (datetime(2024, 3, 1), Decimal("111"), Decimal("52")), + ] + + +def test_metrics_postgres_quoted_identifiers(metric_project): + assert _query( + metric_project, + 'SELECT s."Org" AS "Organization", METRIC("QUOTED_AMOUNT") AS "Value" ' + 'FROM __semantic.__table s WHERE s."Org"=\'A\' GROUP BY s."Org"', + ) == [("A", Decimal("1.25"))] + + +def test_metrics_postgres_distinct_on(metric_project): + assert _query( + metric_project, + "SELECT DISTINCT ON (s.channel) s.org, METRIC(a_sum) " + "FROM __semantic.__table s GROUP BY s.org,s.channel " + "ORDER BY s.channel,s.org NULLS FIRST", + ) == [("A", Decimal("20.5")), ("A", Decimal("10.5")), (None, Decimal("11"))] diff --git a/tests/core/metric/test_rewriter.py b/tests/core/metric/test_rewriter.py index 9fe8169508..4cab6885e7 100644 --- a/tests/core/metric/test_rewriter.py +++ b/tests/core/metric/test_rewriter.py @@ -1,224 +1,479 @@ -from sqlglot import parse_one +import pytest +from sqlglot.optimizer.qualify import qualify -from sqlmesh.core.metric import rewrite +from sqlmesh.core import dialect as d +from sqlmesh.core.metric import expand_metrics, load_metric_ddl, rewrite +from sqlmesh.core.metric.rewriter import Rewriter +from sqlmesh.core.model import load_sql_based_model from sqlmesh.core.reference import ReferenceGraph +from sqlmesh.utils import UniqueKeyDict +from sqlmesh.utils.errors import ConfigError, SQLMeshError -def test_rewrite(sushi_context_pre_scheduling, assert_exp_eq): - context = sushi_context_pre_scheduling - graph = ReferenceGraph(context.models.values()) +def _model(name, columns, grain, references=""): + return load_sql_based_model( + d.parse( + f""" + MODEL ( + name {name}, + dialect duckdb, + columns ({columns}), + grain {grain} + {f", references ({references})" if references else ""} + ); + SELECT {", ".join(column.split()[0] for column in columns.split(","))} + FROM raw.source + """ + ) + ) - query = rewrite( + +@pytest.fixture +def metrics_runtime(): + import duckdb + + columns = "org VARCHAR, channel VARCHAR, amount INT, status VARCHAR" + models = [ + _model(f"facts.{name}", f"{name}_id INT, {columns}", f"{name}_id", "org") + for name in ("a", "b", "c") + ] + models.append( + _model("dims.organizations", "org VARCHAR, region VARCHAR, status VARCHAR", "org") + ) + graph = ReferenceGraph(models) + metas = UniqueKeyDict("metrics") + for expression in d.parse( """ - SELECT - c.customer_id, - METRIC(total_orders), - FROM sushi.customers AS c - GROUP BY c.customer_id + METRIC(name a_sum, expression SUM(facts.a.amount)); + METRIC(name b_sum, expression SUM(facts.b.amount)); + METRIC(name c_sum, expression SUM(facts.c.amount)); + METRIC(name ratio, expression a_sum / b_sum); + """ + ): + meta = load_metric_ddl(expression, dialect="duckdb") + metas[meta.name] = meta + metrics = expand_metrics(metas) + with duckdb.connect(":memory:") as connection: + connection.execute("CREATE SCHEMA facts; CREATE SCHEMA dims") + for name in ("a", "b", "c"): + connection.execute(f"CREATE TABLE facts.{name} ({name}_id INT, {columns})") + connection.execute( + """ + INSERT INTO facts.a VALUES + (1, 'A', 'web', 10, 'fact'), + (2, 'A', 'app', 20, 'fact'), + (3, 'B', 'web', 100, 'fact'); + INSERT INTO facts.b VALUES + (11, 'A', 'web', 2, 'fact'), + (12, 'A', 'app', 3, 'fact'), + (13, 'A', 'app', 5, 'fact'), + (14, 'B', 'web', 50, 'fact'); + INSERT INTO facts.c VALUES + (21, 'B', 'web', 4, 'fact'), + (22, 'C', 'app', 7, 'fact'), + (23, 'D', 'web', 8, 'fact'); + CREATE TABLE dims.organizations (org VARCHAR, region VARCHAR, status VARCHAR); + INSERT INTO dims.organizations VALUES + ('A', 'north', 'ACTIVE'), + ('B', 'south', 'INACTIVE'), + ('C', 'north', 'ACTIVE'), + ('D', 'south', 'INACTIVE'); + """ + ) + yield connection, graph, metrics + + +def _execute(runtime, sql, join_type=None): + connection, graph, metrics = runtime + if join_type is None: + query = rewrite(sql, graph=graph, metrics=metrics, dialect="duckdb") + else: + query = Rewriter( + graph=graph, metrics=metrics, dialect="duckdb", join_type=join_type + ).rewrite(qualify(d.parse_one(sql, dialect="duckdb"), dialect="duckdb")) + return connection.execute(query.sql("duckdb")).fetchall() + + +def test_rewrite_single_source(metrics_runtime): + assert _execute( + metrics_runtime, + """ + SELECT s.org, METRIC(a_sum) + FROM __semantic.__table s + WHERE s.org = 'A' + GROUP BY s.org """, - graph=graph, - metrics=context.metrics, + ) == [("A", 30)] + + +def test_rewrite_physical_base(metrics_runtime): + assert _execute( + metrics_runtime, + """ + SELECT o.org, METRIC(a_sum) + FROM dims.organizations o + GROUP BY o.org + ORDER BY o.org + """, + ) == [("A", 30), ("B", 100), ("C", None), ("D", None)] + + +@pytest.mark.parametrize("grouped", [False, True]) +def test_rewrite_shared_filter(metrics_runtime, grouped): + query = ( + "SELECT s.org, METRIC(a_sum), METRIC(b_sum), METRIC(ratio) " + if grouped + else "SELECT METRIC(a_sum), METRIC(b_sum), METRIC(ratio) " + ) + query += "FROM __semantic.__table s WHERE s.org = 'A'" + if grouped: + query += " GROUP BY s.org" + assert _execute(metrics_runtime, query) == ( + [("A", 30, 10, 3.0)] if grouped else [(30, 10, 3.0)] + ) + + +def test_rewrite_filter_not_in_grain(metrics_runtime): + assert _execute( + metrics_runtime, + """ + SELECT s.org, METRIC(a_sum), METRIC(b_sum), METRIC(ratio) + FROM __semantic.__table s + WHERE s.channel = 'web' + GROUP BY s.org + ORDER BY s.org + """, + ) == [("A", 10, 2, 5.0), ("B", 100, 50, 2.0)] + + +def test_rewrite_boolean_filter(metrics_runtime): + connection, _, _ = metrics_runtime + connection.execute( + """ + INSERT INTO facts.a VALUES (4, 'B', 'app', 21, 'fact'); + INSERT INTO facts.b VALUES (15, 'B', 'app', 7, 'fact'); + """ + ) + assert _execute( + metrics_runtime, + """ + SELECT s.org, METRIC(a_sum), METRIC(b_sum) + FROM __semantic.__table s + WHERE (s.org = 'A' AND s.channel = 'web') + OR (s.org = 'B' AND s.channel = 'app') + GROUP BY s.org + ORDER BY s.org + """, + ) == [("A", 10, 2), ("B", 21, 7)] + + +def test_rewrite_related_dimension_filter(metrics_runtime): + assert _execute( + metrics_runtime, + """ + SELECT s.org, METRIC(a_sum), METRIC(b_sum), METRIC(ratio) + FROM __semantic.__table s + WHERE s.region = 'north' + GROUP BY s.org + """, + ) == [("A", 30, 10, 3.0)] + + +def test_rewrite_dimension_with_multiple_references(metrics_runtime): + connection, graph, _ = metrics_runtime + graph.add_model( + _model( + "dims.labels", + "label_id INT, org VARCHAR, label VARCHAR", + "label_id", + "org", + ) + ) + connection.execute( + """ + CREATE TABLE dims.labels (label_id INT, org VARCHAR, label VARCHAR); + INSERT INTO dims.labels VALUES (101, 'A', 'included'), (102, 'B', 'excluded'); + """ + ) + assert _execute( + metrics_runtime, + """ + SELECT s.org, METRIC(a_sum), METRIC(b_sum) + FROM __semantic.__table s + WHERE s.label = 'included' + GROUP BY s.org + """, + ) == [("A", 30, 10)] + + +def test_rewrite_explicit_dimension_filter(metrics_runtime): + assert _execute( + metrics_runtime, + """ + SELECT s.org, METRIC(a_sum), METRIC(b_sum) + FROM __semantic.__table s + LEFT JOIN dims.organizations o ON s.org = o.org + WHERE o.status = 'ACTIVE' + GROUP BY s.org + """, + ) == [("A", 30, 10)] + + +@pytest.mark.parametrize("predicate", ["s.unknown = 1", "missing.region = 'north'"]) +def test_rewrite_unknown_filter(metrics_runtime, predicate): + with pytest.raises((ConfigError, SQLMeshError)): + _execute( + metrics_runtime, + f"SELECT METRIC(ratio) FROM __semantic.__table s WHERE {predicate}", + ) + + +def test_rewrite_ambiguous_filter(metrics_runtime): + _, graph, _ = metrics_runtime + graph.add_model(_model("dims.audience", "org VARCHAR, region VARCHAR", "org")) + with pytest.raises(ConfigError, match="(?i)ambiguous.*region"): + _execute( + metrics_runtime, + "SELECT METRIC(ratio) FROM __semantic.__table s WHERE s.region = 'north'", + ) + + +def test_rewrite_explicit_dimension_disambiguates_filter(metrics_runtime): + connection, graph, _ = metrics_runtime + graph.add_model(_model("dims.audience", "org VARCHAR, region VARCHAR", "org")) + connection.execute( + """ + CREATE TABLE dims.audience (org VARCHAR, region VARCHAR); + INSERT INTO dims.audience VALUES ('A', 'south'), ('B', 'north'); + """ ) - assert_exp_eq( - query, - parse_one( + assert _execute( + metrics_runtime, + """ + SELECT s.org, METRIC(a_sum), METRIC(b_sum) + FROM __semantic.__table s + LEFT JOIN dims.organizations o ON s.org = o.org + WHERE o.region = 'north' + GROUP BY s.org + """, + ) == [("A", 30, 10)] + + +def test_rewrite_unreachable_filter(metrics_runtime): + _, graph, _ = metrics_runtime + graph.add_model(_model("dims.remote", "remote_id INT, restricted VARCHAR", "remote_id")) + with pytest.raises((ConfigError, SQLMeshError)): + _execute( + metrics_runtime, + "SELECT METRIC(ratio) FROM __semantic.__table s WHERE s.restricted = 'yes'", + ) + + +@pytest.mark.parametrize( + "predicate", + [ + "s.org IN (SELECT org FROM dims.organizations)", + "EXISTS (SELECT 1 FROM dims.organizations o WHERE o.org = s.org)", + ], +) +def test_rewrite_rejects_subquery_filter(metrics_runtime, predicate): + with pytest.raises(ConfigError, match="(?i)subquer"): + _execute( + metrics_runtime, + f"SELECT METRIC(ratio) FROM __semantic.__table s WHERE {predicate}", + ) + + +def test_rewrite_right_only_group(metrics_runtime): + connection, _, _ = metrics_runtime + connection.execute("INSERT INTO facts.b VALUES (15, 'C', 'app', 7, 'fact')") + assert _execute( + metrics_runtime, + """ + SELECT s.org AS organization, METRIC(a_sum) AS total_a, METRIC(b_sum) AS total_b + FROM __semantic.__table s + GROUP BY s.org + ORDER BY organization + """, + ) == [("A", 30, 10), ("B", 100, 50), ("C", None, 7)] + + +def test_rewrite_three_facts(metrics_runtime): + connection, _, _ = metrics_runtime + connection.execute("INSERT INTO facts.b VALUES (15, 'C', 'app', 7, 'fact')") + assert _execute( + metrics_runtime, + """ + SELECT s.org, METRIC(a_sum), METRIC(b_sum), METRIC(c_sum) + FROM __semantic.__table s + GROUP BY s.org + ORDER BY s.org + """, + ) == [("A", 30, 10, None), ("B", 100, 50, 4), ("C", None, 7, 7), ("D", None, None, 8)] + + +def test_rewrite_composite_null_grain(metrics_runtime): + connection, _, _ = metrics_runtime + connection.execute( + """ + INSERT INTO facts.a VALUES + (4, NULL, NULL, 11, 'fact'), (5, 'A', NULL, 13, 'fact'), + (6, NULL, 'web', 17, 'fact'); + INSERT INTO facts.b VALUES + (15, NULL, NULL, 2, 'fact'), (16, 'A', NULL, 3, 'fact'), + (17, NULL, 'app', 5, 'fact'); + INSERT INTO facts.c VALUES + (24, NULL, NULL, 7, 'fact'), (25, NULL, 'app', 9, 'fact'), + (26, 'A', NULL, 11, 'fact'); + """ + ) + assert _execute( + metrics_runtime, + """ + SELECT s.org, s.channel, METRIC(a_sum), METRIC(b_sum), METRIC(c_sum) + FROM __semantic.__table s + GROUP BY s.org, s.channel + ORDER BY s.org NULLS FIRST, s.channel NULLS FIRST + """, + ) == [ + (None, None, 11, 2, 7), + (None, "app", None, 5, 9), + (None, "web", 17, None, None), + ("A", None, 13, 3, 11), + ("A", "app", 20, 8, None), + ("A", "web", 10, 2, None), + ("B", "web", 100, 50, 4), + ("C", "app", None, None, 7), + ("D", "web", None, None, 8), + ] + + +def test_rewrite_computed_grain_and_output_aliases(metrics_runtime): + connection, _, _ = metrics_runtime + connection.execute("INSERT INTO facts.b VALUES (15, 'C', 'app', 7, 'fact')") + assert _execute( + metrics_runtime, + """ + SELECT bucket, total_a, total_b + FROM ( + SELECT CASE WHEN s.org = 'A' THEN 'first' ELSE 'other' END AS bucket, + METRIC(a_sum) AS total_a, METRIC(b_sum) AS total_b + FROM __semantic.__table s + GROUP BY CASE WHEN s.org = 'A' THEN 'first' ELSE 'other' END + ) summary + ORDER BY bucket DESC + """, + ) == [("other", 100, 57), ("first", 30, 10)] + + +@pytest.mark.parametrize( + "join_type, expected", + [ + ("FULL", [("A", 30, 10), ("B", 100, 50), ("C", None, 7), ("D", 9, None)]), + ("LEFT", [("A", 30, 10), ("B", 100, 50), ("D", 9, None)]), + ("INNER", [("A", 30, 10), ("B", 100, 50)]), + ("RIGHT", [("A", 30, 10), ("B", 100, 50), ("C", None, 7)]), + ], +) +def test_rewrite_join_type(metrics_runtime, join_type, expected): + connection, _, _ = metrics_runtime + connection.execute( + """ + INSERT INTO facts.a VALUES (4, 'D', 'web', 9, 'fact'); + INSERT INTO facts.b VALUES (15, 'C', 'app', 7, 'fact'); + """ + ) + assert ( + _execute( + metrics_runtime, """ - SELECT - c.customer_id AS customer_id, - total_orders AS total_orders - FROM ( - SELECT - sushi__customers.customer_id - FROM sushi.customers AS sushi__customers - GROUP BY - sushi__customers.customer_id - ) AS c - FULL JOIN ( - SELECT - sushi__orders.customer_id, - COUNT(sushi__orders.id) AS total_orders - FROM sushi.orders AS sushi__orders - GROUP BY - sushi__orders.customer_id - ) AS sushi__orders - ON c.customer_id = sushi__orders.customer_id - """, - dialect=context.config.dialect, + SELECT s.org, METRIC(a_sum), METRIC(b_sum) + FROM __semantic.__table s + GROUP BY s.org + ORDER BY s.org + """, + join_type=join_type, + ) + == expected + ) + + +def test_rewrite_preserves_non_metric_nested_query(metrics_runtime): + assert _execute( + metrics_runtime, + """ + SELECT nested.org, nested.amount + (SELECT 1) AS incremented + FROM (SELECT org, amount FROM facts.a WHERE amount >= 20) nested + ORDER BY nested.org, nested.amount + """, + ) == [("A", 21), ("B", 101)] + + +def test_rewrite_nested_metric_scopes(metrics_runtime): + assert _execute( + metrics_runtime, + """ + SELECT METRIC(a_sum), + (SELECT METRIC(b_sum) FROM __semantic.__table WHERE org = 'A') AS b + FROM __semantic.__table + WHERE org = 'B' + """, + ) == [(100, 10)] + + +def test_rewrite_conditional_metric_with_dimension_alias(metrics_runtime): + _, _, metrics = metrics_runtime + meta = load_metric_ddl( + d.parse_one( + "METRIC(name active_amount, expression " + "SUM(IF(dims.organizations.status = 'ACTIVE', facts.a.amount, 0)))" ), + dialect="duckdb", ) + metrics.update(expand_metrics({meta.name: meta})) + assert _execute( + metrics_runtime, + """ + SELECT s.org, METRIC(active_amount) + FROM __semantic.__table s + LEFT JOIN dims.organizations o ON s.org = o.org + GROUP BY s.org + ORDER BY s.org + """, + ) == [("A", 30), ("B", 0)] + + +def test_rewrite_case_insensitive_query_metric(metrics_runtime): + connection, graph, metrics = metrics_runtime + query = rewrite( + 'SELECT METRIC("A_SUM") FROM __semantic.__table', + graph=graph, + metrics=metrics, + dialect="snowflake", + ) + assert connection.execute(query.sql("duckdb")).fetchall() == [(130,)] - # query = rewrite( - # """ - # SELECT - # event_date, - # METRIC(total_orders), - # METRIC(items_per_order), - # METRIC(total_orders_from_active_customers), - # FROM __semantic.__table - # GROUP BY event_date - # """, - # graph=graph, - # metrics=context.metrics, - # ) - - # assert_exp_eq( - # query, - # parse_one( - # """ - # SELECT - # __table.event_date AS event_date, - # total_orders AS total_orders, - # total_ordered_items / total_orders AS items_per_order, - # total_orders_from_active_customers AS total_orders_from_active_customers - # FROM ( - # SELECT - # sushi__orders.event_date, - # COUNT(CASE WHEN sushi__customers.status = 'ACTIVE' THEN sushi__orders.id ELSE NULL END) AS total_orders_from_active_customers, - # COUNT(sushi__orders.id) AS total_orders - # FROM sushi.orders AS sushi__orders - # LEFT JOIN sushi.customers AS sushi__customers - # ON sushi__orders.customer_id = sushi__customers.customer_id - # GROUP BY - # sushi__orders.event_date - # ) AS __table - # FULL JOIN ( - # SELECT - # sushi__order_items.event_date, - # SUM(sushi__order_items.quantity) AS total_ordered_items - # FROM sushi.order_items AS sushi__order_items - # GROUP BY - # sushi__order_items.event_date - # ) AS sushi__order_items - # ON __table.event_date = sushi__order_items.event_date - - # """, - # dialect=context.config.dialect, - # ), - # ) - - # query = rewrite( - # """ - # SELECT - # event_date, - # METRIC(total_orders), - # METRIC(items_per_order), - # METRIC(total_orders_from_active_customers), - # FROM sushi.orders - # GROUP BY event_date - # """, - # graph=graph, - # metrics=context.metrics, - # ) - - # assert_exp_eq( - # query, - # parse_one( - # """ - # SELECT - # orders.event_date AS event_date, - # orders.total_orders AS total_orders, - # orders.total_ordered_items / orders.total_orders AS items_per_order, - # orders.total_orders_from_active_customers AS total_orders_from_active_customers - # FROM ( - # SELECT - # sushi__orders.event_date, - # COUNT(IF(sushi__customers.status = 'ACTIVE', sushi__orders.id, NULL)) AS total_orders_from_active_customers, - # COUNT(sushi__orders.id) AS total_orders - # FROM sushi.orders AS sushi__orders - # LEFT JOIN sushi.customers AS sushi__customers - # ON sushi__orders.customer_id = sushi__customers.customer_id - # GROUP BY - # sushi__orders.event_date - # ) AS orders - # FULL JOIN ( - # SELECT - # sushi__order_items.event_date, - # SUM(sushi__order_items.quantity) AS total_ordered_items - # FROM sushi.order_items AS sushi__order_items - # GROUP BY - # sushi__order_items.event_date - # ) AS sushi__order_items - # ON orders.event_date = sushi__order_items.event_date - - # """, - # dialect=context.config.dialect, - # ), - # ) - - # query = rewrite( - # """ - # SELECT - # event_date, - # status, - # METRIC(total_orders), - # FROM __semantic.__table - # GROUP BY event_date, status - # """, - # graph=graph, - # metrics=context.metrics, - # ) - # assert_exp_eq( - # query, - # parse_one( - # """ - # SELECT - # __table.event_date AS event_date, - # __table.status AS status, - # __table.total_orders AS total_orders - # FROM ( - # SELECT - # sushi__orders.event_date, - # sushi__customers.status, - # COUNT(sushi__orders.id) AS total_orders - # FROM sushi.orders AS sushi__orders - # LEFT JOIN sushi.customers AS sushi__customers - # ON sushi__orders.customer_id = sushi__customers.customer_id - # GROUP BY - # sushi__orders.event_date, - # sushi__customers.status - # ) AS __table - # """, - # dialect=context.config.dialect, - # ), - # ) - - # query = rewrite( - # """ - # SELECT - # t.event_date, - # m.status, - # METRIC(t.total_orders), - # FROM __semantic.__table t - # LEFT JOIN sushi.raw_marketing AS m - # WHERE t.event_date > '2022-01-01' - # GROUP BY t.event_date, m.status - # """, - # graph=graph, - # metrics=context.metrics, - # ) - # assert_exp_eq( - # query, - # parse_one( - # """ - # SELECT - # t.event_date AS event_date, - # t.status AS status, - # t.total_orders AS total_orders - # FROM ( - # SELECT - # sushi__orders.event_date, - # m.status, - # COUNT(sushi__orders.id) AS total_orders - # FROM sushi.orders AS sushi__orders - # LEFT JOIN sushi.raw_marketing AS m - # ON sushi__orders.customer_id = m.customer_id - # WHERE - # sushi__orders.event_date > '2022-01-01' - # GROUP BY - # sushi__orders.event_date, - # m.status - # ) AS t - # """, - # dialect=context.config.dialect, - # ), - # ) + +def test_rewrite_distinct_on_group_keys(metrics_runtime): + assert _execute( + metrics_runtime, + """ + SELECT DISTINCT ON (s.channel) s.org, s.channel, METRIC(a_sum) + FROM __semantic.__table s + GROUP BY s.org, s.channel + ORDER BY s.channel, s.org + """, + ) == [("A", "app", 20), ("A", "web", 10)] + + +def test_rewrite_named_window_group_keys(metrics_runtime): + assert _execute( + metrics_runtime, + """ + SELECT s.org, s.channel, METRIC(a_sum), ROW_NUMBER() OVER w AS position + FROM __semantic.__table s + GROUP BY s.org, s.channel + WINDOW w AS (PARTITION BY s.channel ORDER BY s.org) + ORDER BY s.channel, s.org + """, + ) == [("A", "app", 20, 1), ("A", "web", 10, 1), ("B", "web", 100, 2)] diff --git a/tests/core/test_reference.py b/tests/core/test_reference.py index fe08f3f71e..dce89a352d 100644 --- a/tests/core/test_reference.py +++ b/tests/core/test_reference.py @@ -52,3 +52,15 @@ def test_models_for_column(sushi_context_pre_scheduling): "sushi.raw_marketing", ] assert graph.models_for_column("sushi.orders", "event_date") == ["sushi.orders"] + + +def test_target_grain_does_not_override_nonunique_join_reference(make_model): + graph = ReferenceGraph( + [ + make_model("fact", [("org", False)]), + make_model("dimension", [("dimension_id", True), ("org", False)]), + ] + ) + + with pytest.raises(SQLMeshError): + graph.find_path("fact", "dimension")