From 8fffe0c61986b6ccfa854adac61e8d3dca098f2c Mon Sep 17 00:00:00 2001
From: Andrew Liu <159852527+aliu39@users.noreply.github.com>
Date: Thu, 3 Oct 2024 11:42:41 -0700
Subject: [PATCH] Disable sampling in query and add test

---
 src/sentry/tagstore/snuba/backend.py          |  5 ++
 .../api/endpoints/test_organization_tags.py   | 64 ++++++++++++++++++-
 2 files changed, 68 insertions(+), 1 deletion(-)

diff --git a/src/sentry/tagstore/snuba/backend.py b/src/sentry/tagstore/snuba/backend.py
index e7839bd3622f8..bf9922b42b5b9 100644
--- a/src/sentry/tagstore/snuba/backend.py
+++ b/src/sentry/tagstore/snuba/backend.py
@@ -441,6 +441,11 @@ def get_tag_keys_for_projects(
         # So only disable sampling if the timerange is short enough.
         if len(projects) <= max_unsampled_projects and end - start <= timedelta(days=14):
             optimize_kwargs["sample"] = 1
+
+        # Replays doesn't support sampling.
+        if dataset == Dataset.Replays:
+            optimize_kwargs = {}
+
         return self.__get_tag_keys_for_projects(
             projects,
             None,
diff --git a/tests/snuba/api/endpoints/test_organization_tags.py b/tests/snuba/api/endpoints/test_organization_tags.py
index 5431dcf681b74..f3bce93ff3f78 100644
--- a/tests/snuba/api/endpoints/test_organization_tags.py
+++ b/tests/snuba/api/endpoints/test_organization_tags.py
@@ -4,7 +4,8 @@
 from django.urls import reverse
 from rest_framework.exceptions import ErrorDetail
 
-from sentry.testutils.cases import APITestCase, SnubaTestCase
+from sentry.replays.testutils import mock_replay
+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
@@ -385,3 +386,64 @@ def test_different_times_retrieves_cache(self):
             cached_data = response.data
 
             assert original_data == cached_data
+
+
+class ReplayOrganizationTagsTest(APITestCase, ReplaysSnubaTestCase):
+    def test_dataset_replays(self):
+        self.login_as(user=self.user)
+        replay1_id = uuid.uuid4().hex
+        replay2_id = uuid.uuid4().hex
+        replay3_id = uuid.uuid4().hex
+        self.r1_seq0_timestamp = before_now(seconds=22)
+        self.r1_seq1_timestamp = before_now(seconds=15)
+        self.r2_seq0_timestamp = before_now(seconds=10)
+        self.r3_seq0_timestamp = before_now(seconds=10)
+        self.store_replays(
+            mock_replay(
+                self.r1_seq0_timestamp,
+                self.project.id,
+                replay1_id,
+                tags={"fruit": "orange"},
+                segment_id=0,
+            ),
+        )
+        self.store_replays(
+            mock_replay(
+                self.r1_seq1_timestamp,
+                self.project.id,
+                replay1_id,
+                tags={"fruit": "orange"},
+                segment_id=1,
+            ),
+        )
+
+        self.store_replays(
+            mock_replay(
+                self.r2_seq0_timestamp,
+                self.project.id,
+                replay2_id,
+                tags={"fruit": "orange"},
+            )
+        )
+        self.store_replays(
+            mock_replay(
+                self.r3_seq0_timestamp,
+                self.project.id,
+                replay3_id,
+                tags={"fruit": "apple", "drink": "water"},
+            )
+        )
+
+        url = reverse(
+            "sentry-api-0-organization-tags",
+            kwargs={"organization_id_or_slug": self.organization.slug},
+        )
+        response = self.client.get(url, {"statsPeriod": "14d", "dataset": "replays"}, format="json")
+
+        assert response.status_code == 200, response.content
+        data = response.data
+        data.sort(key=lambda val: val["name"])
+        assert data == [
+            {"key": "drink", "name": "Drink", "totalValues": 1},
+            {"key": "fruit", "name": "Fruit", "totalValues": 4},
+        ]