Repository: getsentry/sentry
Pull request: feat(tags): Make dataset a parameter in organization_tags endpoint 
Fixes #73699

This PR replaces the `include_transaction` parameter with a `dataset` parameter in the `organization_tags` endpoint. Previously, the `include_transactions` was used to control whether to pull tags from snuba's `Discover` dataset (if `include_transactions=1`, or `Events` dataset (if `include_transactions=0`). Replacing this "toggle" with an explicit Dataset parameter will enable full control over which dataset is being queried, which is useful for Issues search since it needs Events and IssuePlatform tags, but not Performance tags (included in Discover dataset). 

Problem statement:
Modify organization tags endpoint to accept a dataset parameter
The `OrganizationTagsEndpoint` will query the events dataset with `include_transactions=0` and the Discover dataset when `include_transactions` is enabled. This works for discover and performance issues, because both look at transactions. However, for issues, we need to be searching the errors and issue platform datasets. The way the endpoint works now, there is no way to search the issue platform dataset.

We should add an additional `dataset` param which will override the `dataset` used in the tags request. This param should eventually replace the `include_transactions` parameter. Callers should then be able to send `dataset=errors` or `dataset=issuePlatform` to query the required dataset.

Hints from issue discussion:

Commit message: Add tests for dataset param

Patch to review:
diff --git a/src/sentry/api/endpoints/organization_tags.py b/src/sentry/api/endpoints/organization_tags.py
index b79ac9ffe7d45d..821e7a3bed81ae 100644
--- a/src/sentry/api/endpoints/organization_tags.py
+++ b/src/sentry/api/endpoints/organization_tags.py
@@ -1,4 +1,5 @@
 import sentry_sdk
+from rest_framework.exceptions import ParseError
 from rest_framework.request import Request
 from rest_framework.response import Response
 
@@ -10,6 +11,7 @@
 from sentry.api.bases.organization import OrganizationEndpoint
 from sentry.api.serializers import serialize
 from sentry.api.utils import handle_query_errors
+from sentry.snuba.dataset import Dataset
 from sentry.utils.numbers import format_grouped_length
 from sentry.utils.sdk import set_measurement
 
@@ -27,6 +29,12 @@ def get(self, request: Request, organization) -> Response:
         except NoProjects:
             return Response([])
 
+        if request.GET.get("dataset"):
+            try:
+                Dataset(request.GET.get("dataset"))
+            except ValueError:
+                raise ParseError("Invalid dataset parameter")
+
         with sentry_sdk.start_span(op="tagstore", description="get_tag_keys_for_projects"):
             with handle_query_errors():
                 results = tagstore.backend.get_tag_keys_for_projects(
@@ -35,8 +43,7 @@ def get(self, request: Request, organization) -> Response:
                     filter_params["start"],
                     filter_params["end"],
                     use_cache=request.GET.get("use_cache", "0") == "1",
-                    # Defaults to True, because the frontend caches these tags globally
-                    include_transactions=request.GET.get("include_transactions", "1") == "1",
+                    dataset=Dataset(request.GET.get("dataset", "discover")),
                     tenant_ids={"organization_id": organization.id},
                 )
 
@@ -50,6 +57,7 @@ def get(self, request: Request, organization) -> Response:
                     "custom_tags.count.grouped",
                     format_grouped_length(len(results), [1, 10, 50, 100]),
                 )
+                sentry_sdk.set_tag("dataset_queried", request.GET.get("dataset", "discover"))
                 set_measurement("custom_tags.count", len(results))
 
         return Response(serialize(results, request.user))
diff --git a/src/sentry/tagstore/base.py b/src/sentry/tagstore/base.py
index 90fbef88f7835f..56dca9f08deeac 100644
--- a/src/sentry/tagstore/base.py
+++ b/src/sentry/tagstore/base.py
@@ -1,6 +1,7 @@
 import re
 
 from sentry.constants import TAG_LABELS
+from sentry.snuba.dataset import Dataset
 from sentry.utils.services import Service
 
 # Valid pattern for tag key names
