Coverage for utilities/counters.py: 87%
61 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 18:35 +0000
1from django.apps import apps
2from django.db.models import Count, F, OuterRef, QuerySet, Subquery
3from django.db.models.signals import post_delete, post_save, pre_delete
5from netbox.registry import registry
7from .fields import CounterCacheField
8from .querysets import chunked_update
11def get_counters_for_model(model):
12 """
13 Return field mappings for all counters registered to the given model.
14 """
15 return registry['counter_fields'][model].items()
18def update_counter(model, pk, counter_name, value, using=None):
19 """
20 Increment or decrement a counter field on an object identified by its model and primary key (PK). Positive values
21 will increment; negative values will decrement.
22 """
23 model.objects.using(using).filter(pk=pk).update(
24 **{counter_name: F(counter_name) + value}
25 )
28def update_counts(model, field_name, related_query):
29 """
30 Perform a bulk update for the given model and counter field. For example,
32 update_counts(Device, '_interface_count', 'interfaces')
34 will effectively set
36 Device.objects.update(_interface_count=Count('interfaces'))
37 """
38 subquery = Subquery(
39 model.objects.filter(pk=OuterRef('pk')).annotate(_count=Count(related_query)).values('_count')
40 )
41 return chunked_update(model.objects.all(), **{
42 field_name: subquery
43 })
46#
47# Signal handlers
48#
50def post_save_receiver(sender, instance, created, using=None, **kwargs):
51 """
52 Update counter fields on related objects when a TrackingModelMixin subclass is created or modified.
53 """
54 for field_name, counter_name in get_counters_for_model(sender):
55 parent_model = sender._meta.get_field(field_name).related_model
56 new_pk = getattr(instance, field_name, None)
57 has_old_field = field_name in instance.tracker
58 old_pk = instance.tracker.get(field_name) if has_old_field else None
60 # Update the counters on the old and/or new parents as needed
61 if old_pk is not None: 61 ↛ 62line 61 didn't jump to line 62 because the condition on line 61 was never true
62 update_counter(parent_model, old_pk, counter_name, -1, using=using)
63 if new_pk is not None and (has_old_field or created):
64 update_counter(parent_model, new_pk, counter_name, 1, using=using)
67def _parent_is_being_deleted(origin, parent_model, parent_pk):
68 """
69 Return True if `origin` (the object or queryset that `delete()` was called on) indicates that
70 the parent identified by (parent_model, parent_pk) is itself being deleted as part of the same
71 operation. In that case, decrementing its counter is wasted work: the parent row is going away,
72 so the UPDATE would be a no-op. Skipping it avoids an N+1 storm of pointless UPDATEs when a
73 parent with many tracked children is deleted (e.g. a Device with thousands of Interfaces).
75 Note: only the *direct* parent is detected, since `origin` is just the top-level object/queryset
76 delete() was called on. In a deeper cascade (DeviceType -> Device -> Interface) `origin` stays
77 the DeviceType, so intermediate Devices' interface counters still get the (harmless) no-op
78 UPDATE. Suppressing that would require the full deletion set, which the signals don't expose.
79 """
80 if origin is None: 80 ↛ 81line 80 didn't jump to line 81 because the condition on line 80 was never true
81 return False
82 if isinstance(origin, QuerySet): 82 ↛ 84line 82 didn't jump to line 84 because the condition on line 82 was never true
83 # A bulk delete; every collected child belongs to an object in this queryset by construction
84 return origin.model is parent_model
85 # A single object delete
86 return isinstance(origin, parent_model) and origin.pk == parent_pk
89def pre_delete_receiver(sender, instance, origin, using=None, **kwargs):
90 """
91 Before a tracked object is deleted, check whether its row has already been removed (e.g. by an
92 earlier cascade) and, if so, flag it so post_delete_receiver skips the now-redundant counter
93 update. The existence check is skipped when the tracked parent is itself being deleted, since
94 the counter update would be skipped regardless — this avoids a SELECT per cascaded child.
95 """
96 for field_name, counter_name in get_counters_for_model(sender):
97 parent_model = sender._meta.get_field(field_name).related_model
98 parent_pk = getattr(instance, field_name, None)
99 if parent_pk is None or _parent_is_being_deleted(origin, parent_model, parent_pk):
100 continue
101 # A tracked parent will survive this operation, so the double-delete guard is needed
102 if not sender.objects.using(using).filter(pk=instance.pk).exists(): 102 ↛ 103line 102 didn't jump to line 103 because the condition on line 102 was never true
103 instance._previously_removed = True
104 return
107def post_delete_receiver(sender, instance, origin, using=None, **kwargs):
108 """
109 Update counter fields on related objects when a TrackingModelMixin subclass is deleted.
110 """
111 if hasattr(instance, '_previously_removed'): 111 ↛ 112line 111 didn't jump to line 112 because the condition on line 111 was never true
112 return
114 for field_name, counter_name in get_counters_for_model(sender):
115 parent_model = sender._meta.get_field(field_name).related_model
116 parent_pk = getattr(instance, field_name, None)
118 # Decrement the parent's counter by one, unless the parent is itself being deleted
119 if parent_pk is not None and not _parent_is_being_deleted(origin, parent_model, parent_pk):
120 update_counter(parent_model, parent_pk, counter_name, -1, using=using)
123#
124# Registration
125#
127def connect_counters(*models):
128 """
129 Register counter fields and connect signal handlers for their child models.
130 Ensures exactly one receiver per child (sender), even when multiple counters
131 reference the same sender (e.g., Device).
132 """
133 connected = set() # child models we've already connected
135 for model in models:
136 # Find all CounterCacheFields on the model
137 counter_fields = [field for field in model._meta.get_fields() if isinstance(field, CounterCacheField)]
139 for field in counter_fields:
140 to_model = apps.get_model(field.to_model_name)
142 # Register the counter in the registry
143 change_tracking_fields = registry['counter_fields'][to_model]
144 change_tracking_fields[f'{field.to_field_name}_id'] = field.name
146 # Connect signals once per child model
147 if to_model in connected:
148 continue
150 # Ensure dispatch_uid is unique per model (sender), not per field
151 uid_base = f'countercache.{to_model._meta.label_lower}'
153 # Connect the post_save and post_delete handlers
154 post_save.connect(
155 post_save_receiver,
156 sender=to_model,
157 weak=False,
158 dispatch_uid=f'{uid_base}.post_save',
159 )
160 pre_delete.connect(
161 pre_delete_receiver,
162 sender=to_model,
163 weak=False,
164 dispatch_uid=f'{uid_base}.pre_delete',
165 )
166 post_delete.connect(
167 post_delete_receiver,
168 sender=to_model,
169 weak=False,
170 dispatch_uid=f'{uid_base}.post_delete',
171 )
173 connected.add(to_model)