feat(Replay): Add Function to Query EAP for Replay Details (#102416)

Adding helper function to allow for query EAP for replay details and
added unit test for function. This helper function is not being called
yet (planning on adding a flag to gate)

Relates to: REPLAY-824

Author: cliffordxing

src/sentry/replays/endpoints/organization_replay_details.py

@@ -1,8 +1,10 @@
 import uuid
+from datetime import datetime
 
 from drf_spectacular.utils import extend_schema
 from rest_framework.request import Request
 from rest_framework.response import Response
+from snuba_sdk import Column, Condition, Entity, Function, Granularity, Op, Query
 
 from sentry import features
 from sentry.api.api_owners import ApiOwner
@@ -15,11 +17,98 @@
 from sentry.apidocs.utils import inline_sentry_response_serializer
 from sentry.constants import ALL_ACCESS_PROJECTS
 from sentry.models.organization import Organization
+from sentry.replays.lib.eap import read as eap_read
+from sentry.replays.lib.eap.snuba_transpiler import RequestMeta, Settings
 from sentry.replays.post_process import ReplayDetailsResponse, process_raw_response
 from sentry.replays.query import query_replay_instance
 from sentry.replays.validators import ReplayValidator
 
 
+def query_replay_instance_eap(
+    project_ids: list[int],
+    replay_ids: list[str],
+    start: datetime,
+    end: datetime,
+    organization_id: int,
+    request_user_id: int,
+    referrer: str = "replays.query.details_query",
+):
+    select = [
+        Column("replay_id"),
+        Function("min", parameters=[Column("project_id")], alias="agg_project_id"),
+        Function("min", parameters=[Column("replay_start_timestamp")], alias="started_at"),
+        Function("max", parameters=[Column("timestamp")], alias="finished_at"),
+        Function("count", parameters=[Column("segment_id")], alias="count_segments"),
+        Function("sum", parameters=[Column("count_error_events")], alias="count_errors"),
+        Function("sum", parameters=[Column("count_warning_events")], alias="count_warnings"),
+        Function("sum", parameters=[Column("count_info_events")], alias="count_infos"),
+        Function(
+            "sumIf",
+            parameters=[
+                Column("click_is_dead"),
+                Function(
+                    "greaterOrEquals",
+                    [Column("timestamp"), int(datetime(year=2023, month=7, day=24).timestamp())],
+                ),
+            ],
+            alias="count_dead_clicks",
+        ),
+        Function(
+            "sumIf",
+            parameters=[
+                Column("click_is_rage"),
+                Function(
+                    "greaterOrEquals",
+                    [Column("timestamp"), int(datetime(year=2023, month=7, day=24).timestamp())],
+                ),
+            ],
+            alias="count_rage_clicks",
+        ),
+        Function("max", parameters=[Column("is_archived")], alias="is_archived"),
+    ]
+
+    snuba_query = Query(
+        match=Entity("replays"),
+        select=select,
+        where=[
+            Condition(Column("replay_id"), Op.IN, replay_ids),
+        ],
+        groupby=[Column("replay_id")],
+        granularity=Granularity(3600),
+    )
+
+    settings = Settings(
+        attribute_types={
+            "replay_id": str,
+            "project_id": int,
+            "timestamp": int,
+            "replay_start_timestamp": int,
+            "segment_id": int,
+            "is_archived": int,
+            "count_error_events": int,
+            "count_warning_events": int,
+            "count_info_events": int,
+            "click_is_dead": int,
+            "click_is_rage": int,
+        },
+        default_limit=100,
+        default_offset=0,
+    )
+
+    request_meta = RequestMeta(
+        cogs_category="replays",
+        debug=False,
+        start_datetime=start,
+        end_datetime=end,
+        organization_id=organization_id,
+        project_ids=project_ids,
+        referrer=referrer,
+        request_id=str(uuid.uuid4().hex),
+        trace_item_type="replay",
+    )
+    return eap_read.query(snuba_query, settings, request_meta, [])
+
+
 @region_silo_endpoint
 @extend_schema(tags=["Replays"])
 class OrganizationReplayDetailsEndpoint(OrganizationEndpoint):

src/sentry/replays/usecases/ingest/event_parser.py

@@ -366,6 +366,7 @@ def as_trace_item(
     # Extend the attributes with the replay_id to make it queryable by replay_id after we
     # eventually use the trace_id in its rightful position.
     trace_item_context["attributes"]["replay_id"] = context["replay_id"]
+    trace_item_context["attributes"]["segment_id"] = context["segment_id"]
 
     return new_trace_item(
         {

src/sentry/testutils/cases.py

@@ -9,6 +9,7 @@
 from collections.abc import Generator, Mapping, Sequence
 from contextlib import contextmanager
 from datetime import UTC, datetime, timedelta
+from enum import Enum
 from io import BytesIO
 from typing import Any, Literal, TypedDict, Union
 from unittest import mock
@@ -3786,9 +3787,127 @@ def load_default(self) -> Event:
         )
 
 
+class ReplayBreadcrumbType(Enum):
+    CLICK = "click"
+    DEAD_CLICK = "dead_click"
+    RAGE_CLICK = "rage_click"
+    ERROR = "error"
+    WARNING = "warning"
+    INFO = "info"
+
+
 @pytest.mark.snuba
 @requires_snuba
 @pytest.mark.usefixtures("reset_snuba")
+class ReplayEAPTestCase(BaseTestCase):
+    def create_eap_replay_breadcrumb(
+        self,
+        *,
+        project,
+        replay_id,
+        segment_id,
+        breadcrumb_type,
+        timestamp=None,
+        organization=None,
+        trace_id=None,
+        retention_days=30,
+        **attributes,
+    ):
+        """Create single EAP replay breadcrumb TraceItem."""
+        from datetime import datetime, timezone
+        from uuid import uuid4
+
+        from google.protobuf.timestamp_pb2 import Timestamp
+        from sentry_protos.snuba.v1.request_common_pb2 import TraceItemType
+        from sentry_protos.snuba.v1.trace_item_pb2 import TraceItem
+
+        if organization is None:
+            organization = self.organization
+        if timestamp is None:
+            timestamp = datetime.now(timezone.utc)
+        if trace_id is None:
+            trace_id = replay_id
+
+        if breadcrumb_type == ReplayBreadcrumbType.CLICK:
+            category = "ui.click"
+            click_is_dead = 0
+            click_is_rage = 0
+        elif breadcrumb_type == ReplayBreadcrumbType.DEAD_CLICK:
+            category = "ui.click"
+            click_is_dead = 1
+            click_is_rage = 0
+        elif breadcrumb_type == ReplayBreadcrumbType.RAGE_CLICK:
+            category = "ui.click"
+            click_is_dead = 1
+            click_is_rage = 1
+        elif breadcrumb_type == ReplayBreadcrumbType.ERROR:
+            category = "error"
+            click_is_dead = None
+            click_is_rage = None
+        elif breadcrumb_type == ReplayBreadcrumbType.WARNING:
+            category = "warning"
+            click_is_dead = None
+            click_is_rage = None
+        elif breadcrumb_type == ReplayBreadcrumbType.INFO:
+            category = "info"
+            click_is_dead = None
+            click_is_rage = None
+        else:
+            raise ValueError(f"Unknown breadcrumb type: {breadcrumb_type}")
+
+        breadcrumb_data = {
+            "replay_id": replay_id,
+            "segment_id": segment_id,
+            "project_id": project.id,
+            "timestamp": int(timestamp.timestamp()),
+            "category": category,
+        }
+
+        if category == "ui.click":
+            breadcrumb_data["click_is_dead"] = click_is_dead
+            breadcrumb_data["click_is_rage"] = click_is_rage
+
+        breadcrumb_data.update(attributes)
+
+        attributes_proto = {}
+        for k, v in breadcrumb_data.items():
+            if v is not None:
+                attributes_proto[k] = scalar_to_any_value(v)
+
+        timestamp_proto = Timestamp()
+        timestamp_proto.FromDatetime(timestamp)
+
+        return TraceItem(
+            organization_id=organization.id,
+            project_id=project.id,
+            item_type=TraceItemType.TRACE_ITEM_TYPE_REPLAY,
+            timestamp=timestamp_proto,
+            trace_id=trace_id,
+            item_id=uuid4().bytes,
+            received=timestamp_proto,
+            retention_days=retention_days,
+            attributes=attributes_proto,
+            client_sample_rate=1.0,
+            server_sample_rate=1.0,
+        )
+
+    def store_replays_eap(self, replays):
+        import requests
+        from django.conf import settings
+
+        files = {f"replay_{i}": replay.SerializeToString() for i, replay in enumerate(replays)}
+        response = requests.post(
+            settings.SENTRY_SNUBA + "/tests/entities/eap_items/insert_bytes",
+            files=files,
+        )
+        assert response.status_code == 200
+
+        for replay in replays:
+            # Reverse the ids here since these are stored in little endian in Clickhouse
+            # and end up reversed.
+            replay.item_id = replay.item_id[::-1]
+
+
 class UptimeResultEAPTestCase(BaseTestCase):
     """Test case for creating and storing EAP uptime results."""
 

tests/sentry/replays/endpoints/test_query_replay_instance_eap.py

@@ -0,0 +1,124 @@
+import datetime
+from uuid import uuid4
+
+from sentry.replays.endpoints.organization_replay_details import query_replay_instance_eap
+from sentry.testutils.cases import ReplayBreadcrumbType, ReplayEAPTestCase, TestCase
+
+
+class TestQueryReplayInstanceEAP(TestCase, ReplayEAPTestCase):
+    def test_eap_replay_query(self) -> None:
+        replay_id1 = uuid4().hex
+        replay_id2 = uuid4().hex
+        now = datetime.datetime.now(datetime.UTC)
+
+        replay1_breadcrumbs = [
+            # Dead clicks
+            self.create_eap_replay_breadcrumb(
+                project=self.project,
+                replay_id=replay_id1,
+                segment_id=0,
+                breadcrumb_type=ReplayBreadcrumbType.DEAD_CLICK,
+                timestamp=now,
+            ),
+            self.create_eap_replay_breadcrumb(
+                project=self.project,
+                replay_id=replay_id1,
+                segment_id=0,
+                breadcrumb_type=ReplayBreadcrumbType.DEAD_CLICK,
+                timestamp=now,
+            ),
+            # Rage clicks
+            self.create_eap_replay_breadcrumb(
+                project=self.project,
+                replay_id=replay_id1,
+                segment_id=0,
+                breadcrumb_type=ReplayBreadcrumbType.RAGE_CLICK,
+                timestamp=now,
+            ),
+            # Regular click
+            self.create_eap_replay_breadcrumb(
+                project=self.project,
+                replay_id=replay_id1,
+                segment_id=0,
+                breadcrumb_type=ReplayBreadcrumbType.CLICK,
+                timestamp=now,
+            ),
+        ]
+
+        replay2_breadcrumbs = [
+            # Dead clicks
+            self.create_eap_replay_breadcrumb(
+                project=self.project,
+                replay_id=replay_id2,
+                segment_id=0,
+                breadcrumb_type=ReplayBreadcrumbType.DEAD_CLICK,
+                timestamp=now,
+            ),
+            # Rage clicks
+            self.create_eap_replay_breadcrumb(
+                project=self.project,
+                replay_id=replay_id2,
+                segment_id=0,
+                breadcrumb_type=ReplayBreadcrumbType.RAGE_CLICK,
+                timestamp=now,
+            ),
+            self.create_eap_replay_breadcrumb(
+                project=self.project,
+                replay_id=replay_id2,
+                segment_id=0,
+                breadcrumb_type=ReplayBreadcrumbType.RAGE_CLICK,
+                timestamp=now,
+            ),
+        ]
+
+        self.store_replays_eap(replay1_breadcrumbs + replay2_breadcrumbs)
+
+        start = now - datetime.timedelta(minutes=5)
+        end = now + datetime.timedelta(minutes=5)
+        organization_id = self.organization.id
+        project_ids = [self.project.id]
+
+        res1 = query_replay_instance_eap(
+            project_ids=project_ids,
+            replay_ids=[replay_id1],
+            start=start,
+            end=end,
+            request_user_id=self.user.id,
+            organization_id=organization_id,
+        )
+        res2 = query_replay_instance_eap(
+            project_ids=project_ids,
+            replay_ids=[replay_id2],
+            start=start,
+            end=end,
+            request_user_id=self.user.id,
+            organization_id=organization_id,
+        )
+
+        assert isinstance(res1, dict)
+        assert isinstance(res2, dict)
+        assert res1.get("data") is not None
+        assert res2.get("data") is not None
+
+        assert len(res1["data"]) == 1, f"Expected 1 row for replay_id1, got {len(res1['data'])}"
+        assert len(res2["data"]) == 1, f"Expected 1 row for replay_id2, got {len(res2['data'])}"
+
+        assert res1["data"][0]["replay_id"] == replay_id1
+        assert res2["data"][0]["replay_id"] == replay_id2
+
+        replay1_data = res1["data"][0]
+        assert "count_segments" in replay1_data
+        assert "count_errors" in replay1_data
+        assert "count_warnings" in replay1_data
+        assert "count_dead_clicks" in replay1_data
+        assert "count_rage_clicks" in replay1_data
+        assert "is_archived" in replay1_data
+        assert "started_at" in replay1_data
+        assert "finished_at" in replay1_data
+
+        assert replay1_data["count_dead_clicks"] == 3, "2 DEAD_CLICK + 1 RAGE_CLICK = 3 dead"
+        assert replay1_data["count_rage_clicks"] == 1, "1 RAGE_CLICK = 1 rage"
+
+        replay2_data = res2["data"][0]
+        assert replay2_data["count_dead_clicks"] == 3, "1 DEAD_CLICK + 2 RAGE_CLICK = 3 dead"
+        assert replay2_data["count_rage_clicks"] == 2, "2 RAGE_CLICK = 2 rage"