feat(get_traces): use cross item query path for single item query (#7656)

## Summary

Unifies single-item and cross-item query execution paths by using SQL
subqueries instead of fetching trace IDs in separate queries. This
simplifies the codebase and improves performance by letting ClickHouse
optimize the subquery internally.

## Changes

- **Removed** `get_trace_ids_for_cross_item_query()` function that
fetched trace IDs in a separate query
- **Enhanced** `get_trace_ids_sql_for_cross_item_query()` to return both
SQL string and query result for metadata extraction
- **Refactored** `EndpointGetTraces` to use subquery approach via
`_execute_with_subquery_optimization()`
- **Updated** time series and trace item table resolvers to inject SQL
subqueries using `DangerousRawSQL`
- **Added** feature flag `use_cross_item_path_for_single_item_queries`
to enable unified path for all queries (not just cross-item)
- **Added** test class that runs all existing tests with feature flag
enabled to verify parity

## Benefits

- **Simplified codebase**: Net -101 lines (221 deletions, 120 additions)
- **Unified execution path**: Single code path for both query types
- **Better performance**: ClickHouse can optimize the subquery
internally instead of two round-trips

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Author: volokluev

snuba/web/__init__.py

@@ -1,14 +1,14 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Any, Mapping, TypedDict, cast
+from typing import Any, Dict, Mapping, TypedDict, cast
 
 from snuba.reader import Column, Result, Row, transform_rows
 from snuba.utils.serializable_exception import JsonSerializable, SerializableException
 
 
 class QueryExtraData(TypedDict):
-    stats: Mapping[str, Any]
+    stats: Dict[str, Any]
     sql: str
     experiments: Mapping[str, Any]
 

snuba/web/rpc/v1/endpoint_get_traces.py

@@ -54,7 +54,6 @@
 from snuba.web.rpc.common.exceptions import BadSnubaRPCRequestException
 from snuba.web.rpc.v1.resolvers.common.cross_item_queries import (
     convert_trace_filters_to_trace_item_filter_with_type,
-    get_trace_ids_for_cross_item_query,
     get_trace_ids_sql_for_cross_item_query,
 )
 
@@ -473,11 +472,12 @@ def _execute_with_subquery_optimization(self, in_msg: GetTracesRequest) -> GetTr
         Execute cross-item query using subquery optimization.
         Gets SQL from trace IDs query and uses it as a subquery in metadata query.
         """
-        # Get SQL for trace IDs query (dry run)
-        trace_ids_sql = get_trace_ids_sql_for_cross_item_query(
+        # Get SQL for trace IDs query (dry run) and its query result
+        trace_ids_sql, trace_ids_query_result = get_trace_ids_sql_for_cross_item_query(
             in_msg,
             in_msg.meta,
             convert_trace_filters_to_trace_item_filter_with_type(list(in_msg.filters)),
+            self.routing_decision.tier,
             self._timer,
         )
 
@@ -486,13 +486,12 @@ def _execute_with_subquery_optimization(self, in_msg: GetTracesRequest) -> GetTr
             request=in_msg,
             trace_ids_sql=trace_ids_sql,
         )
-
-        # Build response
+        # Build response - include both query results for proper metadata extraction
         response_meta = extract_response_meta(
             in_msg.meta.request_id,
             in_msg.meta.debug,
-            [metadata_query_result],
-            [self._timer],
+            [trace_ids_query_result, metadata_query_result],
+            [self._timer, self._timer],
         )
 
         return GetTracesResponse(
@@ -503,25 +502,18 @@ def _execute_with_subquery_optimization(self, in_msg: GetTracesRequest) -> GetTr
 
     def _execute(self, in_msg: GetTracesRequest) -> GetTracesResponse:
         _validate_order_by(in_msg)
-        # Feature flag: Use subquery optimization for cross-item queries
-        if self._is_cross_event_query(in_msg.filters) and state.get_config(
-            "enable_cross_item_subquery_optimization", True
-        ):
-            return self._execute_with_subquery_optimization(in_msg)
+
+        # Feature flag: Use cross-item query path for all queries (single-item and cross-item)
+        use_cross_item_path = self._is_cross_event_query(in_msg.filters) or state.get_config(
+            "use_cross_item_path_for_single_item_queries", False
+        )
 
         # Original code path (unchanged)
         query_results: list[Any] = []
 
         # Get a dict of trace IDs and timestamps.
-        if self._is_cross_event_query(in_msg.filters):
-            trace_ids, trace_ids_query_results = get_trace_ids_for_cross_item_query(
-                in_msg,
-                in_msg.meta,
-                convert_trace_filters_to_trace_item_filter_with_type(list(in_msg.filters)),
-                self._timer,
-                return_query_results=True,
-            )
-            query_results.extend(trace_ids_query_results)
+        if use_cross_item_path:
+            return self._execute_with_subquery_optimization(in_msg)
         else:
             trace_ids, trace_ids_query_result = self._get_trace_ids_for_single_item_query(
                 request=in_msg

snuba/web/rpc/v1/resolvers/R_eap_items/resolver_time_series.py

@@ -27,11 +27,12 @@
 from snuba.datasets.entities.entity_key import EntityKey
 from snuba.datasets.entities.factory import get_entity
 from snuba.datasets.pluggable_dataset import PluggableDataset
+from snuba.downsampled_storage_tiers import Tier
 from snuba.query import OrderBy, OrderByDirection, SelectedExpression
 from snuba.query.data_source.simple import Entity
 from snuba.query.dsl import Functions as f
-from snuba.query.dsl import column, in_cond, literal, literals_array
-from snuba.query.expressions import Expression
+from snuba.query.dsl import column, in_cond, literal
+from snuba.query.expressions import DangerousRawSQL, Expression
 from snuba.query.logical import Query
 from snuba.query.query_settings import HTTPQuerySettings
 from snuba.request import Request as SnubaRequest
@@ -62,7 +63,7 @@
     get_count_column,
 )
 from snuba.web.rpc.v1.resolvers.common.cross_item_queries import (
-    get_trace_ids_for_cross_item_query,
+    get_trace_ids_sql_for_cross_item_query,
 )
 from snuba.web.rpc.v1.resolvers.common.formula_reliability import (
     FormulaReliabilityCalculator,
@@ -317,7 +318,9 @@ def _proto_expression_to_ast_expression(
             raise ValueError(f"Unknown expression type: {default}")
 
 
-def build_query(request: TimeSeriesRequest, timer: Optional[Timer] = None) -> Query:
+def build_query(
+    request: TimeSeriesRequest, sampling_tier: Optional[Tier] = None, timer: Optional[Timer] = None
+) -> Query:
     entity = Entity(
         key=EntityKey("eap_items"),
         schema=get_entity(EntityKey("eap_items")).get_data_model(),
@@ -347,14 +350,14 @@ def build_query(request: TimeSeriesRequest, timer: Optional[Timer] = None) -> Qu
 
     # Handle cross item queries by first getting trace IDs
     additional_conditions = []
-    if request.trace_filters and timer is not None:
-        trace_ids = get_trace_ids_for_cross_item_query(
-            request, request.meta, list(request.trace_filters), timer
+    if request.trace_filters and timer is not None and sampling_tier is not None:
+        trace_ids_sql, _ = get_trace_ids_sql_for_cross_item_query(
+            request, request.meta, list(request.trace_filters), sampling_tier, timer
         )
         additional_conditions.append(
             in_cond(
                 column("trace_id"),
-                literals_array(None, [literal(trace_id) for trace_id in trace_ids]),
+                DangerousRawSQL(None, f"({trace_ids_sql})"),
             )
         )
 
@@ -416,7 +419,10 @@ def build_query(request: TimeSeriesRequest, timer: Optional[Timer] = None) -> Qu
 
 
 def _build_snuba_request(
-    request: TimeSeriesRequest, query_settings: HTTPQuerySettings, timer: Optional[Timer] = None
+    request: TimeSeriesRequest,
+    query_settings: HTTPQuerySettings,
+    sampling_tier: Optional[Tier] = None,
+    timer: Optional[Timer] = None,
 ) -> SnubaRequest:
     if request.meta.trace_item_type == TraceItemType.TRACE_ITEM_TYPE_LOG:
         team = "ourlogs"
@@ -430,7 +436,7 @@ def _build_snuba_request(
     return SnubaRequest(
         id=uuid.UUID(request.meta.request_id),
         original_body=MessageToDict(request),
-        query=build_query(request, timer),
+        query=build_query(request, sampling_tier, timer),
         query_settings=query_settings,
         attribution_info=AttributionInfo(
             referrer=request.meta.referrer,
@@ -467,7 +473,9 @@ def resolve(
         except Exception as e:
             sentry_sdk.capture_message(f"Error merging clickhouse settings: {e}")
 
-        snuba_request = _build_snuba_request(in_msg, query_settings, self._timer)
+        snuba_request = _build_snuba_request(
+            in_msg, query_settings, routing_decision.tier, self._timer
+        )
         res = run_query(
             dataset=PluggableDataset(name="eap", all_entities=[]),
             request=snuba_request,

snuba/web/rpc/v1/resolvers/R_eap_items/resolver_trace_item_table.py

@@ -29,13 +29,14 @@
 from snuba.datasets.entities.entity_key import EntityKey
 from snuba.datasets.entities.factory import get_entity
 from snuba.datasets.pluggable_dataset import PluggableDataset
+from snuba.downsampled_storage_tiers import Tier
 from snuba.protos.common import NORMALIZED_COLUMNS_EAP_ITEMS
 from snuba.query import OrderBy, OrderByDirection, SelectedExpression
 from snuba.query.data_source.simple import Entity
 from snuba.query.dsl import Functions as f
 from snuba.query.dsl import and_cond, in_cond, literal, literals_array, or_cond
 from snuba.query.dsl import column as snuba_column
-from snuba.query.expressions import Expression, SubscriptableReference
+from snuba.query.expressions import DangerousRawSQL, Expression, SubscriptableReference
 from snuba.query.logical import Query
 from snuba.query.query_settings import HTTPQuerySettings
 from snuba.request import Request as SnubaRequest
@@ -67,7 +68,7 @@
     get_count_column,
 )
 from snuba.web.rpc.v1.resolvers.common.cross_item_queries import (
-    get_trace_ids_for_cross_item_query,
+    get_trace_ids_sql_for_cross_item_query,
 )
 from snuba.web.rpc.v1.resolvers.common.trace_item_table import convert_results
 
@@ -404,6 +405,7 @@ def _get_offset_from_page_token(page_token: PageToken | None) -> int:
 def build_query(
     request: TraceItemTableRequest,
     time_window: TimeWindow | None = None,
+    sampling_tier: Optional[Tier] = None,
     timer: Optional[Timer] = None,
 ) -> Query:
     entity = Entity(
@@ -429,15 +431,12 @@ def build_query(
 
     # Handle cross item queries by first getting trace IDs
     additional_conditions: List[Expression] = []
-    if request.trace_filters and timer is not None:
-        trace_ids = get_trace_ids_for_cross_item_query(
-            request, request.meta, list(request.trace_filters), timer
+    if request.trace_filters and timer is not None and sampling_tier is not None:
+        trace_ids_sql, _ = get_trace_ids_sql_for_cross_item_query(
+            request, request.meta, list(request.trace_filters), sampling_tier, timer
         )
         additional_conditions.append(
-            in_cond(
-                snuba_column("trace_id"),
-                literals_array(None, [literal(trace_id) for trace_id in trace_ids]),
-            )
+            in_cond(snuba_column("trace_id"), DangerousRawSQL(None, f"({trace_ids_sql})"))
         )
     if time_window is not None:
         additional_conditions.append(
@@ -528,6 +527,7 @@ def _build_snuba_request(
     request: TraceItemTableRequest,
     query_settings: HTTPQuerySettings,
     time_window: TimeWindow | None = None,
+    sampling_tier: Optional[Tier] = None,
     timer: Optional[Timer] = None,
 ) -> SnubaRequest:
     if request.meta.trace_item_type == TraceItemType.TRACE_ITEM_TYPE_LOG:
@@ -542,7 +542,7 @@ def _build_snuba_request(
     return SnubaRequest(
         id=uuid.UUID(request.meta.request_id),
         original_body=MessageToDict(request),
-        query=build_query(request, time_window, timer),
+        query=build_query(request, time_window, sampling_tier, timer),
         query_settings=query_settings,
         attribution_info=AttributionInfo(
             referrer=request.meta.referrer,
@@ -582,7 +582,7 @@ def resolve(
             start_timestamp=in_msg.meta.start_timestamp, end_timestamp=in_msg.meta.end_timestamp
         )
         snuba_request = _build_snuba_request(
-            in_msg, query_settings, routing_decision.time_window, self._timer
+            in_msg, query_settings, routing_decision.time_window, routing_decision.tier, self._timer
         )
         res = run_query(
             dataset=PluggableDataset(name="eap", all_entities=[]),