Cache upstream and downstream power ports for each outlet

netbox/dcim/migrations/0099_outlet_related_powerports.py

@@ -0,0 +1,76 @@
+import sys
+
+from django.db import migrations, models
+
+
+def update_related_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    Updating power outlets with related 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))
+
+        # Copy of PowerOutlet.calculate_upstream_powerports
+        upstream_powerports = PowerPort.objects.none()
+
+        if poweroutlet.power_port:
+            next_powerports = PowerPort.objects.filter(pk=poweroutlet.power_port.pk)
+
+            while next_powerports.exists():
+                upstream_powerports |= next_powerports
+
+                # Prevent loops by excluding those already matched
+                next_powerports = PowerPort.objects.exclude(
+                    pk__in=upstream_powerports,
+                ).filter(
+                    poweroutlets__connected_endpoint__in=upstream_powerports,
+                )
+
+        # Copy of PowerOutlet.calculate_downstream_powerports
+        downstream_powerports = PowerPort.objects.none()
+
+        if hasattr(poweroutlet, 'connected_endpoint'):
+            next_powerports = PowerPort.objects.filter(pk=poweroutlet.connected_endpoint.pk)
+
+            while next_powerports.exists():
+                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.upstream_powerports.set(upstream_powerports)
+        poweroutlet.downstream_powerports.set(downstream_powerports)
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('dcim', '0098_devicetype_images'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='poweroutlet',
+            name='downstream_powerports',
+            field=models.ManyToManyField(blank=True, related_name='upstream_poweroutlets', to='dcim.PowerPort'),
+        ),
+        migrations.AddField(
+            model_name='poweroutlet',
+            name='upstream_powerports',
+            field=models.ManyToManyField(blank=True, related_name='downstream_poweroutlets', to='dcim.PowerPort'),
+        ),
+        migrations.RunPython(
+            code=update_related_powerports,
+            reverse_code=migrations.RunPython.noop,
+        ),
+    ]

netbox/dcim/models/device_components.py

@@ -470,6 +470,16 @@ class PowerOutlet(CableTermination, ComponentModel):
         choices=CONNECTION_STATUS_CHOICES,
         blank=True
     )
+    downstream_powerports = models.ManyToManyField(
+        to='dcim.PowerPort',
+        related_name='upstream_poweroutlets',
+        blank=True
+    )
+    upstream_powerports = models.ManyToManyField(
+        to='dcim.PowerPort',
+        related_name='downstream_poweroutlets',
+        blank=True
+    )
     tags = TaggableManager(through=TaggedItem)
 
     csv_headers = ['device', 'name', 'type', 'power_port', 'feed_leg', 'description']
@@ -502,6 +512,66 @@ class PowerOutlet(CableTermination, ComponentModel):
                 "Parent power port ({}) must belong to the same device".format(self.power_port)
             )
 
+    def calculate_downstream_powerports(self):
+        """
+        Return a queryset of the power ports indirectly drawing power from this outlet.
+        """
+        powerports = PowerPort.objects.none()
+
+        if hasattr(self, 'connected_endpoint'):
+            next_powerports = PowerPort.objects.filter(pk=self.connected_endpoint.pk)
+
+            while next_powerports.exists():
+                powerports |= next_powerports
+
+                # Prevent loops by excluding those already matched
+                next_powerports = PowerPort.objects.exclude(
+                    pk__in=powerports,
+                ).filter(
+                    _connected_poweroutlet__power_port__in=powerports,
+                )
+
+        return powerports
+
+    def calculate_upstream_powerports(self):
+        """
+        Return a queryset of the power ports indirectly supplying power to this outlet.
+        """
+        powerports = PowerPort.objects.none()
+
+        if self.power_port:
+            next_powerports = PowerPort.objects.filter(pk=self.power_port.pk)
+
+            while next_powerports.exists():
+                powerports |= next_powerports
+
+                # Prevent loops by excluding those already matched
+                next_powerports = PowerPort.objects.exclude(
+                    pk__in=powerports,
+                ).filter(
+                    poweroutlets__connected_endpoint__in=powerports,
+                )
+
+        return powerports
+
+    def update_related_powerports(self):
+        """
+        Update the downstream and upstream power ports. This bubbles as needed to any parent power outlets, existing
+        and new.
+        """
+        upstream_powerports = self.calculate_upstream_powerports()
+        downstream_powerports = self.calculate_downstream_powerports()
+
+        old_parents = PowerOutlet.objects.filter(connected_endpoint__in=self.upstream_powerports.all())
+        new_parents = PowerOutlet.objects.filter(connected_endpoint__in=upstream_powerports)
+
+        for outlet in old_parents | new_parents:
+            outlet.upstream_powerports.set(outlet.calculate_upstream_powerports())
+            outlet.downstream_powerports.set(outlet.calculate_downstream_powerports())
+
+        self.upstream_powerports.set(self.calculate_upstream_powerports())
+        self.downstream_powerports.set(self.calculate_downstream_powerports())
+
 
 #
 # Interfaces

netbox/dcim/signals.py

@@ -1,7 +1,7 @@
 from django.db.models.signals import post_save, pre_delete
 from django.dispatch import receiver
 
-from .models import Cable, Device, VirtualChassis
+from .models import Cable, Device, PowerOutlet, VirtualChassis
 
 
 @receiver(post_save, sender=VirtualChassis)
@@ -53,6 +53,12 @@ def update_connected_endpoints(instance, **kwargs):
         endpoint_b.connection_status = path_status
         endpoint_b.save()
 
+    # The cached fields for a power outlet need to be updated after a topology change (both endpoints changed)
+    if isinstance(endpoint_a, PowerOutlet):
+        endpoint_a.update_related_powerports()
+    elif isinstance(endpoint_b, PowerOutlet):
+        endpoint_b.update_related_powerports()
+
 
 @receiver(pre_delete, sender=Cable)
 def nullify_connected_endpoints(instance, **kwargs):
@@ -77,3 +83,9 @@ def nullify_connected_endpoints(instance, **kwargs):
         endpoint_b.connected_endpoint = None
         endpoint_b.connection_status = None
         endpoint_b.save()
+
+    # The cached fields for a power outlet need to be updated after a topology change (both endpoints changed)
+    if isinstance(endpoint_a, PowerOutlet):
+        endpoint_a.update_related_powerports()
+    elif isinstance(endpoint_b, PowerOutlet):
+        endpoint_b.update_related_powerports()