feat(replays): Add issue replay count endpoint (#41996)

## What this PR does

We create an endpoint that provides accurate counts of replay_ids
associated with issue ids within a timeframe. The endpoint will return a
format like so:
```
{
    issue_id_1: 1,
    issue_id_2: 20,
    etc. etc.
}
```

### Constraints 
- Between 1 and 25 issue_ids may be passed in (25 being the current
pagination size of the issue stream)
- We will only count up to 50 replay_ids, 50 chosen somewhat arbitrarily
but seems a reasonable size.
- We will only retrieve up to 100 replay_ids per issue_id. this means
that if over half of replay_ids are sampled, it's possible that we may
undercount replays as we will miss some. This is a small edge case and
is acceptable.

## Modifications:

- We modify discover.py to allow for private_use of the `groupArray`
clickhouse function
- In our endpoint we use QueryBuilder to hit the discover dataset, then
query the replays dataset with the replay_ids returned
- We then count up each replay_id confirmed to exist by its associated
issue_id, and return the resulting dict


### Why are we doing this?
Because of sampling / dropped data, events can have replay_ids that
don't exist. this is nominally fine, although it results in missing
counts / the product not seeming trustworthy in places.


Fixes https://github.com/getsentry/replay-backend/issues/190

mypy.ini

@@ -105,6 +105,7 @@ files = fixtures/mypy-stubs,
         src/sentry/release_health/,
         src/sentry/replays/consumers/,
         src/sentry/replays/usecases/,
+        src/sentry/replays/endpoints/organization_issue_replay_count.py,
         src/sentry/roles/manager.py,
         src/sentry/rules/,
         src/sentry/runner/commands/devserver.py,

src/sentry/api/urls.py

@@ -66,6 +66,9 @@ from sentry.incidents.endpoints.project_alert_rule_index import (
 from sentry.incidents.endpoints.project_alert_rule_task_details import (
     ProjectAlertRuleTaskDetailsEndpoint,
 )
+from sentry.replays.endpoints.organization_issue_replay_count import (
+    OrganizationIssueReplayCountEndpoint,
+)
 from sentry.replays.endpoints.organization_replay_events_meta import (
     OrganizationReplayEventsMetaEndpoint,
 )
@@ -1564,6 +1567,11 @@ urlpatterns = [
                     OrganizationReplayIndexEndpoint.as_view(),
                     name="sentry-api-0-organization-replay-index",
                 ),
+                url(
+                    r"^(?P<organization_slug>[^\/]+)/issue-replay-count/$",
+                    OrganizationIssueReplayCountEndpoint.as_view(),
+                    name="sentry-api-0-organization-issue-replay-count",
+                ),
                 url(
                     r"^(?P<organization_slug>[^\/]+)/replays-events-meta/$",
                     OrganizationReplayEventsMetaEndpoint.as_view(),

src/sentry/replays/endpoints/organization_issue_replay_count.py

@@ -0,0 +1,102 @@
+from __future__ import annotations
+
+from collections import defaultdict
+
+from rest_framework import status
+from rest_framework.response import Response
+from snuba_sdk import Column, Condition, Op, Request
+
+from sentry import features
+from sentry.api.base import region_silo_endpoint
+from sentry.api.bases import NoProjects
+from sentry.api.bases.organization_events import OrganizationEventsV2EndpointBase
+from sentry.models import Organization
+from sentry.replays.query import query_replays_dataset
+from sentry.search.events.builder import QueryBuilder
+from sentry.search.events.types import ParamsType, SnubaParams
+from sentry.utils.snuba import Dataset
+
+
+@region_silo_endpoint
+class OrganizationIssueReplayCountEndpoint(OrganizationEventsV2EndpointBase):
+    """
+    Get all the replay ids associated with a set of issues in discover,
+    then verify that they exist in the replays dataset, and return the count.
+    """
+
+    private = True
+
+    def get(self, request: Request, organization: Organization) -> Response:
+        if not features.has("organizations:session-replay", organization, actor=request.user):
+            return Response(status=404)
+
+        try:
+            snuba_params, params = self.get_snuba_dataclass(
+                request, organization, check_global_views=False
+            )
+        except NoProjects:
+            return Response({})
+
+        try:
+            replay_id_to_issue_map = _query_discover_for_replay_ids(request, params, snuba_params)
+        except ValueError as e:
+            return Response({"detail": str(e)}, status=status.HTTP_400_BAD_REQUEST)
+
+        replay_results = query_replays_dataset(
+            project_ids=[p.id for p in snuba_params.projects],
+            start=snuba_params.start,
+            end=snuba_params.end,
+            sorting=[],
+            where=[
+                Condition(Column("replay_id"), Op.IN, list(replay_id_to_issue_map.keys())),
+            ],
+            search_filters=[],
+            pagination=None,
+        )
+
+        issue_id_counts: dict[int, int] = defaultdict(int)
+
+        for row in replay_results["data"]:
+            issue_ids = replay_id_to_issue_map[row["replay_id"]]
+            for issue_id in issue_ids:
+                # cap count at 50
+                issue_id_counts[issue_id] = min(issue_id_counts[issue_id] + 1, 50)
+
+        return self.respond(issue_id_counts)
+
+
+def _query_discover_for_replay_ids(
+    request: Request, params: ParamsType, snuba_params: SnubaParams
+) -> dict[str, list[int]]:
+    builder = QueryBuilder(
+        dataset=Dataset.Discover,
+        params=params,
+        snuba_params=snuba_params,
+        selected_columns=["group_array(100,replayId)", "issue.id"],
+        query=request.GET.get("query"),
+        limit=25,
+        offset=0,
+        functions_acl=["group_array"],
+    )
+    _validate_params(builder)
+
+    discover_results = builder.run_query("api.organization-issue-replay-count")
+
+    replay_id_to_issue_map = defaultdict(list)
+
+    for row in discover_results["data"]:
+        for replay_id in row["group_array_100_replayId"]:
+            replay_id_to_issue_map[replay_id].append(row["issue.id"])
+
+    return replay_id_to_issue_map
+
+
+def _validate_params(builder: QueryBuilder) -> None:
+    issue_where_condition = next(
+        (cond for cond in builder.where if cond.lhs.name == "group_id"), None
+    )
+
+    if issue_where_condition is None:
+        raise ValueError("Must provide at least one issue id")
+    if len(issue_where_condition.rhs) > 25:
+        raise ValueError("Too many issues ids provided")

src/sentry/search/events/datasets/discover.py

@@ -5,7 +5,17 @@ from typing import Callable, Mapping, Optional, Union
 import sentry_sdk
 from django.utils.functional import cached_property
 from sentry_relay.consts import SPAN_STATUS_NAME_TO_CODE
-from snuba_sdk import Column, Condition, Direction, Function, Identifier, Lambda, Op, OrderBy
+from snuba_sdk import (
+    Column,
+    Condition,
+    CurriedFunction,
+    Direction,
+    Function,
+    Identifier,
+    Lambda,
+    Op,
+    OrderBy,
+)
 
 from sentry.api.event_search import SearchFilter, SearchKey, SearchValue
 from sentry.exceptions import InvalidSearchQuery
@@ -256,6 +266,18 @@ class DiscoverDatasetConfig(DatasetConfig):
                     ),
                     default_result_type="percentage",
                 ),
+                SnQLFunction(
+                    "group_array",
+                    required_args=[NumberRange("max_size", 0, 101), ColumnTagArg("column")],
+                    snql_aggregate=lambda args, alias: CurriedFunction(
+                        "groupArray",
+                        [int(args["max_size"])],
+                        [args["column"]],
+                        alias,
+                    ),
+                    default_result_type="string",  # TODO: support array type
+                    private=True,
+                ),
                 SnQLFunction(
                     "percentile",
                     required_args=[

tests/sentry/replays/test_organization_issue_replay_count.py

@@ -0,0 +1,191 @@
+import datetime
+import uuid
+
+import pytest
+from django.urls import reverse
+
+from sentry.replays.testutils import mock_replay
+from sentry.testutils import APITestCase, SnubaTestCase
+from sentry.testutils.cases import ReplaysSnubaTestCase
+from sentry.testutils.helpers.datetime import before_now, iso_format
+from sentry.testutils.silo import region_silo_test
+
+pytestmark = pytest.mark.sentry_metrics
+
+
+@region_silo_test
+class OrganizationIssueReplayCountEndpoint(APITestCase, SnubaTestCase, ReplaysSnubaTestCase):
+    def setUp(self):
+        super().setUp()
+        self.min_ago = before_now(minutes=1)
+        self.login_as(user=self.user)
+        self.url = reverse(
+            "sentry-api-0-organization-issue-replay-count",
+            kwargs={"organization_slug": self.project.organization.slug},
+        )
+        self.features = {"organizations:session-replay": True}
+
+    def test_simple(self):
+        event_id_a = "a" * 32
+        event_id_b = "b" * 32
+        replay1_id = uuid.uuid4().hex
+        replay2_id = uuid.uuid4().hex
+        replay3_id = uuid.uuid4().hex
+
+        self.store_replays(
+            mock_replay(
+                datetime.datetime.now() - datetime.timedelta(seconds=22),
+                self.project.id,
+                replay1_id,
+            )
+        )
+        self.store_replays(
+            mock_replay(
+                datetime.datetime.now() - datetime.timedelta(seconds=22),
+                self.project.id,
+                replay2_id,
+            )
+        )
+        self.store_replays(
+            mock_replay(
+                datetime.datetime.now() - datetime.timedelta(seconds=22),
+                self.project.id,
+                replay3_id,
+            )
+        )
+        event_a = self.store_event(
+            data={
+                "event_id": event_id_a,
+                "timestamp": iso_format(self.min_ago),
+                "tags": {"replayId": replay1_id},
+                "fingerprint": ["group-1"],
+            },
+            project_id=self.project.id,
+        )
+        self.store_event(
+            data={
+                "event_id": uuid.uuid4().hex,
+                "timestamp": iso_format(self.min_ago),
+                "tags": {"replayId": replay2_id},
+                "fingerprint": ["group-1"],
+            },
+            project_id=self.project.id,
+        )
+        self.store_event(
+            data={
+                "event_id": event_id_b,
+                "timestamp": iso_format(self.min_ago),
+                "tags": {"replayId": "z" * 32},  # a replay id that doesn't exist
+                "fingerprint": ["group-1"],
+            },
+            project_id=self.project.id,
+        )
+        event_c = self.store_event(
+            data={
+                "event_id": event_id_b,
+                "timestamp": iso_format(self.min_ago),
+                "tags": {"replayId": replay3_id},
+                "fingerprint": ["group-2"],
+            },
+            project_id=self.project.id,
+        )
+
+        query = {"query": f"issue.id:[{event_a.group.id}, {event_c.group.id}]"}
+        with self.feature(self.features):
+            response = self.client.get(self.url, query, format="json")
+
+        expected = {
+            event_a.group.id: 2,
+            event_c.group.id: 1,
+        }
+        assert response.status_code == 200, response.content
+        assert response.data == expected
+
+    def test_one_replay_multiple_issues(self):
+        event_id_a = "a" * 32
+        event_id_b = "b" * 32
+        replay1_id = uuid.uuid4().hex
+
+        self.store_replays(
+            mock_replay(
+                datetime.datetime.now() - datetime.timedelta(seconds=22),
+                self.project.id,
+                replay1_id,
+            )
+        )
+        event_a = self.store_event(
+            data={
+                "event_id": event_id_a,
+                "timestamp": iso_format(self.min_ago),
+                "tags": {"replayId": replay1_id},
+                "fingerprint": ["group-1"],
+            },
+            project_id=self.project.id,
+        )
+        event_b = self.store_event(
+            data={
+                "event_id": event_id_b,
+                "timestamp": iso_format(self.min_ago),
+                "tags": {"replayId": replay1_id},
+                "fingerprint": ["group-2"],
+            },
+            project_id=self.project.id,
+        )
+
+        query = {"query": f"issue.id:[{event_a.group.id}, {event_b.group.id}]"}
+        with self.feature(self.features):
+            response = self.client.get(self.url, query, format="json")
+
+        expected = {
+            event_a.group.id: 1,
+            event_b.group.id: 1,
+        }
+        assert response.status_code == 200, response.content
+        assert response.data == expected
+
+    def test_max_50(self):
+        replay_ids = [uuid.uuid4().hex for _ in range(100)]
+        for replay_id in replay_ids:
+            self.store_replays(
+                mock_replay(
+                    datetime.datetime.now() - datetime.timedelta(seconds=22),
+                    self.project.id,
+                    replay_id,
+                )
+            )
+            event_a = self.store_event(
+                data={
+                    "event_id": uuid.uuid4().hex,
+                    "timestamp": iso_format(self.min_ago),
+                    "tags": {"replayId": replay_id},
+                    "fingerprint": ["group-1"],
+                },
+                project_id=self.project.id,
+            )
+
+        query = {"query": f"issue.id:[{event_a.group.id}]"}
+        with self.feature(self.features):
+            response = self.client.get(self.url, query, format="json")
+
+        expected = {
+            event_a.group.id: 50,
+        }
+        assert response.status_code == 200, response.content
+        assert response.data == expected
+
+    def test_invalid_params_need_one_issue_id(self):
+        query = {"query": ""}
+        with self.feature(self.features):
+            response = self.client.get(self.url, query, format="json")
+            assert response.status_code == 400
+            assert response.data["detail"] == "Must provide at least one issue id"
+
+    def test_invalid_params_max_issue_id(self):
+        issue_ids = ",".join(str(i) for i in range(26))
+
+        query = {"query": f"issue.id:[{issue_ids}]"}
+
+        with self.feature(self.features):
+            response = self.client.get(self.url, query, format="json")
+            assert response.status_code == 400
+            assert response.data["detail"] == "Too many issues ids provided"