@@ -139,9 +140,9 @@ def get_tag_keys_for_projects(
         environments,
         start,
         end,
+        dataset: Dataset = Dataset.Events,
         status=TagKeyStatus.ACTIVE,
         use_cache: bool = False,
-        include_transactions: bool = False,
         tenant_ids=None,
     ):
         """
diff --git a/src/sentry/tagstore/snuba/backend.py b/src/sentry/tagstore/snuba/backend.py
index 56ca46094c3b6c..58f78fef0d2fd1 100644
--- a/src/sentry/tagstore/snuba/backend.py
+++ b/src/sentry/tagstore/snuba/backend.py
@@ -218,7 +218,7 @@ def __get_tag_keys(
         limit=1000,
         keys=None,
         include_values_seen=True,
-        include_transactions=False,
+        dataset: Dataset = Dataset.Events,
         denylist=None,
         tenant_ids=None,
         **kwargs,
@@ -232,7 +232,7 @@ def __get_tag_keys(
             limit,
             keys,
             include_values_seen=include_values_seen,
-            include_transactions=include_transactions,
+            dataset=dataset,
             denylist=denylist,
             tenant_ids=tenant_ids,
         )
@@ -248,8 +248,8 @@ def __get_tag_keys_for_projects(
         keys=None,
         include_values_seen=True,
         use_cache=False,
-        include_transactions=False,
         denylist=None,
+        dataset: Dataset = Dataset.Discover,
         **kwargs,
     ):
         """Query snuba for tag keys based on projects
