[jeremystretch]

Closes: #23426

This adds a plugin hook for supplying the primary key of a new object before it is inserted. The motivating case is a clone-based netbox-branching backend, which needs every object written to a branch database to draw its ID from main's sequences. Each allocator is told which database the object is being written to. When no allocator is registered, or every allocator returns None, nothing changes: PostgreSQL assigns the ID from the table's sequence as it does today. Objects saved with a primary key already set behave exactly as before, whether or not an allocator is registered.

Changes

• register_pk_allocator() and allocate_pks() (netbox/utils.py), backed by a new pk_allocators registry store.
  • Allocators are called in registration order, with (model, count, using), where using is the alias of the database the new objects will be written to. Each returns a list of count keys, or None to defer to the next.
  • allocate_pks(model, count=1, using=None) is public, so code that inserts with raw SQL or COPY can reserve IDs the same way the ORM does. If using is omitted, it assumes the router's choice for writing the model.
  • allocate_pks() itself works out which model the allocators receive: the one that owns the primary key column, following multi-table inheritance parent links. It returns None without consulting them unless that column is an auto field. So raw-SQL callers and the ORM wrappers present allocators with the same model.
  • Registering an allocator that's already registered has no effect.
• save_table_with_pk_allocation() (utilities/pk_allocation.py) wraps Django's Model._save_table().
  • When the object has no key, it consults the allocators with the save's using. If one supplies a key, it sets the key and passes force_insert=True.
  • Without that, Django would run UPDATE ... WHERE id=<allocated> before the INSERT. That costs a wasted query, and if the key were already in use it would silently overwrite the existing row instead of raising IntegrityError.
  • If the INSERT fails for any reason, including an interruption such as KeyboardInterrupt, the allocated key is withdrawn from the object, so saving it again consults the allocators again instead of attempting an UPDATE of the allocated key. Under multi-table inheritance, a failed INSERT of the child also withdraws the key allocated for its parent (whose row is rolled back), along with the parent links Django copied it into.
  • Objects that already have a key, and forced updates (force_update/update_fields), pass straight through to Django.
• prepare_for_bulk_create_with_pk_allocation() (utilities/pk_allocation.py) wraps QuerySet._prepare_for_bulk_create().
  • It requests keys for every object without one in a single call, passing the queryset's database. Django then inserts those objects together with any whose key was already set.
  • Allocators are consulted only after Django has validated the objects, so a failed validation (such as an unsaved related object) consumes no keys.
  • Many-to-many add() and set() create through rows with bulk_create(), so they're covered too.
• bulk_create_with_pk_allocation() (utilities/pk_allocation.py) wraps QuerySet.bulk_create(), only to withdraw allocated keys from objects which might not have been inserted.
  • If the operation fails, the keys are withdrawn, as for save().
  • With ignore_conflicts=True, there's no telling which rows were inserted, so the keys are withdrawn from every object. This matches Django, which doesn't set primary keys in this mode.
  • Only objects which received an allocated key are touched. The _prepare_for_bulk_create() wrapper records their prior state on the queryset, so a bulk_create() for which every allocator declines does no extra work.
• Model coverage: the wrappers apply to any model whose primary key is an auto field. That includes NetBox's models, plugin models, auto-created through models, other apps' models such as ContentType, and the historical models used by RunPython data migrations. A proxy model's allocator call receives the concrete model, and a multi-table inheritance child's receives the parent. No model fields change, so there are no migrations and no changes to the REST API or GraphQL schema.
• Wiring: core/apps.py installs all three wrappers at import time, next to the existing custom_deconstruct() patch.
• Docs: a "Primary Key Allocators" section in the plugin development models guide, and the new registry store in the application registry docs. The section covers:
  • when allocators are consulted, and that allocated keys are always inserted as new rows
  • that allocated keys are withdrawn if the insert fails, and how bulk_create()'s ignore_conflicts and update_conflicts options change that
  • what an allocator is responsible for: declining promptly for models and databases it doesn't handle, recognizing historical models by table rather than by class, never returning a key already in use, and keeping the table's sequence ahead of the keys it supplies
  • that save() requests one key at a time, while bulk_create() requests all of its keys at once
  • inserting rows outside the ORM with allocate_pks()

