Merge pull request #1485 from dandi/1469-asset-locking

dandiapi/api/asset_paths.py

@@ -2,7 +2,7 @@
 
 import os
 
-from django.db import transaction
+from django.db import IntegrityError, transaction
 from django.db.models import Count, F, QuerySet, Sum
 from django.db.models.functions import Coalesce
 from tqdm import tqdm
@@ -101,8 +101,23 @@ def search_asset_paths(query: str, version: Version) -> QuerySet[AssetPath] | No
 
 def insert_asset_paths(asset: Asset, version: Version):
     """Add all intermediate paths from an asset and link them together."""
-    # Get or create leaf path
-    leaf, created = AssetPath.objects.get_or_create(path=asset.path, asset=asset, version=version)
+    try:
+        # Get or create leaf path
+        leaf, created = AssetPath.objects.get_or_create(
+            path=asset.path, asset=asset, version=version
+        )
+    except IntegrityError as e:
+        from dandiapi.api.services.asset.exceptions import AssetAlreadyExists
+
+        # If there are simultaneous requests to create the same asset, this check constraint can
+        # fail, and should be handled directly, rather than be allowed to bubble up
+        if 'unique-version-path' in str(e):
+            raise AssetAlreadyExists()
+
+        # Re-raise original exception otherwise
+        raise
+
+    # If the asset was not created, return early, as the work is already done
     if not created:
         return leaf
 

dandiapi/api/tests/test_asset_paths.py

@@ -11,6 +11,7 @@
 )
 from dandiapi.api.models import Asset, AssetPath, Version
 from dandiapi.api.models.asset_paths import AssetPathRelation
+from dandiapi.api.services.asset.exceptions import AssetAlreadyExists
 from dandiapi.api.tasks import publish_dandiset_task
 
 
@@ -96,6 +97,27 @@ def test_asset_path_add_asset_idempotent(draft_version_factory, asset_factory):
     assert path.aggregate_size == asset.size
 
 
+@pytest.mark.django_db
+def test_asset_path_add_asset_conflicting_path(draft_version_factory, asset_factory):
+    # Create asset with version
+    asset1: Asset = asset_factory()
+    asset2: Asset = asset_factory(path=asset1.path)
+    version: Version = draft_version_factory()
+    version.assets.add(asset1)
+    version.assets.add(asset2)
+
+    # Add asset1 paths
+    add_asset_paths(asset1, version)
+    assert version.asset_paths.filter(asset__isnull=False).count() == 1
+
+    # Ensure that adding asset2 raises the correct exception
+    with pytest.raises(AssetAlreadyExists):
+        add_asset_paths(asset2, version)
+
+    # Ensure that there no new asset paths created
+    assert version.asset_paths.filter(asset__isnull=False).count() == 1
+
+
 @pytest.mark.django_db
 def test_asset_path_add_version_asset_paths(draft_version_factory, asset_factory):
     # Create asset with version

dandiapi/api/views/asset.py

@@ -25,6 +25,7 @@
 import os.path
 
 from django.conf import settings
+from django.db import transaction
 from django.db.models import QuerySet
 from django.http import HttpResponseRedirect
 from django.utils.decorators import method_decorator
@@ -316,14 +317,19 @@ def update(self, request, versions__dandiset__pk, versions__version, **kwargs):
         serializer = AssetRequestSerializer(data=self.request.data)
         serializer.is_valid(raise_exception=True)
 
-        asset, _ = change_asset(
-            user=request.user,
-            asset=self.get_object(),
-            version=version,
-            new_asset_blob=serializer.get_blob(),
-            new_zarr_archive=serializer.get_zarr_archive(),
-            new_metadata=serializer.validated_data['metadata'],
-        )
+        # Lock asset for update
+        with transaction.atomic():
+            locked_asset = get_object_or_404(
+                version.assets.select_for_update(), id=self.get_object().id
+            )
+            asset, _ = change_asset(
+                user=request.user,
+                asset=locked_asset,
+                version=version,
+                new_asset_blob=serializer.get_blob(),
+                new_zarr_archive=serializer.get_zarr_archive(),
+                new_metadata=serializer.validated_data['metadata'],
+            )
 
         serializer = AssetDetailSerializer(instance=asset)
         return Response(serializer.data, status=status.HTTP_200_OK)
@@ -338,14 +344,18 @@ def update(self, request, versions__dandiset__pk, versions__version, **kwargs):
                                Only draft versions can be modified.',
     )
     def destroy(self, request, versions__dandiset__pk, versions__version, **kwargs):
-        asset = self.get_object()
         version = get_object_or_404(
             Version,
             dandiset__pk=versions__dandiset__pk,
             version=versions__version,
         )
 
-        remove_asset_from_version(user=request.user, asset=asset, version=version)
+        # Lock asset for delete
+        with transaction.atomic():
+            locked_asset = get_object_or_404(
+                version.assets.select_for_update(), id=self.get_object().id
+            )
+            remove_asset_from_version(user=request.user, asset=locked_asset, version=version)
 
         return Response(None, status=status.HTTP_204_NO_CONTENT)