feat(eap): Add downsampling to the api (#88023)

- This updates the events/ api to accept a sampling param so we can make
a preflight before trying the new best_effort
- Tests are not really ready, but opening this for now so FE can start
sorta testing at least

Author: wmak

src/sentry/api/endpoints/organization_events.py

@@ -421,6 +421,7 @@ def get(self, request: Request, organization) -> Response:
 
         dataset = self.get_dataset(request)
         metrics_enhanced = dataset in {metrics_performance, metrics_enhanced_performance}
+        sampling_mode = request.GET.get("sampling")
 
         sentry_sdk.set_tag("performance.metrics_enhanced", metrics_enhanced)
         allow_metric_aggregates = request.GET.get("preventMetricAggregates") != "1"
@@ -465,6 +466,7 @@ def _data_fn(
                         auto_fields=True,
                         use_aggregate_conditions=use_aggregate_conditions,
                     ),
+                    sampling_mode=sampling_mode,
                 )
             query_source = self.get_request_source(request)
             return dataset_query(

src/sentry/incidents/logic.py

@@ -404,6 +404,7 @@ def get_metric_issue_aggregates(
                 offset=0,
                 limit=1,
                 referrer=Referrer.API_ALERTS_ALERT_RULE_CHART.value,
+                sampling_mode=None,
                 config=SearchResolverConfig(
                     auto_fields=True,
                 ),

src/sentry/profiles/flamegraph.py

@@ -568,6 +568,7 @@ def get_spans_based_candidates(self, query: str | None, limit: int) -> EAPRespon
             offset=0,
             limit=limit,
             referrer=Referrer.API_TRACE_EXPLORER_TRACE_SPANS_CANDIDATES_FLAMEGRAPH.value,
+            sampling_mode=None,
             config=SearchResolverConfig(
                 auto_fields=True,
             ),

src/sentry/search/eap/constants.py

@@ -1,5 +1,6 @@
 from typing import Literal
 
+from sentry_protos.snuba.v1.downsampled_storage_pb2 import DownsampledStorageConfig
 from sentry_protos.snuba.v1.endpoint_trace_item_table_pb2 import AggregationComparisonFilter
 from sentry_protos.snuba.v1.request_common_pb2 import TraceItemType
 from sentry_protos.snuba.v1.trace_item_attribute_pb2 import AttributeKey
@@ -162,3 +163,8 @@
     ],
     5: ["500", "501", "502", "503", "504", "505", "506", "507", "508", "509", "510", "511"],
 }
+
+SAMPLING_MODES = {
+    "BEST_EFFORT": DownsampledStorageConfig.MODE_BEST_EFFORT,
+    "PREFLIGHT": DownsampledStorageConfig.MODE_PREFLIGHT,
+}

src/sentry/search/eap/resolver.py

@@ -47,6 +47,7 @@
     VirtualColumnDefinition,
 )
 from sentry.search.eap.types import SearchResolverConfig
+from sentry.search.eap.utils import validate_sampling
 from sentry.search.events import constants as qb_constants
 from sentry.search.events import fields
 from sentry.search.events import filter as event_filter
@@ -76,7 +77,7 @@ class SearchResolver:
     ] = field(default_factory=dict)
 
     @sentry_sdk.trace
-    def resolve_meta(self, referrer: str) -> RequestMeta:
+    def resolve_meta(self, referrer: str, sampling_mode: str | None = None) -> RequestMeta:
         if self.params.organization_id is None:
             raise Exception("An organization is required to resolve queries")
         span = sentry_sdk.get_current_span()
@@ -89,6 +90,7 @@ def resolve_meta(self, referrer: str) -> RequestMeta:
             start_timestamp=self.params.rpc_start_date,
             end_timestamp=self.params.rpc_end_date,
             trace_item_type=self.definitions.trace_item_type,
+            downsampled_storage_config=validate_sampling(sampling_mode),
         )
 
     @sentry_sdk.trace

src/sentry/search/eap/spans/formulas.py

@@ -422,6 +422,7 @@ def time_spent_percentage(
         orderby=None,
         offset=0,
         limit=1,
+        sampling_mode=None,
         config=SearchResolverConfig(),
     )
 

src/sentry/search/eap/utils.py

@@ -3,11 +3,13 @@
 from typing import Any, Literal
 
 from google.protobuf.timestamp_pb2 import Timestamp
+from sentry_protos.snuba.v1.downsampled_storage_pb2 import DownsampledStorageConfig
 from sentry_protos.snuba.v1.endpoint_time_series_pb2 import Expression, TimeSeriesRequest
 from sentry_protos.snuba.v1.endpoint_trace_item_table_pb2 import Column
 from sentry_protos.snuba.v1.trace_item_attribute_pb2 import Function
 
 from sentry.exceptions import InvalidSearchQuery
+from sentry.search.eap.constants import SAMPLING_MODES
 from sentry.search.eap.ourlogs.attributes import LOGS_INTERNAL_TO_PUBLIC_ALIAS_MAPPINGS
 from sentry.search.eap.spans.attributes import SPANS_INTERNAL_TO_PUBLIC_ALIAS_MAPPINGS
 from sentry.search.eap.types import SupportedTraceItemType
@@ -81,6 +83,16 @@ def transform_column_to_expression(column: Column) -> Expression:
     )
 
 
