feat(discover-snql): Use the snql query builder for events-meta (#29096)

- This adds a feature flag so we can try the new query builder for the
  events-meta endpoint
- This also adds a function_alias_map so we can correctly convert the
  meta results on the discover endpoints

src/sentry/api/endpoints/organization_events_meta.py

@@ -4,7 +4,7 @@
 from rest_framework.exceptions import ParseError
 from rest_framework.response import Response
 
-from sentry import search
+from sentry import features, search
 from sentry.api.base import EnvironmentMixin
 from sentry.api.bases import NoProjects, OrganizationEventsEndpointBase
 from sentry.api.event_search import parse_search_query
@@ -28,6 +28,9 @@ def get(self, request, organization):
                 params=params,
                 query=request.query_params.get("query"),
                 referrer="api.organization-events-meta",
+                use_snql=features.has(
+                    "organizations:discover-use-snql", organization, actor=request.user
+                ),
             )
 
         return Response({"count": result["data"][0]["count"]})

src/sentry/features/__init__.py

@@ -71,6 +71,7 @@
 default_manager.add("organizations:discover-top-events", OrganizationFeature, True)
 default_manager.add("organizations:discover-basic", OrganizationFeature)
 default_manager.add("organizations:discover-query", OrganizationFeature)
+default_manager.add("organizations:discover-use-snql", OrganizationFeature, True)
 default_manager.add("organizations:enterprise-perf", OrganizationFeature)
 default_manager.add("organizations:event-attachments", OrganizationFeature)
 default_manager.add("organizations:event-attachments-viewer", OrganizationFeature)

src/sentry/search/events/fields.py

@@ -2,7 +2,7 @@
 from collections import defaultdict, namedtuple
 from copy import deepcopy
 from datetime import datetime
-from typing import Any, Callable, List, Mapping, Match, Optional, Sequence, Tuple, Union
+from typing import Any, Callable, Dict, List, Mapping, Match, Optional, Sequence, Tuple, Union
 
 import sentry_sdk
 from sentry_relay.consts import SPAN_STATUS_NAME_TO_CODE
@@ -2171,6 +2171,7 @@ def __init__(
     ):
         super().__init__(dataset, params, functions_acl)
 
+        self.function_alias_map: Dict[str, FunctionDetails] = {}
         self.field_alias_converter: Mapping[str, Callable[[str], SelectType]] = {
             # NOTE: `ISSUE_ALIAS` simply maps to the id, meaning that post processing
             # is required to insert the true issue short id into the response.
@@ -2261,6 +2262,7 @@ def __init__(
                 ),
                 SnQLFunction(
                     "count",
+                    optional_args=[NullColumn("column")],
                     snql_aggregate=lambda _, alias: Function(
                         "count",
                         [],
@@ -2783,6 +2785,8 @@ def resolve_function(self, function: str, match: Optional[Match[str]] = None) ->
             raise InvalidSearchQuery(f"{snql_function.name}: no access to private function")
 
         arguments = snql_function.format_as_arguments(name, arguments, self.params)
+        self.function_alias_map[alias] = FunctionDetails(function, snql_function, arguments.copy())
+
         for arg in snql_function.args:
             if isinstance(arg, ColumnArg):
                 arguments[arg.name] = self.resolve_column(arguments[arg.name])

src/sentry/snuba/discover.py

@@ -172,8 +172,8 @@ def get_row(row):
 
     result["data"] = [get_row(row) for row in result["data"]]
 
-    rollup = snuba_filter.rollup
-    if rollup and rollup > 0:
+    if snuba_filter and snuba_filter.rollup and snuba_filter.rollup > 0:
+        rollup = snuba_filter.rollup
         with sentry_sdk.start_span(
             op="discover.discover", description="transform_results.zerofill"
         ) as span:
@@ -231,7 +231,10 @@ def query(
     if not selected_columns:
         raise InvalidSearchQuery("No columns selected")
 
+    sentry_sdk.set_tag("discover.use_snql", use_snql)
     if use_snql:
+        # temporarily add snql to referrer
+        referrer = f"{referrer}.wip-snql"
         builder = QueryBuilder(
             Dataset.Discover,
             params,
@@ -246,8 +249,13 @@ def query(
         )
         snql_query = builder.get_snql_query()
 
-        results = raw_snql_query(snql_query, referrer)
-        return results
+        result = raw_snql_query(snql_query, referrer)
+        with sentry_sdk.start_span(
+            op="discover.discover", description="query.transform_results"
+        ) as span:
+            span.set_data("result_count", len(result.get("data", [])))
+            result = transform_results(result, builder.function_alias_map, {}, None)
+        return result
 
     # We clobber this value throughout this code, so copy the value
     selected_columns = selected_columns[:]

tests/snuba/api/endpoints/test_organization_events_meta.py

@@ -17,12 +17,14 @@ def setUp(self):
             "sentry-api-0-organization-events-meta",
             kwargs={"organization_slug": self.project.organization.slug},
         )
+        self.features = {"organizations:discover-basic": True}
 
     def test_simple(self):
 
         self.store_event(data={"timestamp": iso_format(self.min_ago)}, project_id=self.project.id)
 
-        response = self.client.get(self.url, format="json")
+        with self.feature(self.features):
+            response = self.client.get(self.url, format="json")
 
         assert response.status_code == 200, response.content
         assert response.data["count"] == 1
@@ -37,7 +39,8 @@ def test_multiple_projects(self):
 
         assert response.status_code == 400, response.content
 
-        with self.feature("organizations:global-views"):
+        self.features["organizations:global-views"] = True
+        with self.feature(self.features):
             response = self.client.get(self.url, format="json")
 
         assert response.status_code == 200, response.content
@@ -53,13 +56,15 @@ def test_search(self):
             project_id=self.project.id,
         )
 
-        response = self.client.get(self.url, {"query": "delete"}, format="json")
+        with self.feature(self.features):
+            response = self.client.get(self.url, {"query": "delete"}, format="json")
 
         assert response.status_code == 200, response.content
         assert response.data["count"] == 1
 
     def test_invalid_query(self):
-        response = self.client.get(self.url, {"query": "is:unresolved"}, format="json")
+        with self.feature(self.features):
+            response = self.client.get(self.url, {"query": "is:unresolved"}, format="json")
 
         assert response.status_code == 400, response.content
 
@@ -70,7 +75,8 @@ def test_no_projects(self):
             "sentry-api-0-organization-events-meta",
             kwargs={"organization_slug": no_project_org.slug},
         )
-        response = self.client.get(url, format="json")
+        with self.feature(self.features):
+            response = self.client.get(url, format="json")
 
         assert response.status_code == 200, response.content
         assert response.data["count"] == 0
@@ -91,7 +97,8 @@ def test_transaction_event(self):
             "sentry-api-0-organization-events-meta",
             kwargs={"organization_slug": self.project.organization.slug},
         )
-        response = self.client.get(url, {"query": "transaction.duration:>1"}, format="json")
+        with self.feature(self.features):
+            response = self.client.get(url, {"query": "transaction.duration:>1"}, format="json")
 
         assert response.status_code == 200, response.content
         assert response.data["count"] == 1
@@ -108,37 +115,41 @@ def test_transaction_event_with_last_seen(self):
             "start_timestamp": iso_format(before_now(minutes=1, seconds=3)),
         }
         self.store_event(data=data, project_id=self.project.id)
