Define PK allocator functions in utilities.pk_allocation and re-export from netbox.models

netbox/netbox/models/__init__.py

@@ -2,8 +2,7 @@ from django.conf import settings
 from django.contrib.contenttypes.fields import GenericForeignKey
 from django.core.exceptions import ObjectDoesNotExist
 from django.core.validators import ValidationError
-from django.db import models, router
-from django.db.models.fields import AutoFieldMixin
+from django.db import models
 from django.urls import reverse
 from django.utils.translation import gettext_lazy as _
 from mptt.models import MPTTModel, TreeForeignKey
@@ -11,8 +10,8 @@ from mptt.models import MPTTModel, TreeForeignKey
 from netbox.models.features import *
 from netbox.models.ltree import LtreeManager, LtreeModel, SortPathField
 from netbox.models.mixins import OwnerMixin
-from netbox.registry import registry
 from utilities.mptt import TreeManager
+from utilities.pk_allocation import allocate_pks, register_pk_allocator
 from utilities.querysets import RestrictedQuerySet
 
 __all__ = (
@@ -324,57 +323,3 @@ class AdminModel(
 
     class Meta:
         abstract = True
-
-
-#
-# Primary key allocation
-#
-
-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

@@ -1,11 +1,13 @@
-from django.db import models
+from django.db import models, router
 from django.db.models.fields import AutoFieldMixin
 
 from netbox.registry import registry
 
 __all__ = (
+    'allocate_pks',
     'bulk_create_with_pk_allocation',
     'prepare_for_bulk_create_with_pk_allocation',
+    'register_pk_allocator',
     'save_table_with_pk_allocation',
 )
 
@@ -16,6 +18,56 @@ _prepare_for_bulk_create = models.QuerySet._prepare_for_bulk_create
 _save_table = models.Model._save_table
 
 
+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
+
+
 def _can_allocate_pks(model):
     """
     Return True if primary keys for new instances of the model may be obtained from the registered allocators.
@@ -39,9 +91,6 @@ def save_table_with_pk_allocation(
     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.
     """
-    # This module is imported before the app registry is ready, when netbox.models cannot yet be imported
-    from netbox.models import allocate_pks
-
     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
@@ -86,9 +135,6 @@ def prepare_for_bulk_create_with_pk_allocation(queryset, objs):
     The prior state of each object which receives one is recorded on the queryset, for bulk_create_with_pk_allocation()
     to restore.
     """
-    # This module is imported before the app registry is ready, when netbox.models cannot yet be imported
-    from netbox.models import allocate_pks
-
     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