Cache the downstream power ports for each power outlet

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
+        ),
+    ]

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
+            # Local outlets associated with this power port
-            visited_power_ports = PowerPort.objects.none()
+            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
+            @cached_as(self, extra=outlets)
-            power_outlets = PowerOutlet.objects.filter(power_port=self)
+            def _stats():
-            if leg is not None:
+                return PowerPort.objects.filter(
-                power_outlets = power_outlets.filter(feed_leg=leg)
+                    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.
+            # Power ports drawing power from the local outlets
-            connected_power_ports = PowerPort.objects.filter(_connected_poweroutlet__in=power_outlets).nocache()
+            stats = _stats()
 
-            # Only count the local outlets (i.e. ignore non-immediate ones)
+            return outlets.count(), stats.get('allocated_draw__sum') or 0, stats.get('maximum_draw__sum') or 0
-            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
 
         # 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

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
+        # Power from device 4 to device 1
-        self.power_outlet43 = PowerOutlet.objects.create(
+        self.power_outlet41 = PowerOutlet.objects.create(
             device=self.device4,
             name='Power Outlet 43',
             power_port=self.power_port43,
         )
-        self.power_port34 = PowerPort.objects.create(
+        self.power_feed_cable.delete()
-            device=self.device3,
+        loop_cable = Cable.objects.create(termination_a=self.power_outlet41, termination_b=self.power_port1)
-            name='Power Port 34',
-            maximum_draw=2,
-            allocated_draw=1,
-        )
-        Cable.objects.create(termination_a=self.power_outlet43, termination_b=self.power_port34)
 
         stats = self.power_port1.get_power_draw()
 
-        self.assertEqual(stats['maximum'], 25 + 7 + 4 + 2)
+        self.assertEqual(stats['maximum'], 25 + 7 + 4)
-        self.assertEqual(stats['allocated'], 10 + 5 + 3 + 1)
+        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)