-        response = self.client.get(
-            self.url, {"query": "event.type:transaction last_seen():>2012-12-31"}, format="json"
-        )
+        with self.feature(self.features):
+            response = self.client.get(
+                self.url, {"query": "event.type:transaction last_seen():>2012-12-31"}, format="json"
+            )
 
         assert response.status_code == 200, response.content
         assert response.data["count"] == 1
 
     def test_out_of_retention(self):
-        with self.options({"system.event-retention-days": 10}):
-            response = self.client.get(
-                self.url,
-                format="json",
-                data={
-                    "start": iso_format(before_now(days=20)),
-                    "end": iso_format(before_now(days=15)),
-                },
-            )
+        with self.feature(self.features):
+            with self.options({"system.event-retention-days": 10}):
+                response = self.client.get(
+                    self.url,
+                    format="json",
+                    data={
+                        "start": iso_format(before_now(days=20)),
+                        "end": iso_format(before_now(days=15)),
+                    },
+                )
         assert response.status_code == 400
 
     @mock.patch("sentry.snuba.discover.raw_query")
-    def test_handling_snuba_errors(self, mock_query):
+    @mock.patch("sentry.snuba.discover.raw_snql_query")
+    def test_handling_snuba_errors(self, mock_query, mock_snql_query):
         mock_query.side_effect = ParseError("test")
-        with self.feature("organizations:discover-basic"):
+        mock_snql_query.side_effect = ParseError("test")
+        with self.feature(self.features):
             response = self.client.get(self.url, format="json")
 
         assert response.status_code == 400, response.content
 
     @mock.patch("sentry.utils.snuba.quantize_time")
     def test_quantize_dates(self, mock_quantize):
         mock_quantize.return_value = before_now(days=1).replace(tzinfo=utc)
-        with self.feature("organizations:discover-basic"):
+        with self.feature(self.features):
             # Don't quantize short time periods
             self.client.get(
                 self.url,
@@ -169,6 +180,12 @@ def test_quantize_dates(self, mock_quantize):
             assert len(mock_quantize.mock_calls) == 2
 
 
+class OrganizationEventsMetaEndpointWithSnql(OrganizationEventsMetaEndpoint):
+    def setUp(self):
+        super().setUp()
+        self.features["organizations:discover-use-snql"] = True
+
+
 class OrganizationEventBaselineEndpoint(APITestCase, SnubaTestCase):
     def setUp(self):
         super().setUp()