Add embargo_end_date field to Dandiset model

dandiapi/api/migrations/0031_dandiset_embargo_end_date.py

@@ -0,0 +1,52 @@
+from __future__ import annotations
+
+from datetime import date
+
+from django.conf import settings
+from django.db import migrations, models
+from tqdm import tqdm
+
+
+def populate_embargo_end_date(apps, schema_editor):
+    Dandiset = apps.get_model('api', 'Dandiset')
+    Version = apps.get_model('api', 'Version')
+
+    embargoed_dandisets = Dandiset.objects.filter(
+        embargo_status='EMBARGOED', versions__metadata__access__0__embargoedUntil__isnull=False
+    )
+    for dandiset in tqdm(embargoed_dandisets.iterator(), total=embargoed_dandisets.count()):
+        draft_version = Version.objects.get(dandiset=dandiset, version='draft')
+
+        # embargoedUntil is stored as an ISO 8601 date string, e.g. "2025-01-01"
+        embargoed_until = draft_version.metadata['access'][0]['embargoedUntil']
+
+        try:
+            dandiset.embargo_end_date = date.fromisoformat(embargoed_until)
+            dandiset.save(update_fields=['embargo_end_date'])
+        except (TypeError, ValueError):
+            pass
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ('api', '0030_alter_asset_path'),
+        migrations.swappable_dependency(settings.AUTH_USER_MODEL),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='dandiset',
+            name='embargo_end_date',
+            field=models.DateField(blank=True, default=None, null=True),
+        ),
+        migrations.RunPython(populate_embargo_end_date, migrations.RunPython.noop),
+        migrations.AddConstraint(
+            model_name='dandiset',
+            constraint=models.CheckConstraint(
+                condition=models.Q(
+                    ('embargo_end_date__isnull', False), ('embargo_status', 'OPEN'), _connector='OR'
+                ),
+                name='embargoed_dandiset_has_embargo_end_date',
+            ),
+        ),
+    ]

dandiapi/api/models/asset.py

@@ -13,8 +13,7 @@
 from django.core.exceptions import ValidationError
 from django.core.validators import RegexValidator
 from django.db import models
-from django.db.models import CharField, DateField, Min, Q
-from django.db.models.functions import Cast
+from django.db.models import Min, Q
 from django.urls import reverse
 from django_extensions.db.models import TimeStampedModel
 
@@ -260,13 +259,7 @@ def access_metadata(self):
         # For zarr assets, is_embargoed is based on the dandiset directly, and a zarr can
         # only be associated with one dandiset, so we know this dandiset is embargoed
         if self.zarr is not None:
-            draft_version = self.zarr.dandiset.draft_version
-
-            # EmbargoedUntil isn't guaranteed to be set
-            embargo_end_date: str | None = draft_version.metadata['access'][0].get('embargoedUntil')
-            if embargo_end_date is not None:
-                access['embargoedUntil'] = embargo_end_date
-
+            access['embargoedUntil'] = self.zarr.dandiset.embargo_end_date.isoformat()
             return access
 
         # In the blob case, we need to consider all dandisets this blob might be associated with,
@@ -282,17 +275,16 @@ def access_metadata(self):
             )
 
         # Retrieve the minimum embargo end date, across all dandisets
-        embargo_end_date: datetime.date | None = (
-            self.blob.assets.filter(versions__isnull=False)
-            .annotate(
-                embargo_end_date=Cast(
-                    Cast('versions__metadata__access__0__embargoedUntil', output_field=CharField()),
-                    output_field=DateField(),
-                )
-            )
-            .aggregate(min_embargo_end_date=Min('embargo_end_date'))['min_embargo_end_date']
-        )
+        embargo_end_date: datetime.date | None = self.blob.assets.filter(
+            versions__isnull=False
+        ).aggregate(min_embargo_end_date=Min('versions__dandiset__embargo_end_date'))[
+            'min_embargo_end_date'
+        ]
 
