From d2ab41abfba04036603c3958226e23e992d38dcd Mon Sep 17 00:00:00 2001
From: Jeremy Stretch <stretch@packetlife.net>
Date: Thu, 26 Sep 2019 17:17:12 -0400
Subject: [PATCH] Cache A/B termination devices on the Cable model

---
 netbox/dcim/migrations/0075_cable_devices.py | 46 ++++++++++++++++++++
 netbox/dcim/models.py                        | 22 ++++++++++
 2 files changed, 68 insertions(+)
 create mode 100644 netbox/dcim/migrations/0075_cable_devices.py

diff --git a/netbox/dcim/migrations/0075_cable_devices.py b/netbox/dcim/migrations/0075_cable_devices.py
new file mode 100644
index 000000000..d93f1ab7c
--- /dev/null
+++ b/netbox/dcim/migrations/0075_cable_devices.py
@@ -0,0 +1,46 @@
+from django.db import migrations, models
+import django.db.models.deletion
+
+
+def cache_cable_devices(apps, schema_editor):
+    Cable = apps.get_model('dcim', 'Cable')
+
+    # Cache A/B termination devices on all existing Cables. Note that the custom save() method on Cable is not
+    # available during a migration, so we replicate its logic here.
+    for cable in Cable.objects.all():
+
+        termination_a_model = apps.get_model(cable.termination_a_type.app_label, cable.termination_a_type.model)
+        if hasattr(termination_a_model, 'device'):
+            termination_a = termination_a_model.objects.get(pk=cable.termination_a_id)
+            cable._termination_a_device = termination_a.device
+
+        termination_b_model = apps.get_model(cable.termination_b_type.app_label, cable.termination_b_type.model)
+        if hasattr(termination_b_model, 'device'):
+            termination_b = termination_b_model.objects.get(pk=cable.termination_b_id)
+            cable._termination_b_device = termination_b.device
+
+        cable.save()
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('dcim', '0074_increase_field_length_platform_name_slug'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='cable',
+            name='_termination_a_device',
+            field=models.ForeignKey(blank=True, null=True, on_delete=django.db.models.deletion.CASCADE, related_name='+', to='dcim.Device'),
+        ),
+        migrations.AddField(
+            model_name='cable',
+            name='_termination_b_device',
+            field=models.ForeignKey(blank=True, null=True, on_delete=django.db.models.deletion.CASCADE, related_name='+', to='dcim.Device'),
+        ),
+        migrations.RunPython(
+            code=cache_cable_devices,
+            reverse_code=migrations.RunPython.noop
+        ),
+    ]
diff --git a/netbox/dcim/models.py b/netbox/dcim/models.py
index 5ffd15842..2c554c158 100644
--- a/netbox/dcim/models.py
+++ b/netbox/dcim/models.py
@@ -2781,6 +2781,22 @@ class Cable(ChangeLoggedModel):
         blank=True,
         null=True
     )
+    # Cache the associated device (where applicable) for the A and B terminations. This enables filtering of Cables by
+    # their associated Devices.
+    _termination_a_device = models.ForeignKey(
+        to=Device,
+        on_delete=models.CASCADE,
+        related_name='+',
+        blank=True,
+        null=True
+    )
+    _termination_b_device = models.ForeignKey(
+        to=Device,
+        on_delete=models.CASCADE,
+        related_name='+',
+        blank=True,
+        null=True
+    )
 
     csv_headers = [
         'termination_a_type', 'termination_a_id', 'termination_b_type', 'termination_b_id', 'type', 'status', 'label',
@@ -2895,6 +2911,12 @@ class Cable(ChangeLoggedModel):
         if self.length and self.length_unit:
             self._abs_length = to_meters(self.length, self.length_unit)
 
+        # Store the parent Device for the A and B terminations (if applicable) to enable filtering
+        if hasattr(self.termination_a, 'device'):
+            self._termination_a_device = self.termination_a.device
+        if hasattr(self.termination_b, 'device'):
+            self._termination_b_device = self.termination_b.device
+
         super().save(*args, **kwargs)
 
     def to_csv(self):