ref(tsdb): allow conditions for get_sums + get_distinct_count_totals (#85072)

src/sentry/rules/conditions/event_frequency.py

@@ -569,7 +569,7 @@ class EventUniqueUserFrequencyConditionWithConditions(EventUniqueUserFrequencyCo
                 conditions.append(snuba_condition)
 
         total = self.get_chunked_result(
-            tsdb_function=self.tsdb.get_distinct_counts_totals_with_conditions,
+            tsdb_function=self.tsdb.get_distinct_counts_totals,
             model=get_issue_tsdb_user_group_model(GroupCategory.ERROR),
             organization_id=event.group.project.organization_id,
             group_ids=[event.group.id],
@@ -632,7 +632,7 @@ class EventUniqueUserFrequencyConditionWithConditions(EventUniqueUserFrequencyCo
         )
         if error_issue_ids and organization_id:
             error_totals = self.get_chunked_result(
-                tsdb_function=self.tsdb.get_distinct_counts_totals_with_conditions,
+                tsdb_function=self.tsdb.get_distinct_counts_totals,
                 model=get_issue_tsdb_user_group_model(GroupCategory.ERROR),
                 group_ids=error_issue_ids,
                 organization_id=organization_id,
@@ -646,7 +646,7 @@ class EventUniqueUserFrequencyConditionWithConditions(EventUniqueUserFrequencyCo
 
         if generic_issue_ids and organization_id:
             error_totals = self.get_chunked_result(
-                tsdb_function=self.tsdb.get_distinct_counts_totals_with_conditions,
+                tsdb_function=self.tsdb.get_distinct_counts_totals,
                 model=get_issue_tsdb_user_group_model(GroupCategory.PERFORMANCE),
                 group_ids=generic_issue_ids,
                 organization_id=organization_id,

src/sentry/tsdb/base.py

@@ -118,7 +118,6 @@ class BaseTSDB(Service):
             "get_distinct_counts_series",
             "get_distinct_counts_totals",
             "get_frequency_series",
-            "get_distinct_counts_totals_with_conditions",
         ]
     )
 
@@ -447,6 +446,7 @@ class BaseTSDB(Service):
         jitter_value: int | None = None,
         tenant_ids: dict[str, str | int] | None = None,
         referrer_suffix: str | None = None,
+        conditions: list[tuple[str, str, str]] | None = None,
     ) -> dict[int, int]:
         range_set = self.get_range(
             model,
@@ -459,6 +459,7 @@ class BaseTSDB(Service):
             jitter_value=jitter_value,
             tenant_ids=tenant_ids,
             referrer_suffix=referrer_suffix,
+            conditions=conditions,
         )
         sum_set = {key: sum(p for _, p in points) for (key, points) in range_set.items()}
         return sum_set
@@ -545,28 +546,10 @@ class BaseTSDB(Service):
         jitter_value: int | None = None,
         tenant_ids: dict[str, int | str] | None = None,
         referrer_suffix: str | None = None,
-    ) -> dict[int, Any]:
-        """
-        Count distinct items during a time range.
-        """
-        raise NotImplementedError
-
-    def get_distinct_counts_totals_with_conditions(
-        self,
-        model: TSDBModel,
-        keys: Sequence[int],
-        start: datetime,
-        end: datetime | None = None,
-        rollup: int | None = None,
-        environment_id: int | None = None,
-        use_cache: bool = False,
-        jitter_value: int | None = None,
-        tenant_ids: dict[str, int | str] | None = None,
-        referrer_suffix: str | None = None,
         conditions: list[tuple[str, str, str]] | None = None,
     ) -> dict[int, Any]:
         """
-        Count distinct items during a time range with conditions.
+        Count distinct items during a time range with optional conditions
         """
         raise NotImplementedError
 

src/sentry/tsdb/dummy.py

@@ -72,26 +72,10 @@ class DummyTSDB(BaseTSDB):
         jitter_value=None,
         tenant_ids=None,
         referrer_suffix=None,
-    ):
-        self.validate_arguments([model], [environment_id])
-        return {k: 0 for k in keys}
-
-    def get_distinct_counts_totals_with_conditions(
-        self,
-        model,
-        keys: Sequence[int],
-        start,
-        end=None,
-        rollup=None,
-        environment_id=None,
-        use_cache=False,
-        jitter_value=None,
-        tenant_ids=None,
-        referrer_suffix=None,
         conditions=None,
     ):
         self.validate_arguments([model], [environment_id])
-        return 0
+        return {k: 0 for k in keys}
 
     def merge_distinct_counts(
         self, model, destination, sources, timestamp=None, environment_ids=None

src/sentry/tsdb/redis.py

@@ -530,6 +530,7 @@ class RedisTSDB(BaseTSDB):
         jitter_value: int | None = None,
         tenant_ids: dict[str, int | str] | None = None,
         referrer_suffix: str | None = None,
+        conditions: list[tuple[str, str, str]] | None = None,
     ) -> dict[int, Any]:
         """
         Count distinct items during a time range.

src/sentry/tsdb/redissnuba.py

@@ -30,7 +30,6 @@ method_specifications = {
     "get_sums": (READ, single_model_argument),
     "get_distinct_counts_series": (READ, single_model_argument),
     "get_distinct_counts_totals": (READ, single_model_argument),
-    "get_distinct_counts_totals_with_conditions": (READ, single_model_argument),
     "get_frequency_series": (READ, single_model_argument),
     "incr": (WRITE, single_model_argument),
     "incr_multi": (WRITE, lambda callargs: {item[0] for item in callargs["items"]}),

src/sentry/tsdb/snuba.py

@@ -785,38 +785,8 @@ class SnubaTSDB(BaseTSDB):
         jitter_value=None,
         tenant_ids=None,
         referrer_suffix=None,
+        conditions=None,
     ):
-        return self.get_data(
-            model,
-            keys,
-            start,
-            end,
-            rollup,
-            [environment_id] if environment_id is not None else None,
-            aggregation="uniq",
-            use_cache=use_cache,
-            jitter_value=jitter_value,
-            tenant_ids=tenant_ids,
-            referrer_suffix=referrer_suffix,
-        )
-
-    def get_distinct_counts_totals_with_conditions(
-        self,
-        model: TSDBModel,
-        keys: Sequence[int],
-        start: datetime,
-        end: datetime | None = None,
-        rollup: int | None = None,
-        environment_id: int | None = None,
-        use_cache: bool = False,
-        jitter_value: int | None = None,
-        tenant_ids: dict[str, int | str] | None = None,
-        referrer_suffix: str | None = None,
-        conditions: list[tuple[str, str, str]] | None = None,
-    ) -> dict[int, Any]:
-        """
-        Count distinct items during a time range with conditions.
-        """
         return self.get_data(
             model,
             keys,

tests/snuba/tsdb/test_tsdb_backend.py

@@ -439,8 +439,8 @@ class SnubaTSDBTest(TestCase, SnubaTestCase):
             == {}
         )
 
-    def test_get_distinct_counts_totals_users_with_conditions(self):
-        assert self.db.get_distinct_counts_totals_with_conditions(
+    def test_get_distinct_counts_totals_users__with_conditions(self):
+        assert self.db.get_distinct_counts_totals(
             TSDBModel.users_affected_by_group,
             [self.proj1group1.id],
             self.now,
@@ -451,7 +451,7 @@ class SnubaTSDBTest(TestCase, SnubaTestCase):
         ) == {
             self.proj1group1.id: 2  # 5 unique users with US tag
         }
-        assert self.db.get_distinct_counts_totals_with_conditions(
+        assert self.db.get_distinct_counts_totals(
             TSDBModel.users_affected_by_group,
             [self.proj1group1.id],
             self.now,
@@ -462,7 +462,7 @@ class SnubaTSDBTest(TestCase, SnubaTestCase):
         ) == {
             self.proj1group1.id: 3  # 3 unique users with EU tag
         }
-        assert self.db.get_distinct_counts_totals_with_conditions(
+        assert self.db.get_distinct_counts_totals(
             TSDBModel.users_affected_by_group,
             [self.proj1group1.id],
             self.now,
@@ -472,7 +472,7 @@ class SnubaTSDBTest(TestCase, SnubaTestCase):
             conditions=[("tags[region]", "=", "MARS")],
         ) == {self.proj1group1.id: 0}
         assert (
-            self.db.get_distinct_counts_totals_with_conditions(
+            self.db.get_distinct_counts_totals(
                 TSDBModel.users_affected_by_group,
                 [],
                 self.now,
@@ -837,6 +837,16 @@ class SnubaTSDBGroupProfilingTest(TestCase, SnubaTestCase, SearchIssueTestMixin)
             tenant_ids={"referrer": "test", "organization_id": 1},
         ) == {self.proj1group1.id: 12, self.proj1group2.id: 12}
 
+    def test_get_sums__with_conditions(self):
+        assert self.db.get_sums(
+            model=TSDBModel.group_generic,
+            keys=[self.proj1group1.id, self.proj1group2.id],
+            start=self.now,
+            end=self.now + timedelta(hours=4),
+            tenant_ids={"referrer": "test", "organization_id": 1},
+            conditions=[("tags[environment]", "=", str(self.env1.name))],
+        ) == {self.proj1group1.id: 6, self.proj1group2.id: 6}
+
     def test_get_data_or_conditions_parsed(self):
         """
         Verify parsing the legacy format with nested OR conditions works