Преглед изворни кода

Fixes #22922: Honor the saving database connection in scope propagation signals (#22928)

Jeremy Stretch пре 3 дана
родитељ
комит
c2d39b12d8

+ 3 - 2
netbox/dcim/models/device_components.py

@@ -4,7 +4,7 @@ from django.contrib.contenttypes.fields import GenericForeignKey, GenericRelatio
 from django.contrib.postgres.fields import ArrayField
 from django.contrib.postgres.fields import ArrayField
 from django.core.exceptions import ObjectDoesNotExist, ValidationError
 from django.core.exceptions import ObjectDoesNotExist, ValidationError
 from django.core.validators import MaxValueValidator, MinValueValidator
 from django.core.validators import MaxValueValidator, MinValueValidator
-from django.db import models
+from django.db import models, router
 from django.utils.translation import gettext_lazy as _
 from django.utils.translation import gettext_lazy as _
 from mptt.models import MPTTModel, TreeForeignKey
 from mptt.models import MPTTModel, TreeForeignKey
 
 
@@ -1432,7 +1432,8 @@ class ModuleBay(ModularComponentModel, TrackingModelMixin, MPTTModel):
         # root insert (NB-2800). Children still go through MPTT, which keeps
         # root insert (NB-2800). Children still go through MPTT, which keeps
         # siblings in name order via the same order_insertion_by setting.
         # siblings in name order via the same order_insertion_by setting.
         if self._state.adding and self.parent_id is None and not self.lft and not self.rght:
         if self._state.adding and self.parent_id is None and not self.lft and not self.rght:
-            max_tree_id = ModuleBay._objects_raw.aggregate(
+            using = kwargs.get('using') or router.db_for_write(ModuleBay, instance=self)
+            max_tree_id = ModuleBay._objects_raw.using(using).aggregate(
                 models.Max('tree_id')
                 models.Max('tree_id')
             )['tree_id__max'] or 0
             )['tree_id__max'] or 0
             self.tree_id = max_tree_id + 1
             self.tree_id = max_tree_id + 1

+ 1 - 1
netbox/dcim/models/devices.py

