Merge pull request #241 from netboxlabs/37-branching-2

37 branching 2

Author: bctiemann

netbox_custom_objects/models.py

@@ -544,19 +544,22 @@ def get_model_with_serializer(self):
         return model
 
     def create_model(self):
+        from netbox_custom_objects.api.serializers import get_serializer_class
         # Get the model and ensure it's registered
         model = self.get_model()
 
         # Ensure the ContentType exists and is immediately available
         features = get_model_features(model)
-        self.object_type.features = features + ['branching']
-        self.object_type.public = True
+        if 'branching' in features:
+            features.remove('branching')
         self.object_type.features = features
+        self.object_type.public = True
         self.object_type.save()
 
         with connection.schema_editor() as schema_editor:
             schema_editor.create_model(model)
 
+        get_serializer_class(model)
         self.register_custom_object_search_index(model)
 
     def save(self, *args, **kwargs):

netbox_custom_objects/templates/netbox_custom_objects/customobject_edit.html

@@ -13,6 +13,10 @@
 {% endblock title %}
 
 {% block form %}
+  {% if branch_warning %}
+    {% include 'netbox_custom_objects/inc/branch_warning.html' %}
+  {% endif %}
+
   {# Render hidden fields #}
   {% for field in form.hidden_fields %}
     {{ field }}

netbox_custom_objects/templates/netbox_custom_objects/inc/branch_warning.html

@@ -0,0 +1,12 @@
+{% load i18n %}
+
+<div class="alert alert-warning d-flex align-items-center mb-4" role="alert">
+  <span class="text-warning fs-1">
+    <i class="mdi mdi-alert"></i>
+  </span>
+  <span class="flex-fill">
+  {% blocktrans trimmed  %}
+    This object has fields that reference objects in other apps and you have a branch active. Care must be taken to not reference an object that only exists in another branch.
+  {% endblocktrans %}
+  </span>
+</div>

netbox_custom_objects/utilities.py

@@ -8,6 +8,7 @@
     "AppsProxy",
     "generate_model",
     "get_viewname",
+    "is_in_branch",
 )
 
 
@@ -108,3 +109,18 @@ def generate_model(*args, **kwargs):
             apps.clear_cache = apps.clear_cache
 
     return model
+
+
+def is_in_branch():
+    """
+    Check if currently operating within a branch.
+
+    Returns:
+        bool: True if currently in a branch, False otherwise.
+    """
+    try:
+        from netbox_branching.contextvars import active_branch
+        return active_branch.get() is not None
+    except ImportError:
+        # Branching plugin not installed
+        return False

netbox_custom_objects/views.py

@@ -29,6 +29,7 @@
 from .models import CustomObject, CustomObjectType, CustomObjectTypeField
 from extras.choices import CustomFieldTypeChoices
 from netbox_custom_objects.constants import APP_LABEL
+from netbox_custom_objects.utilities import is_in_branch
 
 logger = logging.getLogger("netbox_custom_objects.views")
 
@@ -572,6 +573,23 @@ def custom_save(self, commit=True):
 
         return form_class
 
+    def get_extra_context(self, request, obj):
+
+        # Check if we're in a branch and if there are external object pointers
+        has_external_object_pointers = False
+        if is_in_branch():
+            # Check all fields in the custom object type
+            for field in self.object.custom_object_type.fields.all():
+                if field.type in [CustomFieldTypeChoices.TYPE_OBJECT, CustomFieldTypeChoices.TYPE_MULTIOBJECT]:
+                    # Check if the related object type is not from the current app
+                    if field.related_object_type.app_label != APP_LABEL:
+                        has_external_object_pointers = True
+                        break
+
+        return {
+            'branch_warning': has_external_object_pointers,
+        }
+
 
 @register_model_view(CustomObject, "delete")
 class CustomObjectDeleteView(generic.ObjectDeleteView):