Closes #23426: Add a plugin API for allocating primary keys (#23443)

Allow plugins to supply primary keys for ORM and raw SQL inserts,
passing the target database and the model that owns the primary key.
Fall back to database-generated keys when no allocator supplies them.

Support batch allocation for bulk creates and force inserts to avoid
overwriting existing rows. Withdraw allocated keys after failed inserts
or ignored bulk conflicts.

docs/development/application-registry.md

@@ -33,6 +33,10 @@ Core model features are listed in the [features matrix](./models.md#features-mat
 
 This key lists all models which have been registered in NetBox which are not designated for private use. (Setting `_netbox_private` to True on a model excludes it from this list.) As with individual features under `model_features`, models are organized by app label.
 
+### `pk_allocators`
+
+A list of callables which may supply primary keys for new objects, consulted in order. Primary key allocators are registered using the `register_pk_allocator()` decorator in `netbox.utils`. (See the [plugin development documentation](../plugins/development/models.md#primary-key-allocators) for details.)
+
 ### `plugins`
 
 This store maintains all registered items for plugins, such as navigation menus, template extensions, etc.

docs/plugins/development/models.md

@@ -17,7 +17,7 @@ class MyModel(models.Model):
         return f'{self.foo} {self.bar}'
 ```
 
-Every model includes by default a numeric primary key. This value is generated automatically by the database, and can be referenced as `pk` or `id`.
+Every model includes by default a numeric primary key. This value is generated automatically by the database (unless supplied by a [primary key allocator](#primary-key-allocators)), and can be referenced as `pk` or `id`.
 
 !!! note
     Model names should adhere to [PEP8](https://www.python.org/dev/peps/pep-0008/#class-names) standards and be CapWords (no underscores).  Using underscores in model names will result in problems with permissions.
@@ -210,6 +210,56 @@ Exceptions raised by an implementation propagate to the caller unchanged: `Cable
 
 The caller is responsible for calling the method only once every related object is in place: `Cable.update_dependent_objects()` retraces the cable's paths, which requires its `CableTermination` objects to exist.
 
+## Primary Key Allocators
+
+!!! info "This feature was introduced in NetBox v4.8."
+
+By default, the database assigns each new object's primary key from its table's sequence. A plugin can instead supply primary keys by registering a primary key allocator with the `register_pk_allocator()` decorator from `netbox.utils`. An allocator is a callable which accepts a model class, a count, and the alias of the database to which the new objects will be written. It returns either a list of exactly that many primary keys for the model, or `None` to defer to the next allocator.
+
+```python
+from netbox.utils import register_pk_allocator
+
+@register_pk_allocator
+def my_allocator(model, count, using):
+    if not handles_database(using) or not handles_model(model):
+        return None
+    return reserve_pks(model, count)
+```
+
+Allocators are consulted in the order in which they were registered. If no allocator supplies primary keys, the database assigns them as usual. Registering an allocator which is already registered has no effect.
+
+NetBox consults allocators whenever an object without a primary key is created via the ORM, from both `save()` and `bulk_create()`, for any model with an automatically assigned (auto field) primary key. This includes NetBox's own models, plugin models, the auto-created through models of many-to-many relationships, models belonging to other Django apps (such as `ContentType`), and the historical models used by data migrations (`RunPython` operations). An object which already has a primary key is saved as usual, without consulting allocators.
+
+When an allocator supplies a primary key, the object is inserted as a new row: if that primary key is already in use, `save()` or `bulk_create()` fails rather than overwriting the existing row. If the insert fails for any reason, the allocated primary key is withdrawn from the object, so that saving it again obtains a new primary key from the allocators. Two options of `bulk_create()` change this behavior:
+
+* With `ignore_conflicts=True`, a row whose primary key is already in use is skipped silently, like any other conflicting row. Since there's no telling which rows were inserted, the allocated primary keys are then withdrawn from every object, as `bulk_create()` doesn't set primary keys in this mode.
+* With `update_conflicts=True` and the primary key among `unique_fields`, the existing row with that primary key is updated instead.
+
+A few things to keep in mind when writing an allocator:
+
+* It is called for every model with an auto field primary key, and for every database. It should return `None` promptly for any model or database it doesn't handle.
+* The model it receives may be a historical model, built from migration state while a data migration runs. Identify models by their attributes (such as `model._meta.label` or `model._meta.db_table`) rather than by comparing classes.
+* Under multi-table inheritance, it receives the parent model, which owns the primary key column. For a proxy model, it receives the concrete model.
+* `save()` requests one primary key at a time, whereas `bulk_create()` requests all of the primary keys it needs in a single call. An allocator which obtains primary keys from a remote source may wish to fetch them in blocks and hand them out from a local cache.
+* It must never return a primary key which is already in use. Some operations, such as adding many-to-many relationships or tags, silently skip rows which conflict with existing ones, so a duplicate primary key won't always raise an error.
+* Primary keys it supplies don't advance the table's sequence. If it supplies values which the sequence could later reach, it is responsible for advancing the sequence past them (for example, using PostgreSQL's `setval()`). Otherwise, once the allocator stops supplying primary keys, the database may assign primary keys which are already in use.
+
+!!! tip
+    Consider registering allocators inside your PluginConfig's `ready()` method.
+
+### Inserting Rows Outside the ORM
+
+Code which inserts rows directly, for example using raw SQL (including `RunSQL` migration operations) or PostgreSQL's `COPY`, bypasses the ORM and with it any registered allocators. Such code should call `allocate_pks()` from `netbox.utils` to obtain primary keys for the new rows, and insert them explicitly. This function returns a list of primary keys from the first allocator which supplies them, or `None` if the database should assign them. As with `save()`, allocators receive the model which owns the primary key column, and are consulted only if it's an auto field. Pass the alias of the database being written to as `using`; if it's omitted, the database router's choice for writing the model is assumed.
+
+```python
+from netbox.utils import allocate_pks
+
+pks = allocate_pks(MyModel, count=len(rows), using=connection.alias)
+if pks is None:
+    # Omit the primary key column and let the database assign values
+    ...
+```
+
 ## Choice Sets
 
 For model fields which support the selection of one or more values from a predefined list of choices, NetBox provides the `ChoiceSet` utility class. This can be used in place of a regular choices tuple to provide enhanced functionality, namely dynamic configuration and colorization. (See [Django's documentation](https://docs.djangoproject.com/en/stable/ref/models/fields/#choices) on the `choices` parameter for supported model fields.)

netbox/core/apps.py

@@ -9,6 +9,11 @@ from PIL import Image
 from core.events import *
 from netbox.events import EVENT_TYPE_KIND_DANGER, EVENT_TYPE_KIND_SUCCESS, EVENT_TYPE_KIND_WARNING, EventType
 from utilities.migration import custom_deconstruct
+from utilities.pk_allocation import (
+    bulk_create_with_pk_allocation,
+    prepare_for_bulk_create_with_pk_allocation,
+    save_table_with_pk_allocation,
+)
 
 # Ignore verbose_name & verbose_name_plural Meta options when calculating model migrations
 AlterModelOptions.ALTER_OPTION_KEYS.remove('verbose_name')
@@ -17,6 +22,11 @@ AlterModelOptions.ALTER_OPTION_KEYS.remove('verbose_name_plural')
 # Use our custom destructor to ignore certain attributes when calculating field migrations
 models.Field.deconstruct = custom_deconstruct
 
+# Consult any registered primary key allocators when creating new objects, via either save() or bulk_create()
+models.Model._save_table = save_table_with_pk_allocation
+models.QuerySet._prepare_for_bulk_create = prepare_for_bulk_create_with_pk_allocation
+models.QuerySet.bulk_create = bulk_create_with_pk_allocation
+
 # Cap the maximum size of an image Pillow will decode, to mitigate decompression-bomb DoS attacks. Pillow raises a
 # DecompressionBombError when an image's declared dimensions exceed 2x this value, before allocating pixel buffers.
 # Django's & DRF's ImageField already convert that exception into a validation error, so no additional handling is

netbox/netbox/registry.py

@@ -30,6 +30,7 @@ registry = Registry({
     'filtersets': dict(),
     'model_actions': collections.defaultdict(set),
     'model_features': dict(),
+    'pk_allocators': list(),
     'plugins': dict(),
     'request_processors': list(),
     'search': dict(),

netbox/netbox/tests/test_pk_allocators.py

@@ -0,0 +1,491 @@
+import inspect
+from unittest.mock import patch
+
+from django.apps import apps
+from django.contrib.contenttypes.models import ContentType
+from django.db import DEFAULT_DB_ALIAS, IntegrityError, connection, models, router, transaction
+from django.db.migrations.state import ProjectState
+from django.test import TestCase
+from django.test.utils import CaptureQueriesContext
+
+from account.models import UserToken
+from core.models import ObjectType
+from dcim.models import Manufacturer, Site
+from extras.models import CachedValue, Tag
+from netbox.registry import registry
+from netbox.utils import allocate_pks, register_pk_allocator
+from tenancy.models import Tenant
+from users.models import Token, User
+from utilities import pk_allocation
+
+
+def get_sequential_allocator(model, start=1_000_000):
+    """
+    Return an allocator which supplies sequential primary keys for `model` only, along with a list recording the
+    arguments of each call it receives. Models are matched by table rather than by class, so that the historical
+    models used by data migrations match too.
+    """
+    calls = []
+    next_pk = start
+
+    def allocator(model_, count, using):
+        nonlocal next_pk
+        calls.append((model_, count, using))
+        if model_._meta.db_table != model._meta.db_table:
+            return None
+        pks = list(range(next_pk, next_pk + count))
+        next_pk += count
+        return pks
+
+    return allocator, calls
+
+
+class PKAllocatorTestCase(TestCase):
+
+    def setUp(self):
+        patcher = patch.dict(registry, {'pk_allocators': []})
+        patcher.start()
+        self.addCleanup(patcher.stop)
+
+    def get_table_writes(self, queries, model):
+        """
+        Return the INSERT and UPDATE statements issued against the given model's table.
+        """
+        table = connection.ops.quote_name(model._meta.db_table)
+        sql = [q['sql'] for q in queries]
+        inserts = [s for s in sql if s.startswith(f'INSERT INTO {table}')]
+        updates = [s for s in sql if s.startswith(f'UPDATE {table}')]
+        return inserts, updates
+
+    def test_register_pk_allocator(self):
+        def allocator(model, count, using):
+            return None
+
+        self.assertIs(register_pk_allocator(allocator), allocator)
+        self.assertEqual(registry['pk_allocators'], [allocator])
+
+        register_pk_allocator(allocator)
+        self.assertEqual(registry['pk_allocators'], [allocator])
+
+    def test_allocate_pks_without_allocators(self):
+        self.assertIsNone(allocate_pks(Site))
+        self.assertIsNone(allocate_pks(Site, 5))
+
+    def test_allocate_pks_consults_allocators_in_order(self):
+        calls = []
+
+        @register_pk_allocator
+        def first(model, count, using):
+            # Returns None to defer to the next allocator
+            calls.append('first')
+
+        @register_pk_allocator
+        def second(model, count, using):
+            calls.append('second')
+            return [100 + i for i in range(count)]
+
+        @register_pk_allocator
+        def third(model, count, using):
+            calls.append('third')
+            return [200 + i for i in range(count)]
+
+        self.assertEqual(allocate_pks(Site, 3), [100, 101, 102])
+        self.assertEqual(calls, ['first', 'second'])
+
+    def test_allocate_pks_using(self):
+        allocator, calls = get_sequential_allocator(Site)
+        register_pk_allocator(allocator)
+
+        # The router is asked about the model passed, rather than the model which owns the primary key column
+        with patch.object(router, 'db_for_write', side_effect=lambda model, **hints: model._meta.label):
+            allocate_pks(Site)
+            allocate_pks(Site, using='other')
+            allocate_pks(UserToken)
+            allocate_pks(ObjectType)
+        self.assertEqual(calls, [
+            (Site, 1, 'dcim.Site'),
+            (Site, 1, 'other'),
+            (Token, 1, 'account.UserToken'),
+            (ContentType, 1, 'core.ObjectType'),
+        ])
+
+    def test_allocate_pks_iterable(self):
+        register_pk_allocator(lambda model, count, using: (100 + i for i in range(count)))
+
+        self.assertEqual(allocate_pks(Site, 2), [100, 101])
+
+    def test_allocate_pks_wrong_count(self):
+        register_pk_allocator(lambda model, count, using: [1])
+
+        with self.assertRaises(ValueError):
+            allocate_pks(Site, 2)
+
+    def test_allocate_pks_model(self):
+        # Allocators receive the model which owns the primary key column, and only if it's an auto field
+        calls = []
+        register_pk_allocator(lambda model, count, using: calls.append(model))
+
+        for model in (ObjectType, UserToken, CachedValue):
+            self.assertIsNone(allocate_pks(model))
+        self.assertEqual(calls, [ContentType, Token])
+
+    def test_save_without_allocator(self):
+        manufacturer = Manufacturer(name='Manufacturer 1', slug='manufacturer-1')
+        self.assertIsNone(manufacturer.pk)
+
+        with CaptureQueriesContext(connection) as ctx:
+            manufacturer.save()
+
+        self.assertIsNotNone(manufacturer.pk)
+        inserts, updates = self.get_table_writes(ctx.captured_queries, Manufacturer)
+        self.assertEqual(len(inserts), 1)
+        self.assertEqual(len(updates), 0)
+
+    def test_save_with_allocator(self):
+        allocator, calls = get_sequential_allocator(Manufacturer)
+        register_pk_allocator(allocator)
+        manufacturer = Manufacturer(name='Manufacturer 1', slug='manufacturer-1')
+
+        with CaptureQueriesContext(connection) as ctx:
+            manufacturer.save()
+
+        self.assertEqual(manufacturer.pk, 1_000_000)
+        self.assertEqual(calls.count((Manufacturer, 1, DEFAULT_DB_ALIAS)), 1)
+        inserts, updates = self.get_table_writes(ctx.captured_queries, Manufacturer)
+        self.assertEqual(len(inserts), 1)
+        self.assertEqual(len(updates), 0)
+        self.assertTrue(Manufacturer.objects.filter(pk=1_000_000, name='Manufacturer 1').exists())
+
+    def test_save_cloned_instance_with_allocator(self):
+        source = Manufacturer.objects.create(name='Manufacturer 1', slug='manufacturer-1')
+        allocator, calls = get_sequential_allocator(Manufacturer)
+        register_pk_allocator(allocator)
+        clone = Manufacturer.objects.get(pk=source.pk)
+        clone.pk = None
+        clone.name, clone.slug = 'Manufacturer 2', 'manufacturer-2'
+
+        with CaptureQueriesContext(connection) as ctx:
+            clone.save()
+
+        self.assertEqual(clone.pk, 1_000_000)
+        inserts, updates = self.get_table_writes(ctx.captured_queries, Manufacturer)
+        self.assertEqual(len(inserts), 1)
+        self.assertEqual(len(updates), 0)
+        self.assertEqual(Manufacturer.objects.count(), 2)
+
+    def test_save_with_colliding_allocator(self):
+        # An allocated primary key which is already in use fails on INSERT, rather than overwriting the existing row
+        source = Manufacturer.objects.create(name='Manufacturer 1', slug='manufacturer-1')
+        register_pk_allocator(lambda model, count, using: [source.pk] if model is Manufacturer else None)
+        clone = Manufacturer.objects.get(pk=source.pk)
+        clone.pk = None
+        clone.name, clone.slug = 'Manufacturer 2', 'manufacturer-2'
+
+        with self.assertRaises(IntegrityError), transaction.atomic():
+            clone.save()
+        source.refresh_from_db()
+        self.assertEqual(source.name, 'Manufacturer 1')
+
+    def test_save_retry_after_colliding_allocator(self):
+        # A failed INSERT withdraws the allocated primary key, so that saving the object again consults the allocators
+        # again, rather than updating the existing row which has the allocated primary key
+        source = Manufacturer.objects.create(name='Manufacturer 1', slug='manufacturer-1')
+        pks = iter([source.pk, 1_000_000])
+        register_pk_allocator(lambda model, count, using: [next(pks)] if model is Manufacturer else None)
+        manufacturer = Manufacturer(name='Manufacturer 2', slug='manufacturer-2')
+
+        with self.assertRaises(IntegrityError), transaction.atomic():
+            manufacturer.save()
+        self.assertIsNone(manufacturer.pk)
+
+        with CaptureQueriesContext(connection) as ctx:
+            manufacturer.save()
+
+        self.assertEqual(manufacturer.pk, 1_000_000)
+        inserts, updates = self.get_table_writes(ctx.captured_queries, Manufacturer)
+        self.assertEqual(len(inserts), 1)
+        self.assertEqual(len(updates), 0)
+        source.refresh_from_db()
+        self.assertEqual(source.name, 'Manufacturer 1')
+
+    def test_save_with_declining_allocator(self):
+        # The allocator supplies primary keys only for Site
+        allocator, calls = get_sequential_allocator(Site)
+        register_pk_allocator(allocator)
+        manufacturer = Manufacturer(name='Manufacturer 1', slug='manufacturer-1')
+        manufacturer.save()
+
+        self.assertEqual(calls.count((Manufacturer, 1, DEFAULT_DB_ALIAS)), 1)
+        self.assertIsNotNone(manufacturer.pk)
+        self.assertLess(manufacturer.pk, 1_000_000)
+
+    def test_bulk_create_with_allocator(self):
+        allocator, calls = get_sequential_allocator(Manufacturer)
+        register_pk_allocator(allocator)
+        manufacturers = [
+            Manufacturer(name=f'Manufacturer {i}', slug=f'manufacturer-{i}') for i in range(1, 4)
+        ]
+
+        with CaptureQueriesContext(connection) as ctx:
+            Manufacturer.objects.bulk_create(manufacturers)
+
+        pks = [1_000_000, 1_000_001, 1_000_002]
+        self.assertEqual([m.pk for m in manufacturers], pks)
+        self.assertEqual(calls, [(Manufacturer, 3, DEFAULT_DB_ALIAS)])
+        inserts, updates = self.get_table_writes(ctx.captured_queries, Manufacturer)
+        self.assertEqual(len(inserts), 1)
+        self.assertEqual(sorted(Manufacturer.objects.values_list('pk', flat=True)), pks)
+
+    def test_bulk_create_with_some_preset_pks(self):
+        allocator, calls = get_sequential_allocator(Manufacturer)
+        register_pk_allocator(allocator)
+        manufacturers = [
+            Manufacturer(pk=500, name='Manufacturer 1', slug='manufacturer-1'),
+            Manufacturer(name='Manufacturer 2', slug='manufacturer-2'),
+            Manufacturer(name='Manufacturer 3', slug='manufacturer-3'),
+        ]
+
+        with CaptureQueriesContext(connection) as ctx:
+            Manufacturer.objects.bulk_create(manufacturers)
+
+        self.assertEqual(calls, [(Manufacturer, 2, DEFAULT_DB_ALIAS)])
+        self.assertEqual([m.pk for m in manufacturers], [500, 1_000_000, 1_000_001])
+        inserts, updates = self.get_table_writes(ctx.captured_queries, Manufacturer)
+        self.assertEqual(len(inserts), 1)
+
+    def test_bulk_create_with_declining_allocator(self):
+        allocator, calls = get_sequential_allocator(Site)
+        register_pk_allocator(allocator)
+        manufacturers = [
+            Manufacturer(name=f'Manufacturer {i}', slug=f'manufacturer-{i}') for i in range(1, 4)
+        ]
+
+        Manufacturer.objects.bulk_create(manufacturers)
+
+        self.assertEqual(calls, [(Manufacturer, 3, DEFAULT_DB_ALIAS)])
+        self.assertTrue(all(m.pk is not None and m.pk < 1_000_000 for m in manufacturers))
+
+    def test_bulk_create_using(self):
+        allocator, calls = get_sequential_allocator(Manufacturer)
+        register_pk_allocator(allocator)
+        manufacturers = [
+            Manufacturer(name=f'Manufacturer {i}', slug=f'manufacturer-{i}') for i in range(1, 3)
+        ]
+
+        # The allocator is told the database the queryset writes to, not the router's choice
+        with patch.object(router, 'db_for_write', return_value='elsewhere'):
+            Manufacturer.objects.using(DEFAULT_DB_ALIAS).bulk_create(manufacturers)
+
+        self.assertEqual(calls, [(Manufacturer, 2, DEFAULT_DB_ALIAS)])
+
+    def test_bulk_create_with_colliding_allocator(self):
+        # The first allocated primary key collides with the existing row
+        source = Manufacturer.objects.create(name='Manufacturer 1', slug='manufacturer-1')
+        register_pk_allocator(lambda model, count, using: [source.pk, 1_000_000] if model is Manufacturer else None)
+        manufacturers = [
+            Manufacturer(name=f'Manufacturer {i}', slug=f'manufacturer-{i}') for i in range(2, 4)
+        ]
+
+        with self.assertRaises(IntegrityError), transaction.atomic():
+            Manufacturer.objects.bulk_create(manufacturers)
+
+        self.assertEqual([m.pk for m in manufacturers], [None, None])
+        self.assertTrue(all(m._state.adding for m in manufacturers))
+
+    def test_interrupted_insert(self):
+        # KeyboardInterrupt derives from BaseException rather than Exception
+        allocator, calls = get_sequential_allocator(Manufacturer)
+        register_pk_allocator(allocator)
+        manufacturer = Manufacturer(name='Manufacturer 1', slug='manufacturer-1')
+        manufacturers = [Manufacturer(name='Manufacturer 2', slug='manufacturer-2')]
+
+        interrupt_save = patch.object(models.Model, '_do_insert', side_effect=KeyboardInterrupt)
+        with self.assertRaises(KeyboardInterrupt), interrupt_save, transaction.atomic():
+            manufacturer.save()
+        interrupt_bulk_create = patch.object(models.QuerySet, '_batched_insert', side_effect=KeyboardInterrupt)
+        with self.assertRaises(KeyboardInterrupt), interrupt_bulk_create, transaction.atomic():
+            Manufacturer.objects.bulk_create(manufacturers)
+
+        self.assertEqual(calls.count((Manufacturer, 1, DEFAULT_DB_ALIAS)), 2)
+        self.assertIsNone(manufacturer.pk)
+        self.assertIsNone(manufacturers[0].pk)
+        self.assertTrue(manufacturers[0]._state.adding)
+
+    def test_bulk_create_ignoring_conflicts(self):
+        # When conflicting rows are ignored, bulk_create() can't tell which objects were inserted, so allocated primary
+        # keys are withdrawn from every object. Otherwise, a skipped object would refer to a row it doesn't represent.
+        source = Manufacturer.objects.create(name='Manufacturer 1', slug='manufacturer-1')
+        register_pk_allocator(lambda model, count, using: [source.pk, 1_000_000] if model is Manufacturer else None)
+        manufacturers = [
+            Manufacturer(name=f'Manufacturer {i}', slug=f'manufacturer-{i}') for i in range(2, 4)
+        ]
+
+        Manufacturer.objects.bulk_create(manufacturers, ignore_conflicts=True)
+
+        # The first object conflicts with the existing row and is skipped, while the second is inserted
+        self.assertEqual([m.pk for m in manufacturers], [None, None])
+        self.assertTrue(all(m._state.adding for m in manufacturers))
+        self.assertEqual(
+            sorted(Manufacturer.objects.values_list('pk', 'name')),
+            [(source.pk, 'Manufacturer 1'), (1_000_000, 'Manufacturer 3')]
+        )
+
+    def test_bulk_create_reusing_queryset(self):
+        # A failed bulk_create() doesn't withdraw the primary keys allocated by an earlier one on the same queryset
+        allocator, calls = get_sequential_allocator(Manufacturer)
+        register_pk_allocator(allocator)
+        queryset = Manufacturer.objects.all()
+        manufacturer = Manufacturer(name='Manufacturer 1', slug='manufacturer-1')
+        queryset.bulk_create([manufacturer])
+
+        # These options are mutually exclusive, so the second call fails before allocating any primary keys
+        with self.assertRaises(ValueError):
+            queryset.bulk_create(
+                [Manufacturer(name='Manufacturer 2', slug='manufacturer-2')], ignore_conflicts=True,
+                update_conflicts=True
+            )
+
+        self.assertEqual(manufacturer.pk, 1_000_000)
+        self.assertFalse(manufacturer._state.adding)
+
+    def test_bulk_create_with_unsaved_related_object(self):
+        # Allocators aren't consulted until the objects have been validated, so a failed validation consumes no keys
+        allocator, calls = get_sequential_allocator(Site)
+        register_pk_allocator(allocator)
+        site = Site(name='Site 1', slug='site-1', tenant=Tenant(name='Tenant 1', slug='tenant-1'))
+
+        with self.assertRaisesRegex(ValueError, 'unsaved related object'):
+            Site.objects.bulk_create([site])
+
+        self.assertEqual(calls, [])
+        self.assertIsNone(site.pk)
+
+    def test_many_to_many_with_allocator(self):
+        through_model = Tag.object_types.through
+        allocator, calls = get_sequential_allocator(through_model)
+        register_pk_allocator(allocator)
+        tag = Tag.objects.create(name='Tag 1', slug='tag-1')
+        tag.object_types.add(ContentType.objects.get_for_model(Site))
+
+        self.assertIn((through_model, 1, DEFAULT_DB_ALIAS), calls)
+        self.assertEqual(list(through_model.objects.filter(tag=tag).values_list('pk', flat=True)), [1_000_000])
+
+    def test_third_party_model_with_allocator(self):
+        # ContentType's primary key is an AutoField, rather than NetBox's default BigAutoField
+        allocator, calls = get_sequential_allocator(ContentType)
+        register_pk_allocator(allocator)
+        content_type = ContentType.objects.create(app_label='test', model='test')
+
+        self.assertEqual(content_type.pk, 1_000_000)
+
+    def test_multi_table_inheritance_with_allocator(self):
+        # The allocator receives the parent model, which owns the primary key column
+        allocator, calls = get_sequential_allocator(ContentType)
+        register_pk_allocator(allocator)
+        object_type = ObjectType(app_label='test', model='test')
+        object_type.save()
+
+        self.assertEqual(
+            [call for call in calls if call[0] in (ContentType, ObjectType)],
+            [(ContentType, 1, DEFAULT_DB_ALIAS)]
+        )
+        self.assertEqual(object_type.pk, 1_000_000)
+        self.assertTrue(ContentType.objects.filter(pk=1_000_000, app_label='test', model='test').exists())
+        self.assertTrue(ObjectType.objects.filter(pk=1_000_000).exists())
+
+    def test_multi_table_inheritance_retry_after_failed_child_insert(self):
+        # If the child's INSERT fails, the parent's row is rolled back, so its allocated primary key is withdrawn along
+        # with the child's parent link
+        allocator, calls = get_sequential_allocator(ContentType)
+        register_pk_allocator(allocator)
+        # A null value for the child's non-nullable field fails only the child's INSERT
+        object_type = ObjectType(app_label='test', model='test', public=None)
+
+        with self.assertRaises(IntegrityError), transaction.atomic():
+            object_type.save()
+        self.assertIsNone(object_type.id)
+        self.assertIsNone(object_type.contenttype_ptr_id)
+
+        object_type.public = False
+        with CaptureQueriesContext(connection) as ctx:
+            object_type.save()
+
+        self.assertEqual(object_type.pk, 1_000_001)
+        for model in (ContentType, ObjectType):
+            inserts, updates = self.get_table_writes(ctx.captured_queries, model)
+            self.assertEqual(len(inserts), 1)
+            self.assertEqual(len(updates), 0)
+        self.assertFalse(ContentType.objects.filter(pk=1_000_000).exists())
+
+        object_type.public = None
+        with self.assertRaises(IntegrityError), transaction.atomic():
+            object_type.save()
+        self.assertEqual(object_type.pk, 1_000_001)
+
+    def test_proxy_model_with_allocator(self):
+        # The allocator receives the concrete model of a proxy model
+        user = User.objects.create_user(username='user1')
+        allocator, calls = get_sequential_allocator(Token)
+        register_pk_allocator(allocator)
+        token = UserToken.objects.create(user=user)
+
+        self.assertEqual(
+            [call for call in calls if call[0] in (Token, UserToken)],
+            [(Token, 1, DEFAULT_DB_ALIAS)]
+        )
+        self.assertEqual(token.pk, 1_000_000)
+
+    def test_historical_model_with_allocator(self):
+        HistoricalManufacturer = ProjectState.from_apps(apps).apps.get_model('dcim', 'Manufacturer')
+        self.assertIsNot(HistoricalManufacturer, Manufacturer)
+        allocator, calls = get_sequential_allocator(Manufacturer)
+        register_pk_allocator(allocator)
+
+        manufacturer = HistoricalManufacturer.objects.create(name='Manufacturer 1', slug='manufacturer-1')
+        manufacturers = HistoricalManufacturer.objects.bulk_create([
+            HistoricalManufacturer(name='Manufacturer 2', slug='manufacturer-2'),
+        ])
+
+        self.assertEqual(manufacturer.pk, 1_000_000)
+        self.assertEqual(manufacturers[0].pk, 1_000_001)
+        self.assertEqual(sorted(Manufacturer.objects.values_list('pk', flat=True)), [1_000_000, 1_000_001])
+
+    def test_save_new_instance_with_existing_pk(self):
+        manufacturer = Manufacturer.objects.create(name='Manufacturer 1', slug='manufacturer-1')
+        allocator, calls = get_sequential_allocator(Manufacturer)
+        register_pk_allocator(allocator)
+
+        with CaptureQueriesContext(connection) as ctx:
+            Manufacturer(pk=manufacturer.pk, name='Manufacturer 2', slug='manufacturer-2').save()
+
+        self.assertNotIn((Manufacturer, 1, DEFAULT_DB_ALIAS), calls)
+        inserts, updates = self.get_table_writes(ctx.captured_queries, Manufacturer)
+        self.assertEqual(len(inserts), 0)
+        self.assertEqual(len(updates), 1)
+        manufacturer.refresh_from_db()
+        self.assertEqual(manufacturer.name, 'Manufacturer 2')
+
+    def test_forced_update_without_pk(self):
+        allocator, calls = get_sequential_allocator(Manufacturer)
+        register_pk_allocator(allocator)
+
+        for kwargs in ({'force_update': True}, {'update_fields': ['name']}):
+            with self.subTest(**kwargs), self.assertRaises(ValueError), transaction.atomic():
+                Manufacturer(name='Manufacturer 1', slug='manufacturer-1').save(**kwargs)
+        self.assertNotIn((Manufacturer, 1, DEFAULT_DB_ALIAS), calls)
+
+    def test_replaced_method_signatures(self):
+        # Each replacement must accept the same arguments as the stock Django method it replaces. If this fails after
+        # upgrading Django, update the replacement in utilities.pk_allocation to match the method's new signature.
+        def get_params(func):
+            # Omit the first parameter, which the replacements name differently (e.g. "instance" rather than "self")
+            return [(p.name, p.kind, p.default) for p in list(inspect.signature(func).parameters.values())[1:]]
+
+        replacements = (
+            (pk_allocation._save_table, pk_allocation.save_table_with_pk_allocation),
+            (pk_allocation._prepare_for_bulk_create, pk_allocation.prepare_for_bulk_create_with_pk_allocation),
+            (pk_allocation._bulk_create, pk_allocation.bulk_create_with_pk_allocation),
+        )
+        for method, replacement in replacements:
+            with self.subTest(method=method.__qualname__):
+                self.assertEqual(get_params(replacement), get_params(method))

netbox/netbox/utils.py

@@ -1,10 +1,15 @@
+from django.db import router
+from django.db.models.fields import AutoFieldMixin
+
 from netbox.registry import registry
 from utilities.choices import Choice
 
 __all__ = (
+    'allocate_pks',
     'get_data_backend_choices',
     'register_data_backend',
     'register_model_feature',
+    'register_pk_allocator',
     'register_request_processor',
 )
 
@@ -65,3 +70,53 @@ def register_request_processor(func):
     registry['request_processors'].append(func)
 
     return func
+
+
+def register_pk_allocator(func):
+    """
+    Decorator for registering a primary key allocator: a callable which may supply primary keys for new objects.
+
+    The callable must accept a model class, a count, and the alias of the database to which the new objects will be
+    written (`using`), and return either a list of exactly `count` primary keys for that model, or None to defer to the
+    next allocator (and ultimately to the database). Allocators are consulted in the order in which they were
+    registered. Registering an allocator which is already registered has no effect.
+
+        @register_pk_allocator
+        def my_allocator(model, count, using):
+            ...
+    """
+    if func not in registry['pk_allocators']:
+        registry['pk_allocators'].append(func)
+
+    return func
+
+
+def allocate_pks(model, count=1, using=None):
+    """
+    Return a list of `count` primary keys for new instances of `model` from the first registered allocator which
+    supplies them, or None if the database should assign them. `using` is the alias of the database to which the new
+    instances will be written; if omitted, the database router's choice for writing `model` is assumed. Code which
+    inserts rows outside the ORM (e.g. using raw SQL or COPY) can call this to obtain primary keys in the same manner
+    as Model.save() and QuerySet.bulk_create().
+    """
+    if using is None:
+        using = router.db_for_write(model)
+    # Allocators receive the model which owns the primary key column (the concrete parent, for a proxy or multi-table
+    # inheritance child), and only if it's an auto field
+    pk = model._meta.pk
+    while pk.remote_field and pk.remote_field.parent_link:
+        pk = pk.target_field
+    if not isinstance(pk, AutoFieldMixin):
+        return None
+    model = pk.model
+    for allocator in registry['pk_allocators']:
+        if (pks := allocator(model, count, using)) is not None:
+            pks = list(pks)
+            if len(pks) != count:
+                name = getattr(allocator, '__qualname__', allocator)
+                raise ValueError(
+                    f'Primary key allocator {name} returned {len(pks)} primary keys for {model._meta.label} '
+                    f'(expected {count}).'
+                )
+            return pks
+    return None

netbox/utilities/pk_allocation.py

@@ -0,0 +1,143 @@
+from django.db import models
+from django.db.models.fields import AutoFieldMixin
+
+from netbox.registry import registry
+from netbox.utils import allocate_pks
+
+__all__ = (
+    'bulk_create_with_pk_allocation',
+    'prepare_for_bulk_create_with_pk_allocation',
+    'save_table_with_pk_allocation',
+)
+
+# The stock methods replaced by the functions below. Callers may pass arguments to these positionally, so each
+# function's signature must match that of the method it replaces; a test verifies this.
+_bulk_create = models.QuerySet.bulk_create
+_prepare_for_bulk_create = models.QuerySet._prepare_for_bulk_create
+_save_table = models.Model._save_table
+
+
+def _can_allocate_pks(model):
+    """
+    Return True if primary keys for new instances of the model may be obtained from the registered allocators.
+    """
+    return bool(registry['pk_allocators']) and isinstance(model._meta.pk, AutoFieldMixin)
+
+
+def save_table_with_pk_allocation(
+    instance, raw=False, cls=None, force_insert=False, force_update=False, using=None, update_fields=None
+):
+    """
+    Imitate the behavior of the stock Model._save_table() method, but first consult any registered primary key
+    allocators for a new object without a primary key. If an allocator supplies one, insert the object directly:
+    Django would otherwise attempt an UPDATE using the allocated primary key first, overwriting any existing row which
+    shares it. Objects saved with a primary key already set, and forced updates (which require one), are unaffected.
+
+    If the INSERT fails, or that of a multi-table inheritance child following it, the allocated primary key is withdrawn
+    from the object, so that saving it again consults the allocators again, rather than attempting an UPDATE of
+    whichever row has the allocated primary key.
+
+    Allocators receive the model which owns the primary key column (the concrete parent, for a proxy or multi-table
+    inheritance child), and the alias of the database being written to.
+    """
+    pk = cls._meta.pk
+    concrete_model = instance._meta.concrete_model
+    # The values to restore if the save fails, carried over from the insertion of any multi-table inheritance parents
+    unallocated = vars(instance).pop('_unallocated_pks', {})
+    try:
+        if (
+            _can_allocate_pks(cls)
+            and not (force_update or update_fields)
+            and not instance._is_pk_set(cls._meta)
+            and (pks := allocate_pks(cls, using=using))
+        ):
+            # Parent links are set from the allocated primary key once the parent has been inserted
+            parent_links = [
+                field.attname
+                for model in (concrete_model, *concrete_model._meta.get_parent_list())
+                for field in model._meta.parents.values() if field
+            ]
+            for attname in (pk.attname, *parent_links):
+                unallocated.setdefault(attname, getattr(instance, attname))
+            setattr(instance, pk.attname, pks[0])
+            force_insert = True
+        updated = _save_table(
+            instance, raw=raw, cls=cls, force_insert=force_insert, force_update=force_update, using=using,
+            update_fields=update_fields
+        )
+    except BaseException:
+        for attname, value in unallocated.items():
+            setattr(instance, attname, value)
+        raise
+    if unallocated and cls is not concrete_model:
+        # This is a multi-table inheritance parent, whose child has yet to be inserted
+        instance._unallocated_pks = unallocated
+    return updated
+
+
+def prepare_for_bulk_create_with_pk_allocation(queryset, objs):
+    """
+    Imitate the behavior of the stock QuerySet._prepare_for_bulk_create() method, but then consult any registered
+    primary key allocators for the objects without a primary key, requesting all of their primary keys in a single
+    call. Objects which receive one are inserted along with any whose primary key was already set. Allocators are
+    consulted only once the objects have been validated, so that a failed validation doesn't consume primary keys.
+    The prior state of each object which receives one is recorded on the queryset, for bulk_create_with_pk_allocation()
+    to restore.
+    """
+    objs_with_pk, objs_without_pk = _prepare_for_bulk_create(queryset, objs)
+    pk = queryset.model._meta.pk
+    # QuerySet.db resolves to the database being written to, since bulk_create() has already set _for_write
+    if (
+        objs_without_pk
+        and _can_allocate_pks(queryset.model)
+        and (pks := allocate_pks(queryset.model, len(objs_without_pk), using=queryset.db))
+    ):
+        queryset._unallocated_pks = [
+            (obj, getattr(obj, pk.attname), obj._state.adding, obj._state.db) for obj in objs_without_pk
+        ]
+        for obj, pk_value in zip(objs_without_pk, pks):
+            setattr(obj, pk.attname, pk_value)
+        # Every object now has a primary key, so insert them all together, in their original order
+        return list(objs), []
+    return objs_with_pk, objs_without_pk
+
+
+def bulk_create_with_pk_allocation(
+    queryset, objs, batch_size=None, ignore_conflicts=False, update_conflicts=False, update_fields=None,
+    unique_fields=None
+):
+    """
+    Wrap the stock QuerySet.bulk_create() method to withdraw allocated primary keys from objects which might not have
+    been inserted, leaving them as bulk_create() would without an allocator:
+
+      * If the operation fails, so that saving the objects again consults the allocators again, rather than attempting
+        an UPDATE of whichever rows have the allocated primary keys.
+      * If conflicting rows are ignored, since there's no telling which objects were inserted. (bulk_create() never
+        sets primary keys in this case.)
+    """
+    pk = queryset.model._meta.pk
+
+    def withdraw_allocated_pks():
+        # Restore any objects to which prepare_for_bulk_create_with_pk_allocation() allocated primary keys
+        for obj, pk_value, adding, db in vars(queryset).pop('_unallocated_pks', ()):
+            setattr(obj, pk.attname, pk_value)
+            obj._state.adding = adding
+            obj._state.db = db
+
+    try:
+        objs = _bulk_create(
+            queryset, objs, batch_size=batch_size, ignore_conflicts=ignore_conflicts,
+            update_conflicts=update_conflicts, update_fields=update_fields, unique_fields=unique_fields
+        )
+    except BaseException:
+        withdraw_allocated_pks()
+        raise
+    if ignore_conflicts:
+        withdraw_allocated_pks()
+    else:
+        vars(queryset).pop('_unallocated_pks', None)
+    return objs
+
+
+# Like the stock method, prevent templates from calling this
+bulk_create_with_pk_allocation.alters_data = True