37 add branching support when creating COT

Author: arthanson

netbox_custom_objects/field_types.py

@@ -31,7 +31,7 @@
 from netbox.tables.columns import BooleanColumn
 
 from netbox_custom_objects.constants import APP_LABEL
-from netbox_custom_objects.utilities import generate_model
+from netbox_custom_objects.utilities import generate_model, get_schema_connection
 
 
 class LazyForeignKey(ForeignKey):

netbox_custom_objects/models.py

@@ -12,7 +12,7 @@
 # from django.contrib.contenttypes.management import create_contenttypes
 from django.contrib.contenttypes.models import ContentType
 from django.core.validators import RegexValidator, ValidationError
-from django.db import connection, IntegrityError, models, transaction
+from django.db import connection, connections, IntegrityError, models, transaction
 from django.db.models import Q
 from django.db.models.functions import Lower
 from django.db.models.signals import pre_delete, post_save
@@ -52,7 +52,7 @@
 
 from netbox_custom_objects.constants import APP_LABEL, RESERVED_FIELD_NAMES
 from netbox_custom_objects.field_types import FIELD_TYPE_CLASS
-from netbox_custom_objects.utilities import generate_model
+from netbox_custom_objects.utilities import generate_model, get_schema_connection
 
 
 class UniquenessConstraintTestError(Exception):
@@ -549,12 +549,11 @@ def create_model(self):
 
         # Ensure the ContentType exists and is immediately available
         features = get_model_features(model)
-        self.object_type.features = features + ['branching']
-        self.object_type.public = True
         self.object_type.features = features
+        self.object_type.public = True
         self.object_type.save()
 
-        with connection.schema_editor() as schema_editor:
+        with get_schema_connection().schema_editor() as schema_editor:
             schema_editor.create_model(model)
 
         self.register_custom_object_search_index(model)
@@ -588,7 +587,7 @@ def delete(self, *args, **kwargs):
         # TODO: Remove this disconnect/reconnect after ObjectType has been exempted from handle_deleted_object
         pre_delete.disconnect(handle_deleted_object)
         object_type.delete()
-        with connection.schema_editor() as schema_editor:
+        with get_schema_connection().schema_editor() as schema_editor:
             schema_editor.delete_model(model)
         pre_delete.connect(handle_deleted_object)
 
@@ -936,7 +935,7 @@ def clean(self):
 
             try:
                 with transaction.atomic():
-                    with connection.schema_editor() as test_schema_editor:
+                    with get_schema_connection().schema_editor() as test_schema_editor:
                         test_schema_editor.alter_field(model, old_field, model_field)
                         # If we get here, the constraint was applied successfully
                         # Now raise a custom exception to rollback the test transaction
@@ -1371,7 +1370,7 @@ def save(self, *args, **kwargs):
         model = self.custom_object_type.get_model()
         model_field.contribute_to_class(model, self.name)
 
-        with connection.schema_editor() as schema_editor:
+        with get_schema_connection().schema_editor() as schema_editor:
             if self._state.adding:
                 schema_editor.add_field(model, model_field)
                 if self.type == CustomFieldTypeChoices.TYPE_MULTIOBJECT:
@@ -1489,7 +1488,7 @@ def delete(self, *args, **kwargs):
         model = self.custom_object_type.get_model()
         model_field.contribute_to_class(model, self.name)
 
-        with connection.schema_editor() as schema_editor:
+        with get_schema_connection().schema_editor() as schema_editor:
             if self.type == CustomFieldTypeChoices.TYPE_MULTIOBJECT:
                 apps = model._meta.apps
                 through_model = apps.get_model(APP_LABEL, self.through_model_name)

netbox_custom_objects/utilities.py

@@ -1,12 +1,22 @@
 import warnings
 
 from django.apps import apps
+from django.db import connection, connections
 
 from netbox_custom_objects.constants import APP_LABEL
 
+# Import branching components if available
+try:
+    from netbox_branching.contextvars import active_branch
+    from netbox_branching.database import BranchAwareRouter
+    BRANCHING_AVAILABLE = True
+except ImportError:
+    BRANCHING_AVAILABLE = False
+
 __all__ = (
     "AppsProxy",
     "generate_model",
+    "get_schema_connection",
     "get_viewname",
 )
 
@@ -59,6 +69,20 @@ def __getattr__(self, attr):
         return getattr(apps, attr)
 
 
+def get_schema_connection():
+    """
+    Get the appropriate database connection for schema operations.
+    Uses BranchAwareRouter logic if branching is available and there's an active branch.
+    """
+    if BRANCHING_AVAILABLE:
+        if branch := active_branch.get():
+            router = BranchAwareRouter()
+            branch_connection_name = router._get_connection(branch)
+            return connections[branch_connection_name]
+
+    return connection
+
+
 def get_viewname(model, action=None, rest_api=False):
     """
     Return the view name for the given model and action, if valid.