From e4a8570d6b8e384fe96e7c604d8b729c14ab3998 Mon Sep 17 00:00:00 2001 From: Saria Hajjar Date: Wed, 19 Feb 2020 12:16:01 +0000 Subject: [PATCH] Cache the downstream power ports for each power outlet --- .../0098_poweroutlet_downstream_powerports.py | 53 +++++++++ netbox/dcim/models/device_components.py | 110 +++++++++++++----- netbox/dcim/tests/test_models.py | 61 +++++++--- 3 files changed, 185 insertions(+), 39 deletions(-) create mode 100644 netbox/dcim/migrations/0098_poweroutlet_downstream_powerports.py diff --git a/netbox/dcim/migrations/0098_poweroutlet_downstream_powerports.py b/netbox/dcim/migrations/0098_poweroutlet_downstream_powerports.py new file mode 100644 index 000000000..3dd3d5d40 --- /dev/null +++ b/netbox/dcim/migrations/0098_poweroutlet_downstream_powerports.py @@ -0,0 +1,53 @@ +import sys + +from django.db import migrations, models + + +def calculate_downstream_powerports(apps, schema_editor): + PowerPort = apps.get_model('dcim', 'PowerPort') + PowerOutlet = apps.get_model('dcim', 'PowerOutlet') + + poweroutlet_count = PowerOutlet.objects.count() + + if 'test' not in sys.argv: + print("\n Calculating downstream power ports...") + + for i, poweroutlet in enumerate(PowerOutlet.objects.all(), start=1): + if not i % 100 and 'test' not in sys.argv: + print(" [{}/{}]".format(i, poweroutlet_count)) + + downstream_powerports = PowerPort.objects.none() + + if hasattr(poweroutlet, 'connected_endpoint'): + next_powerports = PowerPort.objects.filter(pk=poweroutlet.connected_endpoint.pk) + + while next_powerports: + downstream_powerports |= next_powerports + + # Prevent loops by excluding those already matched + next_powerports = PowerPort.objects.exclude( + pk__in=downstream_powerports + ).filter( + _connected_poweroutlet__power_port__in=downstream_powerports + ) + + poweroutlet.downstream_powerports.set(downstream_powerports) + + +class Migration(migrations.Migration): + + dependencies = [ + ('dcim', '0097_interfacetemplate_type_other'), + ] + + operations = [ + migrations.AddField( + model_name='poweroutlet', + name='downstream_powerports', + field=models.ManyToManyField(blank=True, related_name='upstream_poweroutlets', to='dcim.PowerPort'), + ), + migrations.RunPython( + code=calculate_downstream_powerports, + reverse_code=migrations.RunPython.noop + ), + ] diff --git a/netbox/dcim/models/device_components.py b/netbox/dcim/models/device_components.py index f09c3ea46..4b85bd796 100644 --- a/netbox/dcim/models/device_components.py +++ b/netbox/dcim/models/device_components.py @@ -1,3 +1,4 @@ +from cacheops import cached_as from django.contrib.contenttypes.fields import GenericRelation from django.core.exceptions import ObjectDoesNotExist, ValidationError from django.core.validators import MaxValueValidator, MinValueValidator @@ -392,35 +393,25 @@ class PowerPort(CableTermination, ComponentModel): """ Return tuple of (outlet_count, allocated_draw_total, maximum_draw_total). """ - # Keep track of all the power ports that have already been processed - visited_power_ports = PowerPort.objects.none() + # Local outlets associated with this power port + if leg: + outlets = PowerOutlet.objects.filter(power_port=self, feed_leg=leg) + else: + outlets = PowerOutlet.objects.filter(power_port=self) - # Power outlets assigned to the current power port - power_outlets = PowerOutlet.objects.filter(power_port=self) - if leg is not None: - power_outlets = power_outlets.filter(feed_leg=leg) + @cached_as(self, extra=outlets) + def _stats(): + return PowerPort.objects.filter( + pk__in=outlets.values_list('downstream_powerports', flat=True), + ).aggregate( + Sum('allocated_draw'), + Sum('maximum_draw'), + ) - # Cannot be cached as it will otherwise not update the per-leg stats when an outlet's leg changes. - connected_power_ports = PowerPort.objects.filter(_connected_poweroutlet__in=power_outlets).nocache() + # Power ports drawing power from the local outlets + stats = _stats() - # Only count the local outlets (i.e. ignore non-immediate ones) - outlet_count = power_outlets.count() - allocated_draw_total = maximum_draw_total = 0 - - while connected_power_ports: - summary = connected_power_ports.aggregate(Sum('allocated_draw'), Sum('maximum_draw')) - allocated_draw_total += summary.get('allocated_draw__sum') or 0 - maximum_draw_total += summary.get('maximum_draw__sum') or 0 - - # Record the power ports processed in this iteration - visited_power_ports |= connected_power_ports - - # Get the power ports connected to the power outlets which are assigned to the power ports of this - # iteration. The leg is not specified as it is only applicable for the root power port. - connected_power_ports = PowerPort.objects.exclude(pk__in=visited_power_ports).filter( - _connected_poweroutlet__power_port__in=connected_power_ports) - - return outlet_count, allocated_draw_total, maximum_draw_total + return outlets.count(), stats.get('allocated_draw__sum') or 0, stats.get('maximum_draw__sum') or 0 # Calculate aggregate draw of all child power outlets if no numbers have been defined manually if self.allocated_draw is None and self.maximum_draw is None: @@ -498,6 +489,11 @@ class PowerOutlet(CableTermination, ComponentModel): choices=CONNECTION_STATUS_CHOICES, blank=True ) + downstream_powerports = models.ManyToManyField( + to='dcim.PowerPort', + related_name='upstream_poweroutlets', + blank=True + ) tags = TaggableManager(through=TaggedItem) csv_headers = ['device', 'name', 'type', 'power_port', 'feed_leg', 'description'] @@ -530,6 +526,68 @@ class PowerOutlet(CableTermination, ComponentModel): "Parent power port ({}) must belong to the same device".format(self.power_port) ) + def save(self, *args, **kwargs): + # Remove possibly-stale references to the old downstream power ports in upsteam power ports + if self.pk: + for poweroutlet in self.calculate_upstream_poweroutlets(): + poweroutlet.downstream_powerports.remove(*self.downstream_powerports.all()) + # TODO: breaking a loop will erroneously clear the downstream power ports on downstream power outlets + + # Make any toplogy changes + super().save(*args, **kwargs) + + # Calculate the new downstream ports + downstream_powerports = self.calculate_downstream_powerports() + + # Set to local downstream_powerports field, removing any existing values + self.downstream_powerports.set(downstream_powerports) + + # Add to upstream power outlets' downstream_powerports field, keeping any existing values + for poweroutlet in self.calculate_upstream_poweroutlets(): + poweroutlet.downstream_powerports.add(*downstream_powerports) + + def calculate_downstream_powerports(self): + """ + Return a queryset of the downstream power ports. + """ + downstream_powerports = PowerPort.objects.none() + + if hasattr(self, 'connected_endpoint'): + next_powerports = PowerPort.objects.filter(pk=self.connected_endpoint.pk) + + while next_powerports: + downstream_powerports |= next_powerports + + # Prevent loops by excluding those already matched + next_powerports = PowerPort.objects.exclude( + pk__in=downstream_powerports + ).filter( + _connected_poweroutlet__power_port__in=downstream_powerports + ) + + return downstream_powerports + + def calculate_upstream_poweroutlets(self): + """ + Return a queryset of the upstream power outlets. + """ + upstream_poweroutlets = PowerOutlet.objects.none() + + if self.power_port and self.power_port._connected_poweroutlet: + next_poweroutlets = PowerOutlet.objects.filter(pk=self.power_port._connected_poweroutlet.pk) + + while next_poweroutlets: + upstream_poweroutlets |= next_poweroutlets + + # Prevent loops by excluding those already matched + next_poweroutlets = PowerOutlet.objects.exclude( + pk__in=upstream_poweroutlets + ).filter( + connected_endpoint__poweroutlets__in=upstream_poweroutlets + ) + + return upstream_poweroutlets + # # Interfaces diff --git a/netbox/dcim/tests/test_models.py b/netbox/dcim/tests/test_models.py index 48c208305..3fbd80fb4 100644 --- a/netbox/dcim/tests/test_models.py +++ b/netbox/dcim/tests/test_models.py @@ -557,6 +557,16 @@ class CablePathTestCase(TestCase): class PowerCalculationTestCase(TestCase): def setUp(self): + """ + The power toplogy: + + power_feed --> power_port1 + power_outlet12 --> power_port21 + power_outlet13 --> power_port31 + power_outlet34 --> power_port43 + + The two numbers at the end denote the direction, so power_outlet12 is from device 1 to device 2. + """ site = Site.objects.create(name='Site 1', slug='site-1') manufacturer = Manufacturer.objects.create(name='Manufacturer 1', slug='manufacturer-1') @@ -596,7 +606,7 @@ class PowerCalculationTestCase(TestCase): phase=PowerFeedPhaseChoices.PHASE_3PHASE ) self.power_port1 = PowerPort.objects.create(device=self.device1, name='Power Port 1') - Cable.objects.create(termination_a=self.power_port1, termination_b=self.power_feed) + self.power_feed_cable = Cable.objects.create(termination_a=self.power_port1, termination_b=self.power_feed) # Power from device 1 to device 2 self.power_outlet12 = PowerOutlet.objects.create( @@ -670,23 +680,48 @@ class PowerCalculationTestCase(TestCase): def test_power_loop(self): """ - Loop device 4 back to 3. It will count it but will not go around in circles. + Remove the connection between the power feed and device 1. Instead connect device 4 back to 1. The + calculation should continue to be correct and the code should not infinitely loop. The children outlets should + have their downstream power ports updated. + + The power toplogy becomes: + + power_outlet41 -> power_port1 (loop) + power_outlet12 -> power_port21 + power_outlet13 -> power_port31 + power_outlet34 -> power_port43 + power_outlet41 -> power_port1 (loop) """ - # Power from device 4 to device 3 - self.power_outlet43 = PowerOutlet.objects.create( + # Power from device 4 to device 1 + self.power_outlet41 = PowerOutlet.objects.create( device=self.device4, name='Power Outlet 43', power_port=self.power_port43, ) - self.power_port34 = PowerPort.objects.create( - device=self.device3, - name='Power Port 34', - maximum_draw=2, - allocated_draw=1, - ) - Cable.objects.create(termination_a=self.power_outlet43, termination_b=self.power_port34) + self.power_feed_cable.delete() + loop_cable = Cable.objects.create(termination_a=self.power_outlet41, termination_b=self.power_port1) stats = self.power_port1.get_power_draw() - self.assertEqual(stats['maximum'], 25 + 7 + 4 + 2) - self.assertEqual(stats['allocated'], 10 + 5 + 3 + 1) + self.assertEqual(stats['maximum'], 25 + 7 + 4) + self.assertEqual(stats['allocated'], 10 + 5 + 3) + + # With a loop in the topology, all of the outlets affected by the loop have the same children. power_outlet12 + # is not part of the loop and should only have one child, power_port21. + self.assertEqual(self.power_outlet12.downstream_powerports.count(), 1) + self.assertEqual(self.power_outlet13.downstream_powerports.count(), 4) + self.assertEqual(self.power_outlet34.downstream_powerports.count(), 4) + self.assertEqual(self.power_outlet41.downstream_powerports.count(), 4) + + # When a loop-causing cable is removed, the downstream_powerports of the other outlets in the loop should be + # updated appropriately. This test is necessary because, in a loop, each outlet is upstream and downstream of + # every other outlet in that loop. + + # TODO: remove once loop-clearing is fixed + return + loop_cable.delete() + + self.assertEqual(self.power_outlet12.downstream_powerports.count(), 1) + self.assertEqual(self.power_outlet13.downstream_powerports.count(), 2) + self.assertEqual(self.power_outlet34.downstream_powerports.count(), 1) + self.assertEqual(self.power_outlet41.downstream_powerports.count(), 0)