+def validate_sampling(sampling_mode: str | None) -> DownsampledStorageConfig:
+    if sampling_mode is None:
+        return DownsampledStorageConfig(mode=DownsampledStorageConfig.MODE_UNSPECIFIED)
+    sampling_mode = sampling_mode.upper()
+    if sampling_mode not in SAMPLING_MODES:
+        raise InvalidSearchQuery(f"sampling mode: {sampling_mode} is not supported")
+    else:
+        return DownsampledStorageConfig(mode=SAMPLING_MODES[sampling_mode])
+
+
 INTERNAL_TO_PUBLIC_ALIAS_MAPPINGS: dict[
     SupportedTraceItemType, dict[Literal["string", "number"], dict[str, str]]
 ] = {

src/sentry/snuba/ourlogs.py

@@ -58,6 +58,7 @@ def query(
         offset=offset or 0,
         limit=limit,
         referrer=referrer or "referrer unset",
+        sampling_mode=None,
         resolver=get_resolver(
             params=snuba_params,
             config=SearchResolverConfig(

src/sentry/snuba/rpc_dataset_common.py

@@ -43,11 +43,13 @@ def run_table_query(
     offset: int,
     limit: int,
     referrer: str,
+    sampling_mode: str | None,
     resolver: SearchResolver,
     debug: bool = False,
 ) -> EAPResponse:
     """Make the query"""
-    meta = resolver.resolve_meta(referrer=referrer)
+    sentry_sdk.set_tag("query.sampling_mode", sampling_mode)
+    meta = resolver.resolve_meta(referrer=referrer, sampling_mode=sampling_mode)
     where, having, query_contexts = resolver.resolve_query(query_string)
     columns, column_contexts = resolver.resolve_columns(selected_columns)
     contexts = resolver.resolve_contexts(query_contexts + column_contexts)
@@ -98,6 +100,7 @@ def run_table_query(
         virtual_column_contexts=[context for context in contexts if context is not None],
     )
     rpc_response = snuba_rpc.table_rpc([rpc_request])[0]
+    sentry_sdk.set_tag("query.storage_meta.tier", rpc_response.meta.downsampled_storage_meta.tier)
 
     """Process the results"""
     final_data: SnubaData = []

src/sentry/snuba/spans_rpc.py

@@ -84,6 +84,7 @@ def run_table_query(
     limit: int,
     referrer: str,
     config: SearchResolverConfig,
+    sampling_mode: str | None,
     search_resolver: SearchResolver | None = None,
     debug: bool = False,
 ) -> EAPResponse:
@@ -94,6 +95,7 @@ def run_table_query(
         offset,
         limit,
         referrer,
+        sampling_mode,
         search_resolver or get_resolver(params, config),
         debug,
     )
@@ -298,6 +300,7 @@ def run_top_events_timeseries_query(
         limit,
         referrer,
         config,
+        None,
         search_resolver,
     )
     if len(top_events["data"]) == 0:

src/sentry/snuba/uptime_checks.py

@@ -57,6 +57,7 @@ def query(
         offset=offset or 0,
         limit=limit,
         referrer=referrer or "referrer unset",
+        sampling_mode=None,
         resolver=get_resolver(
             params=snuba_params,
             config=SearchResolverConfig(

tests/snuba/api/endpoints/test_organization_events_span_indexed.py

@@ -8,6 +8,11 @@
 from sentry.testutils.helpers import parse_link_header
 from tests.snuba.api.endpoints.test_organization_events import OrganizationEventsEndpointTestBase
 
+# Downsampling is deterministic, so unless the algorithm changes we can find a known id that will appear in the
+# preflight and it will always show up
+# If we need to get a new ID just query for event ids after loading 100s of events and use any of the ids that come back
+KNOWN_PREFLIGHT_ID = "ca056dd858a24299"
+
 
 class OrganizationEventsSpanIndexedEndpointTest(OrganizationEventsEndpointTestBase):
     is_eap = False
@@ -3522,3 +3527,66 @@ def test_filtering_null_numeric_attr(self):
             },
         ]
         assert meta["dataset"] == self.dataset
+
+    def test_preflight_request(self):
+        span = self.create_span(
+            {"description": "foo", "sentry_tags": {"status": "success"}},
+            start_ts=self.ten_mins_ago,
+        )
+        span["span_id"] = KNOWN_PREFLIGHT_ID
+        span2 = self.create_span(
+            {"description": "zoo", "sentry_tags": {"status": "success"}},
+            start_ts=self.ten_mins_ago,
+        )
+        span2["span_id"] = "b" * 16
+        self.store_spans(
+            [span, span2],
+            is_eap=self.is_eap,
+        )
+        response = self.do_request(
+            {
+                "field": ["id", "description", "count()"],
+                "query": "",
+                "orderby": "description",
+                "project": self.project.id,
+                "dataset": self.dataset,
+                "statsPeriod": "1h",
+                "sampling": "PREFLIGHT",
+            }
+        )
+
+        assert response.status_code == 200, response.content
+        assert len(response.data["data"]) == 1
+        assert response.data["data"][0]["id"] == KNOWN_PREFLIGHT_ID
+
+    def test_best_effort_request(self):
+        span = self.create_span(
+            {"description": "foo", "sentry_tags": {"status": "success"}},
+            start_ts=self.ten_mins_ago,
+        )
+        span["span_id"] = KNOWN_PREFLIGHT_ID
+        span2 = self.create_span(
+            {"description": "zoo", "sentry_tags": {"status": "success"}},
+            start_ts=self.ten_mins_ago,
+        )
+        span2["span_id"] = "b" * 16
+        self.store_spans(
+            [span, span2],
+            is_eap=self.is_eap,
+        )
+        response = self.do_request(
+            {
+                "field": ["id", "description", "count()"],
+                "query": "",
+                "orderby": "description",
+                "project": self.project.id,
+                "dataset": self.dataset,
+                "statsPeriod": "1h",
+                "sampling": "BEST_EFFORT",
+            }
+        )
+
+        assert response.status_code == 200, response.content
+        assert len(response.data["data"]) == 2
+        assert response.data["data"][0]["id"] == KNOWN_PREFLIGHT_ID
+        assert response.data["data"][1]["id"] == "b" * 16