Bläddra i källkod

Fixes #22963: Honor the saving database connection in counter cache signals (#22965)

Jeremy Stretch 4 dagar sedan
förälder
incheckning
a6f8848e3c

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

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

+ 6 - 4
netbox/dcim/models/modules.py

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

+ 9 - 9
netbox/utilities/counters.py

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

+ 98 - 1
netbox/utilities/tests/test_counters.py

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