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

This commit is contained in:
Jeremy Stretch 2026-08-18 16:40:30 -04:00 committed by GitHub
parent 40a9df14a8
commit c2d39b12d8
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
13 changed files with 521 additions and 73 deletions

View File

@ -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

View File

@ -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)

View File

@ -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)

View File

@ -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,
)

View File

@ -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])

View File

@ -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

View File

@ -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
})

View File

@ -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)

View File

@ -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)

View File

@ -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')

View File

@ -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

View File

@ -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)

View File

@ -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)