13745 device type migration (#13747)

* 13745 update migrations to use batch_size

* 13745 update migrations to use subquery update

* 13745 refactor and update other counter migrations

netbox/dcim/migrations/0176_device_component_counters.py

@@ -2,47 +2,22 @@ from django.db import migrations
 from django.db.models import Count
 
 import utilities.fields
+from utilities.counters import update_counts
 
 
 def recalculate_device_counts(apps, schema_editor):
     Device = apps.get_model("dcim", "Device")
-    devices = Device.objects.annotate(
-        _console_port_count=Count('consoleports', distinct=True),
-        _console_server_port_count=Count('consoleserverports', distinct=True),
-        _power_port_count=Count('powerports', distinct=True),
-        _power_outlet_count=Count('poweroutlets', distinct=True),
-        _interface_count=Count('interfaces', distinct=True),
-        _front_port_count=Count('frontports', distinct=True),
-        _rear_port_count=Count('rearports', distinct=True),
-        _device_bay_count=Count('devicebays', distinct=True),
-        _module_bay_count=Count('modulebays', distinct=True),
-        _inventory_item_count=Count('inventoryitems', distinct=True),
-    )
 
-    for device in devices:
-        device.console_port_count = device._console_port_count
-        device.console_server_port_count = device._console_server_port_count
-        device.power_port_count = device._power_port_count
-        device.power_outlet_count = device._power_outlet_count
-        device.interface_count = device._interface_count
-        device.front_port_count = device._front_port_count
-        device.rear_port_count = device._rear_port_count
-        device.device_bay_count = device._device_bay_count
-        device.module_bay_count = device._module_bay_count
-        device.inventory_item_count = device._inventory_item_count
-
-    Device.objects.bulk_update(devices, [
-        'console_port_count',
-        'console_server_port_count',
-        'power_port_count',
-        'power_outlet_count',
-        'interface_count',
-        'front_port_count',
-        'rear_port_count',
-        'device_bay_count',
-        'module_bay_count',
-        'inventory_item_count',
-    ], batch_size=100)
+    update_counts(Device, 'console_port_count', 'consoleports')
+    update_counts(Device, 'console_server_port_count', 'consoleserverports')
+    update_counts(Device, 'power_port_count', 'powerports')
+    update_counts(Device, 'power_outlet_count', 'poweroutlets')
+    update_counts(Device, 'interface_count', 'interfaces')
+    update_counts(Device, 'front_port_count', 'frontports')
+    update_counts(Device, 'rear_port_count', 'rearports')
+    update_counts(Device, 'device_bay_count', 'devicebays')
+    update_counts(Device, 'module_bay_count', 'modulebays')
+    update_counts(Device, 'inventory_item_count', 'inventoryitems')
 
 
 class Migration(migrations.Migration):

netbox/dcim/migrations/0177_devicetype_component_counters.py

@@ -2,47 +2,22 @@ from django.db import migrations
 from django.db.models import Count
 
 import utilities.fields
+from utilities.counters import update_counts
 
 
 def recalculate_devicetype_template_counts(apps, schema_editor):
     DeviceType = apps.get_model("dcim", "DeviceType")
-    device_types = list(DeviceType.objects.all().annotate(
-        _console_port_template_count=Count('consoleporttemplates', distinct=True),
-        _console_server_port_template_count=Count('consoleserverporttemplates', distinct=True),
-        _power_port_template_count=Count('powerporttemplates', distinct=True),
-        _power_outlet_template_count=Count('poweroutlettemplates', distinct=True),
-        _interface_template_count=Count('interfacetemplates', distinct=True),
-        _front_port_template_count=Count('frontporttemplates', distinct=True),
-        _rear_port_template_count=Count('rearporttemplates', distinct=True),
-        _device_bay_template_count=Count('devicebaytemplates', distinct=True),
-        _module_bay_template_count=Count('modulebaytemplates', distinct=True),
-        _inventory_item_template_count=Count('inventoryitemtemplates', distinct=True),
-    ))
 
-    for devicetype in device_types:
-        devicetype.console_port_template_count = devicetype._console_port_template_count
-        devicetype.console_server_port_template_count = devicetype._console_server_port_template_count
-        devicetype.power_port_template_count = devicetype._power_port_template_count
-        devicetype.power_outlet_template_count = devicetype._power_outlet_template_count
-        devicetype.interface_template_count = devicetype._interface_template_count
-        devicetype.front_port_template_count = devicetype._front_port_template_count
-        devicetype.rear_port_template_count = devicetype._rear_port_template_count
-        devicetype.device_bay_template_count = devicetype._device_bay_template_count
-        devicetype.module_bay_template_count = devicetype._module_bay_template_count
-        devicetype.inventory_item_template_count = devicetype._inventory_item_template_count
-
-    DeviceType.objects.bulk_update(device_types, [
-        'console_port_template_count',
-        'console_server_port_template_count',
-        'power_port_template_count',
-        'power_outlet_template_count',
-        'interface_template_count',
-        'front_port_template_count',
-        'rear_port_template_count',
-        'device_bay_template_count',
-        'module_bay_template_count',
-        'inventory_item_template_count',
-    ])
+    update_counts(DeviceType, 'console_port_template_count', 'consoleporttemplates')
+    update_counts(DeviceType, 'console_server_port_template_count', 'consoleserverporttemplates')
+    update_counts(DeviceType, 'power_port_template_count', 'powerporttemplates')
+    update_counts(DeviceType, 'power_outlet_template_count', 'poweroutlettemplates')
+    update_counts(DeviceType, 'interface_template_count', 'interfacetemplates')
+    update_counts(DeviceType, 'front_port_template_count', 'frontporttemplates')
+    update_counts(DeviceType, 'rear_port_template_count', 'rearporttemplates')
+    update_counts(DeviceType, 'device_bay_template_count', 'devicebaytemplates')
+    update_counts(DeviceType, 'module_bay_template_count', 'modulebaytemplates')
+    update_counts(DeviceType, 'inventory_item_template_count', 'inventoryitemtemplates')
 
 
 class Migration(migrations.Migration):

netbox/dcim/migrations/0178_virtual_chassis_member_counter.py

@@ -2,17 +2,13 @@ from django.db import migrations
 from django.db.models import Count
 
 import utilities.fields
+from utilities.counters import update_counts
 
 
 def populate_virtualchassis_members(apps, schema_editor):
     VirtualChassis = apps.get_model('dcim', 'VirtualChassis')
 
-    vcs = VirtualChassis.objects.annotate(_member_count=Count('members', distinct=True))
-
-    for vc in vcs:
-        vc.member_count = vc._member_count
-
-    VirtualChassis.objects.bulk_update(vcs, ['member_count'], batch_size=100)
+    update_counts(VirtualChassis, 'member_count', 'members')
 
 
 class Migration(migrations.Migration):

netbox/utilities/counters.py

@@ -1,5 +1,5 @@
 from django.apps import apps
-from django.db.models import F
+from django.db.models import F, Count, OuterRef, Subquery
 from django.db.models.signals import post_delete, post_save
 
 from netbox.registry import registry
@@ -23,6 +23,24 @@ def update_counter(model, pk, counter_name, value):
     )
 
 
