feat(discover): Expose dataset in discover saved query api (#71540)

Expose dataset field to discover saved query API.
All saved queries currently query the `discover` dataset,
so while the dataset selector feature is still under development,
default dataset to `discover`

src/sentry/api/serializers/models/discoversavedquery.py

@@ -3,7 +3,7 @@ from typing import DefaultDict
 
 from sentry.api.serializers import Serializer, register
 from sentry.constants import ALL_ACCESS_PROJECTS
-from sentry.discover.models import DiscoverSavedQuery
+from sentry.discover.models import DiscoverSavedQuery, DiscoverSavedQueryTypes
 from sentry.services.hybrid_cloud.user.service import user_service
 from sentry.utils.dates import outside_retention_with_modified_start, parse_timestamp
 
@@ -55,6 +55,7 @@ class DiscoverSavedQuerySerializer(Serializer):
             "name": obj.name,
             "projects": [project.id for project in obj.projects.all()],
             "version": obj.version or obj.query.get("version", 1),
+            "queryDataset": DiscoverSavedQueryTypes.get_type_name(obj.dataset),
             "expired": False,
             "dateCreated": obj.date_created,
             "dateUpdated": obj.date_updated,

src/sentry/discover/endpoints/discover_homepage_query.py

@@ -92,6 +92,7 @@ class DiscoverHomepageQueryEndpoint(OrganizationEndpoint):
             name="",
             query=data["query"],
             version=data["version"],
+            dataset=data["query_dataset"],
             created_by_id=request.user.id,
             is_homepage=True,
         )

src/sentry/discover/endpoints/discover_saved_queries.py

@@ -141,6 +141,7 @@ class DiscoverSavedQueriesEndpoint(OrganizationEndpoint):
             name=data["name"],
             query=data["query"],
             version=data["version"],
+            dataset=data["query_dataset"],
             created_by_id=request.user.id if request.user.is_authenticated else None,
         )
 

src/sentry/discover/endpoints/discover_saved_query_detail.py

@@ -85,6 +85,7 @@ class DiscoverSavedQueryDetailEndpoint(OrganizationEndpoint):
             name=data["name"],
             query=data["query"],
             version=data["version"],
+            dataset=data["query_dataset"],
         )
 
         model.set_projects(data["project_ids"])

src/sentry/discover/endpoints/serializers.py

@@ -9,7 +9,11 @@ from sentry.api.fields.empty_integer import EmptyIntegerField
 from sentry.api.utils import get_date_range_from_params
 from sentry.constants import ALL_ACCESS_PROJECTS
 from sentry.discover.arithmetic import ArithmeticError, categorize_columns
-from sentry.discover.models import MAX_TEAM_KEY_TRANSACTIONS, TeamKeyTransaction
+from sentry.discover.models import (
+    MAX_TEAM_KEY_TRANSACTIONS,
+    DiscoverSavedQueryTypes,
+    TeamKeyTransaction,
+)
 from sentry.exceptions import InvalidParams, InvalidSearchQuery
 from sentry.models.team import Team
 from sentry.search.events.builder import QueryBuilder
@@ -142,6 +146,10 @@ class DiscoverQuerySerializer(serializers.Serializer):
 class DiscoverSavedQuerySerializer(serializers.Serializer):
     name = serializers.CharField(required=True, max_length=255)
     projects = ListField(child=serializers.IntegerField(), required=False, default=[])
+    queryDataset = serializers.ChoiceField(
+        choices=DiscoverSavedQueryTypes.as_text_choices(),
+        default=DiscoverSavedQueryTypes.get_type_name(DiscoverSavedQueryTypes.DISCOVER),
+    )
     start = serializers.DateTimeField(required=False, allow_null=True)
     end = serializers.DateTimeField(required=False, allow_null=True)
     range = serializers.CharField(required=False, allow_null=True)
@@ -238,11 +246,14 @@ class DiscoverSavedQuerySerializer(serializers.Serializer):
             except (InvalidSearchQuery, ArithmeticError) as err:
                 raise serializers.ValidationError(f"Cannot save invalid query: {err}")
 
