Skip to content

Commit

Permalink
Set version status to PENDING in ingest_zarr_archive
Browse files Browse the repository at this point in the history
  • Loading branch information
jjnesbitt committed Dec 17, 2024
1 parent 513f80e commit 38a90fb
Show file tree
Hide file tree
Showing 2 changed files with 35 additions and 2 deletions.
13 changes: 11 additions & 2 deletions dandiapi/zarr/tasks/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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
Expand Down Expand Up @@ -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):
Expand Down
24 changes: 24 additions & 0 deletions dandiapi/zarr/tests/test_ingest_zarr_archive.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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()
Expand Down

0 comments on commit 38a90fb

Please sign in to comment.