feat(eap): Return bytes scanned for eap queries (#102075)

This returns bytes scanned for all eap queries.

Author: Zylphrex

src/sentry/api/bases/organization_events.py

@@ -370,6 +370,7 @@ def handle_results_with_meta(
                 isMetricsExtractedData = meta.pop("isMetricsExtractedData", False)
                 discoverSplitDecision = meta.pop("discoverSplitDecision", None)
                 full_scan = meta.pop("full_scan", None)
+                bytes_scanned = meta.pop("bytes_scanned", None)
                 debug_info = meta.pop("debug_info", None)
                 fields, units = self.handle_unit_meta(fields_meta)
                 meta = {
@@ -392,6 +393,9 @@ def handle_results_with_meta(
                     # If this key isn't in meta there wasn't any sampling and we can assume all the data was scanned
                     meta["dataScanned"] = "full"
 
+                if bytes_scanned is not None:
+                    meta["bytesScanned"] = bytes_scanned
+
                 # Only appears in meta when debug is passed to the endpoint
                 if debug_info:
                     meta["debug_info"] = debug_info

src/sentry/search/eap/sampling.py

@@ -2,10 +2,11 @@
     DownsampledStorageConfig,
     DownsampledStorageMeta,
 )
+from sentry_protos.snuba.v1.request_common_pb2 import ResponseMeta
 
 from sentry.exceptions import InvalidSearchQuery
 from sentry.search.eap.constants import SAMPLING_MODE_MAP
-from sentry.search.events.types import SAMPLING_MODES
+from sentry.search.events.types import SAMPLING_MODES, EventsMeta
 
 
 def handle_downsample_meta(meta: DownsampledStorageMeta) -> bool:
@@ -19,3 +20,13 @@ def validate_sampling(sampling_mode: SAMPLING_MODES | None) -> DownsampledStorag
         raise InvalidSearchQuery(f"sampling mode: {sampling_mode} is not supported")
     else:
         return DownsampledStorageConfig(mode=SAMPLING_MODE_MAP[sampling_mode])
+
+
+def events_meta_from_rpc_request_meta(meta: ResponseMeta) -> EventsMeta:
+    return EventsMeta(
+        fields={},
+        full_scan=handle_downsample_meta(meta.downsampled_storage_meta),
+        bytes_scanned=(
+            sum(info.stats.progress_bytes for info in meta.query_info) if meta.query_info else None
+        ),
+    )

src/sentry/search/events/types.py

@@ -74,6 +74,7 @@ class EventsMeta(TypedDict):
     # only returned when debug=True
     debug_info: NotRequired[dict[str, Any]]
     full_scan: NotRequired[bool]
+    bytes_scanned: NotRequired[int | None]
 
 
 class EventsResponse(TypedDict):

src/sentry/snuba/ourlogs.py

@@ -7,7 +7,7 @@
 from sentry.search.eap import constants
 from sentry.search.eap.ourlogs.definitions import OURLOG_DEFINITIONS
 from sentry.search.eap.resolver import SearchResolver
-from sentry.search.eap.sampling import handle_downsample_meta
+from sentry.search.eap.sampling import events_meta_from_rpc_request_meta
 from sentry.search.eap.types import AdditionalQueries, EAPResponse, SearchResolverConfig
 from sentry.search.events.types import SAMPLING_MODES, EventsMeta, SnubaParams
 from sentry.snuba import rpc_dataset_common
@@ -112,10 +112,7 @@ def run_timeseries_query(
 
         """Process the results"""
         result = rpc_dataset_common.ProcessedTimeseries()
-        final_meta: EventsMeta = EventsMeta(
-            fields={},
-            full_scan=handle_downsample_meta(rpc_response.meta.downsampled_storage_meta),
-        )
+        final_meta: EventsMeta = events_meta_from_rpc_request_meta(rpc_response.meta)
         for resolved_field in aggregates + groupbys:
             final_meta["fields"][resolved_field.public_alias] = resolved_field.search_type
 

src/sentry/snuba/rpc_dataset_common.py

@@ -47,7 +47,7 @@
 )
 from sentry.search.eap.constants import DOUBLE, MAX_ROLLUP_POINTS, VALID_GRANULARITIES
 from sentry.search.eap.resolver import SearchResolver
-from sentry.search.eap.sampling import handle_downsample_meta
+from sentry.search.eap.sampling import events_meta_from_rpc_request_meta
 from sentry.search.eap.types import (
     CONFIDENCES,
     AdditionalQueries,
@@ -364,10 +364,7 @@ def process_table_response(
         """Process the results"""
         final_data: SnubaData = []
         final_confidence: ConfidenceData = []
-        final_meta: EventsMeta = EventsMeta(
-            fields={},
-            full_scan=handle_downsample_meta(rpc_response.meta.downsampled_storage_meta),
-        )
+        final_meta: EventsMeta = events_meta_from_rpc_request_meta(rpc_response.meta)
         # Mapping from public alias to resolved column so we know type etc.
         columns_by_name = {col.public_alias: col for col in table_request.columns}
 
@@ -767,10 +764,7 @@ def run_top_events_timeseries_query(
         """Process the results"""
         map_result_key_to_timeseries = defaultdict(list)
 
-        final_meta: EventsMeta = EventsMeta(
-            fields={},
-            full_scan=handle_downsample_meta(rpc_response.meta.downsampled_storage_meta),
-        )
+        final_meta: EventsMeta = events_meta_from_rpc_request_meta(rpc_response.meta)
 
         if params.debug:
             set_debug_meta(final_meta, rpc_response.meta, rpc_request)

src/sentry/snuba/spans_rpc.py

@@ -9,7 +9,7 @@
 from sentry.exceptions import InvalidSearchQuery
 from sentry.search.eap.constants import DOUBLE, INT, STRING
 from sentry.search.eap.resolver import SearchResolver
-from sentry.search.eap.sampling import handle_downsample_meta
+from sentry.search.eap.sampling import events_meta_from_rpc_request_meta
 from sentry.search.eap.spans.definitions import SPAN_DEFINITIONS
 from sentry.search.eap.types import AdditionalQueries, EAPResponse, SearchResolverConfig
 from sentry.search.events.types import SAMPLING_MODES, EventsMeta, SnubaParams
@@ -91,10 +91,7 @@ def run_timeseries_query(
         rpc_response = snuba_rpc.timeseries_rpc([rpc_request])[0]
         """Process the results"""
         result = rpc_dataset_common.ProcessedTimeseries()
-        final_meta: EventsMeta = EventsMeta(
-            fields={},
-            full_scan=handle_downsample_meta(rpc_response.meta.downsampled_storage_meta),
-        )
+        final_meta: EventsMeta = events_meta_from_rpc_request_meta(rpc_response.meta)
         if params.debug:
             set_debug_meta(final_meta, rpc_response.meta, rpc_request)
         for resolved_field in aggregates + groupbys:

src/sentry/snuba/trace_metrics.py

@@ -6,7 +6,7 @@
 
 from sentry.search.eap import constants
 from sentry.search.eap.resolver import SearchResolver
-from sentry.search.eap.sampling import handle_downsample_meta
+from sentry.search.eap.sampling import events_meta_from_rpc_request_meta
 from sentry.search.eap.trace_metrics.definitions import TRACE_METRICS_DEFINITIONS
 from sentry.search.eap.types import AdditionalQueries, EAPResponse, SearchResolverConfig
 from sentry.search.events.types import SAMPLING_MODES, EventsMeta, SnubaParams
@@ -104,10 +104,7 @@ def run_timeseries_query(
 
         """Process the results"""
         result = rpc_dataset_common.ProcessedTimeseries()
-        final_meta: EventsMeta = EventsMeta(
-            fields={},
-            full_scan=handle_downsample_meta(rpc_response.meta.downsampled_storage_meta),
-        )
+        final_meta: EventsMeta = events_meta_from_rpc_request_meta(rpc_response.meta)
         for resolved_field in aggregates + groupbys:
             final_meta["fields"][resolved_field.public_alias] = resolved_field.search_type
 

tests/snuba/api/endpoints/test_organization_events_ourlogs.py

@@ -888,3 +888,17 @@ def test_high_accuracy_flex_time_without_feature_flag(self):
 
         response = self.do_request(request)
         assert response.status_code == 400
+
+    def test_bytes_scanned(self):
+        self.store_ourlogs([self.create_ourlog({"body": "log"}, timestamp=self.ten_mins_ago)])
+
+        request = {
+            "field": ["timestamp", "message"],
+            "orderby": "-timestamp",
+            "project": self.project.id,
+            "dataset": self.dataset,
+        }
+
+        response = self.do_request(request)
+        assert response.status_code == 200
+        assert response.data["meta"]["bytesScanned"] > 0

tests/snuba/api/endpoints/test_organization_events_span_indexed.py

@@ -684,6 +684,7 @@ def _test_simple_measurements(self, keys):
             )
             assert response.status_code == 200, response.content
             expected = {
+                "bytesScanned": mock.ANY,
                 "dataScanned": "full",
                 "dataset": mock.ANY,
                 "datasetReason": "unchanged",
@@ -1902,6 +1903,7 @@ def test_span_data_fields_http_resource(self) -> None:
             },
         ]
         expected = {
+            "bytesScanned": mock.ANY,
             "dataScanned": "full",
             "dataset": mock.ANY,
             "datasetReason": "unchanged",