+        dataset = DiscoverSavedQueryTypes.get_id_for_type_name(data["queryDataset"])
+
         return {
             "name": data["name"],
             "project_ids": data["projects"],
             "query": query,
             "version": version,
+            "query_dataset": dataset,
         }
 
     def validate_version_fields(self, version, query):

tests/snuba/api/endpoints/test_discover_homepage_query.py

@@ -1,7 +1,7 @@
 from django.urls import reverse
 
 from sentry.api.serializers import serialize
-from sentry.discover.models import DiscoverSavedQuery
+from sentry.discover.models import DiscoverSavedQuery, DiscoverSavedQueryTypes
 from tests.snuba.api.endpoints.test_discover_saved_queries import DiscoverSavedQueryBase
 
 FEATURES = ("organizations:discover-query",)
@@ -84,6 +84,7 @@ class DiscoverHomepageQueryTest(DiscoverSavedQueryBase):
         assert response.data == serialize(new_query)
         assert new_query.query["fields"] == homepage_query_payload["fields"]
         assert new_query.query["environment"] == homepage_query_payload["environment"]
+        assert new_query.dataset == DiscoverSavedQueryTypes.get_id_for_type_name("discover")
         assert set(new_query.projects.values_list("id", flat=True)) == set(self.project_ids)
 
     def test_put_responds_with_saved_empty_name_field(self):

tests/snuba/api/endpoints/test_discover_saved_query_detail.py

@@ -1,7 +1,11 @@
 import pytest
 from django.urls import NoReverseMatch, reverse
 
-from sentry.discover.models import DiscoverSavedQuery, DiscoverSavedQueryProject
+from sentry.discover.models import (
+    DiscoverSavedQuery,
+    DiscoverSavedQueryProject,
+    DiscoverSavedQueryTypes,
+)
 from sentry.testutils.cases import APITestCase, SnubaTestCase
 
 
@@ -143,6 +147,68 @@ class DiscoverSavedQueryDetailTest(APITestCase, SnubaTestCase):
         assert response.data["conditions"] == []
         assert response.data["limit"] == 20
 
+    def test_put_dataset(self):
+        with self.feature(self.feature_name):
+            url = reverse(
+                "sentry-api-0-discover-saved-query-detail", args=[self.org.slug, self.query_id]
+            )
+
+            response = self.client.put(
+                url,
+                {
+                    "name": "New query",
+                    "projects": self.project_ids,
+                    "fields": [],
+                    "range": "24h",
+                    "limit": 20,
+                    "conditions": [],
+                    "aggregations": [],
+                    "orderby": "-time",
+                    "queryDataset": "transaction-like",
+                },
+            )
+
+        assert response.status_code == 200, response.content
+        assert response.data["id"] == str(self.query_id)
+        assert set(response.data["projects"]) == set(self.project_ids)
+        assert response.data["fields"] == []
+        assert response.data["conditions"] == []
+        assert response.data["limit"] == 20
+        assert response.data["queryDataset"] == "transaction-like"
+
+    def test_dataset_set_to_discover_on_update(self):
+        query = {"fields": ["event_id"], "query": "event.type:error", "limit": 10, "version": 2}
+        model = DiscoverSavedQuery.objects.create(
+            organization=self.org,
+            created_by_id=self.user.id,
+            name="query",
+            query=query,
+            dataset=DiscoverSavedQueryTypes.TRANSACTION_LIKE,
+        )
+
+        with self.feature(self.feature_name):
+            url = reverse(
+                "sentry-api-0-discover-saved-query-detail", args=[self.org.slug, model.id]
+            )
+
+            response = self.client.put(
+                url,
+                {
+                    "name": "New query",
+                    "projects": self.project_ids,
+                    "fields": [],
+                    "range": "24h",
+                    "limit": 20,
+                    "conditions": [],
+                    "aggregations": [],
+                    "orderby": "-time",
+                },
+            )
+
+        assert response.status_code == 200, response.content
+        assert response.data["id"] == str(model.id)
+        assert response.data["queryDataset"] == "discover"
+
     def test_put_with_interval(self):
         with self.feature(self.feature_name):
             url = reverse(