Implement storage- and relationship-aware cleanup

services/cleanup/cleanup.py

@@ -0,0 +1,32 @@
+from django.db.models.query import QuerySet
+
+from services.cleanup.models import MANUAL_CLEANUP
+from services.cleanup.relations import build_relation_graph
+
+
+def run_cleanup(query: QuerySet) -> tuple[int, int]:
+    """
+    Cleans up all the models and storage files reachable from the given `QuerySet`.
+
+    This deletes all database models in topological sort order, and also removes
+    all the files in storage for any of the models in the relationship graph.
+
+    Returns the number of models and files being cleaned up.
+    """
+    models_to_cleanup = build_relation_graph(query)
+
+    cleaned_models = 0
+    cleaned_files = 0
+
+    for model, query in models_to_cleanup:
+        manual_cleanup = MANUAL_CLEANUP.get(model)
+        if manual_cleanup is not None:
+            res = manual_cleanup(query)
+            cleaned_models += res[0]
+            cleaned_files += res[1]
+
+        else:
+            deleted, _ = query.delete()
+            cleaned_models += deleted
+
+    return (cleaned_models, cleaned_files)

services/cleanup/models.py

@@ -0,0 +1,123 @@
+import dataclasses
+import itertools
+from collections.abc import Callable
+from functools import partial
+
+from django.db.models import Model
+from django.db.models.query import Q, QuerySet
+from shared.api_archive.storage import StorageService
+from shared.config import get_config
+from shared.django_apps.core.models import Commit, Pull
+from shared.django_apps.reports.models import CommitReport, ReportDetails, ReportSession
+
+from services.archive import ArchiveService, MinioEndpoints
+
+DELETE_CHUNKS = 25
+
+
+# This has all the `Repository` fields needed by `get_archive_hash`
+@dataclasses.dataclass
+class FakeRepository:
+    repoid: int
+    service: str
+    service_id: str
+
+
+def cleanup_archivefield(field_name: str, query: QuerySet) -> tuple[int, int]:
+    model_field_name = f"_{field_name}_storage_path"
+    # query for a non-`None` `field_name`
+    storage_query = query.filter(**{f"{model_field_name}__isnull": False}).values_list(
+        model_field_name, flat=True
+    )
+
+    # and then delete all those files from storage
+    storage = StorageService()
+    bucket = get_config("services", "minio", "bucket", default="archive")
+    cleaned_files = 0
+    # TODO: possibly fan out the batches to a thread pool, as the storage requests are IO-bound
+    # TODO: do a limit / range query to avoid loading *all* the paths into memory at once
+    for batched_paths in itertools.batched(storage_query, DELETE_CHUNKS):
+        storage.delete_files(bucket, batched_paths)
+        cleaned_files += len(batched_paths)
+
+    cleaned_models, _ = query.delete()
+
+    return (cleaned_models, cleaned_files)
+
+
+def cleanup_commitreport(query: QuerySet) -> tuple[int, int]:
+    coverage_reports = query.filter(
+        Q(report_type=None) | Q(report_type="coverage")
+    ).values_list(
+        "code",
+        "commit__commitid",
+        "repository__repoid",
+        "repository__owner__service",
+        "repository__service_id",
+    )
+
+    storage = StorageService()
+    bucket = get_config("services", "minio", "bucket", default="archive")
+    repo_hashes: dict[int, str] = {}
+    cleaned_files = 0
+    # TODO: figure out a way to run the deletes in batches
+    # TODO: possibly fan out the batches to a thread pool, as the storage requests are IO-bound
+    # TODO: do a limit / range query to avoid loading *all* the paths into memory at once
+    for (
+        report_code,
+        commit_sha,
+        repoid,
+        repo_service,
+        repo_service_id,
+    ) in coverage_reports:
+        if repoid not in repo_hashes:
+            fake_repo = FakeRepository(
+                repoid=repoid, service=repo_service, service_id=repo_service_id
+            )
+            repo_hashes[repoid] = ArchiveService.get_archive_hash(fake_repo)
+        repo_hash = repo_hashes[repoid]
+
+        chunks_file_name = report_code if report_code is not None else "chunks"
+        path = MinioEndpoints.chunks.get_path(
+            version="v4",
+            repo_hash=repo_hash,
+            commitid=commit_sha,
+            chunks_file_name=chunks_file_name,
+        )
+        storage.delete_file(bucket, path)
+        cleaned_files += 1
+
+    cleaned_models, _ = query.delete()
+
+    return (cleaned_models, cleaned_files)
+
+
+def cleanup_upload(query: QuerySet) -> tuple[int, int]:
+    storage_query = query.values_list("storage_path", flat=True)
+
+    storage = StorageService()
+    bucket = get_config("services", "minio", "bucket", default="archive")
+    cleaned_files = 0
+    # TODO: possibly fan out the batches to a thread pool, as the storage requests are IO-bound
+    # TODO: do a limit / range query to avoid loading *all* the paths into memory at once
+    for batched_paths in itertools.batched(storage_query, DELETE_CHUNKS):
+        storage.delete_files(bucket, batched_paths)
+        cleaned_files += len(batched_paths)
+
+    cleaned_models, _ = query.delete()
+
+    return (cleaned_models, cleaned_files)
+
+
+# All the models that need custom python code for deletions so a bulk `DELETE` query does not work.
+MANUAL_CLEANUP: dict[type[Model], Callable[[QuerySet], tuple[int, int]]] = {
+    Commit: partial(cleanup_archivefield, "report"),
+    Pull: partial(cleanup_archivefield, "flare"),
+    ReportDetails: partial(cleanup_archivefield, "files_array"),
+    CommitReport: cleanup_commitreport,
+    ReportSession: cleanup_upload,
+    # TODO: figure out any other models which have files in storage that are not `ArchiveField`
+    # TODO: TA is also storing files in GCS
+    # TODO: BA is also storing files in GCS
+    # TODO: There is also `CompareCommit.report_storage_path`, but that does not seem to be implemented as Django model?
+}

