feat(replays): Update delete script to query clickhouse less (#92163)

Author: cmanallen

src/sentry/replays/scripts/delete_replays.py

@@ -1,16 +1,18 @@
 from __future__ import annotations
 
-import contextlib
 import logging
 from collections.abc import Sequence
 from datetime import datetime, timezone
 
+from snuba_sdk import Column, Condition, Entity, Function, Granularity, Limit, Offset, Op, Query
+
 from sentry.api.event_search import QueryToken, parse_search_query
 from sentry.models.organization import Organization
 from sentry.replays.lib.kafka import initialize_replays_publisher
-from sentry.replays.post_process import generate_normalized_output
-from sentry.replays.query import query_replays_collection_paginated, replay_url_parser_config
-from sentry.replays.tasks import archive_replay, delete_replay_recording_async
+from sentry.replays.query import replay_url_parser_config
+from sentry.replays.tasks import archive_replay, delete_replays_script_async
+from sentry.replays.usecases.query import execute_query, handle_search_filters
+from sentry.replays.usecases.query.configs.scalar import scalar_search_config
 
 logger = logging.getLogger()
 
@@ -33,21 +35,15 @@ def delete_replays(
 
     has_more = True
     while has_more:
-        response = query_replays_collection_paginated(
-            project_ids=[project_id],
+        replays, has_more = _get_rows_matching_deletion_pattern(
+            project_id=project_id,
             start=start_utc,
             end=end_utc,
-            fields=["id"],
             limit=batch_size,
-            environment=environment,
             offset=offset,
             search_filters=search_filters,
-            sort="started_at",
-            organization=Organization.objects.filter(project__id=project_id).get(),
-            preferred_source="scalar",
+            environment=environment,
         )
-        replays = list(generate_normalized_output(response.response))
-        has_more = response.has_more
 
         # Exit early if no replays were found.
         if not replays:
@@ -58,16 +54,16 @@ def delete_replays(
         if dry_run:
             print(f"Replays to be deleted (dry run): {len(replays)}")  # NOQA
         else:
-            delete_replay_ids(project_id, replay_ids=[r["id"] for r in replays])
+            delete_replay_ids(project_id, replays)
 
 
 def translate_cli_tags_param_to_snuba_tag_param(tags: list[str]) -> Sequence[QueryToken]:
     return parse_search_query(" AND ".join(tags), config=replay_url_parser_config)
 
 
-def delete_replay_ids(project_id: int, replay_ids: list[str]) -> None:
+def delete_replay_ids(project_id: int, rows: list[tuple[int, str, int]]) -> None:
     """Delete a set of replay-ids for a specific project."""
-    logger.info("Archiving %d replays.", len(replay_ids))
+    logger.info("Archiving %d replays.", len(rows))
 
     # Bulk produce archived replay rows to the ingest-replay-events topic before flushing.
     #
@@ -79,30 +75,75 @@ def delete_replay_ids(project_id: int, replay_ids: list[str]) -> None:
     #
     # This also gives us reasonable assurances that if the script ran to completion the customer
     # will not be able to access their deleted data even if the actual deletion takes place some
-    # time later.
-    with _bulk_produce_then_flush() as publisher:
-        for replay_id in replay_ids:
-            archive_replay(publisher, project_id, replay_id)
+    # time later
+    publisher = initialize_replays_publisher(is_async=True)
+    for _, replay_id, _ in rows:
+        archive_replay(publisher, project_id, replay_id)
+    publisher.flush()
 
-    logger.info("Scheduling %d replays for deletion.", len(replay_ids))
+    logger.info("Scheduling %d replays for deletion.", len(rows))
 
     # Asynchronously delete RRWeb recording data.
     #
     # Because this operation could involve millions of requests to the blob storage provider we
     # schedule the tasks to run on a cluster of workers. This allows us to parallelize the work
     # and complete the task as quickly as possible.
-    for replay_id in replay_ids:
-        delete_replay_recording_async.delay(project_id, replay_id)
+    for retention_days, replay_id, max_segment_id in rows:
+        delete_replays_script_async.delay(retention_days, project_id, replay_id, max_segment_id)
 
-    logger.info("%d replays were successfully deleted.", len(replay_ids))
+    logger.info("%d replays were successfully deleted.", len(rows))
     logger.info(
         "The customer will no longer have access to the replays passed to this function. Deletion "
         "of RRWeb data will complete asynchronously."
     )
 
 
-@contextlib.contextmanager
-def _bulk_produce_then_flush():
-    publisher = initialize_replays_publisher(is_async=True)
-    yield publisher
-    publisher.flush()
+def _get_rows_matching_deletion_pattern(
+    project_id: int,
+    limit: int,
+    offset: int,
+    end: datetime,
+    start: datetime,
+    search_filters: Sequence[QueryToken],
+    environment: list[str],
+) -> tuple[list[tuple[int, str, int]], bool]:
+    where = handle_search_filters(scalar_search_config, search_filters)
+
+    if environment:
+        where.append(Condition(Column("environment"), Op.IN, environment))
+
+    query = Query(
+        match=Entity("replays"),
+        select=[
+            Function("any", parameters=[Column("retention_days")], alias="retention_days"),
+            Column("replay_id"),
+            Function("max", parameters=[Column("segment_id")], alias="max_segment_id"),
+        ],
+        where=[
+            Condition(Column("project_id"), Op.EQ, project_id),
+            Condition(Column("timestamp"), Op.LT, end),
+            Condition(Column("timestamp"), Op.GTE, start),
+            *where,
+        ],
+        groupby=[Column("replay_id")],
+        granularity=Granularity(3600),
+        limit=Limit(limit),
+        offset=Offset(offset),
+    )
+
+    response = execute_query(
+        query,
+        {"tenant_id": Organization.objects.filter(project__id=project_id).get().id},
+        "replays.scripts.delete_replays",
+    )
+
+    data = response.get("data", [])
+    has_more = len(data) == limit
+
+    return (
+        [
+            (item["retention_days"], item["replay_id"].replace("-", ""), item["max_segment_id"])
+            for item in data
+        ],
+        has_more,
+    )

src/sentry/replays/tasks.py

@@ -6,7 +6,14 @@
 from google.cloud.exceptions import NotFound
 
 from sentry.replays.lib.kafka import initialize_replays_publisher
-from sentry.replays.lib.storage import filestore, make_video_filename, storage, storage_kv
+from sentry.replays.lib.storage import (
+    RecordingSegmentStorageMeta,
+    filestore,
+    make_recording_filename,
+    make_video_filename,
+    storage,
+    storage_kv,
+)
 from sentry.replays.models import ReplayRecordingSegment
 from sentry.replays.usecases.events import archive_event
 from sentry.replays.usecases.reader import fetch_segments_metadata
@@ -59,6 +66,54 @@ def delete_replay_recording_async(project_id: int, replay_id: str) -> None:
     delete_replay_recording(project_id, replay_id)
 
 
+@instrumented_task(
+    name="sentry.replays.tasks.delete_recording_async",
+    queue="replays.delete_replay",
+    default_retry_delay=5,
+    max_retries=5,
+    silo_mode=SiloMode.REGION,
+    taskworker_config=TaskworkerConfig(
+        namespace=replays_tasks,
+        retry=Retry(
+            times=5,
+            delay=5,
+        ),
+    ),
+)
+def delete_replays_script_async(
+    retention_days: int,
+    project_id: int,
+    replay_id: str,
+    max_segment_id: int,
+) -> None:
+    segments = [
+        RecordingSegmentStorageMeta(
+            project_id=project_id,
+            replay_id=replay_id,
+            segment_id=i,
+            retention_days=retention_days,
+        )
+        for i in range(0, max_segment_id)
+    ]
+
+    rrweb_filenames = []
+    video_filenames = []
+    for segment in segments:
+        video_filenames.append(make_video_filename(segment))
+        rrweb_filenames.append(make_recording_filename(segment))
+
+    with cf.ThreadPoolExecutor(max_workers=100) as pool:
+        pool.map(_delete_if_exists, video_filenames)
+        pool.map(_delete_if_exists, rrweb_filenames)
+
+    # Backwards compatibility. Should be deleted one day.
+    segments_from_django_models = ReplayRecordingSegment.objects.filter(
+        replay_id=replay_id, project_id=project_id
+    ).all()
+    for segment_model in segments_from_django_models:
+        segment_model.delete()
+
+
 def delete_replay_recording(project_id: int, replay_id: str) -> None:
     """Delete all recording-segments associated with a Replay."""
     segments_from_metadata = fetch_segments_metadata(project_id, replay_id, offset=0, limit=10000)

tests/sentry/replays/scripts/test_delete_replays.py

@@ -6,9 +6,8 @@
 from zlib import compress
 
 from sentry.models.file import File
-from sentry.replays.lib.storage import RecordingSegmentStorageMeta, storage
 from sentry.replays.models import ReplayRecordingSegment
-from sentry.replays.scripts.delete_replays import delete_replay_ids, delete_replays
+from sentry.replays.scripts.delete_replays import delete_replays
 from sentry.replays.testutils import (
     mock_replay,
     mock_rrweb_div_helloworld,
@@ -35,7 +34,9 @@ def store_replay_segments(
             tags = {}
 
         self.store_replays(
-            mock_replay(timestamp, project_id, replay_id, environment=environment, tags=tags)
+            mock_replay(
+                timestamp, project_id, replay_id, environment=environment, tags=tags, segment_id=5
+            )
         )
 
         segments = [
@@ -280,49 +281,3 @@ def test_deletion_replays_batch_size_all_deleted(self):
 
         replay_recordings = ReplayRecordingSegment.objects.all()
         assert len(replay_recordings) == 0
-
-    def test_delete_replays_by_id(self):
-        # Deleted.
-        deleted_replay_id = uuid4().hex
-        self.store_replays(
-            mock_replay(
-                datetime.datetime.now() - datetime.timedelta(seconds=10),
-                self.project.id,
-                deleted_replay_id,
-            )
-        )
-
-        metadata1 = RecordingSegmentStorageMeta(
-            project_id=self.project.id,
-            replay_id=deleted_replay_id,
-            segment_id=0,
-            retention_days=30,
-            file_id=None,
-        )
-        storage.set(metadata1, b"hello, world!")
-
-        # Kept
-        kept_replay_id = uuid4().hex
-        self.store_replays(
-            mock_replay(
-                datetime.datetime.now() - datetime.timedelta(seconds=10),
-                self.project.id,
-                kept_replay_id,
-            )
-        )
-
-        metadata2 = RecordingSegmentStorageMeta(
-            project_id=self.project.id,
-            replay_id=kept_replay_id,
-            segment_id=0,
-            retention_days=30,
-            file_id=None,
-        )
-        storage.set(metadata2, b"hello, world!")
-
-        with TaskRunner():
-            delete_replay_ids(project_id=self.project.id, replay_ids=[deleted_replay_id])
-
-        # Assert stored data was deleted.
-        assert storage.get(metadata1) is None
-        assert storage.get(metadata2) is not None