From a909ceda84f030591ca4929dee2cfb77076e76f2 Mon Sep 17 00:00:00 2001 From: jeremystretch Date: Wed, 18 May 2022 15:49:52 -0400 Subject: [PATCH] Simplify assignment of new cable terminations --- netbox/dcim/forms/connections.py | 27 +++----------- netbox/dcim/models/cables.py | 64 +++++++++++++++++++++++++------- netbox/dcim/signals.py | 32 ++++++---------- 3 files changed, 67 insertions(+), 56 deletions(-) diff --git a/netbox/dcim/forms/connections.py b/netbox/dcim/forms/connections.py index 7058c01d6..66a9e9c6e 100644 --- a/netbox/dcim/forms/connections.py +++ b/netbox/dcim/forms/connections.py @@ -29,30 +29,13 @@ class BaseCableConnectionForm(TenancyForm, NetBoxModelForm): disabled_indicator='_occupied' ) - def save(self, commit=True): - instance = super().save(commit=commit) + def save(self, *args, **kwargs): - # Create CableTermination instances - terminations = [] - terminations.extend([ - CableTermination(cable=instance, cable_end='A', termination=termination) - for termination in self.cleaned_data['a_terminations'] - ]) - terminations.extend([ - CableTermination(cable=instance, cable_end='B', termination=termination) - for termination in self.cleaned_data['b_terminations'] - ]) + # Set the A/B terminations on the Cable instance + self.instance.a_terminations = self.cleaned_data['a_terminations'] + self.instance.b_terminations = self.cleaned_data['b_terminations'] - if commit: - for ct in terminations: - ct.save() - else: - instance._terminations = [ - *self.cleaned_data['a_terminations'], - *self.cleaned_data['b_terminations'], - ] - - return instance + return super().save(*args, **kwargs) class ConnectCableToDeviceForm(BaseCableConnectionForm): diff --git a/netbox/dcim/models/cables.py b/netbox/dcim/models/cables.py index dd5599a1d..42c298279 100644 --- a/netbox/dcim/models/cables.py +++ b/netbox/dcim/models/cables.py @@ -5,6 +5,7 @@ from django.contrib.contenttypes.models import ContentType from django.core.exceptions import ValidationError from django.db import models from django.db.models import Sum +from django.dispatch import Signal from django.urls import reverse from dcim.choices import * @@ -27,6 +28,9 @@ __all__ = ( ) +trace_paths = Signal() + + # # Cables # @@ -107,20 +111,10 @@ class Cable(NetBoxModel): self._orig_status = self.status # Assign associated CableTerminations (if any) - terminations = [] - if a_terminations and type(a_terminations) is list: - terminations.extend([ - CableTermination(cable=self, cable_end='A', termination=t) for t in a_terminations - ]) - if b_terminations and type(b_terminations) is list: - terminations.extend([ - CableTermination(cable=self, cable_end='B', termination=t) for t in b_terminations - ]) - if terminations: - assert self.pk is None - self._terminations = terminations - else: - self._terminations = [] + if a_terminations is not None: + self.a_terminations = a_terminations + if b_terminations is not None: + self.b_terminations = b_terminations @classmethod def from_db(cls, db, field_names, values): @@ -164,6 +158,7 @@ class Cable(NetBoxModel): self.length_unit = '' def save(self, *args, **kwargs): + _created = self.pk is None # Store the given length (if any) in meters for use in database ordering if self.length and self.length_unit: @@ -183,6 +178,32 @@ class Cable(NetBoxModel): # Update the private pk used in __str__ in case this is a new object (i.e. just got its pk) self._pk = self.pk + # Retrieve existing A/B terminations for the Cable + a_terminations = {ct.termination: ct for ct in self.terminations.filter(cable_end='A')} + b_terminations = {ct.termination: ct for ct in self.terminations.filter(cable_end='B')} + + # Delete stale CableTerminations + if hasattr(self, 'a_terminations'): + for termination, ct in a_terminations.items(): + if termination not in self.a_terminations: + ct.delete() + if hasattr(self, 'b_terminations'): + for termination, ct in b_terminations.items(): + if termination not in self.b_terminations: + ct.delete() + + # Save new CableTerminations (if any) + if hasattr(self, 'a_terminations'): + for termination in self.a_terminations: + if termination not in a_terminations: + CableTermination(cable=self, cable_end='A', termination=termination).save() + if hasattr(self, 'b_terminations'): + for termination in self.b_terminations: + if termination not in b_terminations: + CableTermination(cable=self, cable_end='B', termination=termination).save() + + trace_paths.send(Cable, instance=self, created=_created) + def get_status_color(self): return LinkStatusChoices.colors.get(self.status) @@ -271,6 +292,21 @@ class CableTermination(models.Model): # f"Incompatible termination types: {self.termination_a_type} and {self.termination_b_type}" # ) + def save(self, *args, **kwargs): + super().save(*args, **kwargs) + + # Set the cable on the terminating object + termination_model = self.termination._meta.model + termination_model.objects.filter(pk=self.termination_id).update(cable=self.cable) + + def delete(self, *args, **kwargs): + + # Delete the cable association on the terminating object + termination_model = self.termination._meta.model + termination_model.objects.filter(pk=self.termination_id).update(cable=None) + + super().delete(*args, **kwargs) + class CablePath(models.Model): """ diff --git a/netbox/dcim/signals.py b/netbox/dcim/signals.py index 48fb2f25c..9ea8dd9f8 100644 --- a/netbox/dcim/signals.py +++ b/netbox/dcim/signals.py @@ -1,12 +1,11 @@ -from collections import defaultdict import logging -from django.contrib.contenttypes.models import ContentType from django.db.models.signals import post_save, post_delete, pre_delete from django.dispatch import receiver from .choices import LinkStatusChoices from .models import Cable, CablePath, CableTermination, Device, PathEndpoint, PowerPanel, Rack, Location, VirtualChassis +from .models.cables import trace_paths from .utils import create_cablepath, rebuild_paths @@ -69,8 +68,7 @@ def clear_virtualchassis_members(instance, **kwargs): # Cables # - -@receiver(post_save, sender=Cable) +@receiver(trace_paths, sender=Cable) def update_connected_endpoints(instance, created, raw=False, **kwargs): """ When a Cable is saved, check for and update its two connected endpoints @@ -80,28 +78,22 @@ def update_connected_endpoints(instance, created, raw=False, **kwargs): logger.debug(f"Skipping endpoint updates for imported cable {instance}") return - # Save any new CableTerminations - CableTermination.objects.bulk_create([ - term for term in instance._terminations if not term.pk - ]) - - # Split terminations into A/B sets and save link assignments # TODO: Update link peers - _terms = defaultdict(list) - for t in instance._terminations: - if t.termination.cable != instance: - t.termination.cable = instance - t.termination.save() - _terms[t.cable_end].append(t.termination) # Create/update cable paths if created: - for terms in _terms.values(): + _terms = { + 'A': [t.termination for t in instance.terminations.filter(cable_end='A')], + 'B': [t.termination for t in instance.terminations.filter(cable_end='B')], + } + for nodes in _terms.values(): # Examine type of first termination to determine object type (all must be the same) - if isinstance(terms[0], PathEndpoint): - create_cablepath(terms) + if not nodes: + continue + if isinstance(nodes[0], PathEndpoint): + create_cablepath(nodes) else: - rebuild_paths(terms) + rebuild_paths(nodes) elif instance.status != instance._orig_status: # We currently don't support modifying either termination of an existing Cable. (This # may change in the future.) However, we do need to capture status changes and update