Classification export: pre-filter 'since' exports by changed alleles in SQL (#1778)

annotation/tests/test_annotation_disk_space.py

@@ -49,6 +49,20 @@
 STANDARD = VariantAnnotationPipelineType.STANDARD
 
 
+def past_vep_kwargs() -> dict:
+    """ Timestamps/counts a run carries once VEP has finished, ie ANNOTATION_COMPLETED - what the import
+        lane checks before it does anything. """
+    now = timezone.now()
+    return {
+        "count": 100,
+        "dump_start": now - timedelta(minutes=5),
+        "dump_end": now - timedelta(minutes=4),
+        "dump_count": 100,
+        "annotation_start": now - timedelta(minutes=4),
+        "annotation_end": now,
+    }
+
+
 @override_settings(**get_fake_annotation_settings_dict(columns_version=2))
 class AnnotationRunCleanupTestCase(TestCase):
     """ Each way into the cleanup module, and the one case that must not reach it. """
@@ -134,7 +148,8 @@ def test_failed_import_keeps_everything(self):
         with tempfile.TemporaryDirectory() as tmp_dir, \
                 override_settings(ANNOTATION_VCF_DUMP_DIR=tmp_dir,
                                   ANNOTATION_DELETE_TEMP_FILES_ON_SUCCESS=True):
-            run, paths, _ = self._run_with_output(tmp_dir)
+            run, paths, _ = self._run_with_output(tmp_dir, **past_vep_kwargs())
+            self.assertEqual(run.status, AnnotationStatus.ANNOTATION_COMPLETED)
 
             with mock.patch.object(VEPRunner, "import_results",
                             side_effect=RuntimeError("import blew up")), \
@@ -204,12 +219,8 @@ def _make_run(self, lo_idx, hi_idx, status=AnnotationStatus.CREATED):
         # count stamped as the count lane would have, so the run is ready for the run lanes
         run = AnnotationRun.objects.create(annotation_range_lock=lock, pipeline_type=STANDARD, count=100)
         if status == AnnotationStatus.ANNOTATION_COMPLETED:
-            now = timezone.now()  # past VEP, waiting on the import lane
-            run.dump_start = now - timedelta(minutes=5)
-            run.dump_end = now - timedelta(minutes=4)
-            run.dump_count = 100
-            run.annotation_start = now - timedelta(minutes=4)
-            run.annotation_end = now
+            for k, v in past_vep_kwargs().items():  # past VEP, waiting on the import lane
+                setattr(run, k, v)
             run.vcf_annotated_filename = "/does/not/need/to/exist.vcf.gz"
             run.save()
             self.assertEqual(run.status, AnnotationStatus.ANNOTATION_COMPLETED)

classification/views/exports/classification_export_filter.py

@@ -597,6 +597,31 @@ def _since_flagged_classification_ids(self) -> set[int]:
             created__gte=self.since
         ))
 
+    @cached_property
+    def _since_changed_allele_ids(self) -> set[int]:
+        """
+        Allele ids with any change since the since date, mirroring the per-record checks in _passes_since
+        (which remains the authority) so the whole allele group is still exported when any member changed.
+        Used to pre-filter cms_qs so an incremental export doesn't fetch every last published record
+        only to discard almost all of them in Python.
+        """
+        allele_ids = set(ClassificationModification.objects.filter(
+            is_last_published=True,
+            modified__gte=self.since
+        ).values_list('classification__allele_info__allele_id', flat=True))
+        allele_ids.update(Classification.objects.filter(
+            modified__gt=self.since
+        ).values_list('allele_info__allele_id', flat=True))
+        if flagged_classification_ids := self._since_flagged_classification_ids:
+            allele_ids.update(Classification.objects.filter(
+                pk__in=flagged_classification_ids
+            ).values_list('allele_info__allele_id', flat=True))
+        allele_ids.update(ImportedAlleleInfo.objects.filter(
+            latest_validation__modified__gt=self.since
+        ).values_list('allele_id', flat=True))
+        allele_ids.discard(None)
+        return allele_ids
+
     def _passes_since(self, allele_data: AlleleData) -> bool:
         """
         Is there anything about this AlleleData that indicates it should be included since the since date
@@ -659,6 +684,13 @@ def cms_qs(self) -> QuerySet[ClassificationModification]:
             # only worry about withdrawn if doing 'since' (as we might need to report the withdrawing (json),
             # or at least be aware of it for changes (mvl))
             cms = cms.exclude(classification__withdrawn=True)
+        else:
+            # records without a matched allele can't be found via _since_changed_allele_ids, keep them
+            # here and let _passes_since decide
+            cms = cms.filter(
+                Q(classification__allele_info__allele_id__in=self._since_changed_allele_ids) |
+                Q(classification__allele_info__allele__isnull=True)
+            )
 
         if labs := self.include_sources:
             cms = cms.filter(classification__lab__in=labs)