Enable queries using project slug as filter and groupby in Metrics API

src/sentry/sentry_metrics/querying/data/api.py

@@ -1,19 +1,29 @@
 from collections.abc import Sequence
 from datetime import datetime
-from typing import cast
 
 from snuba_sdk import MetricsQuery, MetricsScope, Rollup
 
 from sentry import features
 from sentry.models.environment import Environment
 from sentry.models.organization import Organization
 from sentry.models.project import Project
-from sentry.sentry_metrics.querying.data.execution import QueryExecutor, QueryResult
+from sentry.sentry_metrics.querying.data.execution import QueryExecutor
+from sentry.sentry_metrics.querying.data.modulation.modulation_value_map import (
+    QueryModulationValueMap,
+)
+from sentry.sentry_metrics.querying.data.modulation.modulator import (
+    Modulator,
+    Project2ProjectIDModulator,
+)
 from sentry.sentry_metrics.querying.data.parsing import QueryParser
+from sentry.sentry_metrics.querying.data.postprocessing.base import run_postprocessing_steps
+from sentry.sentry_metrics.querying.data.postprocessing.demodulation import QueryDemodulationStep
 from sentry.sentry_metrics.querying.data.preparation.base import (
     IntermediateQuery,
+    PreparationStep,
     run_preparation_steps,
 )
+from sentry.sentry_metrics.querying.data.preparation.modulation import QueryModulationStep
 from sentry.sentry_metrics.querying.data.preparation.units_normalization import (
     UnitsNormalizationStep,
 )
@@ -62,12 +72,17 @@ def run_queries(
             )
         )
 
-    preparation_steps = []
+    preparation_steps: list[PreparationStep] = []
+    modulation_value_map = QueryModulationValueMap()
+    modulators: list[Modulator] = [Project2ProjectIDModulator()]
+
     if features.has(
         "organizations:ddm-metrics-api-unit-normalization", organization=organization, actor=None
     ):
         preparation_steps.append(UnitsNormalizationStep())
 
+    preparation_steps.append(QueryModulationStep(projects, modulators, modulation_value_map))
+
     # We run a series of preparation steps which operate on the entire list of queries.
     intermediate_queries = run_preparation_steps(intermediate_queries, *preparation_steps)
 
