Fix asset update/deletion when associated with zarr

dandiapi/api/asset_paths.py

@@ -6,6 +6,7 @@
 from django.db.models import Count, F, QuerySet
 
 from dandiapi.api.models import Asset, AssetPath, AssetPathRelation, Version
+from dandiapi.zarr.models import ZarrArchive
 
 ####################################################################
 # Dandiset and version deletion will cascade to asset path deletion.
@@ -154,3 +155,19 @@ def add_version_asset_paths(version: Version):
     """Add every asset from a version."""
     for asset in version.assets.iterator():
         add_asset_paths(asset, version)
+
+
+def add_zarr_paths(zarr: ZarrArchive):
+    """Add all asset paths that are associated with a zarr."""
+    # Only act on draft assets/versions
+    for asset in zarr.assets.filter(published=False):
+        for version in asset.versions.filter(version='draft'):
+            add_asset_paths(asset, version)
+
+
+def delete_zarr_paths(zarr: ZarrArchive):
+    """Remove all asset paths that are associated with a zarr."""
+    # Only act on draft assets/versions
+    for asset in zarr.assets.filter(published=False):
+        for version in asset.versions.filter(version='draft'):
+            delete_asset_paths(asset, version)

dandiapi/api/tests/test_asset.py

@@ -13,6 +13,7 @@
 from dandiapi.api.models import Asset, AssetBlob, EmbargoedAssetBlob, Version
 from dandiapi.api.models.asset_paths import AssetPath
 from dandiapi.api.models.dandiset import Dandiset
+from dandiapi.zarr.tasks import ingest_zarr_archive
 
 from .fuzzy import HTTP_URL_RE, TIMESTAMP_RE, URN_RE, UTC_ISO_TIMESTAMP_RE, UUID_RE
 
@@ -902,12 +903,21 @@ def test_asset_rest_update_embargo(api_client, user, draft_version, asset, embar
 
 
 @pytest.mark.django_db
-def test_asset_rest_update_zarr(api_client, user, draft_version, asset, zarr_archive):
+def test_asset_rest_update_zarr(
+    api_client, user, draft_version, draft_asset_factory, zarr_archive, zarr_upload_file_factory
+):
     assign_perm('owner', user, draft_version.dandiset)
     api_client.force_authenticate(user=user)
+
+    asset = draft_asset_factory(blob=None, embargoed_blob=None, zarr=zarr_archive)
     draft_version.assets.add(asset)
     add_asset_paths(asset=asset, version=draft_version)
 
+    # Upload file and perform ingest
+    zarr_upload_file_factory(zarr_archive=zarr_archive)
+    ingest_zarr_archive(zarr_archive.zarr_id)
+    zarr_archive.refresh_from_db()
+
     new_path = 'test/asset/rest/update.txt'
     new_metadata = {
         'encodingFormat': 'application/x-zarr',
@@ -916,7 +926,6 @@ def test_asset_rest_update_zarr(api_client, user, draft_version, asset, zarr_arc
         'num': 123,
         'list': ['a', 'b', 'c'],
     }
-
     resp = api_client.put(
         f'/api/dandisets/{draft_version.dandiset.identifier}/'
         f'versions/{draft_version.version}/assets/{asset.asset_id}/',
@@ -929,7 +938,7 @@ def test_asset_rest_update_zarr(api_client, user, draft_version, asset, zarr_arc
         'path': new_path,
         'size': zarr_archive.size,
         'blob': None,
-        'zarr': zarr_archive.zarr_id,
+        'zarr': str(zarr_archive.zarr_id),
         'created': TIMESTAMP_RE,
         'modified': TIMESTAMP_RE,
         'metadata': new_asset.metadata,

dandiapi/zarr/tasks/__init__.py

@@ -10,6 +10,7 @@
 from django.db import transaction
 from django.db.transaction import atomic
 
+from dandiapi.api.asset_paths import add_zarr_paths, delete_zarr_paths
 from dandiapi.api.storage import get_boto_client, yield_files
 from dandiapi.zarr.checksums import (
     ZarrChecksum,
@@ -119,6 +120,9 @@ def ingest_zarr_archive(
     with transaction.atomic():
         zarr = ZarrArchive.objects.select_for_update().get(zarr_id=zarr_id)
 
+        # Remove all asset paths associated with this zarr before ingest
+        delete_zarr_paths(zarr)
+
         # Reset before compute
         if not no_size:
             zarr.size = 0
@@ -178,6 +182,9 @@ def ingest_zarr_archive(
         for asset in zarr.assets.iterator():
             asset.save()
 
+        # Add asset paths after ingest is finished
+        add_zarr_paths(zarr)
+
 
 def ingest_dandiset_zarrs(dandiset_id: int, **kwargs):
     for zarr in ZarrArchive.objects.filter(dandiset__id=dandiset_id):