feat(discover-snql): Add SnQL support for project thresholds (#27412)

Add support for the project transaction threshold configuration
column. We don't ever filter on this col, so not adding
a special filter for it.

src/sentry/search/events/base.py

@@ -1,20 +1,32 @@
-from typing import Callable, List, Mapping, Set
+from typing import Callable, List, Mapping, Optional, Set
 
+import sentry_sdk
 from django.utils.functional import cached_property
 from snuba_sdk.column import Column
 from snuba_sdk.function import CurriedFunction, Function
 from snuba_sdk.orderby import OrderBy
 
+from sentry.exceptions import InvalidSearchQuery
 from sentry.models import Project
+from sentry.models.transaction_threshold import (
+    TRANSACTION_METRICS,
+    ProjectTransactionThreshold,
+    ProjectTransactionThresholdOverride,
+)
 from sentry.search.events.constants import (
+    DEFAULT_PROJECT_THRESHOLD,
+    DEFAULT_PROJECT_THRESHOLD_METRIC,
     ERROR_HANDLED_ALIAS,
     ERROR_UNHANDLED_ALIAS,
     ISSUE_ALIAS,
     ISSUE_ID_ALIAS,
     KEY_TRANSACTION_ALIAS,
+    MAX_QUERYABLE_TRANSACTION_THRESHOLDS,
     PROJECT_ALIAS,
     PROJECT_NAME_ALIAS,
     PROJECT_THRESHOLD_CONFIG_ALIAS,
+    PROJECT_THRESHOLD_CONFIG_INDEX_ALIAS,
+    PROJECT_THRESHOLD_OVERRIDE_CONFIG_INDEX_ALIAS,
     SNQL_FIELD_ALLOWLIST,
     TEAM_KEY_TRANSACTION_ALIAS,
     TIMESTAMP_TO_DAY_ALIAS,
@@ -23,6 +35,7 @@ from sentry.search.events.constants import (
     USER_DISPLAY_ALIAS,
 )
 from sentry.search.events.types import ParamsType, SelectType, WhereType
+from sentry.utils.numbers import format_grouped_length
 from sentry.utils.snuba import Dataset, resolve_column
 
 
@@ -54,12 +67,12 @@ class QueryBase:
             TIMESTAMP_TO_DAY_ALIAS: self._resolve_timestamp_to_day_alias,
             USER_DISPLAY_ALIAS: self._resolve_user_display_alias,
             TRANSACTION_STATUS_ALIAS: self._resolve_transaction_status,
+            PROJECT_THRESHOLD_CONFIG_ALIAS: self._resolve_project_threshold_config,
             ERROR_UNHANDLED_ALIAS: self._resolve_error_unhandled_alias,
             ERROR_HANDLED_ALIAS: self._resolve_error_handled_alias,
             # TODO: implement these
             KEY_TRANSACTION_ALIAS: self._resolve_unimplemented_alias,
             TEAM_KEY_TRANSACTION_ALIAS: self._resolve_unimplemented_alias,
-            PROJECT_THRESHOLD_CONFIG_ALIAS: self._resolve_unimplemented_alias,
         }
 
     @cached_property
@@ -133,6 +146,140 @@ class QueryBase:
             "toUInt8", [self.column(TRANSACTION_STATUS_ALIAS)], TRANSACTION_STATUS_ALIAS
         )
 
