Repository: getsentry/sentry
Pull request: feat(ddm): Hide blocked metrics of projects
This PR adds new logic to automatically remove `projects` from a query, if the metric(s) in that query are blocked for at least one of the supplied projects.

The new logic works in the following way:
1. Before execution, an efficient internal representation of blocked metrics for projects is loaded from project settings given all the metrics queried.
2. Before running a query, we determine a new set of projects ids that are not blocked for the supplied metrics. We start from a set of all the supplied projects and iteratively reduce it when a metric of the query has a different project blocked.
3. We determine if the newly updated query has no more projects, in that case we return a response with empty groups, signaling that the metric was not even queried. For now, we distinguish between metric not queried and metric queried without data.

Closes: https://github.com/getsentry/sentry/issues/64059

Problem statement:
Hide projects when querying a metric that is blocked for those projects
We would like to add some logic to the `metrics/api` that removes project ids of the projects that have a certain metric blocked. The idea is that we don't want to return the data of project `x` if metric `y` for `x` is blocked.

Hints from issue discussion:

Commit message: Add metric

Patch to review:
diff --git a/src/sentry/sentry_metrics/querying/data/api.py b/src/sentry/sentry_metrics/querying/data/api.py
index 2f2d3bb049b364..346e0beac8a711 100644
--- a/src/sentry/sentry_metrics/querying/data/api.py
+++ b/src/sentry/sentry_metrics/querying/data/api.py
@@ -38,7 +38,7 @@ def run_metrics_query(
     )
 
     # We prepare the executor, that will be responsible for scheduling the execution multiple queries.
-    executor = QueryExecutor(organization=organization, referrer=referrer)
+    executor = QueryExecutor(organization=organization, projects=projects, referrer=referrer)
 
     # We parse the input and iterating over each timeseries.
     parser = QueryParser(projects=projects, fields=fields, query=query, group_bys=group_bys)
diff --git a/src/sentry/sentry_metrics/querying/data/execution.py b/src/sentry/sentry_metrics/querying/data/execution.py
index 9fda6c037c3153..049e3bed2a0d20 100644
--- a/src/sentry/sentry_metrics/querying/data/execution.py
+++ b/src/sentry/sentry_metrics/querying/data/execution.py
@@ -1,19 +1,22 @@
 import math
 from dataclasses import dataclass, replace
 from datetime import datetime
-from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple, cast
+from typing import Any, Dict, List, Mapping, Optional, Sequence, Set, Tuple, cast
 
 import sentry_sdk
-from snuba_sdk import Column, Direction, MetricsQuery, Request
+from snuba_sdk import Column, Direction, MetricsQuery, MetricsScope, Request
 from snuba_sdk.conditions import BooleanCondition, BooleanOp, Condition, Op
 
 from sentry.models.organization import Organization
+from sentry.models.project import Project
 from sentry.sentry_metrics.querying.common import DEFAULT_QUERY_INTERVALS, SNUBA_QUERY_LIMIT
 from sentry.sentry_metrics.querying.errors import (
     InvalidMetricsQueryError,
     MetricsQueryExecutionError,
 )
 from sentry.sentry_metrics.querying.types import GroupKey, GroupsCollection
+from sentry.sentry_metrics.querying.visitors import QueriedMetricsVisitor
+from sentry.sentry_metrics.visibility import get_metrics_blocking_state
 from sentry.snuba.dataset import Dataset
 from sentry.snuba.metrics import to_intervals
 from sentry.snuba.metrics_layer.query import run_query
@@ -102,6 +105,9 @@ class ExecutableQuery:
     order_by: Optional[str]
     limit: Optional[int]
 