+def update_counts(model, field_name, related_query):
+    """
+    Perform a bulk update for the given model and counter field. For example,
+
+        update_counts(Device, '_interface_count', 'interfaces')
+
+    will effectively set
+
+        Device.objects.update(_interface_count=Count('interfaces'))
+    """
+    subquery = Subquery(
+        model.objects.filter(pk=OuterRef('pk')).annotate(_count=Count(related_query)).values('_count')
+    )
+    return model.objects.update(**{
+        field_name: subquery
+    })
+
+
 #
 # Signal handlers
 #

netbox/utilities/management/commands/calculate_cached_counts.py

@@ -4,6 +4,7 @@ from django.core.management.base import BaseCommand
 from django.db.models import Count, OuterRef, Subquery
 
 from netbox.registry import registry
+from utilities.counters import update_counts
 
 
 class Command(BaseCommand):
@@ -26,27 +27,9 @@ class Command(BaseCommand):
 
         return models
 
-    def update_counts(self, model, field_name, related_query):
-        """
-        Perform a bulk update for the given model and counter field. For example,
-
-            update_counts(Device, '_interface_count', 'interfaces')
-
-        will effectively set
-
-            Device.objects.update(_interface_count=Count('interfaces'))
-        """
-        self.stdout.write(f'Updating {model.__name__} {field_name}...')
-        subquery = Subquery(
-            model.objects.filter(pk=OuterRef('pk')).annotate(_count=Count(related_query)).values('_count')
-        )
-        return model.objects.update(**{
-            field_name: subquery
-        })
-
     def handle(self, *model_names, **options):
         for model, mappings in self.collect_models().items():
             for field_name, related_query in mappings.items():
-                self.update_counts(model, field_name, related_query)
+                update_counts(model, field_name, related_query)
 
         self.stdout.write(self.style.SUCCESS('Finished.'))

netbox/virtualization/migrations/0035_virtualmachine_interface_count.py

@@ -2,17 +2,13 @@ from django.db import migrations
 from django.db.models import Count
 
 import utilities.fields
+from utilities.counters import update_counts
 
 
 def populate_virtualmachine_counts(apps, schema_editor):
     VirtualMachine = apps.get_model('virtualization', 'VirtualMachine')
 
-    vms = VirtualMachine.objects.annotate(_interface_count=Count('interfaces', distinct=True))
-
-    for vm in vms:
-        vm.interface_count = vm._interface_count
-
-    VirtualMachine.objects.bulk_update(vms, ['interface_count'], batch_size=100)
+    update_counts(VirtualMachine, 'interface_count', 'interfaces')
 
 
 class Migration(migrations.Migration):