+    def _resolve_project_threshold_config(self, _: str) -> SelectType:
+        org_id = self.params.get("organization_id")
+        project_ids = self.params.get("project_id")
+
+        project_threshold_configs = (
+            ProjectTransactionThreshold.objects.filter(
+                organization_id=org_id,
+                project_id__in=project_ids,
+            )
+            .order_by("project_id")
+            .values_list("project_id", "threshold", "metric")
+        )
+
+        transaction_threshold_configs = (
+            ProjectTransactionThresholdOverride.objects.filter(
+                organization_id=org_id,
+                project_id__in=project_ids,
+            )
+            .order_by("project_id")
+            .values_list("transaction", "project_id", "threshold", "metric")
+        )
+
+        num_project_thresholds = project_threshold_configs.count()
+        sentry_sdk.set_tag("project_threshold.count", num_project_thresholds)
+        sentry_sdk.set_tag(
+            "project_threshold.count.grouped",
+            format_grouped_length(num_project_thresholds, [10, 100, 250, 500]),
+        )
+
+        num_transaction_thresholds = transaction_threshold_configs.count()
+        sentry_sdk.set_tag("txn_threshold.count", num_transaction_thresholds)
+        sentry_sdk.set_tag(
+            "txn_threshold.count.grouped",
+            format_grouped_length(num_transaction_thresholds, [10, 100, 250, 500]),
+        )
+
+        if (
+            num_project_thresholds + num_transaction_thresholds
+            > MAX_QUERYABLE_TRANSACTION_THRESHOLDS
+        ):
+            raise InvalidSearchQuery(
+                f"Exceeded {MAX_QUERYABLE_TRANSACTION_THRESHOLDS} configured transaction thresholds limit, try with fewer Projects."
+            )
+
+        project_threshold_config_keys = []
+        project_threshold_config_values = []
+        for project_id, threshold, metric in project_threshold_configs:
+            project_threshold_config_keys.append(Function("toUInt64", [project_id]))
+            project_threshold_config_values.append((TRANSACTION_METRICS[metric], threshold))
+
+        project_threshold_override_config_keys = []
+        project_threshold_override_config_values = []
+        for transaction, project_id, threshold, metric in transaction_threshold_configs:
+            project_threshold_override_config_keys.append(
+                (Function("toUInt64", [project_id]), transaction)
+            )
+            project_threshold_override_config_values.append(
+                (TRANSACTION_METRICS[metric], threshold)
+            )
+
+        project_threshold_config_index: SelectType = Function(
+            "indexOf",
+            [
+                project_threshold_config_keys,
+                self.column("project_id"),
+            ],
+            PROJECT_THRESHOLD_CONFIG_INDEX_ALIAS,
+        )
+
+        project_threshold_override_config_index: SelectType = Function(
+            "indexOf",
+            [
+                project_threshold_override_config_keys,
+                (self.column("project_id"), self.column("transaction")),
+            ],
+            PROJECT_THRESHOLD_OVERRIDE_CONFIG_INDEX_ALIAS,
+        )
+
+        def _project_threshold_config(alias: Optional[str] = None) -> SelectType:
+            return (
+                Function(
+                    "if",
+                    [
+                        Function(
+                            "equals",
+                            [
+                                project_threshold_config_index,
+                                0,
+                            ],
+                        ),
+                        (DEFAULT_PROJECT_THRESHOLD_METRIC, DEFAULT_PROJECT_THRESHOLD),
+                        Function(
+                            "arrayElement",
+                            [
+                                project_threshold_config_values,
+                                project_threshold_config_index,
+                            ],
+                        ),
+                    ],
+                    alias,
+                )
+                if project_threshold_configs
+                else Function(
+                    "tuple",
+                    [DEFAULT_PROJECT_THRESHOLD_METRIC, DEFAULT_PROJECT_THRESHOLD],
+                    alias,
+                )
+            )
+
+        if transaction_threshold_configs:
+            return Function(
+                "if",
+                [
+                    Function(
+                        "equals",
+                        [
+                            project_threshold_override_config_index,
+                            0,
+                        ],
+                    ),
+                    _project_threshold_config(),
+                    Function(
+                        "arrayElement",
+                        [
+                            project_threshold_override_config_values,
+                            project_threshold_override_config_index,
+                        ],
+                    ),
+                ],
+                PROJECT_THRESHOLD_CONFIG_ALIAS,
+            )
+
+        return _project_threshold_config(PROJECT_THRESHOLD_CONFIG_ALIAS)
+
     def _resolve_error_unhandled_alias(self, _: str) -> SelectType:
         return Function("notHandled", [], ERROR_UNHANDLED_ALIAS)
 

src/sentry/search/events/constants.py

@@ -81,6 +81,7 @@ SEARCH_MAP.update(**DATASETS[Dataset.Discover])
 
 DEFAULT_PROJECT_THRESHOLD_METRIC = "duration"
 DEFAULT_PROJECT_THRESHOLD = 300
+MAX_QUERYABLE_TRANSACTION_THRESHOLDS = 500
 
 # Allow list of fields that are compatible with the Snql Query Builder.
 # Once we reach a certain threshold of fields handled should turn this into a denylist