Departures from the proposal in the issue

• No NetBoxAutoField. The issue proposed a BigAutoField subclass, set as DEFAULT_AUTO_FIELD, whose get_pk_value_on_save() consults the allocators. That hook has three limits, so this wraps _save_table() and _prepare_for_bulk_create() instead:

  • Django calls it before the object's database is known, so allocators can't be told where the object is going.
  • bulk_create() calls it once per object.
  • It only covers models whose primary key is the subclass. That leaves out the historical models used by data migrations, which see a plain BigAutoField, and models of apps that set their own default_auto_field.

  Dropping the subclass also removes the DEFAULT_AUTO_FIELD change, the Tag.id change, the strawberry-django type map registration, and the plugin guidance about which field to use.

• Allocators are told the target database. The issue's signature was (model, count), leaving allocators to infer the database from their own context. That context doesn't always answer the question. In netbox-branching, for example, an active branch doesn't mean a row is going to the branch: change records and many-to-many rows follow their parent there, while some branch bookkeeping is written to main. The wrappers already know using, so they pass it on.

• No has_default()/get_default() overrides. The issue proposed making has_default() return True so allocated keys skip the UPDATE. But has_default() belongs to the field and can't see the instance being saved, so it changes behavior for every save. Saving a never-saved instance whose key matches an existing row (Site(pk=5, ...).save(), including with update_fields) would raise IntegrityError instead of updating the row, on every install. The _save_table() wrapper forces the INSERT only when an allocator supplied the key, which resolves the first open question in the issue.

• The allocator receives the model that owns the primary key column, not instance._meta.model. With multi-table inheritance that is the parent model, and its table and sequence are where the insert goes. For a proxy, it's the concrete model. allocate_pks() does this resolution itself, so a raw-SQL caller passing ObjectType or UserToken presents allocators with ContentType or Token, just as save() does.

• allocate_pks() checks the allocator's output. It converts the result to a list, and raises ValueError if an allocator returns a different number of keys than requested. A short list would otherwise silently drop rows for raw-SQL callers that zip keys with rows.

• Allocated keys are withdrawn from objects which weren't inserted. The issue didn't cover failures. Without this, an object whose INSERT failed would keep its allocated key, and saving it again would skip the allocators and attempt an UPDATE of that key.

Open questions from the issue

• has_default(): neither option. See the third departure above.
• Batching in bulk_create(): yes. Wrapping _prepare_for_bulk_create() requests a whole batch in one call. Because the wrappers are installed on Django's QuerySet rather than as overrides on NetBox's querysets, they also cover auto-created through models, which an override would have missed.

Risks

• Three Django methods, two of them private. Model._save_table(), QuerySet._prepare_for_bulk_create() and the public QuerySet.bulk_create() are replaced at startup. A Django upgrade could change the private methods' signatures or behavior without notice. A test checks that each replacement's signature matches the stock method's, and the tests assert exactly which INSERT and UPDATE statements each case issues, and how often allocators are called, so such a change would fail CI rather than slip through. NetBox already patches Django at startup (core/apps.py replaces models.Field.deconstruct()).
• Code that bypasses the bulk_create() wrapper, such as a QuerySet subclass that overrides bulk_create() without calling super(), still receives allocated keys, since they're assigned in _prepare_for_bulk_create(). Those keys aren't withdrawn if its insert fails.
• State carried between Django's internal calls. The _save_table() wrapper passes the values to restore from a multi-table inheritance parent's INSERT to the child's on the instance (_unallocated_pks). The bulk wrappers pass them from _prepare_for_bulk_create() to bulk_create() on the queryset. Both are removed before save() or bulk_create() returns.
• Allocators see every model, including third-party and historical models, so a careless allocator has a wider reach. The docs ask allocators to decline promptly for anything they don't handle.
• Overhead. With no allocator registered, each _save_table() call (one per table a save writes) and each bulk_create() adds a few attribute and dictionary lookups. With an allocator registered, every new object, or bulk_create() batch, also pays a call to allocate_pks(), which loops over the registered allocators.