@@ -77,6 +92,9 @@ def run_queries(
         executor.schedule(intermediate_query=intermediate_query, query_type=query_type)
 
     results = executor.execute()
+    results = run_postprocessing_steps(
+        results, QueryDemodulationStep(projects, modulators, modulation_value_map)
+    )
 
     # We wrap the result in a class that exposes some utils methods to operate on results.
-    return MQLQueriesResult(cast(list[QueryResult], results))
+    return MQLQueriesResult(results)

src/sentry/sentry_metrics/querying/data/execution.py

@@ -1,5 +1,5 @@
 from collections.abc import Mapping, Sequence
-from dataclasses import dataclass, replace
+from dataclasses import dataclass, field, replace
 from datetime import datetime
 from enum import Enum
 from typing import Any, Union, cast
@@ -11,6 +11,7 @@
 from sentry.models.organization import Organization
 from sentry.models.project import Project
 from sentry.sentry_metrics.querying.constants import SNUBA_QUERY_LIMIT
+from sentry.sentry_metrics.querying.data.modulation.modulator import Modulator
 from sentry.sentry_metrics.querying.data.preparation.base import IntermediateQuery
 from sentry.sentry_metrics.querying.data.utils import adjust_time_bounds_with_interval
 from sentry.sentry_metrics.querying.errors import (
@@ -145,6 +146,7 @@ class ScheduledQuery:
     unit_family: UnitFamily | None = None
     unit: MeasurementUnit | None = None
     scaling_factor: float | None = None
+    modulators: list[Modulator] = field(default_factory=list)
 
     def initialize(
         self,
@@ -318,7 +320,7 @@ def _align_date_range(cls, metrics_query: MetricsQuery) -> tuple[MetricsQuery, i
         return metrics_query, None
 
 
-@dataclass(frozen=True)
+@dataclass
 class QueryResult:
     """
     Represents the result of a ScheduledQuery containing its associated series and totals results.
@@ -445,12 +447,24 @@ def modified_end(self) -> datetime:
 
     @property
     def series(self) -> Sequence[Mapping[str, Any]]:
+        if "series" not in self.result:
+            return []
         return self.result["series"]["data"]
 
+    @series.setter
+    def series(self, value: Sequence[Mapping[str, Any]]) -> None:
+        self.result["series"]["data"] = value
+
     @property
     def totals(self) -> Sequence[Mapping[str, Any]]:
+        if "totals" not in self.result:
+            return []
         return self.result["totals"]["data"]
 
+    @totals.setter
+    def totals(self, value: Sequence[Mapping[str, Any]]) -> None:
+        self.result["totals"]["data"] = value
+
     @property
     def meta(self) -> Sequence[Mapping[str, str]]:
         # By default, we extract the metadata from the totals query, if that is not there we extract from the series
@@ -464,7 +478,11 @@ def group_bys(self) -> list[str]:
         # that we can correctly render groups in case they are not returned from the db because of missing data.
         #
         # Sorting of the groups is done to maintain consistency across function calls.
-        return sorted(UsedGroupBysVisitor().visit(self._any_query().metrics_query.query))
+        scheduled_query = self._any_query()
+        modulators = scheduled_query.modulators
+        return sorted(
+            UsedGroupBysVisitor(modulators=modulators).visit(scheduled_query.metrics_query.query)
+        )
 
     @property
     def interval(self) -> int | None:
@@ -774,7 +792,7 @@ def _execution_loop(self):
         while continue_execution:
             continue_execution = self._bulk_execute()
 
-    def execute(self) -> Sequence[QueryResult]:
+    def execute(self) -> list[QueryResult]:
         """
         Executes the scheduled queries in the execution loop.
 
@@ -798,7 +816,7 @@ def execute(self) -> Sequence[QueryResult]:
                     "Not all queries were executed in the execution loop"
                 )
 
-        return cast(Sequence[QueryResult], self._query_results)
+        return cast(list[QueryResult], self._query_results)
 
     def schedule(self, intermediate_query: IntermediateQuery, query_type: QueryType):
         """
@@ -813,6 +831,7 @@ def schedule(self, intermediate_query: IntermediateQuery, query_type: QueryType)
             unit_family=intermediate_query.unit_family,
             unit=intermediate_query.unit,
             scaling_factor=intermediate_query.scaling_factor,
+            modulators=intermediate_query.modulators,
         )
 
         # In case the user chooses to run also a series query, we will duplicate the query and chain it after totals.

src/sentry/sentry_metrics/querying/data/modulation/modulation_value_map.py

@@ -0,0 +1,13 @@
+class QueryModulationValueMap(dict):
+    def __init__(self, *args, **kwargs):
+        self.update(*args, **kwargs)
+
+    def __getitem__(self, key):
+        return dict.__getitem__(self, key)
+
+    def __setitem__(self, key, value):
+        dict.__setitem__(self, key, value)
+
+    def update(self, *args, **kwargs):
+        for key, value in dict(*args, **kwargs).items():
+            self[key] = value

src/sentry/sentry_metrics/querying/data/modulation/modulator.py

@@ -0,0 +1,84 @@
+import abc
+from collections.abc import Sequence
+
+from snuba_sdk import Formula
+
+from sentry.models.project import Project
+from sentry.sentry_metrics.querying.data.modulation.modulation_value_map import (
+    QueryModulationValueMap,
+)
+
+
+class Modulator(abc.ABC):
+    def __init__(self, from_key: str, to_key: str):
+        self.from_key = from_key
+        self.to_key = to_key
+
+    def __hash__(self):
+        return hash((self.from_key, self.to_key))
+
+    @abc.abstractmethod
+    def modulate(
+        self,
+        projects: Sequence[Project],
+        value_map: QueryModulationValueMap,
+        formula: Formula,
+        **kwargs,
+    ) -> Formula:
+        return formula
+
+    @abc.abstractmethod
+    def demodulate(
+        self,
+        projects: Sequence[Project],
+        value_map: QueryModulationValueMap,
+        formula: Formula,
+        **kwargs,
+    ) -> Formula:
+        return formula
+
+
+class Project2ProjectIDModulator(Modulator):
+    def __init__(self, from_key: str = "project", to_key: str = "project_id"):
+        self.from_key = from_key
+        self.to_key = to_key
+
+    def modulate(
+        self,
+        projects: Sequence[Project],
+        value_map: QueryModulationValueMap,
+        formula: Formula,
+    ) -> Formula:
+        if formula not in value_map:
+            value_map[formula] = None
+            for project in projects:
+                if project.slug == formula:
+                    value_map[formula] = project.id
+        return value_map[formula]
+
+    def demodulate(
+        self,
+        projects: Sequence[Project],
+        value_map: QueryModulationValueMap,
+        formula: Formula,
+    ) -> Formula:
+        if formula not in value_map:
+            for project in projects:
+                if project.id == formula:
+                    value_map[formula] = project.slug
+
+        return value_map[formula]
+
+
+def find_modulator(
+    modulators: Sequence[Modulator], from_key: str | None = None, to_key: str | None = None
+) -> Modulator | None:
+    for modulator in modulators:
+        if from_key:
+            if modulator.from_key == from_key:
+                return modulator
+        if to_key:
+            if modulator.to_key == to_key:
+                return modulator
+
+    return None

src/sentry/sentry_metrics/querying/data/postprocessing/base.py

@@ -0,0 +1,37 @@
+from abc import ABC, abstractmethod
+
+from sentry.sentry_metrics.querying.data.execution import QueryResult
+
+
+class PostProcessingStep(ABC):
+    """
+    Represents an abstract step that post-processes a collection of QueryResult objects.
+
+    The post-processing of these objects might include transforming them or just obtaining some intermediate data that
+    is useful to compute other things before returning the results.
+    """
+
+    @abstractmethod
+    def run(self, query_results: list[QueryResult]) -> list[QueryResult]:
+        """
+        Runs the post-processing steps on a list of query results.
+
+        Returns:
+            A list of post-processed query results.
+        """
+        raise NotImplementedError
+
+
+def run_postprocessing_steps(query_results: list[QueryResult], *steps) -> list[QueryResult]:
+    """
+    Takes a series of query results and steps and runs the post-processing steps one after each other in order they are
+    supplied in.
+
+    Returns:
+        A list of query results after running the post-processing steps.
+    """
+    for step in steps:
+        if isinstance(step, PostProcessingStep):
+            query_results = step.run(query_results=query_results)
+
+    return query_results

src/sentry/sentry_metrics/querying/data/postprocessing/demodulation.py

@@ -0,0 +1,53 @@
+from collections.abc import Mapping, Sequence
+from typing import Any
+
+from sentry.models.project import Project
+from sentry.sentry_metrics.querying.data.execution import QueryResult
+from sentry.sentry_metrics.querying.data.modulation.modulator import Modulator
+from sentry.sentry_metrics.querying.data.postprocessing.base import PostProcessingStep
+
+
+class QueryDemodulationStep(PostProcessingStep):
+    def __init__(
+        self,
+        projects: Sequence[Project],
+        modulators: Sequence[Modulator],
+        value_map: dict[Any, Any],
+    ):
+        self.projects = projects
+        self.modulators = modulators
+        self.value_map = value_map
+
+    def run(self, query_results: list[QueryResult]) -> list[QueryResult]:
+        for query_result in query_results:
+            if query_result.totals:
+                query_result.totals = self._demodulate_data(
+                    query_result.totals, query_result.totals_query.modulators
+                )
+            if query_result.series:
+                query_result.series = self._demodulate_data(
+                    query_result.series, query_result.series_query.modulators
+                )
+
+        return query_results
+
+    def _demodulate_data(
+        self, data: Sequence[Mapping[str, Any]], modulators: list[Modulator]
+    ) -> Sequence[Mapping[str, Any]]:
+        for element in data:
+            updated_element = dict()
+            keys_to_delete = []
+            for result_key in element.keys():
+                for modulator in modulators:
+                    if modulator.to_key == result_key:
+                        original_value = modulator.demodulate(
+                            self.projects, self.value_map, element[result_key]
+                        )
+                        updated_element[modulator.from_key] = original_value
+                        keys_to_delete.append(result_key)
+
+            for key in keys_to_delete:
+                del element[key]
+            element.update(updated_element)
+
+        return data

src/sentry/sentry_metrics/querying/data/preparation/base.py

@@ -1,8 +1,9 @@
 from abc import ABC, abstractmethod
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 
 from snuba_sdk import MetricsQuery
 
+from sentry.sentry_metrics.querying.data.modulation.modulator import Modulator
 from sentry.sentry_metrics.querying.types import QueryOrder
 from sentry.sentry_metrics.querying.units import MeasurementUnit, UnitFamily
 
@@ -27,6 +28,7 @@ class IntermediateQuery:
     unit_family: UnitFamily | None = None
     unit: MeasurementUnit | None = None
     scaling_factor: float | None = None
+    modulators: list[Modulator] = field(default_factory=list)
 
 
 class PreparationStep(ABC):