fix(metrics): infer measure table for aggregates with ORDER BY

sqlmesh/core/metric/definition.py

@@ -199,6 +199,10 @@ def is_measure(node: exp.Expr) -> bool:
             return is_measure(parent)
         if isinstance(parent, (exp.Binary, exp.Paren, exp.Distinct)):
             return is_measure(parent)
+        # AGG(x ORDER BY y) is parsed as AGG(Order(this=x, expressions=[Ordered(y)])).
+        # Only the Order's `this` can be the measure, not the ORDER BY expressions.
+        if isinstance(parent, exp.Order) and node.arg_key == "this":
+            return is_measure(parent)
         return False
 
     for node in expression.walk():

tests/core/metric/test_metric.py

@@ -154,3 +154,11 @@ def test_get_measure_and_dim_tables():
         "a",
         ("b",),
     )
+    assert _get_measure_and_dim_tables(d.parse_one("ARRAY_AGG(a.x ORDER BY a.x)")) == ("a", ())
+    assert _get_measure_and_dim_tables(d.parse_one("ARRAY_AGG(a.x ORDER BY b.y)")) == ("a", ("b",))
+    assert _get_measure_and_dim_tables(
+        d.parse_one("ARRAY_AGG(IF(c.z = 'dim', a.x, NULL) ORDER BY b.y)")
+    ) == ("a", ("b", "c"))
+    # A column that only appears in the ORDER BY clause cannot be the measure.
+    with pytest.raises(ConfigError, match="Could not infer a measures table"):
+        _get_measure_and_dim_tables(d.parse_one("ARRAY_AGG(1 ORDER BY b.y)"))

tests/core/metric/test_rewriter.py

@@ -1,7 +1,9 @@
 from sqlglot import parse_one
 
-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.reference import ReferenceGraph
+from sqlmesh.utils import UniqueKeyDict
 
 
 def test_rewrite(sushi_context_pre_scheduling, assert_exp_eq):
@@ -222,3 +224,27 @@ def test_rewrite(sushi_context_pre_scheduling, assert_exp_eq):
     #        dialect=context.config.dialect,
     #    ),
     # )
+
+
+def test_rewrite_order_by_metric(sushi_context_pre_scheduling):
+    context = sushi_context_pre_scheduling
+    graph = ReferenceGraph(context.models.values())
+    meta = load_metric_ddl(
+        d.parse_one(
+            "METRIC (name sorted_amounts, expression ARRAY_AGG(sushi.orders.amount ORDER BY sushi.orders.amount))",
+            dialect=context.config.dialect,
+        ),
+        dialect=context.config.dialect,
+    )
+    metrics = expand_metrics(UniqueKeyDict("metas", {meta.name: meta}))
+
+    query = rewrite(
+        "SELECT c.customer_id, METRIC(sorted_amounts) FROM sushi.customers AS c GROUP BY c.customer_id",
+        graph=graph,
+        metrics=metrics,
+    )
+    sql = query.sql()
+    # Stop before the generated NULLS LAST so the assertion doesn't depend on null ordering.
+    assert "ARRAY_AGG(sushi__orders.amount ORDER BY sushi__orders.amount" in sql
+    assert "sushi.orders AS sushi__orders" in sql
+    parse_one(sql, dialect=context.config.dialect)