@@ -272,9 +272,6 @@ def __get_tag_keys_for_projects(
         if end is None:
             end = default_end
 
-        dataset = Dataset.Events
-        if include_transactions:
-            dataset = Dataset.Discover
         conditions = []
         aggregations = [["count()", "", "count"]]
 
@@ -282,7 +279,7 @@ def __get_tag_keys_for_projects(
         if environments:
             filters["environment"] = sorted(environments)
         if group is not None:
-            # We override dataset changes from `include_transactions` here. They aren't relevant
+            # We override dataset changes from `dataset` here. They aren't relevant
             # when filtering by a group.
             dataset, conditions, filters = self.apply_group_filters_conditions(
                 group, conditions, filters
@@ -428,9 +425,9 @@ def get_tag_keys_for_projects(
         environments,
         start,
         end,
+        dataset: Dataset = Dataset.Events,
         status=TagKeyStatus.ACTIVE,
         use_cache: bool = False,
-        include_transactions: bool = False,
         tenant_ids=None,
     ):
         max_unsampled_projects = _max_unsampled_projects
@@ -450,9 +447,9 @@ def get_tag_keys_for_projects(
             environments,
             start,
             end,
+            dataset=dataset,
             include_values_seen=False,
             use_cache=use_cache,
-            include_transactions=include_transactions,
             tenant_ids=tenant_ids,
             **optimize_kwargs,
         )
diff --git a/tests/snuba/api/endpoints/test_organization_tags.py b/tests/snuba/api/endpoints/test_organization_tags.py
index 98aa826d51fe4d..a3cd50560c0715 100644
--- a/tests/snuba/api/endpoints/test_organization_tags.py
+++ b/tests/snuba/api/endpoints/test_organization_tags.py
@@ -1,12 +1,15 @@
+import uuid
 from unittest import mock
 
 from django.urls import reverse
 
 from sentry.testutils.cases import APITestCase, 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 OrganizationTagsTest(APITestCase, SnubaTestCase):
+class OrganizationTagsTest(APITestCase, OccurrenceTestMixin, SnubaTestCase):
     def setUp(self):
         super().setUp()
         self.min_ago = iso_format(before_now(minutes=1))
@@ -55,6 +58,112 @@ def test_simple(self):
             {"name": "Some Tag", "key": "some_tag", "totalValues": 1},
         ]
 
+    def test_dataset_events(self):
+        user = self.create_user()
+        org = self.create_organization()
+        team = self.create_team(organization=org)
+        self.create_member(organization=org, user=user, teams=[team])
+
+        self.login_as(user=user)
+
+        project = self.create_project(organization=org, teams=[team])
+        self.store_event(
+            data={"event_id": "a" * 32, "tags": {"berry": "raspberry"}, "timestamp": self.min_ago},
+            project_id=project.id,
+        )
+
+        url = reverse(
+            "sentry-api-0-organization-tags", kwargs={"organization_id_or_slug": org.slug}
+        )
+
+        response = self.client.get(url, {"statsPeriod": "14d", "dataset": "events"}, format="json")
+
+        assert response.status_code == 200, response.content
+        data = response.data
+        data.sort(key=lambda val: val["name"])
+        assert data == [
+            {"name": "Berry", "key": "berry", "totalValues": 1},
+            {"name": "Level", "key": "level", "totalValues": 1},
+        ]
+
+    def test_dataset_discover(self):
+        user = self.create_user()
+        org = self.create_organization()
+        team = self.create_team(organization=org)
+        self.create_member(organization=org, user=user, teams=[team])
+
+        self.login_as(user=user)
+
+        project = self.create_project(organization=org, teams=[team])
+
+        url = reverse(
+            "sentry-api-0-organization-tags", kwargs={"organization_id_or_slug": org.slug}
+        )
+
+        event = load_data("transaction")
+        event["tags"].extend([["apple", "fuji"]])
+        event.update(
+            {
+                "transaction": "example_transaction",
+                "event_id": uuid.uuid4().hex,
+                "start_timestamp": self.min_ago,
+                "timestamp": self.min_ago,
+            }
+        )
+        event["measurements"]["lcp"]["value"] = 5000
+        self.store_event(data=event, project_id=project.id)
+
+        discoverResponse = self.client.get(
+            url,
+            {"statsPeriod": "14d", "dataset": "discover"},
+            format="json",
+        )
+        assert discoverResponse.status_code == 200, discoverResponse.content
+        assert {"name": "Apple", "key": "apple", "totalValues": 1} in discoverResponse.data
+
+    def test_dataset_issue_platform(self):
+        user = self.create_user()
+        org = self.create_organization()
+        team = self.create_team(organization=org)
+        self.create_member(organization=org, user=user, teams=[team])
+
+        self.login_as(user=user)
+
+        project = self.create_project(organization=org, teams=[team])
+
+        self.store_event(
+            data={"event_id": "a" * 32, "tags": {"berry": "raspberry"}, "timestamp": self.min_ago},
+            project_id=project.id,
+        )
+
+        self.process_occurrence(
+            event_id=uuid.uuid4().hex,
+            project_id=project.id,
+            event_data={
+                "title": "some problem",
+                "platform": "python",
+                "tags": {"stone_fruit": "cherry"},
+                "timestamp": before_now(minutes=1).isoformat(),
+                "received": before_now(minutes=1).isoformat(),
+            },
+        )
+
+        url = reverse(
+            "sentry-api-0-organization-tags", kwargs={"organization_id_or_slug": org.slug}
+        )
+
+        response = self.client.get(
+            url, {"statsPeriod": "14d", "dataset": "search_issues"}, format="json"
+        )
+
+        assert response.status_code == 200, response.content
+        data = response.data
+        data.sort(key=lambda val: val["name"])
+        assert data == [
+            {"name": "Level", "key": "level", "totalValues": 1},
+            {"name": "Stone Fruit", "key": "stone_fruit", "totalValues": 1},
+        ]
+
     def test_no_projects(self):
         user = self.create_user()
         org = self.create_organization(owner=user)


Code review comments:
diff_hunk: @@ -27,6 +29,12 @@ def get(self, request: Request, organization) -> Response:
         except NoProjects:
             return Response([])
 
+        if request.GET.get("dataset"):
+            try:
+                Dataset(request.GET.get("dataset"))
+            except ValueError:
+                raise ParseError("Invalid dataset parameter")
path: src/sentry/api/endpoints/organization_tags.py
text: @user1:
```suggestion
        try:
            dataset = Dataset(request.GET.get("dataset", "discover"))
        except ValueError:
            raise ParseError("Invalid dataset parameter")
```
If you pull out the dataset here you should be able to avoid having to keep the fallback "discover" aligned in all the gets everywhere else - line 46 should be able to be just `dataset=dataset,` and line 60 is a lot easier too

@author:
Good call!