Bläddra i källkod

Simplify assignment of new cable terminations

jeremystretch 3 år sedan
förälder
incheckning
a909ceda84
3 ändrade filer med 67 tillägg och 56 borttagningar
  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'
     )
 
-    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):

+ 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.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):
     """

+ 12 - 20
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