@@ -89,6 +90,7 @@ SNQL_FIELD_ALLOWLIST = {
     "environment",
     "message",
     "project",
+    "transaction",
     "project.id",
     "release",
     USER_DISPLAY_ALIAS,

tests/sentry/snuba/test_discover.py

@@ -411,6 +411,116 @@ class QueryIntegrationTest(SnubaTestCase, TestCase):
             assert len(data) == 1, query_fn
             assert [item["user.display"] for item in data] == ["bruce@example.com"]
 
+    def test_snql_wip_project_threshold_config(self):
+        ProjectTransactionThreshold.objects.create(
+            project=self.project,
+            organization=self.project.organization,
+            threshold=100,
+            metric=TransactionMetric.DURATION.value,
+        )
+
+        project2 = self.create_project()
+        ProjectTransactionThreshold.objects.create(
+            project=project2,
+            organization=project2.organization,
+            threshold=600,
+            metric=TransactionMetric.LCP.value,
+        )
+
+        events = [
+            ("a" * 10, 300),
+            ("b" * 10, 300),
+            ("c" * 10, 3000),
+            ("d" * 10, 3000),
+        ]
+        for idx, event in enumerate(events):
+            data = load_data(
+                "transaction",
+                timestamp=before_now(minutes=(3 + idx)),
+                start_timestamp=before_now(minutes=(3 + idx), milliseconds=event[1]),
+            )
+            data["event_id"] = f"{idx}" * 32
+            data["transaction"] = event[0]
+            self.store_event(data, project_id=self.project.id)
+
+            if idx % 2:
+                ProjectTransactionThresholdOverride.objects.create(
+                    transaction=event[0],
+                    project=self.project,
+                    organization=self.organization,
+                    threshold=1000,
+                    metric=TransactionMetric.DURATION.value,
+                )
+
+        data = load_data(
+            "transaction", timestamp=before_now(minutes=3), start_timestamp=before_now(minutes=4)
+        )
+        data["transaction"] = "e" * 10
+        self.store_event(data, project_id=project2.id)
+
+        expected_transaction = ["a" * 10, "b" * 10, "c" * 10, "d" * 10, "e" * 10]
+        expected_project_threshold_config = [
+            ["duration", 100],
+            ["duration", 1000],
+            ["duration", 100],
+            ["duration", 1000],
+            ["lcp", 600],
+        ]
+
+        for query_fn in [discover.query, discover.wip_snql_query]:
+            result = query_fn(
+                selected_columns=["project", "transaction", "project_threshold_config"],
+                query="",
+                params={
+                    "start": before_now(minutes=10),
+                    "end": before_now(minutes=2),
+                    "project_id": [self.project.id, project2.id],
+                    "organization_id": self.organization.id,
+                },
+            )
+
+            assert len(result["data"]) == 5
+            sorted_data = sorted(result["data"], key=lambda k: k["transaction"])
+
+            assert [row["transaction"] for row in sorted_data] == expected_transaction
+            assert [row["project_threshold_config"][0] for row in sorted_data] == [
+                r[0] for r in expected_project_threshold_config
+            ]
+            assert [row["project_threshold_config"][1] for row in sorted_data] == [
+                r[1] for r in expected_project_threshold_config
+            ]
+
+        ProjectTransactionThreshold.objects.filter(
+            project=project2,
+            organization=project2.organization,
+        ).delete()
+
+        expected_transaction = ["e" * 10]
+        expected_project_threshold_config = [["duration", 300]]
+
+        for query_fn in [discover.query, discover.wip_snql_query]:
+            result = query_fn(
+                selected_columns=["project", "transaction", "project_threshold_config"],
+                query="",
+                params={
+                    "start": before_now(minutes=10),
+                    "end": before_now(minutes=2),
+                    "project_id": [project2.id],
+                    "organization_id": self.organization.id,
+                },
+            )
+
+            assert len(result["data"]) == 1
+            sorted_data = sorted(result["data"], key=lambda k: k["transaction"])
+
+            assert [row["transaction"] for row in sorted_data] == expected_transaction
+            assert [row["project_threshold_config"][0] for row in sorted_data] == [
+                r[0] for r in expected_project_threshold_config
+            ]
+            assert [row["project_threshold_config"][1] for row in sorted_data] == [
+                r[1] for r in expected_project_threshold_config
+            ]
+
     def test_transaction_status(self):
         data = load_data("transaction", timestamp=before_now(minutes=1))
         data["transaction"] = "/test_transaction/success"