Atomic IntegerField increment based on another model instance

[Ryan Causey]

Assume the following models for this question:

class Jurisdiction(models.Model):
    name = models.CharField(max_length = 100)

class Account(models.Model):
    account_number = models.PositiveIntegerField(blank = True)
    jurisdiction = models.ForeignKey("Jurisdiction", on_delete = models.CASCADE)

    class Meta:
        """Meta model to define unique_together constraints."""
        #make the account_number unique for a jurisdiction
        unique_together = ("account_number", "jurisdiction")

    def create_account_number(self):
        """Create the account number for this account."""
        # If the account number is already set, do nothing.
        if self.account_number is not None:
            return

        # Get the last account created for the jurisdiction, based on having the largest account number.
        last_account_for_jurisdiction = Account.objects.filter(jurisdiction = self.jurisdiction).order_by("-account_number").first()
        # If the account exists set the account number.
        if last_account_for_jurisdiction is not None:
            self.account_number = last_account_for_jurisdiction.account_number + 1
        # Else there are no other accounts for this jurisdiction, set the first account number.
        else:
            self.account_number = 1

    def save(self, *args, **kwargs):
        # create the account number
        self.create_account_number()
        # call the superclass save to write the account to the database.
        super().save(*args, **kwargs)

Is there a way in Django to be able to atomically set account_number on save based on the largest account number for a given jurisdiction? I was hoping some combination of F expressions, conditional expressions, and subquery expressions would allow me to do this, but they don't appear to allow referencing a different model's fields.

The unique_together will properly prevent duplicate Jurisdiction + account_number pairs from being entered into the DB, but I was hoping there was a way to perform the setting of account_number in an atomic way so I wouldn't have to handle an IntegrityError and try again. I get the feeling what I'm asking for isn't possible, or there's a completely different way I should be going about this. I haven't used AutoField as I need a sequental series of numbers within each Jurisdiction.

Thanks in advance for the help.

[Ryan Causey]

The solution I ended up going with was using a PostgreSQL table lock:

@contextmanager

@transaction.atomic

def lock_table(model):

"""Performs a table lock on the provided model.

Note that this locks the entire table for the given model

and prevents any other operations from occurring on that

model. This should be used sparingly and not be used on

long running processes if possible.

This only works on PostgreSQL databases.

"""

with connection.cursor() as cursor:

# We use the most restrictive lock mode, ACCESS EXCLUSIVE, to prevent any

# other operations from being performed on the table while we hold this lock.

# See: https://www.postgresql.org/docs/current/explicit-locking.html

cursor.execute(f"LOCK TABLE {model._meta.db_table} IN ACCESS EXCLUSIVE MODE")

yield

I found, perhaps accidentally a way of doing it purely in the ORM but it seemed gross, was slower, and probably was a happy accident rather than a real solution. For the morbidly curious, it took making account_number a nullable field and updating the account number after save like so:

# My best guess is this locks the entire table? Without this it doesn't work.
list(Account.objects.select_for_update().all())
# Update ourself
Account.objects.filter(
    pk = self.pk,
).annotate(
    other_accounts_exist = models.Exists(
        Account.objects.filter(
            jurisdiction = models.OuterRef("jurisdiction"),
            account_number__isnull = False,
        ).exclude(
            pk = models.OuterRef("pk"),
        ),
    ),
).update(
    account_number = models.Case(
        models.When(
            other_accounts_exist = True,
            then = models.Subquery(
                Account.objects.filter(
                    jurisdiction = models.OuterRef("jurisdiction"),
                    account_number__isnull = False,
                ).exclude(
                    pk = models.OuterRef("pk"),
                ).order_by(
                    "-account_number",
                ).values(
                    "account_number",
                )[:1]
            ) + models.Value(1),
        ),
        default = 1,
    )
)
# Since we did an update we need to refresh ourselves from the db.
self.refresh_from_db()