Verification

• New tests in netbox/netbox/tests/test_pk_allocators.py (31 tests) cover:

  • registration, including registering the same allocator twice, and allocator order
  • using: passed explicitly, defaulted from the router, and taken from the queryset rather than the router
  • the count check, and allocators that return an iterable other than a list
  • allocate_pks() given a multi-table inheritance child (ObjectType), a proxy (UserToken) and a UUID-keyed model (CachedValue): allocators receive ContentType and Token, and aren't consulted for the UUID-keyed model
  • save() with no allocator, a matching allocator, and a declining one: one INSERT, no UPDATE, and the allocator consulted once
  • saving a cloned instance (primary key cleared): one INSERT, no UPDATE
  • an allocator that returns a key already in use, which raises IntegrityError and leaves the existing row untouched. Saving the object again obtains a new key and issues one INSERT, no UPDATE.
  • a preset key with an allocator registered: the allocator isn't consulted, and the existing row is updated
  • force_update/update_fields on an object with no key: the allocator isn't consulted, and Django's ValueError is raised as usual
  • bulk_create(): one allocator call and one INSERT for the batch, including when some objects already have keys, and when the allocator declines
  • bulk_create() withdrawing keys after a failed insert and with ignore_conflicts=True, consuming no keys when validation fails, and leaving the objects of an earlier call on the same queryset alone
  • an INSERT interrupted by KeyboardInterrupt, via both save() and bulk_create(), which still withdraws the allocated keys
  • M2M add()
  • a third-party model (ContentType)
  • multi-table inheritance (ObjectType): the allocator receives the parent, and a failed child INSERT withdraws the parent's key and parent link, so the retry issues one INSERT per table
  • a proxy model (UserToken): the allocator receives the concrete model
  • a historical model built from migration state, via both save() and bulk_create()
  • each replacement's signature matching the Django method it replaces

  Removing force_insert = True from the _save_table() wrapper fails five of these tests. Not installing the _prepare_for_bulk_create() wrapper fails ten, and not installing the bulk_create() wrapper fails three.

• No existing tests changed.

• Full test suite: manage.py test --parallel 16 ran 14,539 tests (22 skipped), using a copy of configuration_testing.py with its own test database name. The only failure was LoadLdapConfigTest.test_configuration_dir_setting_matches_active_configuration, which fails under any configuration module other than the stock one. It passes under the stock configuration_testing.py.

• Pre-commit hooks pass: ruff, manage.py check, makemigrations --check, and zensical build.

[jeremystretch]

@pheus this functionality is a bit esoteric; it's being driven by a need surfaced in netbox-branching specifically. Happy to walk through it with you if you'd prefer.

[pheus]

Thanks! I have one routing fix to request, plus a couple of suggestions to trim the repeated comments and docstrings.

[pheus]

Could we resolve using from the original model first? A proxy or inherited child can route differently from its parent, so the allocator may receive the wrong database and incorrectly decline.

[pheus]

Could we shorten the three wrapper docstrings? The allocator contract is already covered in the plugin guide. Keeping the explanations for forced inserts and restoring allocated keys would make these easier to follow.

Suggested change

	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.
	"""
	Allocate missing primary keys in one batch after Django prepares the objects.
	Record their prior state on the queryset for bulk_create_with_pk_allocation() to restore.
	"""

[pheus]

Could we trim comments that repeat the test names or assertions throughout this module? This one restates the assertion below it. The notes explaining the collision and inheritance setup are useful.

[pheus]

Suggested change

	"""
	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.
	"""
	"""
	Force an INSERT when allocating a primary key to avoid overwriting an existing row.
	Restore the original keys and parent links if this table's INSERT or a child INSERT raises.
	"""

[pheus]

Suggested change

	"""
	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.)
	"""
	"""
	Withdraw allocated keys and restore object state on failure or when conflicts are ignored.
	Ignored conflicts prevent us from identifying which objects were inserted.
	"""