+        # The only way embargo_end_date can be None here is if asset isn't associated with any
+        # versions (most likely due to being updated). Even so, sometimes these assets are accessed
+        # directly, so we need to handle that case.
+        # TODO: Update once https://github.com/dandi/dandi-archive/issues/2733 is addressed
         if embargo_end_date is not None:
             access['embargoedUntil'] = embargo_end_date.isoformat()
 

dandiapi/api/models/dandiset.py

@@ -6,27 +6,40 @@
 from guardian.models import GroupObjectPermissionBase, UserObjectPermissionBase
 
 
+class DandisetEmbargoStatus(models.TextChoices):
+    EMBARGOED = 'EMBARGOED', 'Embargoed'
+    UNEMBARGOING = 'UNEMBARGOING', 'Unembargoing'
+    OPEN = 'OPEN', 'Open'
+
+
 class Dandiset(TimeStampedModel):
     # Don't add beginning and end markers, so this can be embedded in larger regexes
     IDENTIFIER_REGEX = r'\d{6}'
 
-    class EmbargoStatus(models.TextChoices):
-        EMBARGOED = 'EMBARGOED', 'Embargoed'
-        UNEMBARGOING = 'UNEMBARGOING', 'Unembargoing'
-        OPEN = 'OPEN', 'Open'
+    EmbargoStatus = DandisetEmbargoStatus
 
     embargo_status = models.CharField(
         max_length=max(len(choice[0]) for choice in EmbargoStatus.choices),
         choices=EmbargoStatus.choices,
         default=EmbargoStatus.OPEN,
     )
+    embargo_end_date = models.DateField(null=True, blank=True, default=None)
     starred_users = models.ManyToManyField(
         to=User, through='DandisetStar', related_name='starred_dandisets'
     )
 
     class Meta:
         ordering = ['id']
         permissions = [('owner', 'Owns the dandiset')]
+        constraints = [
+            models.CheckConstraint(
+                name='embargoed_dandiset_has_embargo_end_date',
+                condition=(
+                    models.Q(embargo_end_date__isnull=False)
+                    | models.Q(embargo_status=DandisetEmbargoStatus.OPEN)
+                ),
+            )
+        ]
 
     @property
     def identifier(self) -> str:

dandiapi/api/models/version.py

@@ -224,6 +224,8 @@ def _populate_access_metadata(self):
             if self.dandiset.embargoed
             else AccessType.OpenAccess.value,
         }
+        if self.dandiset.embargo_end_date is not None:
+            access[0]['embargoedUntil'] = self.dandiset.embargo_end_date.isoformat()
 
         return access
 

dandiapi/api/services/dandiset/__init__.py

@@ -5,7 +5,7 @@
 from django.db import transaction
 
 if TYPE_CHECKING:
-    from datetime import datetime
+    from datetime import date
 
     from django.contrib.auth.models import User
 
