Allow unrolling subqueries in cleanup task (#1156)

Author: Swatinem

services/cleanup/models.py

@@ -9,7 +9,12 @@
 from shared.django_apps.compare.models import CommitComparison
 from shared.django_apps.core.models import Commit, Pull, Repository
 from shared.django_apps.profiling.models import ProfilingUpload
-from shared.django_apps.reports.models import CommitReport, ReportDetails
+from shared.django_apps.reports.models import (
+    CommitReport,
+    DailyTestRollup,
+    ReportDetails,
+    TestInstance,
+)
 from shared.django_apps.reports.models import ReportSession as Upload
 from shared.django_apps.staticanalysis.models import StaticAnalysisSingleFileSnapshot
 from shared.django_apps.timeseries.models import Dataset, Measurement
@@ -18,6 +23,7 @@
 from shared.utils.sessions import SessionType
 
 from services.archive import ArchiveService, MinioEndpoints
+from services.cleanup.relations import reverse_filter
 from services.cleanup.utils import CleanupContext, CleanupResult
 
 MANUAL_QUERY_CHUNKSIZE = 1_000
@@ -220,6 +226,21 @@ def cleanup_repository(context: CleanupContext, query: QuerySet) -> CleanupResul
     return CleanupResult(query._raw_delete(query.db))
 
 
+def unroll_subquery(context: CleanupContext, query: QuerySet) -> CleanupResult:
+    reversed_query = reverse_filter(query)
+    if not reversed_query:
+        return CleanupResult(query._raw_delete(query.db))
+    field, subquery = reversed_query
+
+    cleaned_models = 0
+    for parent in subquery:
+        cleaned_models += query.model.objects.filter(**{field: parent.pk})._raw_delete(
+            query.db
+        )
+
+    return CleanupResult(cleaned_models)
+
+
 # All the models that need custom python code for deletions so a bulk `DELETE` query does not work.
 MANUAL_CLEANUP: dict[
     type[Model], Callable[[CleanupContext, QuerySet], CleanupResult]
@@ -235,4 +256,6 @@ def cleanup_repository(context: CleanupContext, query: QuerySet) -> CleanupResul
         cleanup_with_storage_field, "content_location"
     ),
     Repository: cleanup_repository,
+    TestInstance: unroll_subquery,
+    DailyTestRollup: unroll_subquery,
 }

services/cleanup/relations.py

@@ -4,6 +4,7 @@
 
 from django.db.models import Model
 from django.db.models.expressions import Col, Expression
+from django.db.models.fields import Field
 from django.db.models.lookups import Exact, In
 from django.db.models.query import QuerySet
 from shared.django_apps.bundle_analysis.models import CacheConfig
@@ -191,3 +192,28 @@ def simplified_lookup(queryset: QuerySet) -> QuerySet | list[int]:
             return condition.rhs
 
     return queryset
+
+
+def reverse_filter(queryset: QuerySet) -> None | tuple[str, QuerySet]:
+    if queryset.query.is_sliced:
+        return None
+
+    where = queryset.query.where
+    if len(where.children) != 1:
+        return None
+
+    condition = where.children[0]
+    if (
+        not isinstance(condition, In)
+        or not isinstance(condition.lhs, Col)
+        or not isinstance(condition.lhs.target, Field)
+        or condition.rhs_is_direct_value()
+    ):
+        return None
+
+    column = condition.lhs.target.name
+
+    query = condition.rhs
+    model = query.model
+
+    return (column, QuerySet(model=model, query=query))

services/cleanup/tests/test_relations.py

@@ -3,10 +3,16 @@
 from django.db.models.query import QuerySet
 from django.db.models.sql.subqueries import DeleteQuery
 from shared.django_apps.codecov_auth.models import Owner
+from shared.django_apps.codecov_auth.tests.factories import OwnerFactory
 from shared.django_apps.core.models import Repository
-from shared.django_apps.reports.models import ReportDetails
+from shared.django_apps.core.tests.factories import RepositoryFactory
+from shared.django_apps.reports.models import ReportDetails, TestInstance
 
-from services.cleanup.relations import build_relation_graph, simplified_lookup
+from services.cleanup.relations import (
+    build_relation_graph,
+    reverse_filter,
+    simplified_lookup,
+)
 
 
 def dump_delete_queries(queryset: QuerySet) -> str:
