perf(metrics): Infer automatically the correct MRI based on select (#41260)

src/sentry/snuba/metrics/mqb_query_transformer.py

@@ -257,20 +257,42 @@ def _derive_mri_to_apply(project_ids, select, orderby):
                 break
 
     if has_order_by_team_key_transaction:
-        # TODO: add here optimization that gets an entity from the select (either set of distribution) and sets it
-        #  to all tkt in the query.
         entities = set()
-        for orderby_field in orderby:
-            if orderby_field.field.op != TEAM_KEY_TRANSACTION_OP:
-                expr = metric_object_factory(orderby_field.field.op, orderby_field.field.metric_mri)
-                entity = expr.get_entity(project_ids, use_case_id=UseCaseKey.PERFORMANCE)
-                if isinstance(entity, str):
-                    entities.add(entity)
 
-        if len(entities) > 1:
-            raise InvalidParams("The orderby cannot have fields with multiple entities.")
+        if len(orderby) == 1:
+            # If the number of clauses in the order by is equal to 1 and the order by has a team_key_transaction it
+            # means that it must be the only one, therefore we want to infer the MRI type of the team_key_transaction
+            # from one entity in the select in order to save up a query. This is just an optimization for the edge case
+            # in which the select has a different entity than the default entity for the team_key_transaction, which
+            # is the distribution, inferred from TransactionMRI.DURATION.
+            for select_field in select:
+                if select_field.op != TEAM_KEY_TRANSACTION_OP:
+                    expr = metric_object_factory(select_field.op, select_field.metric_mri)
+                    entity = expr.get_entity(project_ids, use_case_id=UseCaseKey.PERFORMANCE)
+                    if isinstance(entity, str):
+                        entities.add(entity)
+        else:
+            # If the number of clauses in the order by is more than 1 it means that together with team_key_transaction
+            # there are other order by conditions and by definition we want all the order by conditions to belong to
+            # the same entity type, therefore we want to check how many entities are there in the other order by
+            # conditions and if there is only one we will infer the MRI type of the team_key_transaction
+            # from that one entity. If, on the other hand, there are multiple entities, then we throw an error because
+            # an order by across multiple entities is not supported.
+            for orderby_field in orderby:
+                if orderby_field.field.op != TEAM_KEY_TRANSACTION_OP:
+                    expr = metric_object_factory(
+                        orderby_field.field.op, orderby_field.field.metric_mri
+                    )
+                    entity = expr.get_entity(project_ids, use_case_id=UseCaseKey.PERFORMANCE)
+                    if isinstance(entity, str):
+                        entities.add(entity)
+
+            if len(entities) > 1:
+                raise InvalidParams("The orderby cannot have fields with multiple entities.")
 
-        if len(entities) != 0:
+        if len(entities) > 0:
+            # Only if entities are found in the clauses we are going to update the MRI to apply, otherwise we will just
+            # resort to the default one.
             mri_to_apply = mri_dictionary[entities.pop()]
 
     return mri_to_apply

tests/sentry/snuba/metrics/test_mqb_query_transformer.py

@@ -546,6 +546,154 @@ VALID_QUERIES_INTEGRATION_TEST_CASES = [
         ),
         id="team_key_transaction transformation with p95 in select",
     ),