@@ -1032,7 +1032,7 @@ class Device(
                 # Set default values for any applicable custom fields
                 # Set default values for any applicable custom fields
                 if cf_defaults := CustomField.objects.get_defaults_for_model(model):
                 if cf_defaults := CustomField.objects.get_defaults_for_model(model):
                     component.custom_field_data = cf_defaults
                     component.custom_field_data = cf_defaults
-                component.save()
+                component.save(using=using)
 
 
     def save(self, *args, **kwargs):
     def save(self, *args, **kwargs):
         is_new = not bool(self.pk)
         is_new = not bool(self.pk)

+ 11 - 8
netbox/dcim/models/modules.py

@@ -334,16 +334,18 @@ class Module(TrackingModelMixin, PrimaryModel):
         old_module_bay_id = None
         old_module_bay_id = None
 
 
         if not is_new:
         if not is_new:
-            old_module_bay_id = Module.objects.filter(pk=self.pk).values_list(
+            old_module_bay_id = Module.objects.using(self._state.db).filter(pk=self.pk).values_list(
                 'module_bay_id', flat=True
                 'module_bay_id', flat=True
             ).first()
             ).first()
 
 
         super().save(*args, **kwargs)
         super().save(*args, **kwargs)
 
 
+        using = self._state.db
+
         if old_module_bay_id is not None and old_module_bay_id != self.module_bay_id:
         if old_module_bay_id is not None and old_module_bay_id != self.module_bay_id:
-            for child_bay in self.modulebays.select_related('module__module_bay'):
+            for child_bay in self.modulebays.db_manager(using).select_related('module__module_bay'):
                 child_bay.snapshot()
                 child_bay.snapshot()
-                child_bay.save()
+                child_bay.save(using=using)
 
 
         adopt_components = getattr(self, '_adopt_components', False)
         adopt_components = getattr(self, '_adopt_components', False)
         disable_replication = getattr(self, '_disable_replication', False)
         disable_replication = getattr(self, '_disable_replication', False)
@@ -353,8 +355,6 @@ class Module(TrackingModelMixin, PrimaryModel):
         if not is_new or (disable_replication and not adopt_components):
         if not is_new or (disable_replication and not adopt_components):
             return
             return
 
 
-        using = self._state.db
-
         # Iterate all component types
         # Iterate all component types
         for templates, component_attribute, component_model in [
         for templates, component_attribute, component_model in [
             ("consoleporttemplates", "consoleports", ConsolePort),
             ("consoleporttemplates", "consoleports", ConsolePort),
@@ -372,7 +372,9 @@ class Module(TrackingModelMixin, PrimaryModel):
             # Prefetch installed components
             # Prefetch installed components
             installed_components = {
             installed_components = {
                 component.name: component
                 component.name: component
-                for component in getattr(self.device, component_attribute).filter(module__isnull=True)
+                for component in getattr(self.device, component_attribute).db_manager(using).filter(
+                    module__isnull=True
+                )
             }
             }
 
 
             # Get the template for the module type.
             # Get the template for the module type.
@@ -420,7 +422,7 @@ class Module(TrackingModelMixin, PrimaryModel):
             else:
             else:
                 # MPTT models must be saved individually to maintain tree structure
                 # MPTT models must be saved individually to maintain tree structure
                 for instance in create_instances:
                 for instance in create_instances:
-                    instance.save()
+                    instance.save(using=using)
 
 
             update_fields = ['module']
             update_fields = ['module']
 
 
@@ -439,7 +441,8 @@ class Module(TrackingModelMixin, PrimaryModel):
 
 
             # Rebuild MPTT tree if needed (bulk_update bypasses model save)
             # Rebuild MPTT tree if needed (bulk_update bypasses model save)
             if issubclass(component_model, MPTTModel) and update_instances:
             if issubclass(component_model, MPTTModel) and update_instances:
-                component_model.objects.rebuild()
+                # db_manager() is used in place of using(), as rebuild() is a manager method
+                component_model.objects.db_manager(using).rebuild()
 
 
         # Replicate any front/rear port mappings from the ModuleType
         # Replicate any front/rear port mappings from the ModuleType
         create_port_mappings(self.device, self.module_type, self)
         create_port_mappings(self.device, self.module_type, self)

+ 77 - 30
netbox/dcim/signals.py

@@ -95,21 +95,28 @@ def handle_location_site_change(instance, created, using=None, **kwargs):
     """
     """
     Update child objects when a Location is saved. All updates are queryset update() calls,
     Update child objects when a Location is saved. All updates are queryset update() calls,
     which fire no signals and generate no change records for the affected objects.
     which fire no signals and generate no change records for the affected objects.
+
+    Each query is pinned to the connection the Location was saved on: on an installation
+    with database routers configured, letting the router pick the alias would both write to
+    a different database than the one being saved and leave the row locks below outside the
+    transaction opened here. For the same reason the new Site is assigned by ID: reading
+    instance.site would fetch the related object over a router-selected connection whenever
+    the save left it uncached (a rename, say).
     """
     """
     if created:
     if created:
         return
         return
-    with transaction.atomic(savepoint=False):
-        instance.get_descendants().update(site=instance.site)
+    with transaction.atomic(using=using, savepoint=False):
+        instance.get_descendants().using(using).update(site_id=instance.site_id)
         # Materialized once so every statement below sees the same membership, even if a
         # Materialized once so every statement below sees the same membership, even if a
         # concurrent commit renumbers the tree mid-handler.
         # concurrent commit renumbers the tree mid-handler.
-        locations = list(instance.get_descendants(include_self=True).values_list('pk', flat=True))
-        Rack.objects.filter(location__in=locations).update(site=instance.site)
-        Device.objects.filter(location__in=locations).update(site=instance.site)
-        PowerPanel.objects.filter(location__in=locations).update(site=instance.site)
-        CableTermination.objects.filter(_location__in=locations).update(_site=instance.site)
+        locations = list(instance.get_descendants(include_self=True).using(using).values_list('pk', flat=True))
+        Rack.objects.using(using).filter(location__in=locations).update(site_id=instance.site_id)
+        Device.objects.using(using).filter(location__in=locations).update(site_id=instance.site_id)
+        PowerPanel.objects.using(using).filter(location__in=locations).update(site_id=instance.site_id)
+        CableTermination.objects.using(using).filter(_location__in=locations).update(_site_id=instance.site_id)
         # Update component models for devices in these locations
         # Update component models for devices in these locations
         for model in COMPONENT_MODELS:
         for model in COMPONENT_MODELS:
-            model.objects.filter(device__location__in=locations).update(_site=instance.site)
+            model.objects.using(using).filter(device__location__in=locations).update(_site_id=instance.site_id)
 
 
         # Objects scoped to descendant Locations receive no post_save of their own from the
         # Objects scoped to descendant Locations receive no post_save of their own from the
         # queryset updates above, so their cached scope fields are updated here whenever the
         # queryset updates above, so their cached scope fields are updated here whenever the
@@ -123,7 +130,8 @@ def handle_location_site_change(instance, created, using=None, **kwargs):
             # a concurrent scope change on that Site serializes against this move; an
             # a concurrent scope change on that Site serializes against this move; an
             # unlocked read could stamp region/group values from before that change.
             # unlocked read could stamp region/group values from before that change.
             site = (
             site = (
-                Site.objects.filter(pk=instance.site_id)
+                Site.objects.using(using)
+                .filter(pk=instance.site_id)
                 .select_for_update(no_key=True)
                 .select_for_update(no_key=True)
                 .values('region_id', 'group_id')
                 .values('region_id', 'group_id')
                 .first()
                 .first()
@@ -132,9 +140,12 @@ def handle_location_site_change(instance, created, using=None, **kwargs):
                 # Select rows through the authoritative scope rather than the cached
                 # Select rows through the authoritative scope rather than the cached
                 # _location, which may itself be stale; scope_id doubles as the correct
                 # _location, which may itself be stale; scope_id doubles as the correct
                 # _location value for Location-scoped rows.
                 # _location value for Location-scoped rows.
-                location_ct = ContentType.objects.get_for_model(Location)
+                # The content type is read on the saving connection as well, since its ID is
+                # fed straight into the pinned filter below; a router-selected read could
+                # return an ID which means something else on that connection.
+                location_ct = ContentType.objects.db_manager(using).get_for_model(Location)
                 for model in (Prefix, Cluster, WirelessLAN):
                 for model in (Prefix, Cluster, WirelessLAN):
-                    model.objects.filter(scope_type=location_ct, scope_id__in=locations).update(
+                    model.objects.using(using).filter(scope_type=location_ct, scope_id__in=locations).update(
                         _location_id=F('scope_id'),
                         _location_id=F('scope_id'),
                         _site_id=instance.site_id,
                         _site_id=instance.site_id,
                         _region_id=site['region_id'],
                         _region_id=site['region_id'],
@@ -160,31 +171,38 @@ def handle_location_site_change(instance, created, using=None, **kwargs):
 
 
 
 
 @receiver(post_save, sender=Rack)
 @receiver(post_save, sender=Rack)
-def handle_rack_site_change(instance, created, **kwargs):
+def handle_rack_site_change(instance, created, using=None, **kwargs):
     """
     """
-    Update child Devices if Site or Location assignment has changed.
+    Update child Devices if Site or Location assignment has changed. Queries are pinned to
+    the connection the Rack was saved on, and the new values are assigned by ID so that no
+    related object is fetched over a router-selected connection.
     """
     """
     if not created:
     if not created:
-        Device.objects.filter(rack=instance).update(site=instance.site, location=instance.location)
+        Device.objects.using(using).filter(rack=instance).update(
+            site_id=instance.site_id,
+            location_id=instance.location_id,
+        )
         # Update component models for devices in this rack
         # Update component models for devices in this rack
         for model in COMPONENT_MODELS:
         for model in COMPONENT_MODELS:
-            model.objects.filter(device__rack=instance).update(
-                _site=instance.site,
-                _location=instance.location,
+            model.objects.using(using).filter(device__rack=instance).update(
+                _site_id=instance.site_id,
+                _location_id=instance.location_id,
             )
             )
 
 
 
 
 @receiver(post_save, sender=Device)
 @receiver(post_save, sender=Device)
-def handle_device_site_change(instance, created, **kwargs):
+def handle_device_site_change(instance, created, using=None, **kwargs):
     """
     """
     Update child components to update the parent Site, Location, and Rack when a Device is saved.
     Update child components to update the parent Site, Location, and Rack when a Device is saved.
+    Queries are pinned to the connection the Device was saved on, and the new values are
+    assigned by ID so that no related object is fetched over a router-selected connection.
     """
     """
     if not created:
     if not created:
         for model in COMPONENT_MODELS:
         for model in COMPONENT_MODELS:
-            model.objects.filter(device=instance).update(
-                _site=instance.site,
-                _location=instance.location,
-                _rack=instance.rack,
+            model.objects.using(using).filter(device=instance).update(
+                _site_id=instance.site_id,
+                _location_id=instance.location_id,
+                _rack_id=instance.rack_id,
             )
             )
 
 
 
 
@@ -323,9 +341,30 @@ def update_mac_address_interface(instance, created, raw, **kwargs):
         instance.primary_mac_address.save()
         instance.primary_mac_address.save()
 
 
 
 
+def _get_scope_object(scope_type_id, scope_id, using):
+    """
+    Return the object referenced by a CachedScopeMixin generic scope, read on the given
+    database connection. The ancestors which cache_related_objects() traverses are selected
+    in the same query, so recomputing the cached fields from the returned object issues no
+    further reads. Returns None if the scope is unset or dangling.
+    """
+    if scope_type_id is None or scope_id is None:
+        return None
+    scope_type = ContentType.objects.db_manager(using).get_for_id(scope_type_id)
+    scope_model = scope_type.model_class()
+    if scope_model is None:
+        return None
+    queryset = scope_model._base_manager.using(using)
+    if scope_model is Location:
+        queryset = queryset.select_related('site__region', 'site__group')
+    elif scope_model is Site:
+        queryset = queryset.select_related('region', 'group')
+    return queryset.filter(pk=scope_id).first()
+
+
 @receiver(post_save, sender=Location)
 @receiver(post_save, sender=Location)
 @receiver(post_save, sender=Site)
 @receiver(post_save, sender=Site)
-def sync_cached_scope_fields(instance, created, **kwargs):
+def sync_cached_scope_fields(instance, created, using=None, **kwargs):
     """
     """
     Rebuild cached scope fields for all CachedScopeMixin-based models
     Rebuild cached scope fields for all CachedScopeMixin-based models
     affected by a change to a Site or Location.
     affected by a change to a Site or Location.
@@ -362,9 +401,9 @@ def sync_cached_scope_fields(instance, created, **kwargs):
 
 
     # These models are explicitly listed because they all subclass CachedScopeMixin
     # These models are explicitly listed because they all subclass CachedScopeMixin
     # and therefore require their cached scope fields to be recomputed.
     # and therefore require their cached scope fields to be recomputed.
-    with transaction.atomic(savepoint=False):
+    with transaction.atomic(using=using, savepoint=False):
         for model in (Prefix, Cluster, WirelessLAN):
         for model in (Prefix, Cluster, WirelessLAN):
-            qs = model.objects.filter(**filters)
+            qs = model.objects.using(using).filter(**filters)
 
 
             # Recompute the cached fields once per distinct scope, then apply each result with a
             # Recompute the cached fields once per distinct scope, then apply each result with a
             # single UPDATE. This avoids loading every object into memory as well as the per-row
             # single UPDATE. This avoids loading every object into memory as well as the per-row
@@ -376,11 +415,19 @@ def sync_cached_scope_fields(instance, created, **kwargs):
             # all-or-nothing outside a request transaction.
             # all-or-nothing outside a request transaction.
             scopes = qs.values_list('scope_type_id', 'scope_id').order_by('scope_type_id', 'scope_id').distinct()
             scopes = qs.values_list('scope_type_id', 'scope_id').order_by('scope_type_id', 'scope_id').distinct()
             for scope_type_id, scope_id in scopes:
             for scope_type_id, scope_id in scopes:
-                ref = model(scope_type_id=scope_type_id, scope_id=scope_id)
+                # Resolve the scope (and the ancestors cache_related_objects() traverses) on
+                # the saving connection, then hand it to a throwaway reference object with
+                # its relations already populated, so that recomputing the cached fields
+                # reads nothing further. Assigning ref._state.db alone would not suffice:
+                # Django consults DATABASE_ROUTERS first for related-object lookups and only
+                # falls back to the instance's recorded database when every router declines.
+                ref = model()
+                ref._state.db = using
+                ref.scope = _get_scope_object(scope_type_id, scope_id, using)
                 ref.cache_related_objects()
                 ref.cache_related_objects()
                 qs.filter(scope_type_id=scope_type_id, scope_id=scope_id).update(
                 qs.filter(scope_type_id=scope_type_id, scope_id=scope_id).update(
-                    _location=ref._location,
-                    _site=ref._site,
-                    _site_group=ref._site_group,
-                    _region=ref._region,
+                    _location_id=ref._location_id,
+                    _site_id=ref._site_id,
+                    _site_group_id=ref._site_group_id,
+                    _region_id=ref._region_id,
                 )
                 )

+ 124 - 1
netbox/dcim/tests/test_models.py

@@ -1,9 +1,10 @@
 from decimal import Decimal
 from decimal import Decimal
+from unittest.mock import patch
 
 
 from django.core.exceptions import ValidationError
 from django.core.exceptions import ValidationError
 from django.db.models import ProtectedError
 from django.db.models import ProtectedError
 from django.db.models.signals import post_save
 from django.db.models.signals import post_save
-from django.test import TestCase, tag
+from django.test import TestCase, override_settings, tag
 
 
 from circuits.models import *
 from circuits.models import *
 from core.models import ObjectType
 from core.models import ObjectType
@@ -15,6 +16,7 @@ from ipam.models import Prefix
 from netbox.choices import WeightUnitChoices
 from netbox.choices import WeightUnitChoices
 from tenancy.models import Tenant
 from tenancy.models import Tenant
 from utilities.data import drange
 from utilities.data import drange
+from utilities.testing import PinnedConnectionRouter
 from virtualization.models import Cluster, ClusterType
 from virtualization.models import Cluster, ClusterType
 
 
 
 
@@ -2923,3 +2925,124 @@ class PowerPortDrawTestCase(TestCase):
         self.assertEqual(legs_by_name['A']['maximum'], 200)
         self.assertEqual(legs_by_name['A']['maximum'], 200)
         self.assertEqual(legs_by_name['B']['allocated'], 0)
         self.assertEqual(legs_by_name['B']['allocated'], 0)
         self.assertEqual(legs_by_name['C']['allocated'], 0)
         self.assertEqual(legs_by_name['C']['allocated'], 0)
+
+
+class ComponentInstantiationConnectionTestCase(TestCase):
+    """
+    Verify that component instantiation issues its queries against the connection the
+    parent object was written to, rather than letting DATABASE_ROUTERS select one. On an
+    installation with routers configured (e.g. netbox_branching), a routed query reads or
+    writes the component in the wrong database.
+
+    Where a path instantiates components, PinnedConnectionRouter cannot be used: Django's
+    own forward-relation descriptor consults the router when a related object is assigned
+    to an unsaved instance. Those paths are checked by capturing the alias handed to the
+    call instead.
+    """
+    @classmethod
+    def setUpTestData(cls):
+        cls.site = Site.objects.create(name='Site 1', slug='site-1')
+        manufacturer = Manufacturer.objects.create(name='Manufacturer 1', slug='manufacturer-1')
+        cls.device_type = DeviceType.objects.create(manufacturer=manufacturer, model='Device Type 1')
+        cls.device_role = DeviceRole.objects.create(name='Device Role 1', slug='device-role-1')
+        cls.module_type = ModuleType.objects.create(manufacturer=manufacturer, model='Module Type 1')
+
+    def _record_module_bay_save_aliases(self):
+        """
+        Patch ModuleBay.save() to record the database alias passed to each call.
+        """
+        aliases = []
+        original_save = ModuleBay.save
+
+        def record_alias(instance, *args, **kwargs):
+            aliases.append(kwargs.get('using'))
+            return original_save(instance, *args, **kwargs)
+
+        return aliases, patch.object(ModuleBay, 'save', record_alias)
+
+    def test_module_bay_tree_id_lookup_pinned_to_saving_connection(self):
+        """
+        Inserting a root ModuleBay looks up the highest existing tree ID, which must be
+        read from the connection the bay is being written to.
+        """
+        device = Device.objects.create(
+            name='Device 1', device_type=self.device_type, role=self.device_role, site=self.site
+        )
+        # Instantiate outside the router, as assigning the Device consults it.
+        module_bay = ModuleBay(device=device, name='Module Bay 1')
+
+        with override_settings(DATABASE_ROUTERS=[PinnedConnectionRouter(ModuleBay)]):
+            module_bay.save(using='default')
+
+        self.assertTrue(ModuleBay.objects.filter(pk=module_bay.pk).exists())
+
+    def test_device_module_bays_receive_saving_connection(self):
+        """
+        ModuleBays are instantiated individually (rather than in bulk) to maintain the MPTT
+        tree, so each save() must be given the Device's connection.
+        """
+        ModuleBayTemplate.objects.create(device_type=self.device_type, name='Module Bay 1')
+
+        device = Device(
+            name='Device 1', device_type=self.device_type, role=self.device_role, site=self.site
+        )
+        aliases, spy = self._record_module_bay_save_aliases()
+        with spy:
+            device.save()
+
+        self.assertEqual(aliases, [device._state.db])
+        self.assertEqual(ModuleBay.objects.filter(device=device).count(), 1)
+
+    def test_module_module_bays_receive_saving_connection(self):
+        """
+        Replicated MPTT components are likewise saved individually, and must be given the
+        Module's connection.
+        """
+        ModuleBayTemplate.objects.create(module_type=self.module_type, name='Module Bay 1')
+
+        device = Device.objects.create(
+            name='Device 1', device_type=self.device_type, role=self.device_role, site=self.site
+        )
+        parent_bay = ModuleBay.objects.create(device=device, name='Parent Bay')
+
+        module = Module(device=device, module_bay=parent_bay, module_type=self.module_type)
+        aliases, spy = self._record_module_bay_save_aliases()
+        with spy:
+            module.save()
+
+        self.assertEqual(aliases, [module._state.db])
+        self.assertEqual(ModuleBay.objects.filter(module=module).count(), 1)
+
+    def test_module_component_rebuild_uses_saving_connection(self):
+        """
+        Adopting existing components assigns them to the Module via bulk_update(), which
+        bypasses save() and so requires an explicit MPTT tree rebuild. That rebuild must
+        run on the Module's connection.
+        """
+        ModuleBayTemplate.objects.create(module_type=self.module_type, name='Module Bay 1')
+
+        device = Device.objects.create(
+            name='Device 1', device_type=self.device_type, role=self.device_role, site=self.site
+        )
+        parent_bay = ModuleBay.objects.create(device=device, name='Parent Bay')
+        child_bay = ModuleBay.objects.create(device=device, name='Module Bay 1')
+
+        aliases = []
+        manager_class = type(ModuleBay.objects)
+        original_rebuild = manager_class.rebuild
+
+        def record_alias(manager, *args, **kwargs):
+            # Manager.db falls back to the router, so the private attribute is the only
+            # indication of whether an alias was set explicitly.
+            aliases.append(manager._db)
+            return original_rebuild(manager, *args, **kwargs)
+
+        module = Module(device=device, module_bay=parent_bay, module_type=self.module_type)
+        module._adopt_components = True
+        module._disable_replication = True
+        with patch.object(manager_class, 'rebuild', record_alias):
+            module.save()
+
+        child_bay.refresh_from_db()
+        self.assertEqual(child_bay.module, module)
+        self.assertEqual(aliases, [module._state.db])

+ 124 - 1
netbox/dcim/tests/test_signals.py

@@ -3,7 +3,7 @@ from unittest.mock import MagicMock, patch
 
 
 from django.contrib.contenttypes.models import ContentType
 from django.contrib.contenttypes.models import ContentType
 from django.db import connection, transaction
 from django.db import connection, transaction
-from django.test import SimpleTestCase, TestCase, TransactionTestCase
+from django.test import SimpleTestCase, TestCase, TransactionTestCase, override_settings
 from django.test.utils import CaptureQueriesContext
 from django.test.utils import CaptureQueriesContext
 
 
 from circuits.models import Circuit, CircuitTermination, CircuitType, Provider
 from circuits.models import Circuit, CircuitTermination, CircuitType, Provider
@@ -12,6 +12,7 @@ from dcim.choices import CableEndChoices, CableProfileChoices, LinkStatusChoices
 from dcim.models import (
 from dcim.models import (
     Cable,
     Cable,
     CablePath,
     CablePath,
+    CableTermination,
     Device,
     Device,
     DeviceRole,
     DeviceRole,
     DeviceType,
     DeviceType,
@@ -30,6 +31,7 @@ from dcim.models import (
     VirtualChassis,
     VirtualChassis,
 )
 )
 from ipam.models import Prefix
 from ipam.models import Prefix
+from utilities.testing import PinnedConnectionRouter
 from virtualization.models import Cluster, ClusterType
 from virtualization.models import Cluster, ClusterType
 from wireless.models import WirelessLAN
 from wireless.models import WirelessLAN
 
 
@@ -162,6 +164,127 @@ class RackSiteChangeSignalTestCase(TestCase):
         self.assertEqual(interface._location, self.location_b)
         self.assertEqual(interface._location, self.location_b)
 
 
 
 
+class ScopeSignalConnectionTestCase(TestCase):
+    """
+    Verify the scope-propagation handlers issue every query against the connection the
+    saved object was written to, rather than letting DATABASE_ROUTERS select one. On an
+    installation with routers configured (e.g. netbox_branching), a routed query both
+    writes to the wrong database and falls outside the transaction opened by the handler,
+    which makes the handler's select_for_update() raise.
+    """
+
+    @classmethod
+    def setUpTestData(cls):
+        cls.site_a = Site.objects.create(name='Site A', slug='site-a')
+        cls.site_b = Site.objects.create(name='Site B', slug='site-b')
+        manufacturer = Manufacturer.objects.create(name='Manufacturer', slug='manufacturer')
+        cls.device_type = DeviceType.objects.create(manufacturer=manufacturer, model='Device Type')
+        cls.device_role = DeviceRole.objects.create(name='Device Role', slug='device-role')
+
+    def test_location_save_pins_queries_to_saving_connection(self):
+        parent = Location.objects.create(name='Parent', slug='parent', site=self.site_a)
+        child = Location.objects.create(name='Child', slug='child', site=self.site_a, parent=parent)
+        rack = Rack.objects.create(name='Rack', site=self.site_a, location=parent)
+        device = Device.objects.create(
+            name='Device',
+            site=self.site_a,
+            location=parent,
+            device_type=self.device_type,
+            role=self.device_role,
+        )
+        interface = Interface.objects.create(device=device, name='Interface 1')
+        power_panel = PowerPanel.objects.create(name='Panel', site=self.site_a, location=parent)
+        cluster_type = ClusterType.objects.create(name='Cluster Type', slug='cluster-type')
+        cluster = Cluster.objects.create(name='Cluster', type=cluster_type, scope=child)
+
+        # Re-fetch and assign the new Site by ID, leaving the site relation uncached: a
+        # handler which reads instance.site rather than instance.site_id would fetch it
+        # over a routed connection, which is what the Site entry below catches.
+        parent = Location.objects.get(pk=parent.pk)
+        parent.site_id = self.site_b.pk
+        router = PinnedConnectionRouter(
+            CableTermination,
+            CircuitTermination,
+            Cluster,
+            Device,
+            Interface,
+            PowerPanel,
+            Prefix,
+            Rack,
+            Site,
+            WirelessLAN,
+        )
+        with override_settings(DATABASE_ROUTERS=[router]):
+            parent.save()
+
+        for obj in (child, rack, device, power_panel):
+            obj.refresh_from_db()
+            self.assertEqual(obj.site, self.site_b)
+        interface.refresh_from_db()
+        self.assertEqual(interface._site, self.site_b)
+        cluster.refresh_from_db()
+        self.assertEqual(cluster._site, self.site_b)
+
+    def test_rack_save_pins_queries_to_saving_connection(self):
+        rack = Rack.objects.create(name='Rack', site=self.site_a)
+        device = Device.objects.create(
+            name='Device',
+            site=self.site_a,
+            rack=rack,
+            device_type=self.device_type,
+            role=self.device_role,
+        )
+        interface = Interface.objects.create(device=device, name='Interface 1')
+
+        rack = Rack.objects.get(pk=rack.pk)
+        rack.site_id = self.site_b.pk
+        router = PinnedConnectionRouter(CableTermination, Device, Interface, Site)
+        with override_settings(DATABASE_ROUTERS=[router]):
+            rack.save()
+
+        device.refresh_from_db()
+        interface.refresh_from_db()
+        self.assertEqual(device.site, self.site_b)
+        self.assertEqual(interface._site, self.site_b)
+
+    def test_device_save_pins_queries_to_saving_connection(self):
+        device = Device.objects.create(
+            name='Device',
+            site=self.site_a,
+            device_type=self.device_type,
+            role=self.device_role,
+        )
+        interface = Interface.objects.create(device=device, name='Interface 1')
+
+        device = Device.objects.get(pk=device.pk)
+        device.site_id = self.site_b.pk
+        with override_settings(DATABASE_ROUTERS=[PinnedConnectionRouter(CableTermination, Interface, Site)]):
+            device.save()
+
+        interface.refresh_from_db()
+        self.assertEqual(interface._site, self.site_b)
+
+    def test_site_save_pins_scope_resync_to_saving_connection(self):
+        region = Region.objects.create(name='Region', slug='region')
+        cluster_type = ClusterType.objects.create(name='Cluster Type', slug='cluster-type')
+        # Scope the Cluster to a Location rather than to the Site itself: the rebuild then
+        # has to resolve the Location behind the object's generic scope, which is the read
+        # that must follow the connection the Site was saved on.
+        location = Location.objects.create(name='Location', slug='location', site=self.site_a)
+        cluster = Cluster.objects.create(name='Cluster', type=cluster_type, scope=location)
+
+        site = Site.objects.get(pk=self.site_a.pk)
+        site.region = region
+        # Region is included to catch the Location's site.region read made while rebuilding
+        # the cached fields; Site itself cannot be, as Django routes the save under test.
+        router = PinnedConnectionRouter(CircuitTermination, Cluster, Location, Prefix, Region, WirelessLAN)
+        with override_settings(DATABASE_ROUTERS=[router]):
+            site.save()
+
+        cluster.refresh_from_db()
+        self.assertEqual(cluster._region, region)
+
+
 class DeviceSiteChangeSignalTestCase(TestCase):
 class DeviceSiteChangeSignalTestCase(TestCase):
     """
     """
     Verify dcim.signals.handle_device_site_change propagates a Device's site/location/rack
     Verify dcim.signals.handle_device_site_change propagates a Device's site/location/rack

+ 2 - 2
netbox/ipam/models/ip.py

@@ -399,7 +399,7 @@ class Prefix(ContactsMixin, GetAvailablePrefixesMixin, CachedScopeMixin, Primary
         """
         """
         lookup = 'net_contains_or_equals' if include_self else 'net_contains'
         lookup = 'net_contains_or_equals' if include_self else 'net_contains'
         return Prefix.objects.filter(**{
         return Prefix.objects.filter(**{
-            'vrf': self.vrf,
+            'vrf_id': self.vrf_id,
             f'prefix__{lookup}': self.prefix
             f'prefix__{lookup}': self.prefix
         })
         })
 
 
@@ -409,7 +409,7 @@ class Prefix(ContactsMixin, GetAvailablePrefixesMixin, CachedScopeMixin, Primary
         """
         """
         lookup = 'net_contained_or_equal' if include_self else 'net_contained'
         lookup = 'net_contained_or_equal' if include_self else 'net_contained'
         return Prefix.objects.filter(**{
         return Prefix.objects.filter(**{
-            'vrf': self.vrf,
+            'vrf_id': self.vrf_id,
             f'prefix__{lookup}': self.prefix
             f'prefix__{lookup}': self.prefix
         })
         })
 
 

+ 17 - 15
netbox/ipam/signals.py

@@ -7,47 +7,49 @@ from virtualization.models import VirtualMachine
 from .models import IPAddress, Prefix
 from .models import IPAddress, Prefix
 
 
 
 
-def update_parents_children(prefix):
+def update_parents_children(prefix, using=None):
     """
     """
     Update depth on prefix & containing prefixes
     Update depth on prefix & containing prefixes
     """
     """
-    parents = prefix.get_parents(include_self=True).annotate_hierarchy()
+    parents = prefix.get_parents(include_self=True).using(using).annotate_hierarchy()
     for parent in parents:
     for parent in parents:
         parent._children = parent.hierarchy_children
         parent._children = parent.hierarchy_children
-    Prefix.objects.bulk_update(parents, ['_children'], batch_size=100)
+    Prefix.objects.using(using).bulk_update(parents, ['_children'], batch_size=100)
 
 
 
 
-def update_children_depth(prefix):
+def update_children_depth(prefix, using=None):
     """
     """
     Update children count on prefix & contained prefixes
     Update children count on prefix & contained prefixes
     """
     """
-    children = prefix.get_children(include_self=True).annotate_hierarchy()
+    children = prefix.get_children(include_self=True).using(using).annotate_hierarchy()
     for child in children:
     for child in children:
         child._depth = child.hierarchy_depth
         child._depth = child.hierarchy_depth
-    Prefix.objects.bulk_update(children, ['_depth'], batch_size=100)
+    Prefix.objects.using(using).bulk_update(children, ['_depth'], batch_size=100)
 
 
 
 
 @receiver(post_save, sender=Prefix)
 @receiver(post_save, sender=Prefix)
-def handle_prefix_saved(instance, created, **kwargs):
-
+def handle_prefix_saved(instance, created, using=None, **kwargs):
+    """
+    Recompute the cached hierarchy counters for the prefixes surrounding this one.
+    """
     # Prefix has changed (or new instance has been created)
     # Prefix has changed (or new instance has been created)
     if created or instance.vrf_id != instance._vrf_id or instance.prefix != instance._prefix:
     if created or instance.vrf_id != instance._vrf_id or instance.prefix != instance._prefix:
 
 
-        update_parents_children(instance)
-        update_children_depth(instance)
+        update_parents_children(instance, using)
+        update_children_depth(instance, using)
 
 
         # If this is not a new prefix, clean up parent/children of previous prefix
         # If this is not a new prefix, clean up parent/children of previous prefix
         if not created:
         if not created:
             old_prefix = Prefix(vrf_id=instance._vrf_id, prefix=instance._prefix)
             old_prefix = Prefix(vrf_id=instance._vrf_id, prefix=instance._prefix)
-            update_parents_children(old_prefix)
-            update_children_depth(old_prefix)
+            update_parents_children(old_prefix, using)
+            update_children_depth(old_prefix, using)
 
 
 
 
 @receiver(post_delete, sender=Prefix)
 @receiver(post_delete, sender=Prefix)
-def handle_prefix_deleted(instance, **kwargs):
+def handle_prefix_deleted(instance, using=None, **kwargs):
 
 
-    update_parents_children(instance)
-    update_children_depth(instance)
+    update_parents_children(instance, using)
+    update_children_depth(instance, using)
 
 
 
 
 @receiver(pre_delete, sender=IPAddress)
 @receiver(pre_delete, sender=IPAddress)

+ 57 - 2
netbox/ipam/tests/test_signals.py

@@ -1,13 +1,15 @@
 import uuid
 import uuid
 
 
 from django.contrib.contenttypes.models import ContentType
 from django.contrib.contenttypes.models import ContentType
-from django.test import RequestFactory, TestCase
+from django.test import RequestFactory, TestCase, override_settings
 
 
 from core.choices import ObjectChangeActionChoices
 from core.choices import ObjectChangeActionChoices
 from core.models import ObjectChange
 from core.models import ObjectChange
-from ipam.models import IPAddress, Prefix
+from ipam import signals
+from ipam.models import VRF, IPAddress, Prefix
 from netbox.context_managers import event_tracking
 from netbox.context_managers import event_tracking
 from users.models import User
 from users.models import User
+from utilities.testing import PinnedConnectionRouter
 from utilities.testing.utils import create_test_device, create_test_virtualmachine
 from utilities.testing.utils import create_test_device, create_test_virtualmachine
 
 
 
 
@@ -229,3 +231,56 @@ class ClearOOBIPSignalTestCase(TestCase):
                 action=ObjectChangeActionChoices.ACTION_UPDATE,
                 action=ObjectChangeActionChoices.ACTION_UPDATE,
             ).exists()
             ).exists()
         )
         )
+
+
+class PrefixHierarchySignalConnectionTestCase(TestCase):
+    """
+    Verify the prefix hierarchy handlers issue every query against the connection the saved
+    Prefix was written to, rather than letting DATABASE_ROUTERS select one. On an
+    installation with routers configured (e.g. netbox_branching), a routed query would
+    recount the hierarchy against one database and write the result to another.
+
+    These handlers are invoked directly rather than through save()/delete(): every query
+    they make is against Prefix, which is also the model being written, so a router which
+    fails routed Prefix queries would trip on the save itself.
+    """
+
+    @classmethod
+    def setUpTestData(cls):
+        cls.vrf = VRF.objects.create(name='VRF 1')
+
+    def test_prefix_saved_handler_pins_queries_to_given_connection(self):
+        parent = Prefix.objects.create(prefix='10.0.0.0/16', vrf=self.vrf)
+        child = Prefix.objects.create(prefix='10.0.1.0/24', vrf=self.vrf)
+
+        # Re-fetch and move the child, leaving the vrf relation uncached: a lookup which
+        # filters on self.vrf rather than self.vrf_id fetches it over a routed connection,
+        # which the VRF entry below catches. The same applies to the throwaway Prefix the
+        # handler builds to clean up the child's previous position. The instance is not
+        # re-fetched after the save, as that would reset the _prefix snapshot the handler
+        # compares against and it would decline to do any work at all.
+        child = Prefix.objects.get(pk=child.pk)
+        child.prefix = '10.0.2.0/24'
+        child.save()
+        self.assertNotEqual(child.prefix, child._prefix)
+
+        router = PinnedConnectionRouter(Prefix, VRF)
+        with override_settings(DATABASE_ROUTERS=[router]):
+            signals.handle_prefix_saved(instance=child, created=False, using='default')
+
+        parent.refresh_from_db()
+        child.refresh_from_db()
+        self.assertEqual(parent._children, 1)
+        self.assertEqual(child._depth, 1)
+
+    def test_prefix_deleted_handler_pins_queries_to_given_connection(self):
+        parent = Prefix.objects.create(prefix='10.0.0.0/16', vrf=self.vrf)
+        child = Prefix.objects.create(prefix='10.0.1.0/24', vrf=self.vrf)
+
+        child = Prefix.objects.get(pk=child.pk)
+        router = PinnedConnectionRouter(Prefix, VRF)
+        with override_settings(DATABASE_ROUTERS=[router]):
+            signals.handle_prefix_deleted(instance=child, using='default')
+
+        parent.refresh_from_db()
+        self.assertEqual(parent._children, 1)

+ 6 - 3
netbox/netbox/denormalized.py

@@ -28,7 +28,7 @@ def register(model, field_name, mappings):
 
 
 
 
 @receiver(post_save)
 @receiver(post_save)
-def update_denormalized_fields(sender, instance, created, raw, **kwargs):
+def update_denormalized_fields(sender, instance, created, raw, using=None, **kwargs):
     """
     """
     Check if the sender has denormalized fields registered, and update them as necessary.
     Check if the sender has denormalized fields registered, and update them as necessary.
     """
     """
@@ -52,6 +52,9 @@ def update_denormalized_fields(sender, instance, created, raw, **kwargs):
         }
         }
 
 
         # TODO: Improve efficiency here by placing conditions on the query?
         # TODO: Improve efficiency here by placing conditions on the query?
-        # Update all the denormalized fields with the triggering object's new values
-        count = model.objects.filter(**filter_params).update(**update_params)
+        # Update all the denormalized fields with the triggering object's new values. The
+        # update is pinned to the connection the instance was saved on: letting a database
+        # router select one could write these values to a different database than the one
+        # holding the change which triggered them.
+        count = model.objects.using(using).filter(**filter_params).update(**update_params)
         logger.debug(f'Updated {count} rows')
         logger.debug(f'Updated {count} rows')

+ 27 - 0
netbox/utilities/testing/utils.py

@@ -197,3 +197,30 @@ def get_random_string(length, charset=None):
     """
     """
     characters = string.ascii_letters + string.digits  # a-z, A-Z, 0-9
     characters = string.ascii_letters + string.digits  # a-z, A-Z, 0-9
     return ''.join(random.choice(characters) for __ in range(length))
     return ''.join(random.choice(characters) for __ in range(length))
+
+
+#
+# Database routing
+#
+
+class UnpinnedQuery(Exception):
+    """Raised when a query which should have been pinned to a connection is routed instead."""
+
+
+class PinnedConnectionRouter:
+    """
+    Fails any read or write of the given models which is not pinned to an explicit database
+    alias. Django consults DATABASE_ROUTERS only for queries which name no connection, so a
+    signal handler which threads through the alias supplied by the signal never reaches
+    this router. Each test leaves out the model being saved, as Django routes that save
+    itself.
+    """
+    def __init__(self, *models):
+        self.models = models
+
+    def _check(self, model, **hints):
+        if model in self.models:
+            raise UnpinnedQuery(f"{model.__name__} query was routed rather than pinned to a connection")
+
+    db_for_read = _check
+    db_for_write = _check

+ 7 - 7
netbox/virtualization/signals.py

@@ -6,20 +6,20 @@ from .models import Cluster, VirtualDisk, VirtualMachine
 
 
 
 
 @receiver((post_delete, post_save), sender=VirtualDisk)
 @receiver((post_delete, post_save), sender=VirtualDisk)
-def update_virtualmachine_disk(instance, **kwargs):
+def update_virtualmachine_disk(instance, using=None, **kwargs):
     """
     """
     When a VirtualDisk has been modified, update the aggregate disk_size value of its VM.
     When a VirtualDisk has been modified, update the aggregate disk_size value of its VM.
     """
     """
-    vm = instance.virtual_machine
-    VirtualMachine.objects.filter(pk=vm.pk).update(
-        disk=vm.virtualdisks.aggregate(Sum('size'))['size__sum']
+    disks = VirtualDisk.objects.using(using).filter(virtual_machine_id=instance.virtual_machine_id)
+    VirtualMachine.objects.using(using).filter(pk=instance.virtual_machine_id).update(
+        disk=disks.aggregate(Sum('size'))['size__sum']
     )
     )
 
 
 
 
 @receiver(post_save, sender=Cluster)
 @receiver(post_save, sender=Cluster)
-def update_virtualmachine_site(instance, **kwargs):
+def update_virtualmachine_site(instance, using=None, **kwargs):
     """
     """
     Update the assigned site for all VMs to match that of the Cluster (if any).
     Update the assigned site for all VMs to match that of the Cluster (if any).
     """
     """
-    if instance._site:
-        VirtualMachine.objects.filter(cluster=instance).update(site=instance._site)
+    if instance._site_id:
+        VirtualMachine.objects.using(using).filter(cluster=instance).update(site_id=instance._site_id)

+ 65 - 1
netbox/virtualization/tests/test_signals.py

@@ -1,7 +1,9 @@
 from django.contrib.contenttypes.models import ContentType
 from django.contrib.contenttypes.models import ContentType
-from django.test import TestCase
+from django.test import TestCase, override_settings
 
 
 from dcim.models import Site
 from dcim.models import Site
+from utilities.testing import PinnedConnectionRouter
+from virtualization import signals
 from virtualization.models import Cluster, ClusterType, VirtualDisk, VirtualMachine
 from virtualization.models import Cluster, ClusterType, VirtualDisk, VirtualMachine
 
 
 
 
@@ -83,3 +85,65 @@ class UpdateVirtualMachineSiteSignalTestCase(TestCase):
 
 
         vm.refresh_from_db()
         vm.refresh_from_db()
         self.assertEqual(vm.site, self.site_a)
         self.assertEqual(vm.site, self.site_a)
+
+
+class VirtualizationSignalConnectionTestCase(TestCase):
+    """
+    Verify the propagation handlers issue every query against the connection the saved
+    object was written to, rather than letting DATABASE_ROUTERS select one. On an
+    installation with routers configured (e.g. netbox_branching), a routed query reads from
+    or writes to a different database than the one being saved.
+    """
+
+    @classmethod
+    def setUpTestData(cls):
+        cls.site_a = Site.objects.create(name='Site A', slug='site-a')
+        cls.site_b = Site.objects.create(name='Site B', slug='site-b')
+        cls.cluster_type = ClusterType.objects.create(name='Cluster Type', slug='cluster-type')
+
+    def test_cluster_save_pins_vm_update_to_saving_connection(self):
+        cluster = Cluster.objects.create(name='Cluster', type=self.cluster_type, scope=self.site_a)
+        vm = VirtualMachine.objects.create(name='VM 1', cluster=cluster)
+
+        # Site is deliberately absent from the router: Cluster.save() resolves its generic
+        # scope through CachedScopeMixin.cache_related_objects(), which is a routed read of
+        # its own and not something this handler controls.
+        cluster = Cluster.objects.get(pk=cluster.pk)
+        cluster.scope_id = self.site_b.pk
+        with override_settings(DATABASE_ROUTERS=[PinnedConnectionRouter(VirtualMachine)]):
+            cluster.save()
+
+        vm.refresh_from_db()
+        self.assertEqual(vm.site, self.site_b)
+
+    def test_virtualdisk_save_pins_vm_update_to_saving_connection(self):
+        cluster = Cluster.objects.create(name='Cluster', type=self.cluster_type)
+        vm = VirtualMachine.objects.create(name='VM 1', cluster=cluster)
+        disk = VirtualDisk.objects.create(virtual_machine=vm, name='disk0', size=50)
+
+        # Re-fetch so the virtual_machine relation is uncached; resolving it to reach the VM
+        # or its disks is itself a routed read.
+        disk = VirtualDisk.objects.get(pk=disk.pk)
+        disk.size = 80
+        with override_settings(DATABASE_ROUTERS=[PinnedConnectionRouter(VirtualMachine)]):
+            disk.save()
+
+        vm.refresh_from_db()
+        self.assertEqual(vm.disk, 80)
+
+    def test_virtualdisk_handler_pins_disk_aggregate_to_given_connection(self):
+        # VirtualDisk cannot be listed in the router above, as Django routes the save of the
+        # disk itself; calling the handler directly leaves the aggregate over the sibling
+        # disks as the only VirtualDisk query in scope.
+        cluster = Cluster.objects.create(name='Cluster', type=self.cluster_type)
+        vm = VirtualMachine.objects.create(name='VM 1', cluster=cluster)
+        disk = VirtualDisk.objects.create(virtual_machine=vm, name='disk0', size=50)
+        VirtualDisk.objects.create(virtual_machine=vm, name='disk1', size=75)
+
+        disk = VirtualDisk.objects.get(pk=disk.pk)
+        router = PinnedConnectionRouter(VirtualDisk, VirtualMachine)
+        with override_settings(DATABASE_ROUTERS=[router]):
+            signals.update_virtualmachine_disk(instance=disk, using='default')
+
+        vm.refresh_from_db()
+        self.assertEqual(vm.disk, 125)