Refactor IsicId

isic/conftest.py

@@ -3,7 +3,7 @@
 import pytest
 from pytest_factoryboy import register
 
-from isic.core.tests.factories import CollectionFactory, ImageFactory
+from isic.core.tests.factories import CollectionFactory, ImageFactory, IsicIdFactory
 from isic.ingest.tests.factories import (
     AccessionFactory,
     AccessionReviewFactory,
@@ -88,6 +88,7 @@ def _eager_celery(settings):
 register(ZipUploadFactory)
 
 # core factories
+register(IsicIdFactory)
 register(ImageFactory)
 register(CollectionFactory)
 

isic/core/migrations/0001_initial.py

@@ -9,8 +9,6 @@
 import oauth2_provider.generators
 import s3_file_field.fields
 
-import isic.core.models.isic_id
-
 
 class Migration(migrations.Migration):
     initial = True
@@ -256,7 +254,6 @@ class Migration(migrations.Migration):
                 (
                     "id",
                     models.CharField(
-                        default=isic.core.models.isic_id._default_id,  # noqa: SLF001
                         max_length=12,
                         primary_key=True,
                         serialize=False,

isic/core/migrations/0002_initial.py

@@ -6,8 +6,6 @@
 import django.db.models.deletion
 import django.db.models.functions.text
 
-import isic.core.models.isic_id
-
 
 class Migration(migrations.Migration):
     initial = True
@@ -39,7 +37,6 @@ class Migration(migrations.Migration):
             model_name="image",
             name="isic",
             field=models.OneToOneField(
-                default=isic.core.models.isic_id.IsicId.safe_create,
                 editable=False,
                 on_delete=django.db.models.deletion.PROTECT,
                 to="core.isicid",

isic/core/models/__init__.py

@@ -8,6 +8,7 @@
 from .girder_image import GirderDataset, GirderImage
 from .image import Image
 from .image_alias import ImageAlias
+from .isic_id import IsicId
 from .segmentation import Segmentation, SegmentationReview
 
 __all__ = [
@@ -18,6 +19,7 @@
     "GirderDataset",
     "GirderImage",
     "Image",
+    "IsicId",
     "ImageAlias",
     "IsicOAuthApplication",
     "Segmentation",

isic/core/models/image.py

@@ -58,7 +58,6 @@ class Meta(CreationSortedTimeStampedModel.Meta):
     isic = models.OneToOneField(
         IsicId,
         on_delete=models.PROTECT,
-        default=IsicId.safe_create,
         editable=False,
         verbose_name="isic id",
     )

isic/core/models/isic_id.py

@@ -7,32 +7,25 @@
 from isic.core.constants import ISIC_ID_REGEX
 
 
-def _default_id():
-    while True:
-        isic_id = f"ISIC_{secrets.randbelow(9999999):07}"
-        # This has a race condition, so the actual creation should be retried
-        if not IsicId.objects.filter(id=isic_id).exists():
-            return isic_id
+class IsicIdManager(models.Manager):
+    @retry(
+        reraise=True,
+        retry=retry_if_exception_type(IntegrityError),
+        stop=stop_after_attempt(10),
+    )
+    def create_random(self):
+        return self.create(id=f"ISIC_{secrets.randbelow(9999999):07}")
 
 
 class IsicId(models.Model):
     id = models.CharField(
         primary_key=True,
-        default=_default_id,
         verbose_name="ISIC ID",
         max_length=12,
         validators=[RegexValidator(f"^{ISIC_ID_REGEX}$")],
     )
 
+    objects = IsicIdManager()
+
     def __str__(self) -> str:
         return self.id
-
-    @classmethod
-    @retry(
-        reraise=True,
-        retry=retry_if_exception_type(IntegrityError),
-        stop=stop_after_attempt(10),
-    )
-    def safe_create(cls):
-        """Safely create an IsicId, without race conditions."""
-        return cls.objects.create()

isic/core/services/image/__init__.py

@@ -4,15 +4,17 @@
 from django.db import transaction
 from django.db.models import QuerySet
 
-from isic.core.models.image import Image
+from isic.core.models import Image, IsicId
 from isic.ingest.models.accession import Accession
 
 
 def image_create(*, creator: User, accession: Accession, public: bool) -> Image:
-    image = Image(creator=creator, accession=accession, public=public)
-    image.full_clean()
-    image.save()
-    return image
+    with transaction.atomic():
+        isic_id = IsicId.objects.create_random()
+        image = Image(isic=isic_id, creator=creator, accession=accession, public=public)
+        image.full_clean()
+        image.save()
+        return image
 
 
 def image_share(

isic/core/tests/factories.py

@@ -2,10 +2,18 @@
 import factory.django
 
 from isic.core.models import Collection, Image
+from isic.core.models.isic_id import IsicId
 from isic.factories import UserFactory
 from isic.ingest.tests.factories import AccessionFactory
 
 
+class IsicIdFactory(factory.django.DjangoModelFactory):
+    class Meta:
+        model = IsicId
+
+    id = factory.Faker("pystr_format", string_format="ISIC_#######")
+
+
 class ImageFactory(factory.django.DjangoModelFactory):
     class Meta:
         model = Image
@@ -14,6 +22,7 @@ class Meta:
     creator = factory.SelfAttribute("accession.creator")
     accession = factory.SubFactory(AccessionFactory)
     public = factory.Faker("boolean")
+    isic = factory.SubFactory(IsicIdFactory)
 
 
 class CollectionFactory(factory.django.DjangoModelFactory):

isic/core/tests/test_isic_id.py

@@ -6,7 +6,7 @@
 
 @pytest.mark.django_db()
 def test_isic_id_safe_create_success():
-    IsicId.safe_create()
+    IsicId.objects.create_random()
 
     assert IsicId.objects.count() == 1
 
@@ -15,26 +15,22 @@ def test_isic_id_safe_create_success():
 @pytest.mark.django_db(transaction=True)
 def test_isic_id_safe_create_retry(mocker):
     # Simulate a race condition where the first default returns a collision
-    id_field = IsicId._meta.get_field("id")
-    # Since the "default" property is cached, mock "get_default"
-    mocker.patch.object(id_field, "get_default", side_effect=["ISIC_0000000", "ISIC_0000001"])
+    mocker.patch("isic.core.models.isic_id.secrets.randbelow", side_effect=[0, 1])
 
     IsicId.objects.create(id="ISIC_0000000")
 
-    IsicId.safe_create()
+    IsicId.objects.create_random()
 
     assert IsicId.objects.filter(id="ISIC_0000001").exists()
 
 
 @pytest.mark.django_db(transaction=True)
 def test_isic_id_safe_create_failure(mocker):
     # Simulate very unlikely race condition where every default returns a collision
-    id_field = IsicId._meta.get_field("id")
-    mock_default = mocker.patch.object(id_field, "get_default", return_value="ISIC_0000000")
-
+    mocked_rand = mocker.patch("isic.core.models.isic_id.secrets.randbelow", return_value=0)
     IsicId.objects.create(id="ISIC_0000000")
 
     with pytest.raises(IntegrityError):
-        IsicId.safe_create()
+        IsicId.objects.create_random()
 
-    assert mock_default.call_count == 10
+    assert mocked_rand.call_count == 10

isic/core/utils/db.py

@@ -0,0 +1,14 @@
+from contextlib import contextmanager
+
+from django.db import models, transaction
+
+
+@contextmanager
+def lock_table(cls: type[models.Model]):
+    with transaction.atomic():
+        cursor = transaction.get_connection().cursor()
+        cursor.execute(f"LOCK TABLE {cls._meta.db_table}")
+        try:
+            yield
+        finally:
+            cursor.close()

isic/ingest/services/cohort/__init__.py

@@ -6,10 +6,12 @@
 from django.db.models import Count
 
 from isic.core.models.collection import Collection
+from isic.core.models.isic_id import IsicId
 from isic.core.services.collection import collection_create, collection_merge_magic_collections
 from isic.core.services.collection.image import collection_add_images
 from isic.core.services.image import image_create
 from isic.core.tasks import sync_elasticsearch_index_task
+from isic.core.utils.db import lock_table
 from isic.ingest.models.accession import Accession
 from isic.ingest.models.cohort import Cohort
 from isic.ingest.models.metadata_file import MetadataFile
@@ -44,20 +46,23 @@ def cohort_publish_initialize(
     publish_cohort_task.delay(cohort.pk, publisher.pk, public=public, collection_ids=collection_ids)
 
 
-@transaction.atomic()
 def cohort_publish(
     *, cohort: Cohort, publisher: User, public: bool, collection_ids: list[int] | None = None
 ) -> None:
     additional_collections = (
         Collection.objects.filter(pk__in=collection_ids) if collection_ids else []
     )
 
-    for accession in cohort.accessions.publishable().iterator():
-        image = image_create(creator=publisher, accession=accession, public=public)
-        collection_add_images(collection=cohort.collection, image=image, ignore_lock=True)
+    # this creates a transaction
+    with lock_table(IsicId):
+        for accession in cohort.accessions.publishable().iterator():
+            image = image_create(creator=publisher, accession=accession, public=public)
+            collection_add_images(collection=cohort.collection, image=image, ignore_lock=True)
 
-        for additional_collection in additional_collections:
-            collection_add_images(collection=additional_collection, image=image, ignore_lock=True)
+            for additional_collection in additional_collections:
+                collection_add_images(
+                    collection=additional_collection, image=image, ignore_lock=True
+                )
 
     sync_elasticsearch_index_task.delay()