diff --git a/netbox/dcim/apps.py b/netbox/dcim/apps.py index 8643f412c..0868f031c 100644 --- a/netbox/dcim/apps.py +++ b/netbox/dcim/apps.py @@ -1,4 +1,4 @@ -from django.apps import AppConfig +from django.apps import AppConfig, apps from netbox import denormalized @@ -14,7 +14,7 @@ class DCIMConfig(AppConfig): Interface, InventoryItem, PowerOutlet, PowerPort, RearPort, ) - from utilities.counter import connect_counter + from utilities.counter import connect_counters # Register denormalized fields denormalized.register(CableTermination, '_device', { @@ -30,12 +30,4 @@ class DCIMConfig(AppConfig): '_site': 'site', }) - connect_counter('_console_port_count', ConsolePort.device) - connect_counter('_console_server_port_count', ConsoleServerPort.device) - connect_counter('_interface_count', Interface.device) - connect_counter('_front_port_count', FrontPort.device) - connect_counter('_rear_port_count', RearPort.device) - connect_counter('_device_bay_count', DeviceBay.device) - connect_counter('_inventory_item_count', InventoryItem.device) - connect_counter('_power_port_count', PowerPort.device) - connect_counter('_power_outlet_count', PowerOutlet.device) + connect_counters(self) diff --git a/netbox/dcim/migrations/0175_device__console_port_count_and_more.py b/netbox/dcim/migrations/0175_device__console_port_count_and_more.py index 68b131149..426902cfe 100644 --- a/netbox/dcim/migrations/0175_device__console_port_count_and_more.py +++ b/netbox/dcim/migrations/0175_device__console_port_count_and_more.py @@ -53,47 +53,47 @@ class Migration(migrations.Migration): migrations.AddField( model_name='device', name='_console_port_count', - field=utilities.fields.CounterCacheField(default=0), + field=utilities.fields.CounterCacheField(default=0, to_field='device', to_model='dcim.ConsolePort'), ), migrations.AddField( model_name='device', name='_console_server_port_count', - field=utilities.fields.CounterCacheField(default=0), + field=utilities.fields.CounterCacheField(default=0, to_field='device', to_model='dcim.ConsoleServerPort'), ), migrations.AddField( model_name='device', name='_device_bay_count', - field=utilities.fields.CounterCacheField(default=0), + field=utilities.fields.CounterCacheField(default=0, to_field='device', to_model='dcim.DeviceBay'), ), migrations.AddField( model_name='device', name='_front_port_count', - field=utilities.fields.CounterCacheField(default=0), + field=utilities.fields.CounterCacheField(default=0, to_field='device', to_model='dcim.FrontPort'), ), migrations.AddField( model_name='device', name='_interface_count', - field=utilities.fields.CounterCacheField(default=0), + field=utilities.fields.CounterCacheField(default=0, to_field='device', to_model='dcim.Interface'), ), migrations.AddField( model_name='device', name='_inventory_item_count', - field=utilities.fields.CounterCacheField(default=0), + field=utilities.fields.CounterCacheField(default=0, to_field='device', to_model='dcim.InventoryItem'), ), migrations.AddField( model_name='device', name='_power_outlet_count', - field=utilities.fields.CounterCacheField(default=0), + field=utilities.fields.CounterCacheField(default=0, to_field='device', to_model='dcim.PowerOutlet'), ), migrations.AddField( model_name='device', name='_power_port_count', - field=utilities.fields.CounterCacheField(default=0), + field=utilities.fields.CounterCacheField(default=0, to_field='device', to_model='dcim.PowerPort'), ), migrations.AddField( model_name='device', name='_rear_port_count', - field=utilities.fields.CounterCacheField(default=0), + field=utilities.fields.CounterCacheField(default=0, to_field='device', to_model='dcim.RearPort'), ), migrations.RunPython( recalculate_device_counts, diff --git a/netbox/dcim/models/devices.py b/netbox/dcim/models/devices.py index a6fca4483..6299bd39f 100644 --- a/netbox/dcim/models/devices.py +++ b/netbox/dcim/models/devices.py @@ -640,15 +640,15 @@ class Device(PrimaryModel, ConfigContextModel, TrackingModelMixin): help_text=_("GPS coordinate in decimal format (xx.yyyyyy)") ) - _console_port_count = CounterCacheField() - _console_server_port_count = CounterCacheField() - _power_port_count = CounterCacheField() - _power_outlet_count = CounterCacheField() - _interface_count = CounterCacheField() - _front_port_count = CounterCacheField() - _rear_port_count = CounterCacheField() - _device_bay_count = CounterCacheField() - _inventory_item_count = CounterCacheField() + _console_port_count = CounterCacheField(to_model='dcim.ConsolePort', to_field='device') + _console_server_port_count = CounterCacheField(to_model='dcim.ConsoleServerPort', to_field='device') + _power_port_count = CounterCacheField(to_model='dcim.PowerPort', to_field='device') + _power_outlet_count = CounterCacheField(to_model='dcim.PowerOutlet', to_field='device') + _interface_count = CounterCacheField(to_model='dcim.Interface', to_field='device') + _front_port_count = CounterCacheField(to_model='dcim.FrontPort', to_field='device') + _rear_port_count = CounterCacheField(to_model='dcim.RearPort', to_field='device') + _device_bay_count = CounterCacheField(to_model='dcim.DeviceBay', to_field='device') + _inventory_item_count = CounterCacheField(to_model='dcim.InventoryItem', to_field='device') # Generic relations contacts = GenericRelation( diff --git a/netbox/netbox/registry.py b/netbox/netbox/registry.py index 23b9ad4cb..546eb4404 100644 --- a/netbox/netbox/registry.py +++ b/netbox/netbox/registry.py @@ -28,4 +28,5 @@ registry = Registry({ 'search': dict(), 'views': collections.defaultdict(dict), 'widgets': dict(), + 'cached_counter_fields': dict(), }) diff --git a/netbox/utilities/counter.py b/netbox/utilities/counter.py index 4dd850ee1..9396f49c8 100644 --- a/netbox/utilities/counter.py +++ b/netbox/utilities/counter.py @@ -1,3 +1,4 @@ +from django.apps import apps from django.db.models import F from django.db.models.signals import post_delete, post_save, pre_save from functools import partial @@ -77,5 +78,13 @@ class Counter: return self.set_counter_field(parent_id, F(self.counter_name) + amount) -def connect_counter(counter_name, foreign_key_field): - return Counter(counter_name, foreign_key_field) +def connect_counters(app_config): + models = app_config.get_models() + for model in models: + if issubclass(model, TrackingModelMixin): + fields = model._meta.get_fields() + for field in fields: + if type(field) is CounterCacheField: + to_model = apps.get_model(field.to_model_name) + to_field = getattr(to_model, field.to_field_name) + Counter(field.name, to_field) diff --git a/netbox/utilities/fields.py b/netbox/utilities/fields.py index e8df8a709..4c17e77d2 100644 --- a/netbox/utilities/fields.py +++ b/netbox/utilities/fields.py @@ -2,6 +2,7 @@ from collections import defaultdict from django.contrib.contenttypes.fields import GenericForeignKey from django.db import models +from django.utils.translation import gettext_lazy as _ from utilities.ordering import naturalize from .forms.widgets import ColorSelect @@ -150,6 +151,35 @@ class CounterCacheField(models.BigIntegerField): """ Counter field to keep track of related model counts. """ - def __init__(self, *args, **kwargs): + def __init__(self, to_model, to_field, *args, **kwargs): + if not isinstance(to_model, str): + raise TypeError( + _("%s(%r) is invalid. to_model parameter to CounterCacheField must be " + "a string in the format 'app.model'") + % ( + self.__class__.__name__, + to_model, + ) + ) + + if not isinstance(to_field, str): + raise TypeError( + _("%s(%r) is invalid. to_field parameter to CounterCacheField must be " + "a string in the format 'field'") + % ( + self.__class__.__name__, + to_field, + ) + ) + + self.to_model_name = to_model + self.to_field_name = to_field kwargs['default'] = kwargs.get('default', 0) + super().__init__(*args, **kwargs) + + def deconstruct(self): + name, path, args, kwargs = super().deconstruct() + kwargs["to_model"] = self.to_model_name + kwargs["to_field"] = self.to_field_name + return name, path, args, kwargs diff --git a/netbox/virtualization/apps.py b/netbox/virtualization/apps.py index 913c8d795..32e895b0b 100644 --- a/netbox/virtualization/apps.py +++ b/netbox/virtualization/apps.py @@ -6,8 +6,6 @@ class VirtualizationConfig(AppConfig): def ready(self): from . import search - from .models import VMInterface + from utilities.counter import connect_counters - from utilities.counter import connect_counter - - connect_counter('_interface_count', VMInterface.virtual_machine) + connect_counters(self) diff --git a/netbox/virtualization/migrations/0035_virtualmachine__interface_count.py b/netbox/virtualization/migrations/0035_virtualmachine__interface_count.py index 7f641984f..78a270f6a 100644 --- a/netbox/virtualization/migrations/0035_virtualmachine__interface_count.py +++ b/netbox/virtualization/migrations/0035_virtualmachine__interface_count.py @@ -25,7 +25,9 @@ class Migration(migrations.Migration): migrations.AddField( model_name='virtualmachine', name='_interface_count', - field=utilities.fields.CounterCacheField(default=0), + field=utilities.fields.CounterCacheField( + default=0, to_field='virtual_machine', to_model='virtualization.VMInterface' + ), ), migrations.RunPython( code=populate_virtualmachine_counts, diff --git a/netbox/virtualization/models/virtualmachines.py b/netbox/virtualization/models/virtualmachines.py index 6e1c9326b..d7847d08e 100644 --- a/netbox/virtualization/models/virtualmachines.py +++ b/netbox/virtualization/models/virtualmachines.py @@ -121,7 +121,7 @@ class VirtualMachine(PrimaryModel, ConfigContextModel, TrackingModelMixin): verbose_name='Disk (GB)' ) - _interface_count = CounterCacheField() + _interface_count = CounterCacheField(to_model='virtualization.VMInterface', to_field='virtual_machine') # Generic relation contacts = GenericRelation(