@@ -60,3 +66,34 @@ def test_can_simplify_queries():
 def test_leaf_table(snapshot):
     query = ReportDetails.objects.all()
     assert dump_delete_queries(query) == snapshot("leaf.txt")
+
+
+@pytest.mark.django_db
+def test_can_reverse_filter():
+    query = TestInstance.objects.filter(repoid=123)
+    assert reverse_filter(query) is None
+
+    query = TestInstance.objects.filter(repoid__in=[123, 234])
+    assert reverse_filter(query) is None
+
+    query = TestInstance.objects.filter(
+        repoid__in=Repository.objects.filter(author=123), branch="foo"
+    )
+    assert reverse_filter(query) is None
+
+    owner = OwnerFactory()
+    r1 = RepositoryFactory(author=owner)
+    r2 = RepositoryFactory(author=owner)
+    r3 = RepositoryFactory(author=owner)
+
+    filtered_qs = Repository.objects.filter(author=owner.ownerid)
+    query = TestInstance.objects.filter(repoid__in=filtered_qs)
+
+    column, reversed_query = reverse_filter(query)
+
+    assert column == "repoid"
+    # ideally, we would assert that the `QuerySet` itself is the same,
+    # but that is not really doable. but in essense we only care that it yields
+    # the same results, which it does:
+    assert set(reversed_query) == set(filtered_qs)
+    assert set(reversed_query) == {r1, r2, r3}

tasks/tests/unit/test_delete_owner.py

@@ -10,6 +10,20 @@
     CommitFactory,
     RepositoryFactory,
 )
+from shared.django_apps.reports.models import (
+    CommitReport,
+    DailyTestRollup,
+    Test,
+    TestInstance,
+)
+from shared.django_apps.reports.models import ReportSession as Upload
+from shared.django_apps.reports.tests.factories import (
+    CommitReportFactory,
+    DailyTestRollupFactory,
+    TestFactory,
+    TestInstanceFactory,
+    UploadFactory,
+)
 
 from services.cleanup.utils import CleanupResult, CleanupSummary
 from tasks.delete_owner import DeleteOwnerTask
@@ -53,26 +67,50 @@ def test_delete_owner_deletes_owner_with_commit_compares(mock_storage):
     compare_commit = CommitFactory(repository=repo, author=user)
     CommitComparisonFactory(base_commit=base_commit, compare_commit=compare_commit)
 
+    report = CommitReportFactory(commit=base_commit)
+    upload = UploadFactory(report=report)
+    test = TestFactory(repository=repo)
+    TestInstanceFactory(test=test, upload=upload)
+    DailyTestRollupFactory(test=test, repoid=repo.repoid)
+
+    # This test factory implicitly creates:
+    # - Test with a Repository and an Owner,
+    # - An Upload with a CommitReport, a Commit that has a different Owner and
+    #   one more Repository with yet another different Owner.
+    # And then also a Branch and a Pull via DB triggers because of the Commit.
+    remaining = TestInstanceFactory()
+
     res = DeleteOwnerTask().run_impl({}, user.ownerid)
 
     assert res == CleanupSummary(
-        CleanupResult(7),
+        CleanupResult(12),
         {
             Branch: CleanupResult(1),
             Commit: CleanupResult(2),
             CommitComparison: CleanupResult(1),
             Owner: CleanupResult(1),
             Pull: CleanupResult(1),
             Repository: CleanupResult(1),
+            CommitReport: CleanupResult(1),
+            Upload: CleanupResult(1),
+            Test: CleanupResult(1),
+            TestInstance: CleanupResult(1),
+            DailyTestRollup: CleanupResult(1),
         },
     )
 
-    assert Branch.objects.count() == 0
-    assert Commit.objects.count() == 0
+    assert list(TestInstance.objects.all()) == [remaining]
+    # See the comment above why we have all of these objects
+    assert Branch.objects.count() == 1
+    assert Commit.objects.count() == 1
     assert CommitComparison.objects.count() == 0
-    assert Owner.objects.count() == 0
-    assert Pull.objects.count() == 0
-    assert Repository.objects.count() == 0
+    assert Owner.objects.count() == 3
+    assert Pull.objects.count() == 1
+    assert Repository.objects.count() == 2
+    assert CommitReport.objects.count() == 1
+    assert Upload.objects.count() == 1
+    assert Test.objects.count() == 1
+    assert DailyTestRollup.objects.count() == 0
 
 
 @pytest.mark.django_db(databases=["timeseries", "default"], transaction=True)