6347 make connect_counter auto-discoverable

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_counters(self)
-        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)

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,

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_port_count = CounterCacheField(to_model='dcim.ConsolePort', to_field='device')
-    _console_server_port_count = CounterCacheField()
+    _console_server_port_count = CounterCacheField(to_model='dcim.ConsoleServerPort', to_field='device')
-    _power_port_count = CounterCacheField()
+    _power_port_count = CounterCacheField(to_model='dcim.PowerPort', to_field='device')
-    _power_outlet_count = CounterCacheField()
+    _power_outlet_count = CounterCacheField(to_model='dcim.PowerOutlet', to_field='device')
-    _interface_count = CounterCacheField()
+    _interface_count = CounterCacheField(to_model='dcim.Interface', to_field='device')
-    _front_port_count = CounterCacheField()
+    _front_port_count = CounterCacheField(to_model='dcim.FrontPort', to_field='device')
-    _rear_port_count = CounterCacheField()
+    _rear_port_count = CounterCacheField(to_model='dcim.RearPort', to_field='device')
-    _device_bay_count = CounterCacheField()
+    _device_bay_count = CounterCacheField(to_model='dcim.DeviceBay', to_field='device')
-    _inventory_item_count = CounterCacheField()
+    _inventory_item_count = CounterCacheField(to_model='dcim.InventoryItem', to_field='device')
 
     # Generic relations
     contacts = GenericRelation(

netbox/netbox/registry.py

@@ -28,4 +28,5 @@ registry = Registry({
     'search': dict(),
     'views': collections.defaultdict(dict),
     'widgets': dict(),
+    'cached_counter_fields': dict(),
 })

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):
+def connect_counters(app_config):
-    return Counter(counter_name, foreign_key_field)
+    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)

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

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_counters(self)
-
-        connect_counter('_interface_count', VMInterface.virtual_machine)

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,

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(