feat(tags): Add dataset param to OrganizationTagKeyValues endpoint (#74525)

This PR adds a dataset parameter to the OrganizationTagKeyValues
endpoint, allowing full control over which dataset will be queried. Our
frontend currently queries the Discover dataset to retrieve tag values
for the issue stream - we need this control to query only the Events and
IssuePlatform datasets.

Author: MichaelSun48

src/sentry/api/endpoints/organization_tagkey_values.py

@@ -10,6 +10,7 @@
 from sentry.api.paginator import SequencePaginator
 from sentry.api.serializers import serialize
 from sentry.api.utils import handle_query_errors
+from sentry.snuba.dataset import Dataset
 from sentry.tagstore.base import TAG_KEY_RE
 
 
@@ -31,6 +32,13 @@ def get(self, request: Request, organization, key) -> Response:
 
         sentry_sdk.set_tag("query.tag_key", key)
 
+        dataset = None
+        if request.GET.get("dataset"):
+            try:
+                dataset = Dataset(request.GET.get("dataset"))
+            except ValueError:
+                raise ParseError(detail="Invalid dataset parameter")
+
         try:
             # still used by events v1 which doesn't require global views
             filter_params = self.get_snuba_params(request, organization, check_global_views=False)
@@ -47,6 +55,7 @@ def get(self, request: Request, organization, key) -> Response:
                     key,
                     filter_params["start"],
                     filter_params["end"],
+                    dataset=dataset,
                     query=request.GET.get("query"),
                     order_by=validate_sort_field(request.GET.get("sort", "-last_seen")),
                     include_transactions=request.GET.get("includeTransactions") == "1",

src/sentry/tagstore/base.py

@@ -232,6 +232,7 @@ def get_tag_value_paginator_for_projects(
         key,
         start=None,
         end=None,
+        dataset: Dataset | None = None,
         query=None,
         order_by="-last_seen",
         include_transactions: bool = False,

src/sentry/tagstore/snuba/backend.py

@@ -1176,6 +1176,7 @@ def get_tag_value_paginator_for_projects(
         key,
         start=None,
         end=None,
+        dataset: Dataset = None,
         query: str | None = None,
         order_by="-last_seen",
         include_transactions: bool = False,
@@ -1193,11 +1194,12 @@ def get_tag_value_paginator_for_projects(
         if order_by == "-count":
             order_by = "-times_seen"
 
-        dataset = Dataset.Events
-        if include_transactions:
-            dataset = Dataset.Discover
-        if include_replays:
-            dataset = Dataset.Replays
+        if not dataset:
+            dataset = Dataset.Events
+            if include_transactions:
+                dataset = Dataset.Discover
+            if include_replays:
+                dataset = Dataset.Replays
 
         snuba_key = snuba.get_snuba_column_name(key, dataset=dataset)
 

tests/snuba/api/endpoints/test_organization_tagkey_values.py

@@ -7,9 +7,11 @@
 
 from sentry.replays.testutils import mock_replay
 from sentry.search.events.constants import RELEASE_ALIAS, SEMVER_ALIAS
+from sentry.snuba.dataset import Dataset
 from sentry.testutils.cases import APITestCase, ReplaysSnubaTestCase, SnubaTestCase
 from sentry.testutils.helpers.datetime import before_now, iso_format
 from sentry.utils.samples import load_data
+from tests.sentry.issues.test_utils import OccurrenceTestMixin
 
 
 class OrganizationTagKeyTestCase(APITestCase, SnubaTestCase):
@@ -808,3 +810,134 @@ def test_schema(self):
             "name",
             "value",
         ]
+
+
+class DatasetParamOrganizationTagKeyValuesTest(OrganizationTagKeyTestCase, OccurrenceTestMixin):
+    def setUp(self):
+        super().setUp()
+
+    def run_dataset_test(self, key, expected, dataset: Dataset, **kwargs):
+        # all tests here require that we search in transactions so make that the default here
+        qs_params = kwargs.get("qs_params", {})
+        qs_params["dataset"] = dataset.value
+        kwargs["qs_params"] = qs_params
+        super().run_test(key, expected, **kwargs)
+
+    def test_dataset_events(self):
+        self.store_event(
+            data={
+                "event_id": "a" * 32,
+                "tags": {"berry": "raspberry"},
+                "timestamp": iso_format(self.min_ago),
+            },
+            project_id=self.project.id,
+        )
+        self.store_event(
+            data={
+                "event_id": "b" * 32,
+                "tags": {"berry": "blueberry"},
+                "timestamp": iso_format(self.min_ago),
+            },
+            project_id=self.project.id,
+        )
+        self.store_event(
+            data={
+                "event_id": "c" * 32,
+                "tags": {"berry": "banana"},
+                "timestamp": iso_format(self.min_ago),
+            },
+            project_id=self.project.id,
+        )
+        self.store_event(
+            data={
+                "event_id": "d" * 32,
+                "tags": {"berry": "banana"},
+                "timestamp": iso_format(self.min_ago),
+            },
+            project_id=self.project.id,
+        )
+        # Should appear in Events and Discover datasets, but not IssuePlatform
+        self.run_dataset_test(
+            "berry",
+            expected=[("raspberry", 1), ("blueberry", 1), ("banana", 2)],
+            dataset=Dataset.Events,
+        )
+        self.run_dataset_test(
+            "berry",
+            expected=[("raspberry", 1), ("blueberry", 1), ("banana", 2)],
+            dataset=Dataset.Discover,
+        )
+        self.run_dataset_test(
+            "berry",
+            expected=[],
+            dataset=Dataset.IssuePlatform,
+        )
+
+    def test_dataset_issue_platform(self):
+        self.store_event(
+            data={
+                "event_id": "a" * 32,
+                "tags": {"stone_fruit": "peach"},
+                "timestamp": iso_format(self.min_ago),
+            },
+            project_id=self.project.id,
+        )
+        self.process_occurrence(
+            event_id=uuid.uuid4().hex,
+            project_id=self.project.id,
+            event_data={
+                "title": "some problem",
+                "platform": "python",
+                "tags": {"stone_fruit": "cherry"},
+                "timestamp": iso_format(self.min_ago),
+                "received": iso_format(self.min_ago),
+            },
+        )
+
+        # (stone_fruit: cherry) should appear in IssuePlatform dataset,
+        # but (sonte_fruit: peach) should not
+        self.run_dataset_test(
+            "stone_fruit",
+            expected=[("cherry", 1)],
+            dataset=Dataset.IssuePlatform,
+        )
+        self.run_dataset_test(
+            "stone_fruit",
+            expected=[("peach", 1)],
+            dataset=Dataset.Events,
+        )
+        self.run_dataset_test(
+            "stone_fruit",
+            expected=[("peach", 1)],
+            dataset=Dataset.Discover,
+        )
+
+    def test_dataset_discover(self):
+        event = load_data("transaction")
+        event["tags"].extend([["fake_fruit", "tomato"]])
+        event.update(
+            {
+                "transaction": "example_transaction",
+                "event_id": uuid.uuid4().hex,
+                "start_timestamp": iso_format(self.min_ago),
+                "timestamp": iso_format(self.min_ago),
+            }
+        )
+        event["measurements"]["lcp"]["value"] = 5000
+        self.store_event(data=event, project_id=self.project.id)
+
+        self.run_dataset_test(
+            "fake_fruit",
+            expected=[],
+            dataset=Dataset.IssuePlatform,
+        )
+        self.run_dataset_test(
+            "fake_fruit",
+            expected=[],
+            dataset=Dataset.Events,
+        )
+        self.run_dataset_test(
+            "fake_fruit",
+            expected=[("tomato", 1)],
+            dataset=Dataset.Discover,
+        )