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

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.core.exceptions import ObjectDoesNotExist, ValidationError
 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 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
         # 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:
-            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')
             )['tree_id__max'] or 0
             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
                 if cf_defaults := CustomField.objects.get_defaults_for_model(model):
                     component.custom_field_data = cf_defaults
-                component.save()
+                component.save(using=using)
 
     def save(self, *args, **kwargs):
         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
 
         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
             ).first()
 
         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:
-            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.save()
+                child_bay.save(using=using)
 
         adopt_components = getattr(self, '_adopt_components', 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):
             return
 
-        using = self._state.db
-
         # Iterate all component types
         for templates, component_attribute, component_model in [
             ("consoleporttemplates", "consoleports", ConsolePort),
@@ -372,7 +372,9 @@ class Module(TrackingModelMixin, PrimaryModel):
             # Prefetch installed components
             installed_components = {
                 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.
@@ -420,7 +422,7 @@ class Module(TrackingModelMixin, PrimaryModel):
             else:
                 # MPTT models must be saved individually to maintain tree structure
                 for instance in create_instances:
-                    instance.save()
+                    instance.save(using=using)
 
             update_fields = ['module']
 
@@ -439,7 +441,8 @@ class Module(TrackingModelMixin, PrimaryModel):
 
             # Rebuild MPTT tree if needed (bulk_update bypasses model save)
             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
         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,
     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:
         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
         # 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
         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
         # 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
             # unlocked read could stamp region/group values from before that change.
             site = (
-                Site.objects.filter(pk=instance.site_id)
+                Site.objects.using(using)
+                .filter(pk=instance.site_id)
                 .select_for_update(no_key=True)
                 .values('region_id', 'group_id')
                 .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
                 # _location, which may itself be stale; scope_id doubles as the correct
                 # _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):
-                    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'),
                         _site_id=instance.site_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)
-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:
-        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
         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)
-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.
+    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:
         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()
 
 
+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=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
     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
     # 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):
-            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
             # 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.
             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:
-                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()
                 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 unittest.mock import patch
 
 from django.core.exceptions import ValidationError
 from django.db.models import ProtectedError
 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 core.models import ObjectType
@@ -15,6 +16,7 @@ from ipam.models import Prefix
 from netbox.choices import WeightUnitChoices
 from tenancy.models import Tenant
 from utilities.data import drange
+from utilities.testing import PinnedConnectionRouter
 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['B']['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.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 circuits.models import Circuit, CircuitTermination, CircuitType, Provider
@@ -12,6 +12,7 @@ from dcim.choices import CableEndChoices, CableProfileChoices, LinkStatusChoices
 from dcim.models import (
     Cable,
     CablePath,
+    CableTermination,
     Device,
     DeviceRole,
     DeviceType,
@@ -30,6 +31,7 @@ from dcim.models import (
     VirtualChassis,
 )
 from ipam.models import Prefix
+from utilities.testing import PinnedConnectionRouter
 from virtualization.models import Cluster, ClusterType
 from wireless.models import WirelessLAN
 
@@ -162,6 +164,127 @@ class RackSiteChangeSignalTestCase(TestCase):
         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):
     """
     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'
         return Prefix.objects.filter(**{
-            'vrf': self.vrf,
+            'vrf_id': self.vrf_id,
             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'
         return Prefix.objects.filter(**{
-            'vrf': self.vrf,
+            'vrf_id': self.vrf_id,
             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
 
 
-def update_parents_children(prefix):
+def update_parents_children(prefix, using=None):
     """
     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:
         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
     """
-    children = prefix.get_children(include_self=True).annotate_hierarchy()
+    children = prefix.get_children(include_self=True).using(using).annotate_hierarchy()
     for child in children:
         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)
-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)
     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 not created:
             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)
-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)

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

@@ -1,13 +1,15 @@
 import uuid
 
 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.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 users.models import User
+from utilities.testing import PinnedConnectionRouter
 from utilities.testing.utils import create_test_device, create_test_virtualmachine
 
 
@@ -229,3 +231,56 @@ class ClearOOBIPSignalTestCase(TestCase):
                 action=ObjectChangeActionChoices.ACTION_UPDATE,
             ).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)
-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.
     """
@@ -52,6 +52,9 @@ def update_denormalized_fields(sender, instance, created, raw, **kwargs):
         }
 
         # 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')

+ 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
     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)
-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.
     """
-    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)
-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).
     """
-    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.test import TestCase
+from django.test import TestCase, override_settings
 
 from dcim.models import Site
+from utilities.testing import PinnedConnectionRouter
+from virtualization import signals
 from virtualization.models import Cluster, ClusterType, VirtualDisk, VirtualMachine
 
 
@@ -83,3 +85,65 @@ class UpdateVirtualMachineSiteSignalTestCase(TestCase):
 
         vm.refresh_from_db()
         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)