From 0f3618028c56dcd3621091b83878aef7e93071ec Mon Sep 17 00:00:00 2001 From: Jeremy Stretch Date: Fri, 1 Nov 2024 08:49:08 -0400 Subject: [PATCH 1/3] Move model registration logic to register_models() --- netbox_branching/__init__.py | 23 ++++------------------- netbox_branching/utilities.py | 32 +++++++++++++++++++++++++++++++- 2 files changed, 35 insertions(+), 20 deletions(-) diff --git a/netbox_branching/__init__.py b/netbox_branching/__init__.py index 2d3054c..6254f33 100644 --- a/netbox_branching/__init__.py +++ b/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 diff --git a/netbox_branching/utilities.py b/netbox_branching/utilities.py index ea38cea..306be9a 100644 --- a/netbox_branching/utilities.py +++ b/netbox_branching/utilities.py @@ -6,7 +6,9 @@ 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, REPLICATE_TABLES from .contextvars import active_branch __all__ = ( @@ -19,6 +21,7 @@ 'get_tables_to_replicate', 'is_api_request', 'record_applied_change', + 'register_models', 'update_object', ) @@ -80,6 +83,33 @@ 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'), + ) + + # Determine which models support branching + 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 + + def get_tables_to_replicate(): """ Return an ordered list of database tables to replicate when provisioning a new schema. From f32a8a50528b64954e7d1fed84444b2666e4dbc9 Mon Sep 17 00:00:00 2001 From: Jeremy Stretch Date: Fri, 1 Nov 2024 08:49:57 -0400 Subject: [PATCH 2/3] Restore BranchAwareRouter logic removed under #98 --- netbox_branching/database.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/netbox_branching/database.py b/netbox_branching/database.py index 3c749fc..842f881 100644 --- a/netbox_branching/database.py +++ b/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}' From d6a7f2099154961f8bb82a4d156cd90b20da5deb Mon Sep 17 00:00:00 2001 From: Jeremy Stretch Date: Fri, 1 Nov 2024 09:56:19 -0400 Subject: [PATCH 3/3] Declare branchin support for CablePath model --- netbox_branching/constants.py | 8 ++++---- netbox_branching/utilities.py | 26 +++++++++++++++----------- 2 files changed, 19 insertions(+), 15 deletions(-) diff --git a/netbox_branching/constants.py b/netbox_branching/constants.py index 938963d..d90e811 100644 --- a/netbox_branching/constants.py +++ b/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 diff --git a/netbox_branching/utilities.py b/netbox_branching/utilities.py index 306be9a..d7293d0 100644 --- a/netbox_branching/utilities.py +++ b/netbox_branching/utilities.py @@ -1,5 +1,6 @@ import datetime import logging +from collections import defaultdict from contextlib import contextmanager from dataclasses import dataclass @@ -8,7 +9,7 @@ from netbox.plugins import get_plugin_config from netbox.registry import registry -from .constants import EXEMPT_MODELS, REPLICATE_TABLES +from .constants import EXEMPT_MODELS, INCLUDE_MODELS from .contextvars import active_branch __all__ = ( @@ -94,27 +95,30 @@ def register_models(): *get_plugin_config('netbox_branching', 'exempt_models'), ) - # Determine which models support branching - branching_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 - models = [ - model for model in models - if f'{app_label}.{model}' not in exempt_models - ] - if models: - branching_models[app_label] = models + for model in models: + if f'{app_label}.{model}' not in exempt_models: + branching_models[app_label].append(model) - registry['model_features']['branching'] = branching_models + # 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()