From 00c99dfd2afd4d59a121dc186ed534731c000f00 Mon Sep 17 00:00:00 2001
From: Nar Saynorath <nar.saynorath@sentry.io>
Date: Fri, 26 Jul 2024 20:40:23 -0400
Subject: [PATCH] feat(discover): Query specific datasets in events

---
 .../api/endpoints/organization_events.py      | 36 ++++--------
 .../endpoints/test_organization_events_mep.py | 55 +++++++++++++++++++
 2 files changed, 66 insertions(+), 25 deletions(-)

diff --git a/src/sentry/api/endpoints/organization_events.py b/src/sentry/api/endpoints/organization_events.py
index c997f1021bd354..74bc501f7b2a9c 100644
--- a/src/sentry/api/endpoints/organization_events.py
+++ b/src/sentry/api/endpoints/organization_events.py
@@ -23,7 +23,13 @@
 from sentry.exceptions import InvalidParams
 from sentry.models.dashboard_widget import DashboardWidget, DashboardWidgetTypes
 from sentry.models.organization import Organization
-from sentry.snuba import discover, metrics_enhanced_performance, metrics_performance
+from sentry.snuba import (
+    discover,
+    errors,
+    metrics_enhanced_performance,
+    metrics_performance,
+    transactions,
+)
 from sentry.snuba.metrics.extraction import MetricSpecType
 from sentry.snuba.referrer import Referrer
 from sentry.snuba.utils import dataset_split_decision_inferred_from_query, get_dataset
@@ -404,33 +410,18 @@ def _dashboards_data_fn(scoped_dataset, offset, limit, scoped_query, dashboard_w
 
                 if does_widget_have_split and not has_override_feature:
                     # This is essentially cached behaviour and we skip the check
-                    split_query = scoped_query
                     if widget.discover_widget_split == DashboardWidgetTypes.ERROR_EVENTS:
-                        split_dataset = discover
-                        split_query = (
-                            f"({scoped_query}) AND !event.type:transaction"
-                            if scoped_query
-                            else "!event.type:transaction"
-                        )
+                        split_dataset = errors
                     elif widget.discover_widget_split == DashboardWidgetTypes.TRANSACTION_LIKE:
                         # We can't add event.type:transaction for now because of on-demand.
                         split_dataset = scoped_dataset
                     else:
                         split_dataset = discover
 
-                    return _data_fn(split_dataset, offset, limit, split_query)
+                    return _data_fn(split_dataset, offset, limit, scoped_query)
 
                 try:
-                    error_results = _data_fn(
-                        discover,
-                        offset,
-                        limit,
-                        (
-                            f"({scoped_query}) AND !event.type:transaction"
-                            if scoped_query
-                            else "!event.type:transaction"
-                        ),
-                    )
+                    error_results = _data_fn(errors, offset, limit, scoped_query)
                     # Widget has not split the discover dataset yet, so we need to check if there are errors etc.
                     has_errors = len(error_results["data"]) > 0
                 except SnubaError:
@@ -452,12 +443,7 @@ def _dashboards_data_fn(scoped_dataset, offset, limit, scoped_query, dashboard_w
                 if has_errors and has_other_data and not using_metrics:
                     # In the case that the original request was not using the metrics dataset, we cannot be certain that other data is solely transactions.
                     sentry_sdk.set_tag("third_split_query", True)
-                    transactions_only_query = (
-                        f"({scoped_query}) AND event.type:transaction"
-                        if scoped_query
-                        else "event.type:transaction"
-                    )
-                    transaction_results = _data_fn(discover, offset, limit, transactions_only_query)
+                    transaction_results = _data_fn(transactions, offset, limit, scoped_query)
                     has_transactions = len(transaction_results["data"]) > 0
 
                 decision = self.save_split_decision(
diff --git a/tests/snuba/api/endpoints/test_organization_events_mep.py b/tests/snuba/api/endpoints/test_organization_events_mep.py
index 8e644620b2a776..438199e3045d44 100644
--- a/tests/snuba/api/endpoints/test_organization_events_mep.py
+++ b/tests/snuba/api/endpoints/test_organization_events_mep.py
@@ -3791,6 +3791,61 @@ def test_split_decision_for_ambiguous_widget_with_data(self):
         assert widget.discover_widget_split == DashboardWidgetTypes.ERROR_EVENTS
         assert widget.dataset_source == DatasetSourcesTypes.FORCED.value
 
+    @mock.patch("sentry.snuba.errors.query")
+    def test_errors_request_made_for_saved_error_dashboard_widget_type(self, mock_errors_query):
+        mock_errors_query.return_value = {
+            "data": [],
+            "meta": {},
+        }
+        _, widget, __ = create_widget(
+            ["count()"], "", self.project, discover_widget_split=DashboardWidgetTypes.ERROR_EVENTS
+        )
+
+        response = self.do_request(
+            {
+                "field": [
+                    "count()",
+                ],
+                "query": "",
+                "dataset": "metricsEnhanced",
+                "per_page": 50,
+                "dashboardWidgetId": widget.id,
+            }
+        )
+
+        assert response.status_code == 200, response.content
+        mock_errors_query.assert_called_once()
+
+    @mock.patch("sentry.snuba.metrics_enhanced_performance.query")
+    def test_metrics_enhanced_request_made_for_saved_transaction_like_dashboard_widget_type(
+        self, mock_mep_query
+    ):
+        mock_mep_query.return_value = {
+            "data": [],
+            "meta": {},
+        }
+        _, widget, __ = create_widget(
+            ["count()"],
+            "",
+            self.project,
+            discover_widget_split=DashboardWidgetTypes.TRANSACTION_LIKE,
+        )
+
+        response = self.do_request(
+            {
+                "field": [
+                    "count()",
+                ],
+                "query": "",
+                "dataset": "metricsEnhanced",
+                "per_page": 50,
+                "dashboardWidgetId": widget.id,
+            }
+        )
+
+        assert response.status_code == 200, response.content
+        mock_mep_query.assert_called_once()
+
 
 class OrganizationEventsMetricsEnhancedPerformanceEndpointTestWithMetricLayer(
     OrganizationEventsMetricsEnhancedPerformanceEndpointTest