diff --git a/netbox/dcim/migrations/0099_outlet_related_powerports.py b/netbox/dcim/migrations/0099_outlet_related_powerports.py new file mode 100644 index 000000000..1bc26776f --- /dev/null +++ b/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, + ), + ] diff --git a/netbox/dcim/models/device_components.py b/netbox/dcim/models/device_components.py index a41eda576..d43b9e933 100644 --- a/netbox/dcim/models/device_components.py +++ b/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 diff --git a/netbox/dcim/signals.py b/netbox/dcim/signals.py index 71ee7ec3c..64c9dce15 100644 --- a/netbox/dcim/signals.py +++ b/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()