+    pytest.param(
+        Query(
+            match=Entity("generic_metrics_distributions"),
+            select=[
+                AliasedExpression(
+                    exp=Column("e:transactions/user_misery@ratio"),
+                    alias="user_misery",
+                ),
+                Function(
+                    function="team_key_transaction",
+                    parameters=[
+                        Column(TransactionMRI.TEAM_KEY_TRANSACTION.value),
+                        [(13, "foo_transaction")],
+                    ],
+                    alias="team_key_transaction",
+                ),
+            ],
+            groupby=[
+                Function(
+                    function="team_key_transaction",
+                    parameters=[
+                        Column(TransactionMRI.TEAM_KEY_TRANSACTION.value),
+                        [(13, "foo_transaction")],
+                    ],
+                    alias="team_key_transaction",
+                ),
+            ],
+            array_join=None,
+            where=[
+                Condition(
+                    lhs=Column(
+                        name="timestamp",
+                    ),
+                    op=Op.GTE,
+                    rhs=datetime.datetime(2022, 3, 24, 11, 11, 36, 75132),
+                ),
+                Condition(
+                    lhs=Column(
+                        name="timestamp",
+                    ),
+                    op=Op.LT,
+                    rhs=datetime.datetime(2022, 6, 22, 11, 11, 36, 75132),
+                ),
+                Condition(
+                    lhs=Column(
+                        name="project_id",
+                    ),
+                    op=Op.IN,
+                    rhs=[13],
+                ),
+                Condition(
+                    lhs=Column(
+                        name="org_id",
+                    ),
+                    op=Op.EQ,
+                    rhs=14,
+                ),
+                Condition(
+                    lhs=Function(
+                        function="team_key_transaction",
+                        parameters=[
+                            Column(TransactionMRI.TEAM_KEY_TRANSACTION.value),
+                            [(13, "foo_transaction")],
+                        ],
+                        alias="team_key_transaction",
+                    ),
+                    op=Op.EQ,
+                    rhs=1,
+                ),
+            ],
+            having=[],
+            orderby=[
+                OrderBy(
+                    exp=Function(
+                        function="team_key_transaction",
+                        parameters=[
+                            Column(TransactionMRI.TEAM_KEY_TRANSACTION.value),
+                            [(13, "foo_transaction")],
+                        ],
+                        alias="team_key_transaction",
+                    ),
+                    direction=Direction.ASC,
+                ),
+            ],
+            limitby=None,
+            limit=Limit(limit=51),
+            offset=Offset(offset=0),
+            granularity=Granularity(granularity=86400),
+            totals=None,
+        ),
+        MetricsQuery(
+            org_id=14,
+            project_ids=[13],
+            select=[
+                MetricField(
+                    op=None,
+                    metric_mri="e:transactions/user_misery@ratio",
+                    alias="user_misery",
+                ),
+                MetricField(
+                    op="team_key_transaction",
+                    metric_mri="s:transactions/user@none",
+                    params={"team_key_condition_rhs": [(13, "foo_transaction")]},
+                    alias="team_key_transaction",
+                ),
+            ],
+            start=datetime.datetime(2022, 3, 24, 11, 11, 36, 75132),
+            end=datetime.datetime(2022, 6, 22, 11, 11, 36, 75132),
+            granularity=Granularity(granularity=86400),
+            where=[
+                MetricConditionField(
+                    lhs=MetricField(
+                        op="team_key_transaction",
+                        metric_mri="s:transactions/user@none",
+                        params={"team_key_condition_rhs": [(13, "foo_transaction")]},
+                        alias="team_key_transaction",
+                    ),
+                    op=Op.EQ,
+                    rhs=1,
+                )
+            ],
+            groupby=[
+                MetricGroupByField(
+                    MetricField(
+                        op="team_key_transaction",
+                        metric_mri="s:transactions/user@none",
+                        params={"team_key_condition_rhs": [(13, "foo_transaction")]},
+                        alias="team_key_transaction",
+                    ),
+                ),
+            ],
+            orderby=[
+                MetricsOrderBy(
+                    field=MetricField(
+                        op="team_key_transaction",
+                        metric_mri="s:transactions/user@none",
+                        params={"team_key_condition_rhs": [(13, "foo_transaction")]},
+                        alias="team_key_transaction",
+                    ),
+                    direction=Direction.ASC,
+                ),
+            ],
+            limit=Limit(limit=51),
+            offset=Offset(offset=0),
+            include_series=False,
+        ),
+        id="team_key_transaction transformation with single order by and user_misery in select",
+    ),
     pytest.param(
         Query(
             match=Entity("generic_metrics_distributions"),