+    def is_empty(self) -> bool:
+        return not self.metrics_query.scope.org_ids or not self.metrics_query.scope.project_ids
+
     def replace_date_range(self, start: datetime, end: datetime) -> "ExecutableQuery":
         return replace(
             self,
@@ -144,8 +150,10 @@ def add_group_filters(
         groups_collection: Optional[GroupsCollection],
     ) -> "ExecutableQuery":
         """
-        Adds a series of filters to the query which will make sure that the results returned only belong to the supplied
-        groups.
+        Returns a new `ExecutableQuery` with a series of filters that ensure that the new query will have the same
+        groups returned. Keep in mind that there is no guarantee that all the groups will be returned, since data might
+        change in the meanwhile, so the guarantee of this method is that the returned groups will all be belonging to
+        `groups_collection`.
 
         The need for this filter arises because when executing multiple queries, we want to have the same groups
         returned, in order to make results consistent. Note that in case queries have different groups, some results
@@ -179,6 +187,34 @@ def add_group_filters(
             ),
         )
 
+    def merge_with_blocked_projects(
+        self,
+        organization: Organization,
+        projects: Set[Project],
+        blocked_metrics_for_projects: Mapping[str, Set[int]],
+    ) -> "ExecutableQuery":
+        """
+        Returns a new `ExecutableQuery` with the projects for which all the queries are not blocked. In case no projects
+        exist, the query will be returned with empty projects, signaling the executor to not run the query.
+        """
+        intersected_projects: Set[int] = {project.id for project in projects}
+
+        for queried_metric in QueriedMetricsVisitor().visit(self.metrics_query.query):
+            blocked_for_projects = blocked_metrics_for_projects.get(queried_metric)
+            if blocked_for_projects:
+                metrics.incr(key="ddm.metrics_api.blocked_metric_queried", amount=1)
+                intersected_projects -= blocked_for_projects
+
+        return replace(
+            self,
+            metrics_query=self.metrics_query.set_scope(
+                MetricsScope(
+                    org_ids=[organization.id],
+                    project_ids=list(intersected_projects),
+                )
+            ),
+        )
+
 
 @dataclass(frozen=True)
 class QueryResult:
@@ -189,6 +225,20 @@ class QueryResult:
     def __post_init__(self):
         assert self.series_executable_query or self.totals_executable_query
 
+    @classmethod
+    def empty_from(cls, executable_query: ExecutableQuery) -> "QueryResult":
+        return QueryResult(
+            series_executable_query=executable_query,
+            totals_executable_query=executable_query,
+            result={
+                "series": {"data": {}, "meta": {}},
+                "totals": {"data": {}, "meta": {}},
+                # We want to honor the date ranges of the supplied query.
+                "modified_start": executable_query.metrics_query.start,
+                "modified_end": executable_query.metrics_query.end,
+            },
+        )
+
     @property
     def query_name(self) -> str:
         if self.series_executable_query:
@@ -286,6 +336,8 @@ def align_with(self, reference_query_result: "QueryResult") -> "QueryResult":
             self.totals, reference_query_result.totals, alignment_keys, indexed_totals
         )
 
+        # We only mutate with the aligned data, only if we have data, since if it's empty it could be that we are
+        # trying to align on a query that has no data, and we want to avoid deleting the data of this query.
         if aligned_series:
             self.result["series"]["data"] = aligned_series
         if aligned_totals:
@@ -324,9 +376,11 @@ def align_series_to_totals(self) -> "QueryResult":
 
 
 class QueryExecutor:
-    def __init__(self, organization: Organization, referrer: str):
+    def __init__(self, organization: Organization, projects: Sequence[Project], referrer: str):
         self._organization = organization
+        self._projects = projects
         self._referrer = referrer
+
         # Ordered list of the intervals that can be chosen by the executor. They are removed when tried, in order
         # to avoid an infinite recursion.
         self._interval_choices = sorted(DEFAULT_QUERY_INTERVALS)
@@ -335,6 +389,26 @@ def __init__(self, organization: Organization, referrer: str):
         # Tracks the number of queries that have been executed (for measuring purposes).
         self._number_of_executed_queries = 0
 
+        # We load the blocked metrics for the supplied projects.
+        self._blocked_metrics_for_projects = self._load_blocked_metrics_for_projects()
+
+    def _load_blocked_metrics_for_projects(self) -> Mapping[str, Set[int]]:
+        """
+        Load the blocked metrics for the supplied projects and stores them in the executor in an efficient way that
+        speeds up the determining of the projects to exclude from the query.
+        """
+        blocked_metrics_for_projects: Dict[str, Set[int]] = {}
+
+        for project_id, metrics_blocking_state in get_metrics_blocking_state(
+            self._projects
+        ).items():
+            for metric_blocking in metrics_blocking_state.metrics.values():
+                blocked_metrics_for_projects.setdefault(metric_blocking.metric_mri, set()).add(
+                    project_id
+                )
+
+        return blocked_metrics_for_projects
+
     def _build_request(self, query: MetricsQuery) -> Request:
         """
         Builds a Snuba request given a MetricsQuery to execute.
@@ -354,10 +428,16 @@ def _execute(
         Executes a query as series and/or totals and returns the result.
         """
         try:
+            # We merge the query with the blocked projects, in order to obtain a new query with only the projects that
+            # all have the queried metrics unblocked.
+            executable_query = executable_query.merge_with_blocked_projects(
+                organization=self._organization,
+                projects=set(self._projects),
+                blocked_metrics_for_projects=self._blocked_metrics_for_projects,
+            )
+
             # We try to determine the interval of the query, which will be used to define clear time bounds for both
             # queries. This is done here since the metrics layer doesn't adjust the time for totals queries.
-            # TODO: maybe we can find a way to tell the layer to use the interval in totals but just to honor the same
-            #   time interval as used in the series query.
             interval = executable_query.metrics_query.rollup.interval
             if interval:
                 modified_start, modified_end, _ = to_intervals(
@@ -370,6 +450,11 @@ def _execute(
                         modified_start, modified_end
                     )
 
+            # If, after merging the query with the blocked projects, the query becomes empty, we will return an empty
+            # result.
+            if executable_query.is_empty():
+                return QueryResult.empty_from(executable_query)
+
             totals_executable_query = executable_query
             totals_result = None
             if executable_query.with_totals:
diff --git a/src/sentry/sentry_metrics/querying/visitors.py b/src/sentry/sentry_metrics/querying/visitors.py
index 8f0efacdafc8ef..2175beb13c4b63 100644
--- a/src/sentry/sentry_metrics/querying/visitors.py
+++ b/src/sentry/sentry_metrics/querying/visitors.py
@@ -1,5 +1,5 @@
 from abc import ABC
-from typing import Generic, Mapping, Optional, Sequence, TypeVar
+from typing import Generic, Mapping, Optional, Sequence, Set, TypeVar
 
 from snuba_sdk import BooleanCondition, BooleanOp, Column, Condition, Formula, Op, Timeseries
 from snuba_sdk.conditions import ConditionGroup
@@ -7,7 +7,10 @@
 from sentry.api.serializers import bulk_fetch_project_latest_releases
 from sentry.models.environment import Environment
 from sentry.models.project import Project
-from sentry.sentry_metrics.querying.errors import LatestReleaseNotFoundError
+from sentry.sentry_metrics.querying.errors import (
+    InvalidMetricsQueryError,
+    LatestReleaseNotFoundError,
+)
 from sentry.sentry_metrics.querying.types import QueryCondition, QueryExpression
 
 TVisited = TypeVar("TVisited")
@@ -228,3 +231,23 @@ def _visit_condition(self, condition: Condition) -> QueryCondition:
             op=condition.op,
             rhs=condition.rhs,
         )
+
+
+class QueriedMetricsVisitor(QueryExpressionVisitor[Set[str]]):
+    """
+    Visitor that recursively computes all the metrics MRI that have been queried.
+    """
+
+    def _visit_formula(self, formula: Formula) -> Set[str]:
+        metrics: Set[str] = set()
+
+        for parameter in formula.parameters:
+            metrics.union(self.visit(parameter))
+
+        return metrics
+
+    def _visit_timeseries(self, timeseries: Timeseries) -> Set[str]:
+        if timeseries.metric.mri is None:
+            raise InvalidMetricsQueryError("Can't determine queried metrics without a MRI")
+
+        return {timeseries.metric.mri}
diff --git a/tests/sentry/sentry_metrics/querying/data/test_api.py b/tests/sentry/sentry_metrics/querying/data/test_api.py
index 1bd17dd616d181..26520b417391aa 100644
--- a/tests/sentry/sentry_metrics/querying/data/test_api.py
+++ b/tests/sentry/sentry_metrics/querying/data/test_api.py
@@ -10,6 +10,7 @@
     MetricsQueryExecutionError,
 )
 from sentry.sentry_metrics.use_case_id_registry import UseCaseID
+from sentry.sentry_metrics.visibility import block_metric
 from sentry.snuba.metrics.naming_layer import SessionMRI, TransactionMRI
 from sentry.testutils.cases import BaseMetricsTestCase, TestCase
 from sentry.testutils.helpers.datetime import freeze_time
@@ -737,3 +738,115 @@ def test_with_sessions(self) -> None:
         assert len(groups) == 1
         assert groups[0]["by"] == {}
         assert groups[0]["series"] == {field: [60.0]}
+
+    def test_query_with_one_metric_blocked_for_one_project(self):
+        mri = "d:custom/page_load@millisecond"
+
+        project_1 = self.create_project()
+        project_2 = self.create_project()
+
+        block_metric(mri, [project_1])
+
+        for project, value in ((project_1, 10.0), (project_2, 15.0)):
+            self.store_metric(
+                self.project.organization.id,
+                project.id,
+                "distribution",
+                mri,
+                {},
+                self.ts(self.now()),
+                value,
+                UseCaseID.CUSTOM,
+            )
+
+        field = f"sum({mri})"
+        results = run_metrics_query(
+            fields=[field],
+            start=self.now() - timedelta(minutes=30),
+            end=self.now() + timedelta(hours=1, minutes=30),
+            interval=3600,
+            organization=self.project.organization,
+            projects=[project_1, project_2],
+            environments=[],
+            referrer="metrics.data.api",
+        )
+        groups = results["groups"]
+        assert len(groups) == 1
+        assert groups[0]["by"] == {}
+        assert groups[0]["series"] == {field: [0.0, 15.0, 0.0]}
+        assert groups[0]["totals"] == {field: 15.0}
+
+    def test_query_with_one_metric_blocked_for_all_projects(self):
+        mri = "d:custom/page_load@millisecond"
+
+        project_1 = self.create_project()
+        project_2 = self.create_project()
+
+        block_metric(mri, [project_1, project_2])
+
+        for project, value in ((project_1, 10.0), (project_2, 15.0)):
+            self.store_metric(
+                self.project.organization.id,
+                project.id,
+                "distribution",
+                mri,
+                {},
+                self.ts(self.now()),
+                value,
+                UseCaseID.CUSTOM,
+            )
+
+        field = f"sum({mri})"
+        results = run_metrics_query(
+            fields=[field],
+            start=self.now() - timedelta(minutes=30),
+            end=self.now() + timedelta(hours=1, minutes=30),
+            interval=3600,
+            organization=self.project.organization,
+            projects=[project_1, project_2],
+            environments=[],
+            referrer="metrics.data.api",
+        )
+        groups = results["groups"]
+        assert len(groups) == 0
+
+    def test_query_with_two_metrics_and_one_blocked_for_a_project(self):
+        mri_1 = "d:custom/page_load@millisecond"
+        mri_2 = "d:custom/app_load@millisecond"
+
+        project_1 = self.create_project()
+        project_2 = self.create_project()
+
+        block_metric(mri_1, [project_1, project_2])
+
+        for project, mri in ((project_1, mri_1), (project_2, mri_2)):
+            self.store_metric(
+                self.project.organization.id,
+                project.id,
+                "distribution",
+                mri,
+                {},
+                self.ts(self.now()),
+                10.0,
+                UseCaseID.CUSTOM,
+            )
+
+        field_1 = f"sum({mri_1})"
+        field_2 = f"sum({mri_2})"
+        results = run_metrics_query(
+            fields=[field_1, field_2],
+            # We test with the order by to make sure alignment doesn't remove data.
+            order_by=field_1,
+            start=self.now() - timedelta(minutes=30),
+            end=self.now() + timedelta(hours=1, minutes=30),
+            interval=3600,
+            organization=self.project.organization,
+            projects=[project_1, project_2],
+            environments=[],
+            referrer="metrics.data.api",
+        )
+        groups = results["groups"]
+        assert len(groups) == 1
+        assert groups[0]["by"] == {}
+        assert groups[0]["series"] == {field_2: [0.0, 10.0, 0.0]}
+        assert groups[0]["totals"] == {field_2: 10.0}


Code review comments:
diff_hunk: @@ -179,6 +187,34 @@ def add_group_filters(
             ),
         )
 
+    def merge_with_blocked_projects(
path: src/sentry/sentry_metrics/querying/data/execution.py
text: @user1:
nit: this is more of a filter (diff) than a merge

@author:
How would you call it? More than happy to change it.

@user1:
`filter_blocked_projects`, `get_non_blocked_projects`, or something along those lines