fix(csv-export): Handle the issues platform (#62702)

- This updates the csv export to handle the issues platform dataset
- This is because if you go from a perf issue to discover and press
export csv it fails since the export assumes the dataset will only ever
be discover
- This moves the get_dataset function to a util file so it can be reused
by the export
- This provides a workaround for #62698 but doesn't fix it IMO since we
still need to handle the fact that the dataset being set in discover is
currently implicit and only visible in the URL

src/sentry/api/bases/organization_events.py

@@ -28,40 +28,15 @@ from sentry.models.team import Team
 from sentry.search.events.constants import DURATION_UNITS, SIZE_UNITS
 from sentry.search.events.fields import get_function_alias
 from sentry.search.events.types import SnubaParams
-from sentry.snuba import (
-    discover,
-    errors,
-    functions,
-    issue_platform,
-    metrics_enhanced_performance,
-    metrics_performance,
-    profiles,
-    spans_indexed,
-    spans_metrics,
-)
+from sentry.snuba import discover
 from sentry.snuba.metrics.extraction import MetricSpecType
+from sentry.snuba.utils import DATASET_LABELS, DATASET_OPTIONS, get_dataset
 from sentry.utils import snuba
 from sentry.utils.cursors import Cursor
 from sentry.utils.dates import get_interval_from_range, get_rollup_from_request, parse_stats_period
 from sentry.utils.http import absolute_uri
 from sentry.utils.snuba import MAX_FIELDS, SnubaTSResult
 
-# Doesn't map 1:1 with real datasets, but rather what we present to users
-# ie. metricsEnhanced is not a real dataset
-DATASET_OPTIONS = {
-    "discover": discover,
-    "errors": errors,
-    "metricsEnhanced": metrics_enhanced_performance,
-    "metrics": metrics_performance,
-    "profiles": profiles,
-    "issuePlatform": issue_platform,
-    "profileFunctions": functions,
-    "spansIndexed": spans_indexed,
-    "spansMetrics": spans_metrics,
-}
-
-DATASET_LABELS = {value: key for key, value in DATASET_OPTIONS.items()}
-
 
 def resolve_axis_column(column: str, index: int = 0) -> str:
     return get_function_alias(column) if not is_equation(column) else f"equation[{index}]"
@@ -103,10 +78,11 @@ class OrganizationEventsEndpointBase(OrganizationEndpoint):
 
     def get_dataset(self, request: Request) -> Any:
         dataset_label = request.GET.get("dataset", "discover")
-        if dataset_label not in DATASET_OPTIONS:
+        result = get_dataset(dataset_label)
+        if result is None:
             raise ParseError(detail=f"dataset must be one of: {', '.join(DATASET_OPTIONS.keys())}")
         sentry_sdk.set_tag("query.dataset", dataset_label)
-        return DATASET_OPTIONS[dataset_label]
+        return result
 
     def get_snuba_dataclass(
         self, request: Request, organization: Organization, check_global_views: bool = True

src/sentry/data_export/endpoints/data_export.py

@@ -24,6 +24,13 @@ from ..models import ExportedData
 from ..processors.discover import DiscoverProcessor
 from ..tasks import assemble_download
 
+# To support more datasets we may need to change the QueryBuilder being used
+# for now only doing issuePlatform since the product is forcing our hand
+SUPPORTED_DATASETS = {
+    "discover": Dataset.Discover,
+    "issuePlatform": Dataset.IssuePlatform,
+}
+
 
 class DataExportQuerySerializer(serializers.Serializer):
     query_type = serializers.ChoiceField(choices=ExportQueryType.as_str_choices(), required=True)
@@ -86,6 +93,9 @@ class DataExportQuerySerializer(serializers.Serializer):
                 del query_info["statsPeriodEnd"]
             query_info["start"] = start.isoformat()
             query_info["end"] = end.isoformat()
+            dataset = query_info.get("dataset", "discover")
+            if dataset not in SUPPORTED_DATASETS:
+                raise serializers.ValidationError(f"{dataset} is not support for csv exports")
 
             # validate the query string by trying to parse it
             processor = DiscoverProcessor(
@@ -94,7 +104,7 @@ class DataExportQuerySerializer(serializers.Serializer):
             )
             try:
                 builder = QueryBuilder(
-                    Dataset.Discover,
+                    SUPPORTED_DATASETS[dataset],
                     processor.params,
                     query=query_info["query"],
                     selected_columns=fields.copy(),

src/sentry/data_export/processors/discover.py

@@ -8,6 +8,7 @@ from sentry.models.group import Group
 from sentry.models.project import Project
 from sentry.search.events.fields import get_function_alias
 from sentry.snuba import discover
+from sentry.snuba.utils import get_dataset
 
 from ..base import ExportError
 
@@ -45,6 +46,7 @@ class DiscoverProcessor:
             query=discover_query["query"],
             params=self.params,
             sort=discover_query.get("sort"),
+            dataset=discover_query.get("dataset"),
         )
 
     @staticmethod
@@ -76,9 +78,13 @@ class DiscoverProcessor:
         return environment_names
 
     @staticmethod
-    def get_data_fn(fields, equations, query, params, sort):
+    def get_data_fn(fields, equations, query, params, sort, dataset):
+        dataset = get_dataset(dataset)
+        if dataset is None:
+            dataset = discover
+
         def data_fn(offset, limit):
-            return discover.query(
+            return dataset.query(
                 selected_columns=fields,
                 equations=equations,
                 query=query,

src/sentry/snuba/utils.py

@@ -0,0 +1,32 @@
+from typing import Any, Optional
+
+from sentry.snuba import (
+    discover,
+    errors,
+    functions,
+    issue_platform,
+    metrics_enhanced_performance,
+    metrics_performance,
+    profiles,
+    spans_indexed,
+    spans_metrics,
+)
+
+# Doesn't map 1:1 with real datasets, but rather what we present to users
+# ie. metricsEnhanced is not a real dataset
+DATASET_OPTIONS = {
+    "discover": discover,
+    "errors": errors,
+    "metricsEnhanced": metrics_enhanced_performance,
+    "metrics": metrics_performance,
+    "profiles": profiles,
+    "issuePlatform": issue_platform,
+    "profileFunctions": functions,
+    "spansIndexed": spans_indexed,
+    "spansMetrics": spans_metrics,
+}
+DATASET_LABELS = {value: key for key, value in DATASET_OPTIONS.items()}
+
+
+def get_dataset(dataset_label: str) -> Optional[Any]:
+    return DATASET_OPTIONS.get(dataset_label)

tests/sentry/data_export/endpoints/test_data_export.py

@@ -314,3 +314,28 @@ class DataExportTest(APITestCase):
         query_info = data_export.query_info
         assert query_info["field"] == ["count()"]
         assert query_info["equations"] == ["count() / 2"]
+
+    def test_valid_dataset(self):
+        """
+        Ensures that equations are handled
+        """
+        payload = self.make_payload(
+            "discover", {"field": ["title", "count()"], "dataset": "issuePlatform"}
+        )
+        with self.feature(["organizations:discover-query"]):
+            response = self.get_success_response(self.org.slug, status_code=201, **payload)
+        data_export = ExportedData.objects.get(id=response.data["id"])
+        query_info = data_export.query_info
+        assert query_info["field"] == ["title", "count()"]
+        assert query_info["dataset"] == "issuePlatform"
+
+    def test_invalid_dataset(self):
+        """
+        Ensures that equations are handled
+        """
+        payload = self.make_payload(
+            "discover", {"field": ["title", "count()"], "dataset": "somefakedataset"}
+        )
+        with self.feature(["organizations:discover-query"]):
+            response = self.get_response(self.org.slug, **payload)
+        assert response.status_code == 400

tests/sentry/data_export/processors/test_discover.py

@@ -3,7 +3,7 @@ from sentry_relay.consts import SPAN_STATUS_NAME_TO_CODE
 
 from sentry.data_export.base import ExportError
 from sentry.data_export.processors.discover import DiscoverProcessor
-from sentry.testutils.cases import SnubaTestCase, TestCase
+from sentry.testutils.cases import PerformanceIssueTestCase, SnubaTestCase, TestCase
 
 
 class DiscoverProcessorTest(TestCase, SnubaTestCase):
@@ -88,3 +88,24 @@ class DiscoverProcessorTest(TestCase, SnubaTestCase):
         assert new_result_list[0] != result_list
         assert new_result_list[0]["count(id) / fake(field)"] == 5
         assert new_result_list[0]["count(id) / 2"] == 8
+
+
+class DiscoverIssuesProcessorTest(TestCase, PerformanceIssueTestCase):
+    def test_handle_dataset(self):
+        query = {
+            "statsPeriod": "14d",
+            "project": [self.project.id],
+            "field": ["count(id)", "fake(field)", "issue"],
+            "query": "",
+        }
+        query["field"] = ["title", "count()"]
+        query["dataset"] = "issuePlatform"
+        self.create_performance_issue()
+        processor = DiscoverProcessor(organization_id=self.organization.id, discover_query=query)
+        assert processor.header_fields == [
+            "title",
+            "count",
+        ]
+        result = processor.data_fn(0, 1)
+        assert len(result["data"]) == 1
+        assert result["data"][0]["title"] == "N+1 Query"

tests/snuba/api/endpoints/test_organization_events_mep.py

@@ -7,7 +7,6 @@ import pytest
 from django.urls import reverse
 from rest_framework.response import Response
 
-from sentry.api.bases.organization_events import DATASET_OPTIONS
 from sentry.discover.models import TeamKeyTransaction
 from sentry.models.projectteam import ProjectTeam
 from sentry.models.transaction_threshold import (
@@ -20,6 +19,7 @@ from sentry.search.utils import map_device_class_level
 from sentry.snuba.metrics.extraction import MetricSpecType, OnDemandMetricSpec
 from sentry.snuba.metrics.naming_layer.mri import TransactionMRI
 from sentry.snuba.metrics.naming_layer.public import TransactionMetricKey
+from sentry.snuba.utils import DATASET_OPTIONS
 from sentry.testutils.cases import MetricsEnhancedPerformanceTestCase
 from sentry.testutils.helpers.datetime import before_now, iso_format
 from sentry.testutils.silo import region_silo_test