This commit is contained in:
parent
560d595ac3
commit
a6f8848e3c
|
|
@ -987,6 +987,7 @@ class Device(
|
|||
(default). Otherwise, save() will be called on each instance individually.
|
||||
"""
|
||||
model = queryset.model.component_model
|
||||
using = self._state.db
|
||||
|
||||
if bulk_create:
|
||||
components = [obj.instantiate(device=self) for obj in queryset]
|
||||
|
|
@ -1005,7 +1006,7 @@ class Device(
|
|||
component._site = self.site
|
||||
component._location = self.location
|
||||
component._rack = self.rack
|
||||
components = model.objects.bulk_create(components)
|
||||
components = model.objects.using(using).bulk_create(components)
|
||||
# Prefetch related objects to minimize queries needed during post_save
|
||||
prefetch_fields = get_prefetchable_fields(model)
|
||||
prefetch_related_objects(components, *prefetch_fields)
|
||||
|
|
@ -1016,7 +1017,7 @@ class Device(
|
|||
instance=component,
|
||||
created=True,
|
||||
raw=False,
|
||||
using='default',
|
||||
using=using,
|
||||
update_fields=None
|
||||
)
|
||||
else:
|
||||
|
|
|
|||
|
|
@ -353,6 +353,8 @@ 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),
|
||||
|
|
@ -404,7 +406,7 @@ class Module(TrackingModelMixin, PrimaryModel):
|
|||
|
||||
# we handle create and update separately - this is for create
|
||||
if not issubclass(component_model, MPTTModel):
|
||||
component_model.objects.bulk_create(create_instances)
|
||||
component_model.objects.using(using).bulk_create(create_instances)
|
||||
# Emit the post_save signal for each newly created object
|
||||
for component in create_instances:
|
||||
post_save.send(
|
||||
|
|
@ -412,7 +414,7 @@ class Module(TrackingModelMixin, PrimaryModel):
|
|||
instance=component,
|
||||
created=True,
|
||||
raw=False,
|
||||
using='default',
|
||||
using=using,
|
||||
update_fields=None
|
||||
)
|
||||
else:
|
||||
|
|
@ -423,7 +425,7 @@ class Module(TrackingModelMixin, PrimaryModel):
|
|||
update_fields = ['module']
|
||||
|
||||
# we handle create and update separately - this is for update
|
||||
component_model.objects.bulk_update(update_instances, update_fields)
|
||||
component_model.objects.using(using).bulk_update(update_instances, update_fields)
|
||||
# Emit the post_save signal for each updated object
|
||||
for component in update_instances:
|
||||
post_save.send(
|
||||
|
|
@ -431,7 +433,7 @@ class Module(TrackingModelMixin, PrimaryModel):
|
|||
instance=component,
|
||||
created=False,
|
||||
raw=False,
|
||||
using='default',
|
||||
using=using,
|
||||
update_fields=update_fields
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -14,12 +14,12 @@ def get_counters_for_model(model):
|
|||
return registry['counter_fields'][model].items()
|
||||
|
||||
|
||||
def update_counter(model, pk, counter_name, value):
|
||||
def update_counter(model, pk, counter_name, value, using=None):
|
||||
"""
|
||||
Increment or decrement a counter field on an object identified by its model and primary key (PK). Positive values
|
||||
will increment; negative values will decrement.
|
||||
"""
|
||||
model.objects.filter(pk=pk).update(
|
||||
model.objects.using(using).filter(pk=pk).update(
|
||||
**{counter_name: F(counter_name) + value}
|
||||
)
|
||||
|
||||
|
|
@ -46,7 +46,7 @@ def update_counts(model, field_name, related_query):
|
|||
# Signal handlers
|
||||
#
|
||||
|
||||
def post_save_receiver(sender, instance, created, **kwargs):
|
||||
def post_save_receiver(sender, instance, created, using=None, **kwargs):
|
||||
"""
|
||||
Update counter fields on related objects when a TrackingModelMixin subclass is created or modified.
|
||||
"""
|
||||
|
|
@ -58,9 +58,9 @@ def post_save_receiver(sender, instance, created, **kwargs):
|
|||
|
||||
# Update the counters on the old and/or new parents as needed
|
||||
if old_pk is not None:
|
||||
update_counter(parent_model, old_pk, counter_name, -1)
|
||||
update_counter(parent_model, old_pk, counter_name, -1, using=using)
|
||||
if new_pk is not None and (has_old_field or created):
|
||||
update_counter(parent_model, new_pk, counter_name, 1)
|
||||
update_counter(parent_model, new_pk, counter_name, 1, using=using)
|
||||
|
||||
|
||||
def _parent_is_being_deleted(origin, parent_model, parent_pk):
|
||||
|
|
@ -85,7 +85,7 @@ def _parent_is_being_deleted(origin, parent_model, parent_pk):
|
|||
return isinstance(origin, parent_model) and origin.pk == parent_pk
|
||||
|
||||
|
||||
def pre_delete_receiver(sender, instance, origin, **kwargs):
|
||||
def pre_delete_receiver(sender, instance, origin, using=None, **kwargs):
|
||||
"""
|
||||
Before a tracked object is deleted, check whether its row has already been removed (e.g. by an
|
||||
earlier cascade) and, if so, flag it so post_delete_receiver skips the now-redundant counter
|
||||
|
|
@ -98,12 +98,12 @@ def pre_delete_receiver(sender, instance, origin, **kwargs):
|
|||
if parent_pk is None or _parent_is_being_deleted(origin, parent_model, parent_pk):
|
||||
continue
|
||||
# A tracked parent will survive this operation, so the double-delete guard is needed
|
||||
if not sender.objects.filter(pk=instance.pk).exists():
|
||||
if not sender.objects.using(using).filter(pk=instance.pk).exists():
|
||||
instance._previously_removed = True
|
||||
return
|
||||
|
||||
|
||||
def post_delete_receiver(sender, instance, origin, **kwargs):
|
||||
def post_delete_receiver(sender, instance, origin, using=None, **kwargs):
|
||||
"""
|
||||
Update counter fields on related objects when a TrackingModelMixin subclass is deleted.
|
||||
"""
|
||||
|
|
@ -116,7 +116,7 @@ def post_delete_receiver(sender, instance, origin, **kwargs):
|
|||
|
||||
# Decrement the parent's counter by one, unless the parent is itself being deleted
|
||||
if parent_pk is not None and not _parent_is_being_deleted(origin, parent_model, parent_pk):
|
||||
update_counter(parent_model, parent_pk, counter_name, -1)
|
||||
update_counter(parent_model, parent_pk, counter_name, -1, using=using)
|
||||
|
||||
|
||||
#
|
||||
|
|
|
|||
|
|
@ -1,9 +1,17 @@
|
|||
from unittest.mock import patch
|
||||
|
||||
from django.db.utils import ConnectionDoesNotExist
|
||||
from django.test import override_settings
|
||||
from django.urls import reverse
|
||||
|
||||
from dcim.models import *
|
||||
from utilities.counters import connect_counters, update_counter
|
||||
from utilities.counters import (
|
||||
connect_counters,
|
||||
post_delete_receiver,
|
||||
post_save_receiver,
|
||||
pre_delete_receiver,
|
||||
update_counter,
|
||||
)
|
||||
from utilities.testing.base import TestCase
|
||||
from utilities.testing.utils import create_test_device
|
||||
|
||||
|
|
@ -192,3 +200,92 @@ class CountersTestCase(TestCase):
|
|||
vc.refresh_from_db()
|
||||
self.assertEqual(device1.device_type.device_count, 2, 'device_count should decrement exactly once')
|
||||
self.assertEqual(vc.member_count, 0, 'member_count should decrement exactly once')
|
||||
|
||||
|
||||
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 written, as Django routes that write 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
|
||||
|
||||
|
||||
class CounterConnectionTestCase(TestCase):
|
||||
"""
|
||||
Validate that the counter cache handlers issue their queries against the connection the
|
||||
triggering object was written to, rather than letting DATABASE_ROUTERS select one. A routed
|
||||
query updates a counter in a different database than the one holding the change which triggered
|
||||
it, leaving the cached count silently wrong.
|
||||
"""
|
||||
@classmethod
|
||||
def setUpTestData(cls):
|
||||
cls.device = create_test_device('Device 1')
|
||||
|
||||
def test_create_pins_counter_increment(self):
|
||||
with override_settings(DATABASE_ROUTERS=[PinnedConnectionRouter(Device)]):
|
||||
Interface.objects.create(device=self.device, name='Interface 1')
|
||||
|
||||
self.device.refresh_from_db()
|
||||
self.assertEqual(self.device.interface_count, 1)
|
||||
|
||||
def test_move_pins_both_counter_updates(self):
|
||||
other = create_test_device('Device 2')
|
||||
interface = Interface.objects.create(device=self.device, name='Interface 1')
|
||||
|
||||
interface = Interface.objects.get(pk=interface.pk)
|
||||
interface.device = other
|
||||
with override_settings(DATABASE_ROUTERS=[PinnedConnectionRouter(Device)]):
|
||||
interface.save()
|
||||
|
||||
self.device.refresh_from_db()
|
||||
other.refresh_from_db()
|
||||
self.assertEqual(self.device.interface_count, 0)
|
||||
self.assertEqual(other.interface_count, 1)
|
||||
|
||||
def test_receivers_use_the_alias_supplied_by_the_signal(self):
|
||||
"""
|
||||
The tests above prove the queries name *an* alias, but with one configured database that
|
||||
alias is always 'default' — they would pass just as well against a hardcoded
|
||||
.using('default'). Invoking each receiver with an alias that does not exist distinguishes
|
||||
"threaded through from the signal" from "happens to be the default".
|
||||
"""
|
||||
interface = Interface.objects.create(device=self.device, name='Interface 1')
|
||||
interface = Interface.objects.get(pk=interface.pk)
|
||||
|
||||
with self.assertRaises(ConnectionDoesNotExist):
|
||||
post_save_receiver(Interface, interface, created=True, using='nonexistent')
|
||||
|
||||
# origin=None: the parent is not itself being deleted, so the existence guard runs
|
||||
with self.assertRaises(ConnectionDoesNotExist):
|
||||
pre_delete_receiver(Interface, interface, origin=None, using='nonexistent')
|
||||
|
||||
with self.assertRaises(ConnectionDoesNotExist):
|
||||
post_delete_receiver(Interface, interface, origin=None, using='nonexistent')
|
||||
|
||||
def test_delete_pins_counter_decrement(self):
|
||||
interface = Interface.objects.create(device=self.device, name='Interface 1')
|
||||
self.device.refresh_from_db()
|
||||
self.assertEqual(self.device.interface_count, 1)
|
||||
|
||||
# The delete itself is pinned, as Django routes an unpinned one; that leaves
|
||||
# pre_delete_receiver's existence guard as the only Interface query in scope, and that read
|
||||
# decides whether the decrement below happens at all.
|
||||
with override_settings(DATABASE_ROUTERS=[PinnedConnectionRouter(Device, Interface)]):
|
||||
Interface.objects.using('default').filter(pk=interface.pk).delete()
|
||||
|
||||
self.device.refresh_from_db()
|
||||
self.assertEqual(self.device.interface_count, 0)
|
||||
|
|
|
|||
Loading…
Reference in New Issue