Fixes #163: Ensure changelog records for non-branching models are created in main schema

netbox_branching/__init__.py

@@ -1,8 +1,8 @@
 from django.conf import settings
 from django.core.exceptions import ImproperlyConfigured
 
-from netbox.plugins import PluginConfig, get_plugin_config
-from netbox.registry import registry
+from netbox.plugins import PluginConfig
+from .utilities import register_models
 
 
 class AppConfig(PluginConfig):
@@ -44,23 +44,8 @@ def ready(self):
                 "netbox_branching: DATABASE_ROUTERS must contain 'netbox_branching.database.BranchAwareRouter'."
             )
 
-        # Record all object types which support branching in the NetBox registry
-        exempt_models = (
-            *constants.EXEMPT_MODELS,
-            *get_plugin_config('netbox_branching', 'exempt_models'),
-        )
-        branching_models = {}
-        for app_label, models in registry['model_features']['change_logging'].items():
-            # Wildcard exclusion for all models in this app
-            if f'{app_label}.*' in exempt_models:
-                continue
-            models = [
-                model for model in models
-                if f'{app_label}.{model}' not in exempt_models
-            ]
-            if models:
-                branching_models[app_label] = models
-        registry['model_features']['branching'] = branching_models
+        # Register models which support branching
+        register_models()
 
 
 config = AppConfig

netbox_branching/constants.py

@@ -10,10 +10,10 @@
 # URL query parameter name
 QUERY_PARAM = '_branch'
 
-# Tables which must be replicated within a branch even though their
-# models don't directly support branching.
-REPLICATE_TABLES = (
-    'dcim_cablepath',
+# Models which do not support change logging, but whose database tables
+# must be replicated for each branch to ensure proper functionality
+INCLUDE_MODELS = (
+    'dcim.cablepath',
 )
 
 # Models for which branching support is explicitly disabled

netbox_branching/database.py

@@ -21,6 +21,11 @@ def _get_db(self, model, **hints):
             warnings.warn(f"Routing database query for {model} before branching support is initialized.")
             return
 
+        # Bail if the model does not support branching
+        app_label, model_name = model._meta.label.lower().split('.')
+        if model_name not in registry['model_features']['branching'].get(app_label, []):
+            return
+
         # Return the schema for the active branch (if any)
         if branch := active_branch.get():
             return f'schema_{branch.schema_name}'

netbox_branching/utilities.py

@@ -1,12 +1,15 @@
 import datetime
 import logging
+from collections import defaultdict
 from contextlib import contextmanager
 from dataclasses import dataclass
 
 from django.db.models import ForeignKey, ManyToManyField
 from django.urls import reverse
 
-from .constants import REPLICATE_TABLES
+from netbox.plugins import get_plugin_config
+from netbox.registry import registry
+from .constants import EXEMPT_MODELS, INCLUDE_MODELS
 from .contextvars import active_branch
 
 __all__ = (
@@ -19,6 +22,7 @@
     'get_tables_to_replicate',
     'is_api_request',
     'record_applied_change',
+    'register_models',
     'update_object',
 )
 
@@ -80,11 +84,41 @@ def get_branchable_object_types():
     return ObjectType.objects.with_feature('branching')
 
 
+def register_models():
+    """
+    Register all models which support branching in the NetBox registry.
+    """
+    # Compile a list of exempt models (those for which change logging may
+    # be enabled, but branching is not supported)
+    exempt_models = (
+        *EXEMPT_MODELS,
+        *get_plugin_config('netbox_branching', 'exempt_models'),
+    )
+
+    # Register all models which support change logging and are not exempt
+    branching_models = defaultdict(list)
+    for app_label, models in registry['model_features']['change_logging'].items():
+        # Wildcard exclusion for all models in this app
+        if f'{app_label}.*' in exempt_models:
+            continue
+        for model in models:
+            if f'{app_label}.{model}' not in exempt_models:
+                branching_models[app_label].append(model)
+
+    # Register additional included models
+    # TODO: Allow plugins to declare additional models?
+    for label in INCLUDE_MODELS:
+        app_label, model = label.split('.')
+        branching_models[app_label].append(model)
+
+    registry['model_features']['branching'] = dict(branching_models)
+
+
 def get_tables_to_replicate():
     """
     Return an ordered list of database tables to replicate when provisioning a new schema.
     """
-    tables = set(REPLICATE_TABLES)
+    tables = set()
 
     branch_aware_models = [
         ot.model_class() for ot in get_branchable_object_types()