6347 make connect_counter auto-discoverable

This commit is contained in:
Arthur 2023-07-19 19:04:35 +07:00
parent b98fecd188
commit b6e1abc629
9 changed files with 70 additions and 38 deletions

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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