@@ -90,7 +90,7 @@ def create_embargoed_dandiset(  # noqa: PLR0913
     version_metadata: dict,
     funding_source: str | None,
     award_number: str | None,
-    embargo_end_date: datetime,
+    embargo_end_date: date,
 ) -> tuple[Dandiset, Version]:
     with transaction.atomic():
         dandiset, draft_version = _create_dandiset(
@@ -101,6 +101,7 @@ def create_embargoed_dandiset(  # noqa: PLR0913
         )
 
         dandiset.embargo_status = Dandiset.EmbargoStatus.EMBARGOED
+        dandiset.embargo_end_date = embargo_end_date
         dandiset.full_clean()
         dandiset.save()
 

dandiapi/api/services/embargo/__init__.py

@@ -4,6 +4,7 @@
 from typing import TYPE_CHECKING
 
 from django.db import transaction
+from django.utils import timezone
 
 from dandiapi.api.mail import send_dandiset_unembargoed_message
 from dandiapi.api.models import AssetBlob, Dandiset, Version
@@ -57,8 +58,11 @@ def unembargo_dandiset(ds: Dandiset, user: User):
     logger.info('Set %s assets to PENDING', updated_assets)
     logger.info('Updated %s asset blobs', updated_blobs)
 
-    # Set status to OPEN
-    Dandiset.objects.filter(pk=ds.pk).update(embargo_status=Dandiset.EmbargoStatus.OPEN)
+    # Set status to OPEN, update embargo end date
+    Dandiset.objects.filter(pk=ds.pk).update(
+        embargo_status=Dandiset.EmbargoStatus.OPEN,
+        embargo_end_date=timezone.now().date(),
+    )
     logger.info('Dandiset embargo status updated')
 
     # Fetch version to ensure changed embargo_status is included

dandiapi/api/tests/factories.py

@@ -8,6 +8,7 @@
 from dandischema.consts import DANDI_SCHEMA_VERSION
 from dandischema.models import AccessType
 from django.contrib.auth.models import User
+from django.utils import timezone
 import factory
 import faker
 
@@ -76,6 +77,15 @@ class Meta:
         model = Dandiset
         skip_postgeneration_save = True
 
+    embargo_status = Dandiset.EmbargoStatus.OPEN
+    embargo_end_date = factory.LazyAttribute(
+        lambda self: (
+            timezone.now().date() + datetime.timedelta(days=365 * 2)
+            if self.embargo_status != Dandiset.EmbargoStatus.OPEN
+            else None
+        )
+    )
+
     @factory.post_generation
     def owners(self, create: bool, extracted: list[User] | None) -> None:  # noqa: FBT001
         if not create:

dandiapi/api/tests/test_asset_access_metadata.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from datetime import date
 from typing import TYPE_CHECKING
 
 from dandischema.consts import DANDI_SCHEMA_VERSION
@@ -39,19 +40,15 @@ def test_asset_full_metadata_access(
     )
 
     # Test that access is correctly inferred from embargo status
-    assert embargoed_zarr_asset.full_metadata['access'] == [
-        {'schemaKey': 'AccessRequirements', 'status': AccessType.EmbargoedAccess.value}
-    ]
-    assert embargoed_blob_asset.full_metadata['access'] == [
-        {'schemaKey': 'AccessRequirements', 'status': AccessType.EmbargoedAccess.value}
-    ]
+    for embargoed_asset in [embargoed_zarr_asset, embargoed_blob_asset]:
+        assert embargoed_asset.full_metadata['access'][0]['schemaKey'] == 'AccessRequirements'
+        assert (
+            embargoed_asset.full_metadata['access'][0]['status'] == AccessType.EmbargoedAccess.value
+        )
 
-    assert open_zarr_asset.full_metadata['access'] == [
-        {'schemaKey': 'AccessRequirements', 'status': AccessType.OpenAccess.value}
-    ]
-    assert open_blob_asset.full_metadata['access'] == [
-        {'schemaKey': 'AccessRequirements', 'status': AccessType.OpenAccess.value}
-    ]
+    for open_asset in [open_zarr_asset, open_blob_asset]:
+        assert open_asset.full_metadata['access'][0]['schemaKey'] == 'AccessRequirements'
+        assert open_asset.full_metadata['access'][0]['status'] == AccessType.OpenAccess.value
 
 
 @pytest.mark.django_db
@@ -80,10 +77,10 @@ def test_access_metadata_embargoed_zarr_with_embargoed_until(
     embargoed_zarr_archive_factory, draft_asset_factory
 ):
     """Embargoed zarr asset returns embargoedUntil from draft version."""
-    zarr = embargoed_zarr_archive_factory()
+    zarr = embargoed_zarr_archive_factory(
+        dandiset__embargo_end_date=date.fromisoformat('2026-06-15')
+    )
     draft_version = zarr.dandiset.draft_version
-    draft_version.metadata['access'][0]['embargoedUntil'] = '2026-06-15'
-    draft_version.save()
 
     asset = draft_asset_factory(zarr=zarr, blob=None)
     draft_version.assets.add(asset)
@@ -95,28 +92,15 @@ def test_access_metadata_embargoed_zarr_with_embargoed_until(
     }
 
 
-@pytest.mark.django_db
-def test_access_metadata_embargoed_zarr_without_embargoed_until(
-    embargoed_zarr_archive_factory, draft_asset_factory
-):
-    """Embargoed zarr asset without embargoedUntil on version has no embargoedUntil in access."""
-    zarr = embargoed_zarr_archive_factory()
-    asset = draft_asset_factory(zarr=zarr, blob=None)
-    zarr.dandiset.draft_version.assets.add(asset)
-
-    assert 'embargoedUntil' not in asset.access_metadata()
-
-
 @pytest.mark.django_db
 def test_access_metadata_embargoed_blob_with_embargoed_until(
     embargoed_asset_blob, draft_asset_factory
 ):
     """Embargoed blob asset returns embargoedUntil from version."""
     draft_version = DraftVersionFactory.create(
-        dandiset__embargo_status=Dandiset.EmbargoStatus.EMBARGOED
+        dandiset__embargo_status=Dandiset.EmbargoStatus.EMBARGOED,
+        dandiset__embargo_end_date=date.fromisoformat('2026-06-15'),
     )
-    draft_version.metadata['access'][0]['embargoedUntil'] = '2026-06-15'
-    draft_version.save()
 
     asset = draft_asset_factory(blob=embargoed_asset_blob)
     draft_version.assets.add(asset)
@@ -134,18 +118,16 @@ def test_access_metadata_embargoed_blob_shared_across_embargoed_dandisets(
 ):
     """Blob shared by multiple embargoed dandisets returns minimum embargo end date."""
     version_a = DraftVersionFactory.create(
-        dandiset__embargo_status=Dandiset.EmbargoStatus.EMBARGOED
+        dandiset__embargo_status=Dandiset.EmbargoStatus.EMBARGOED,
+        dandiset__embargo_end_date=date.fromisoformat('2026-08-01'),
     )
-    version_a.metadata['access'][0]['embargoedUntil'] = '2026-08-01'
-    version_a.save()
     asset_a = draft_asset_factory(blob=embargoed_asset_blob)
     version_a.assets.add(asset_a)
 
     version_b = DraftVersionFactory.create(
-        dandiset__embargo_status=Dandiset.EmbargoStatus.EMBARGOED
+        dandiset__embargo_status=Dandiset.EmbargoStatus.EMBARGOED,
+        dandiset__embargo_end_date=date.fromisoformat('2018-10-25'),
     )
-    version_b.metadata['access'][0]['embargoedUntil'] = '2018-10-25'
-    version_b.save()
     asset_b = draft_asset_factory(blob=embargoed_asset_blob)
     version_b.assets.add(asset_b)
 
@@ -165,21 +147,3 @@ def test_access_metadata_embargoed_blob_in_open_dandiset_raises(
 
     with pytest.raises(EmbargoedAssetWithinOpenDandisetError):
         asset.access_metadata()
-
-
-@pytest.mark.django_db
-def test_access_metadata_embargoed_blob_no_embargoed_until(
-    embargoed_asset_blob, draft_asset_factory
-):
-    """Embargoed blob with no embargoedUntil on any version has no embargoedUntil in access."""
-    assets = []
-    for _ in range(5):
-        draft_version = DraftVersionFactory.create(
-            dandiset__embargo_status=Dandiset.EmbargoStatus.EMBARGOED
-        )
-        asset = draft_asset_factory(blob=embargoed_asset_blob)
-        draft_version.assets.add(asset)
-        assets.append(asset)
-
-    for asset in assets:
-        assert 'embargoedUntil' not in asset.access_metadata()