Merge pull request #2111 from dandi/ingest-zarr-dandiset-metadata

Set version status to PENDING in ingest_zarr_archive

dandiapi/zarr/tasks/__init__.py

@@ -3,10 +3,12 @@
 from celery import shared_task
 from celery.utils.log import get_task_logger
 from django.db import transaction
+from django.utils import timezone
 from zarr_checksum import compute_zarr_checksum
 from zarr_checksum.generators import S3ClientOptions, yield_files_s3
 
 from dandiapi.api.asset_paths import add_zarr_paths, delete_zarr_paths
+from dandiapi.api.models.version import Version
 from dandiapi.api.storage import get_storage_params
 from dandiapi.zarr.models import ZarrArchive, ZarrArchiveStatus
 
@@ -29,8 +31,10 @@ def ingest_zarr_archive(zarr_id: str, *, force: bool = False):
 
     # Zarr is in correct state, lock until ingestion finishes
     with transaction.atomic():
-        zarr = ZarrArchive.objects.select_for_update().get(
-            zarr_id=zarr_id, status=ZarrArchiveStatus.INGESTING
+        zarr = (
+            ZarrArchive.objects.select_related('dandiset')
+            .select_for_update()
+            .get(zarr_id=zarr_id, status=ZarrArchiveStatus.INGESTING)
         )
 
         # Remove all asset paths associated with this zarr before ingest
@@ -62,6 +66,11 @@ def ingest_zarr_archive(zarr_id: str, *, force: bool = False):
         # Add asset paths after ingest is finished
         add_zarr_paths(zarr)
 
+        # Set version status back to PENDING, and update modified.
+        Version.objects.filter(id=zarr.dandiset.draft_version.id).update(
+            status=Version.Status.PENDING, modified=timezone.now()
+        )
+
 
 def ingest_dandiset_zarrs(dandiset_id: int, **kwargs):
     for zarr in ZarrArchive.objects.filter(dandiset__id=dandiset_id):

dandiapi/zarr/tests/factories.py

@@ -14,6 +14,15 @@ class Meta:
     name = factory.Faker('catch_phrase')
     dandiset = factory.SubFactory(DandisetFactory)
 
+    @factory.post_generation
+    def ensure_draft_version(obj: ZarrArchive, *args, **kwargs):  # type: ignore  # noqa: N805, PGH003
+        from dandiapi.api.tests.factories import DraftVersionFactory
+
+        if obj.dandiset.versions.filter(version='draft').exists():
+            return
+
+        DraftVersionFactory(dandiset=obj.dandiset)
+
 
 class EmbargoedZarrArchiveFactory(ZarrArchiveFactory):
     embargoed = True

dandiapi/zarr/tests/test_ingest_zarr_archive.py

@@ -6,6 +6,7 @@
 from zarr_checksum.checksum import EMPTY_CHECKSUM
 
 from dandiapi.api.models import AssetPath
+from dandiapi.api.models.version import Version
 from dandiapi.api.services.asset import add_asset_to_version
 from dandiapi.zarr.models import ZarrArchive, ZarrArchiveStatus
 from dandiapi.zarr.tasks import ingest_dandiset_zarrs, ingest_zarr_archive
@@ -144,6 +145,29 @@ def test_ingest_zarr_archive_modified(user, draft_version, zarr_archive_factory,
     ingest_zarr_archive(zarr_archive.zarr_id)
 
 
+@pytest.mark.django_db(transaction=True)
+def test_ingest_zarr_archive_sets_version_pending(
+    draft_version_factory, zarr_archive_factory, zarr_file_factory
+):
+    """Ensure that when a zarr is ingested, it sets the version back to PENDING."""
+    draft_version = draft_version_factory(status=Version.Status.VALID)
+    assert draft_version.status == Version.Status.VALID
+
+    # Ensure zarr has non-zero size
+    zarr_archive = zarr_archive_factory(
+        dandiset=draft_version.dandiset, status=ZarrArchiveStatus.UPLOADED
+    )
+    zarr_file_factory(zarr_archive=zarr_archive, size=100)
+
+    # Kick off ingest
+    ingest_zarr_archive(zarr_archive.zarr_id)
+    zarr_archive.refresh_from_db()
+    draft_version.refresh_from_db()
+
+    # Check that version is now `PENDING`, instead of VALID
+    assert draft_version.status == Version.Status.PENDING
+
+
 @pytest.mark.django_db(transaction=True)
 def test_ingest_dandiset_zarrs(dandiset_factory, zarr_archive_factory, zarr_file_factory):
     dandiset = dandiset_factory()