|
1 | 1 | import datetime |
2 | 2 | import logging |
| 3 | +from collections import defaultdict |
3 | 4 | from contextlib import contextmanager |
4 | 5 | from dataclasses import dataclass |
5 | 6 |
|
6 | 7 | from django.db.models import ForeignKey, ManyToManyField |
7 | 8 | from django.urls import reverse |
8 | 9 |
|
9 | | -from .constants import REPLICATE_TABLES |
| 10 | +from netbox.plugins import get_plugin_config |
| 11 | +from netbox.registry import registry |
| 12 | +from .constants import EXEMPT_MODELS, INCLUDE_MODELS |
10 | 13 | from .contextvars import active_branch |
11 | 14 |
|
12 | 15 | __all__ = ( |
|
19 | 22 | 'get_tables_to_replicate', |
20 | 23 | 'is_api_request', |
21 | 24 | 'record_applied_change', |
| 25 | + 'register_models', |
22 | 26 | 'update_object', |
23 | 27 | ) |
24 | 28 |
|
@@ -80,11 +84,41 @@ def get_branchable_object_types(): |
80 | 84 | return ObjectType.objects.with_feature('branching') |
81 | 85 |
|
82 | 86 |
|
| 87 | +def register_models(): |
| 88 | + """ |
| 89 | + Register all models which support branching in the NetBox registry. |
| 90 | + """ |
| 91 | + # Compile a list of exempt models (those for which change logging may |
| 92 | + # be enabled, but branching is not supported) |
| 93 | + exempt_models = ( |
| 94 | + *EXEMPT_MODELS, |
| 95 | + *get_plugin_config('netbox_branching', 'exempt_models'), |
| 96 | + ) |
| 97 | + |
| 98 | + # Register all models which support change logging and are not exempt |
| 99 | + branching_models = defaultdict(list) |
| 100 | + for app_label, models in registry['model_features']['change_logging'].items(): |
| 101 | + # Wildcard exclusion for all models in this app |
| 102 | + if f'{app_label}.*' in exempt_models: |
| 103 | + continue |
| 104 | + for model in models: |
| 105 | + if f'{app_label}.{model}' not in exempt_models: |
| 106 | + branching_models[app_label].append(model) |
| 107 | + |
| 108 | + # Register additional included models |
| 109 | + # TODO: Allow plugins to declare additional models? |
| 110 | + for label in INCLUDE_MODELS: |
| 111 | + app_label, model = label.split('.') |
| 112 | + branching_models[app_label].append(model) |
| 113 | + |
| 114 | + registry['model_features']['branching'] = dict(branching_models) |
| 115 | + |
| 116 | + |
83 | 117 | def get_tables_to_replicate(): |
84 | 118 | """ |
85 | 119 | Return an ordered list of database tables to replicate when provisioning a new schema. |
86 | 120 | """ |
87 | | - tables = set(REPLICATE_TABLES) |
| 121 | + tables = set() |
88 | 122 |
|
89 | 123 | branch_aware_models = [ |
90 | 124 | ot.model_class() for ot in get_branchable_object_types() |
|
0 commit comments