fix(replays): use keyset pagination in delete_replays finder so it stops timing out Snuba

src/sentry/replays/scripts/delete_replays.py

@@ -4,7 +4,18 @@
 from collections.abc import Sequence
 from datetime import datetime, timezone
 
-from snuba_sdk import Column, Condition, Entity, Function, Granularity, Limit, Offset, Op, Query
+from snuba_sdk import (
+    Column,
+    Condition,
+    Direction,
+    Entity,
+    Function,
+    Granularity,
+    Limit,
+    Op,
+    OrderBy,
+    Query,
+)
 
 from sentry.api.event_search import QueryToken, parse_search_query
 from sentry.models.organization import Organization
@@ -27,19 +38,21 @@ def delete_replays(
 ) -> None:
     """Delete a set of replays from a query."""
     search_filters = translate_cli_tags_param_to_snuba_tag_param(tags)
-    offset = 0
 
     start_utc = start_utc.replace(tzinfo=timezone.utc)
     end_utc = end_utc.replace(tzinfo=timezone.utc)
 
+    # Keyset (seek) pagination cursor - page by the last replay_id we saw
+    last_replay_id = None
+
     has_more = True
     while has_more:
-        replays, has_more = _get_rows_matching_deletion_pattern(
+        replays, has_more, last_replay_id = _get_rows_matching_deletion_pattern(
             project_id=project_id,
             start=start_utc,
             end=end_utc,
             limit=batch_size,
-            offset=offset,
+            after_replay_id=last_replay_id,
             search_filters=search_filters,
             environment=environment,
         )
@@ -48,16 +61,13 @@ def delete_replays(
         if not replays:
             return None
 
-        offset += len(replays)
-
         logging_context = {
             "project_id": project_id,
             "dry_run": dry_run,
             "batch_size": batch_size,
             "tags": tags,
             "start_utc": start_utc,
             "end_utc": end_utc,
-            "offset": offset,
             "has_more": has_more,
         }
         if dry_run:
@@ -113,17 +123,21 @@ def delete_replay_ids(project_id: int, rows: list[tuple[int, str, int]]) -> None
 def _get_rows_matching_deletion_pattern(
     project_id: int,
     limit: int,
-    offset: int,
+    after_replay_id: str | None,
     end: datetime,
     start: datetime,
     search_filters: Sequence[QueryToken],
     environment: list[str],
-) -> tuple[list[tuple[int, str, int]], bool]:
+) -> tuple[list[tuple[int, str, int]], bool, str | None]:
     where = handle_search_filters(scalar_search_config, search_filters)
 
     if environment:
         where.append(Condition(Column("environment"), Op.IN, environment))
 
+    # Keyset cursor in order to walk result set in fixed-cost pages. Need the ORDER BY to keep it deterministic
+    if after_replay_id is not None:
+        where.append(Condition(Column("replay_id"), Op.GT, after_replay_id))
+
     query = Query(
         match=Entity("replays"),
         select=[
@@ -138,9 +152,9 @@ def _get_rows_matching_deletion_pattern(
             *where,
         ],
         groupby=[Column("replay_id")],
+        orderby=[OrderBy(Column("replay_id"), Direction.ASC)],
         granularity=Granularity(3600),
         limit=Limit(limit),
-        offset=Offset(offset),
     )
 
     response = execute_query(
@@ -152,10 +166,14 @@ def _get_rows_matching_deletion_pattern(
     data = response.get("data", [])
     has_more = len(data) == limit
 
+    # The next page seeks past the last raw replay_id in this page
+    next_cursor = data[-1]["replay_id"] if data else None
+
     return (
         [
             (item["retention_days"], item["replay_id"].replace("-", ""), item["max_segment_id"])
             for item in data
         ],
         has_more,
+        next_cursor,
     )

tests/sentry/replays/scripts/test_delete_replays.py

@@ -281,3 +281,47 @@ def test_deletion_replays_batch_size_all_deleted(self) -> None:
 
         replay_recordings = ReplayRecordingSegment.objects.all()
         assert len(replay_recordings) == 0
+
+    def test_deletion_replays_multi_page_keyset_pagination(self) -> None:
+        # Store several full pages worth of deletable replays so the keyset cursor has to walk
+        # past multiple pages. This guards the property that seek pagination never skips or
+        # double-processes a replay across page boundaries.
+        num_pages = 3
+        to_delete = [uuid4().hex for _ in range(self.small_batch_size * num_pages + 1)]
+        for replay_id in to_delete:
+            self.store_replay_segments(
+                replay_id=replay_id,
+                project_id=self.project.id,
+                timestamp=datetime.datetime.now() - datetime.timedelta(seconds=10),
+            )
+
+        # Keepers that fall inside the id space but must not be touched: a replay in another
+        # project and a replay outside the deletion time range.
+        replay_id_other_project = uuid4().hex
+        self.store_replay_segments(
+            replay_id_other_project,
+            self.other_project.id,
+            datetime.datetime.now() - datetime.timedelta(seconds=10),
+        )
+        replay_id_outside_timerange = uuid4().hex
+        self.store_replay_segments(
+            replay_id_outside_timerange,
+            self.project.id,
+            datetime.datetime.now() + datetime.timedelta(seconds=10),
+        )
+
+        with TaskRunner():
+            delete_replays(
+                project_id=self.project.id,
+                batch_size=self.small_batch_size,
+                tags=[],
+                start_utc=self.default_start_time,
+                end_utc=self.default_end_time,
+                dry_run=False,
+                environment=[],
+            )
+
+        for replay_id in to_delete:
+            self.assert_recording_deleted(replay_id)
+        self.assert_recording_not_deleted(replay_id_other_project)
+        self.assert_recording_not_deleted(replay_id_outside_timerange)