feat(cross-item-queres): add cross item support in timeseries endpoint (#7400)

Author: davidtsuk

setup.cfg

@@ -7,6 +7,7 @@ markers =
     clickhouse_db: Use clickhouse
     redis_db: Use redis
     ci_only: Only run in CI
+    eap: Use clickhouse with EAP migrations only
 
 [flake8]
 # tests/state/test_state.py:19:36: E712 comparison to True should be 'if cond is True:' or 'if cond:'

snuba/subscriptions/data.py

@@ -121,18 +121,12 @@ class _SubscriptionData(ABC, Generic[TRequest]):
 
     def validate(self) -> None:
         if self.time_window_sec < 60:
-            raise InvalidSubscriptionError(
-                "Time window must be greater than or equal to 1 minute"
-            )
+            raise InvalidSubscriptionError("Time window must be greater than or equal to 1 minute")
         elif self.time_window_sec > 60 * 60 * 24:
-            raise InvalidSubscriptionError(
-                "Time window must be less than or equal to 24 hours"
-            )
+            raise InvalidSubscriptionError("Time window must be less than or equal to 24 hours")
 
         if self.resolution_sec < 60:
-            raise InvalidSubscriptionError(
-                "Resolution must be greater than or equal to 1 minute"
-            )
+            raise InvalidSubscriptionError("Resolution must be greater than or equal to 1 minute")
 
     @abstractmethod
     def build_request(
@@ -238,9 +232,7 @@ def build_request(
         request_class.ParseFromString(base64.b64decode(self.time_series_request))
 
         start_time_proto = Timestamp()
-        start_time_proto.FromDatetime(
-            timestamp - timedelta(seconds=self.time_window_sec)
-        )
+        start_time_proto.FromDatetime(timestamp - timedelta(seconds=self.time_window_sec))
         end_time_proto = Timestamp()
         end_time_proto.FromDatetime(timestamp)
         request_class.meta.start_timestamp.CopyFrom(start_time_proto)
@@ -261,28 +253,24 @@ def run_query(
         concurrent_queries_gauge: Optional[Gauge] = None,
     ) -> QueryResult:
         response = EndpointTimeSeries().execute(request)
-        if not response.result_timeseries:
+        if not response.result_timeseries or not any(
+            dp.data_present for dp in response.result_timeseries[0].data_points
+        ):
             result: Result = {
                 "meta": [],
                 "data": [{request.expressions[0].label: None}],
                 "trace_output": "",
             }
-            return QueryResult(
-                result=result, extra={"stats": {}, "sql": "", "experiments": {}}
-            )
+            return QueryResult(result=result, extra={"stats": {}, "sql": "", "experiments": {}})
 
         timeseries = response.result_timeseries[0]
         data = [{timeseries.label: timeseries.data_points[0].data}]
 
         result = {"meta": [], "data": data, "trace_output": ""}
-        return QueryResult(
-            result=result, extra={"stats": {}, "sql": "", "experiments": {}}
-        )
+        return QueryResult(result=result, extra={"stats": {}, "sql": "", "experiments": {}})
 
     @classmethod
-    def from_dict(
-        cls, data: Mapping[str, Any], entity_key: EntityKey
-    ) -> RPCSubscriptionData:
+    def from_dict(cls, data: Mapping[str, Any], entity_key: EntityKey) -> RPCSubscriptionData:
         entity: Entity = get_entity(entity_key)
         metadata = {}
         for key in data.keys():
@@ -370,9 +358,7 @@ def add_conditions(
         elif isinstance(from_clause, EntityDS):
             entities = [(None, get_entity(from_clause.key))]
         else:
-            raise InvalidSubscriptionError(
-                "Only simple queries and join queries are supported"
-            )
+            raise InvalidSubscriptionError("Only simple queries and join queries are supported")
         for entity_alias, entity in entities:
             conditions_to_add: List[Expression] = [
                 binary_condition(
@@ -406,9 +392,7 @@ def add_conditions(
             new_condition = combine_and_conditions(conditions_to_add)
             condition = query.get_condition()
             if condition:
-                new_condition = binary_condition(
-                    BooleanFunctions.AND, condition, new_condition
-                )
+                new_condition = binary_condition(BooleanFunctions.AND, condition, new_condition)
 
             query.set_ast_condition(new_condition)
 
@@ -480,9 +464,7 @@ def run_query(
         )
 
     @classmethod
-    def from_dict(
-        cls, data: Mapping[str, Any], entity_key: EntityKey
-    ) -> SnQLSubscriptionData:
+    def from_dict(cls, data: Mapping[str, Any], entity_key: EntityKey) -> SnQLSubscriptionData:
         entity: Entity = get_entity(entity_key)
 
         metadata = {}

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

@@ -2,7 +2,7 @@
 from collections import defaultdict
 from dataclasses import replace
 from datetime import datetime
-from typing import Any, Callable, Dict, Iterable
+from typing import Any, Callable, Dict, Iterable, Optional
 
 import sentry_sdk
 from google.protobuf.json_format import MessageToDict
@@ -30,13 +30,14 @@
 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, literal
+from snuba.query.dsl import column, in_cond, literal, literals_array
 from snuba.query.expressions import Expression
 from snuba.query.logical import Query
 from snuba.query.query_settings import HTTPQuerySettings
 from snuba.request import Request as SnubaRequest
 from snuba.settings import ENABLE_FORMULA_RELIABILITY_DEFAULT
 from snuba.state import get_int_config
+from snuba.utils.metrics.timer import Timer
 from snuba.web.query import run_query
 from snuba.web.rpc.common.common import (
     add_existence_check_to_subscriptable_references,
@@ -61,6 +62,9 @@
     get_confidence_interval_column,
     get_count_column,
 )
+from snuba.web.rpc.v1.resolvers.common.cross_item_queries import (
+    get_trace_ids_for_cross_item_query,
+)
 from snuba.web.rpc.v1.resolvers.common.formula_reliability import (
     FormulaReliabilityCalculator,
 )
@@ -141,13 +145,11 @@ def _convert_result_timeseries(
     #       time_converted_to_integer_timestamp: row_data_for_that_time_bucket
     #   }
     # }
-    result_timeseries_timestamp_to_row: defaultdict[
-        tuple[str, str], dict[int, Dict[str, Any]]
-    ] = defaultdict(dict)
-
-    query_duration = (
-        request.meta.end_timestamp.seconds - request.meta.start_timestamp.seconds
+    result_timeseries_timestamp_to_row: defaultdict[tuple[str, str], dict[int, Dict[str, Any]]] = (
+        defaultdict(dict)
     )
+
+    query_duration = request.meta.end_timestamp.seconds - request.meta.start_timestamp.seconds
     time_buckets = [
         Timestamp(seconds=(request.meta.start_timestamp.seconds) + secs)
         for secs in range(0, query_duration, request.granularity_secs)
@@ -183,9 +185,7 @@ def _convert_result_timeseries(
             if not row_data:
                 timeseries.data_points.append(DataPoint(data=0, data_present=False))
             else:
-                extrapolation_context = ExtrapolationContext.from_row(
-                    timeseries.label, row_data
-                )
+                extrapolation_context = ExtrapolationContext.from_row(timeseries.label, row_data)
                 if row_data.get(timeseries.label, None) is not None:
                     timeseries.data_points.append(
                         DataPoint(
@@ -199,9 +199,7 @@ def _convert_result_timeseries(
                 else:
                     timeseries.data_points.append(DataPoint(data=0, data_present=False))
 
-    if get_int_config(
-        "enable_formula_reliability_ts", ENABLE_FORMULA_RELIABILITY_DEFAULT
-    ):
+    if get_int_config("enable_formula_reliability_ts", ENABLE_FORMULA_RELIABILITY_DEFAULT):
         frc = FormulaReliabilityCalculator(request, data, time_buckets)
         for timeseries in result_timeseries.values():
             if timeseries.label in frc:
@@ -240,10 +238,7 @@ def _get_reliability_context_columns(
         which_oneof = expr.WhichOneof("expression")
         assert which_oneof in ["conditional_aggregation", "aggregation"]
         aggregation = getattr(expr, which_oneof)
-        if (
-            aggregation.extrapolation_mode
-            == ExtrapolationMode.EXTRAPOLATION_MODE_SAMPLE_WEIGHTED
-        ):
+        if aggregation.extrapolation_mode == ExtrapolationMode.EXTRAPOLATION_MODE_SAMPLE_WEIGHTED:
             confidence_interval_column = get_confidence_interval_column(
                 aggregation, _get_attribute_key_to_expression_function(request_meta)
             )
@@ -271,9 +266,7 @@ def _get_reliability_context_columns(
             SelectedExpression(name=count_column.alias, expression=count_column)
         )
     elif expr.WhichOneof("expression") == "formula":
-        if not get_int_config(
-            "enable_formula_reliability_ts", ENABLE_FORMULA_RELIABILITY_DEFAULT
-        ):
+        if not get_int_config("enable_formula_reliability_ts", ENABLE_FORMULA_RELIABILITY_DEFAULT):
             return []
         # also query for the left and right parts of the formula separately
         # this will be used later to calculate the reliability of the formula
@@ -286,9 +279,7 @@ def _get_reliability_context_columns(
                         expression=_proto_expression_to_ast_expression(e, request_meta),
                     )
                 )
-            additional_context_columns.extend(
-                _get_reliability_context_columns(e, request_meta)
-            )
+            additional_context_columns.extend(_get_reliability_context_columns(e, request_meta))
     return additional_context_columns
 
 
@@ -311,13 +302,9 @@ def _proto_expression_to_ast_expression(
                 case None:
                     pass
                 case "default_value_double":
-                    formula_expr = f.coalesce(
-                        formula_expr, expr.formula.default_value_double
-                    )
+                    formula_expr = f.coalesce(formula_expr, expr.formula.default_value_double)
                 case "default_value_int64":
-                    formula_expr = f.coalesce(
-                        formula_expr, expr.formula.default_value_int64
-                    )
+                    formula_expr = f.coalesce(formula_expr, expr.formula.default_value_int64)
                 case default:
                     raise BadSnubaRPCRequestException(
                         f"Unknown default_value in formula. Expected default_value_double or default_value_int64 but got {default}"
@@ -329,7 +316,7 @@ def _proto_expression_to_ast_expression(
             raise ValueError(f"Unknown expression type: {default}")
 
 
-def build_query(request: TimeSeriesRequest) -> Query:
+def build_query(request: TimeSeriesRequest, timer: Optional[Timer] = None) -> Query:
     entity = Entity(
         key=EntityKey("eap_items"),
         schema=get_entity(EntityKey("eap_items")).get_data_model(),
@@ -346,21 +333,30 @@ def build_query(request: TimeSeriesRequest) -> Query:
 
     additional_context_columns = []
     for expr in request.expressions:
-        additional_context_columns.extend(
-            _get_reliability_context_columns(expr, request.meta)
-        )
+        additional_context_columns.extend(_get_reliability_context_columns(expr, request.meta))
 
     groupby_columns = [
         SelectedExpression(
             name=attr_key.name,
-            expression=_get_attribute_key_to_expression_function(request.meta)(
-                attr_key
-            ),
+            expression=_get_attribute_key_to_expression_function(request.meta)(attr_key),
         )
         for attr_key in request.group_by
     ]
     item_type_conds = [f.equals(column("item_type"), request.meta.trace_item_type)]
 
+    # 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
+        )
+        additional_conditions.append(
+            in_cond(
+                column("trace_id"),
+                literals_array(None, [literal(trace_id) for trace_id in trace_ids]),
+            )
+        )
+
     res = Query(
         from_clause=entity,
         selected_columns=[
@@ -402,6 +398,7 @@ def build_query(request: TimeSeriesRequest) -> Query:
                 request.filter, _get_attribute_key_to_expression_function(request.meta)
             ),
             *item_type_conds,
+            *additional_conditions,
         ),
         groupby=[
             column("time_slot"),
@@ -410,17 +407,15 @@ def build_query(request: TimeSeriesRequest) -> Query:
                 for attr_key in request.group_by
             ],
         ],
-        order_by=[
-            OrderBy(expression=column("time_slot"), direction=OrderByDirection.ASC)
-        ],
+        order_by=[OrderBy(expression=column("time_slot"), direction=OrderByDirection.ASC)],
     )
     treeify_or_and_conditions(res)
     add_existence_check_to_subscriptable_references(res)
     return res
 
 
 def _build_snuba_request(
-    request: TimeSeriesRequest, query_settings: HTTPQuerySettings
+    request: TimeSeriesRequest, query_settings: HTTPQuerySettings, timer: Optional[Timer] = None
 ) -> SnubaRequest:
     if request.meta.trace_item_type == TraceItemType.TRACE_ITEM_TYPE_LOG:
         team = "ourlogs"
@@ -434,7 +429,7 @@ def _build_snuba_request(
     return SnubaRequest(
         id=uuid.UUID(request.meta.request_id),
         original_body=MessageToDict(request),
-        query=build_query(request),
+        query=build_query(request, timer),
         query_settings=query_settings,
         attribution_info=AttributionInfo(
             referrer=request.meta.referrer,
@@ -464,18 +459,14 @@ def resolve(
         # if the user passes it in
         assert len(in_msg.aggregations) == 0
 
-        query_settings = (
-            setup_trace_query_settings() if in_msg.meta.debug else HTTPQuerySettings()
-        )
+        query_settings = setup_trace_query_settings() if in_msg.meta.debug else HTTPQuerySettings()
         try:
-            routing_decision.strategy.merge_clickhouse_settings(
-                routing_decision, query_settings
-            )
+            routing_decision.strategy.merge_clickhouse_settings(routing_decision, query_settings)
             query_settings.set_sampling_tier(routing_decision.tier)
         except Exception as e:
             sentry_sdk.capture_message(f"Error merging clickhouse settings: {e}")
 
-        snuba_request = _build_snuba_request(in_msg, query_settings)
+        snuba_request = _build_snuba_request(in_msg, query_settings, self._timer)
         res = run_query(
             dataset=PluggableDataset(name="eap", all_entities=[]),
             request=snuba_request,

tests/web/rpc/v1/conftest.py

@@ -0,0 +1,46 @@
+from typing import Generator
+
+import pytest
+
+from snuba.datasets.factory import reset_dataset_factory
+from snuba.migrations.groups import MigrationGroup
+from snuba.migrations.runner import Runner
+
+
+@pytest.fixture
+def eap(request: pytest.FixtureRequest, create_databases: None) -> Generator[None, None, None]:
+    """
+    A custom ClickHouse fixture that only runs EAP (Events Analytics Platform) migrations and Outcomes migrations (for storage routing).
+    This is much faster than running all migrations for tests that only need EAP tables.
+
+    Use this with @pytest.mark.eap marker.
+    """
+    if not request.node.get_closest_marker("eap"):
+        pytest.fail("Need to use eap marker if eap fixture is used")
+
+    try:
+        reset_dataset_factory()
+        # Run only SYSTEM migrations (required for migrations table) and EAP migrations
+        runner = Runner()
+        runner.run_all(group=MigrationGroup.EVENTS_ANALYTICS_PLATFORM, force=True)
+        runner.run_all(group=MigrationGroup.OUTCOMES, force=True)
+        yield
+    finally:
+        # Import here to avoid circular imports
+        from tests.conftest import _clear_db
+
+        _clear_db()
+
+
+# Hook to modify test collection
+def pytest_runtest_setup(item: pytest.Item) -> None:
+    """Custom setup to handle eap marker."""
+    if item.get_closest_marker("eap"):
+        # Remove block_clickhouse_db if it was added by parent conftest
+        fixturenames = getattr(item, "fixturenames", None)
+        if fixturenames is not None:
+            if "block_clickhouse_db" in fixturenames:
+                fixturenames.remove("block_clickhouse_db")
+            # Add our custom fixture if not already present
+            if "eap" not in fixturenames:
+                fixturenames.append("eap")