services/cleanup/relations.py

@@ -0,0 +1,77 @@
+# from pprint import pprint
+
+from graphlib import TopologicalSorter
+
+from django.db.models import Model
+from django.db.models.query import QuerySet
+
+
+def build_relation_graph(query: QuerySet) -> list[tuple[type[Model], QuerySet]]:
+    """
+    This takes as input a django `QuerySet`, like `Repository.objects.filter(repoid=123)`.
+
+    It then walks the django relation graph, resolving all the models that have a relationship **to** the input model,
+    returning those models along with a `QuerySet` that allows either querying or deleting those models.
+
+    The returned list is in topological sorting order, so related models are always sorted before models they depend on.
+    """
+    graph: TopologicalSorter[type[Model]] = TopologicalSorter()
+    querysets: dict[type[Model], QuerySet] = {}
+
+    def process_model(model: type[Model], query: QuerySet):
+        if model in querysets:
+            return
+        querysets[model] = query
+
+        if not (meta := model._meta):
+            return
+
+        for field in meta.get_fields(include_hidden=True):
+            if not field.is_relation:
+                continue
+
+            if field.one_to_many or field.one_to_one:
+                # Most likely the reverse of a `ForeignKey`
+                # <https://docs.djangoproject.com/en/5.1/ref/models/fields/#django.db.models.Field.one_to_many>
+
+                if not hasattr(field, "field"):
+                    # I believe this is the actual *forward* definition of a `OneToOne`
+                    continue
+
+                # this should be the actual `ForeignKey` definition:
+                actual_field = field.field
+                if actual_field.model == model:
+                    # this field goes from *this* model to another, but we are interested in the reverse actually
+                    continue
+
+                related_model = actual_field.model
+                related_model_field = actual_field.name
+                related_query = related_model.objects.filter(
+                    **{f"{related_model_field}__in": query}
+                )
+                graph.add(model, related_model)
+                process_model(related_model, related_query)
+
+            elif field.many_to_many:
+                if not hasattr(field, "through"):
+                    # we want to delete all related records on the join table
+                    continue
+
+                related_model = field.through
+                join_meta = related_model._meta
+                for field in join_meta.get_fields(include_hidden=True):
+                    if not field.is_relation or field.model != model:
+                        continue
+
+                    related_model_field = actual_field.name
+                    related_query = related_model.objects.filter(
+                        **{f"{related_model_field}__in": query}
+                    )
+                    graph.add(model, related_model)
+                    process_model(related_model, related_query)
+
+                # pprint(vars(field.through._meta))
+
+    process_model(query.model, query)
+
+    return [(model, querysets[model]) for model in graph.static_order()]

services/cleanup/tests/test_relations.py

@@ -0,0 +1,10 @@
+from shared.django_apps.core.models import Repository
+
+from services.cleanup.relations import build_relation_graph
+
+
+def test_builds_relation_graph(db):
+    print()
+    relations = build_relation_graph(Repository.objects.filter(repoid=123))
+    for model, query in relations:
+        print(model, str(query.query))