Browse Source

Simplify assignment of new cable terminations

jeremystretch 3 years ago
parent
commit
a909ceda84
3 changed files with 67 additions and 56 deletions
  1. 5 22
      netbox/dcim/forms/connections.py
  2. 50 14
      netbox/dcim/models/cables.py
  3. 12 20
      netbox/dcim/signals.py

+ 5 - 22
netbox/dcim/forms/connections.py

@@ -29,30 +29,13 @@ class BaseCableConnectionForm(TenancyForm, NetBoxModelForm):
         disabled_indicator='_occupied'
         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):
 class ConnectCableToDeviceForm(BaseCableConnectionForm):

+ 50 - 14
netbox/dcim/models/cables.py

@@ -5,6 +5,7 @@ from django.contrib.contenttypes.models import ContentType
 from django.core.exceptions import ValidationError
 from django.core.exceptions import ValidationError
 from django.db import models
 from django.db import models
 from django.db.models import Sum
 from django.db.models import Sum
+from django.dispatch import Signal
 from django.urls import reverse
 from django.urls import reverse
 
 
 from dcim.choices import *
 from dcim.choices import *
@@ -27,6 +28,9 @@ __all__ = (
 )
 )
 
 
 
 
+trace_paths = Signal()
+
+
 #
 #
 # Cables
 # Cables
 #
 #
@@ -107,20 +111,10 @@ class Cable(NetBoxModel):
         self._orig_status = self.status
         self._orig_status = self.status
 
 
         # Assign associated CableTerminations (if any)
         # 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
     @classmethod
     def from_db(cls, db, field_names, values):
     def from_db(cls, db, field_names, values):
@@ -164,6 +158,7 @@ class Cable(NetBoxModel):
             self.length_unit = ''
             self.length_unit = ''
 
 
     def save(self, *args, **kwargs):
     def save(self, *args, **kwargs):
+        _created = self.pk is None
 
 
         # Store the given length (if any) in meters for use in database ordering
         # Store the given length (if any) in meters for use in database ordering
         if self.length and self.length_unit:
         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)
         # Update the private pk used in __str__ in case this is a new object (i.e. just got its pk)
         self._pk = self.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):
     def get_status_color(self):
         return LinkStatusChoices.colors.get(self.status)
         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}"
         #         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):
 class CablePath(models.Model):
     """
     """

+ 12 - 20
netbox/dcim/signals.py

@@ -1,12 +1,11 @@
-from collections import defaultdict
 import logging
 import logging
 
 
-from django.contrib.contenttypes.models import ContentType
 from django.db.models.signals import post_save, post_delete, pre_delete
 from django.db.models.signals import post_save, post_delete, pre_delete
 from django.dispatch import receiver
 from django.dispatch import receiver
 
 
 from .choices import LinkStatusChoices
 from .choices import LinkStatusChoices
 from .models import Cable, CablePath, CableTermination, Device, PathEndpoint, PowerPanel, Rack, Location, VirtualChassis
 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
 from .utils import create_cablepath, rebuild_paths
 
 
 
 
@@ -69,8 +68,7 @@ def clear_virtualchassis_members(instance, **kwargs):
 # Cables
 # Cables
 #
 #
 
 
-
-@receiver(post_save, sender=Cable)
+@receiver(trace_paths, sender=Cable)
 def update_connected_endpoints(instance, created, raw=False, **kwargs):
 def update_connected_endpoints(instance, created, raw=False, **kwargs):
     """
     """
     When a Cable is saved, check for and update its two connected endpoints
     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}")
         logger.debug(f"Skipping endpoint updates for imported cable {instance}")
         return
         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
     # 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
     # Create/update cable paths
     if created:
     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)
             # 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:
             else:
-                rebuild_paths(terms)
+                rebuild_paths(nodes)
     elif instance.status != instance._orig_status:
     elif instance.status != instance._orig_status:
         # We currently don't support modifying either termination of an existing Cable. (This
         # 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
         # may change in the future.) However, we do need to capture status changes and update