Develop triggers for setting parents

netbox/ipam/models/ip.py

@@ -1,4 +1,5 @@
 import netaddr
+import pgtrigger
 from django.contrib.contenttypes.fields import GenericForeignKey
 from django.contrib.contenttypes.models import ContentType
 from django.contrib.postgres.indexes import GistIndex
@@ -17,6 +18,7 @@ from ipam.fields import IPNetworkField, IPAddressField
 from ipam.lookups import Host
 from ipam.managers import IPAddressManager
 from ipam.querysets import PrefixQuerySet
+from ipam.triggers import ipam_prefix_delete_adjust_prefix_parent, ipam_prefix_insert_adjust_prefix_parent
 from ipam.validators import DNSValidator
 from netbox.config import get_config
 from netbox.models import OrganizationalModel, PrimaryModel
@@ -186,25 +188,6 @@ class Aggregate(ContactsMixin, GetAvailablePrefixesMixin, PrimaryModel):
         return min(utilization, 100)
 
 
-class Role(OrganizationalModel):
-    """
-    A Role represents the functional role of a Prefix or VLAN; for example, "Customer," "Infrastructure," or
-    "Management."
-    """
-    weight = models.PositiveSmallIntegerField(
-        verbose_name=_('weight'),
-        default=1000
-    )
-
-    class Meta:
-        ordering = ('weight', 'name')
-        verbose_name = _('role')
-        verbose_name_plural = _('roles')
-
-    def __str__(self):
-        return self.name
-
-
 class Prefix(ContactsMixin, GetAvailablePrefixesMixin, CachedScopeMixin, PrimaryModel):
     """
     A Prefix represents an IPv4 or IPv6 network, including mask length. Prefixes can optionally be scoped to certain
@@ -306,6 +289,20 @@ class Prefix(ContactsMixin, GetAvailablePrefixesMixin, CachedScopeMixin, Primary
                 opclasses=['inet_ops'],
             ),
         ]
+        triggers = [
+            pgtrigger.Trigger(
+                name='ipam_prefix_delete',
+                operation=pgtrigger.Delete,
+                when=pgtrigger.Before,
+                func=ipam_prefix_delete_adjust_prefix_parent,
+            ),
+            pgtrigger.Trigger(
+                name='ipam_prefix_insert',
+                operation=pgtrigger.Insert,
+                when=pgtrigger.After,
+                func=ipam_prefix_insert_adjust_prefix_parent,
+            ),
+        ]
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
@@ -546,6 +543,25 @@ class Prefix(ContactsMixin, GetAvailablePrefixesMixin, CachedScopeMixin, Primary
         return prefixes.last()
 
 
+class Role(OrganizationalModel):
+    """
+    A Role represents the functional role of a Prefix or VLAN; for example, "Customer," "Infrastructure," or
+    "Management."
+    """
+    weight = models.PositiveSmallIntegerField(
+        verbose_name=_('weight'),
+        default=1000
+    )
+
+    class Meta:
+        ordering = ('weight', 'name')
+        verbose_name = _('role')
+        verbose_name_plural = _('roles')
+
+    def __str__(self):
+        return self.name
+
+
 class IPRange(ContactsMixin, PrimaryModel):
     """
     A range of IP addresses, defined by start and end addresses.

netbox/ipam/signals.py

@@ -29,123 +29,12 @@ def update_children_depth(prefix):
     Prefix.objects.bulk_update(children, ['_depth'], batch_size=100)
 
 
-def update_object_prefix(prefix, child_model=IPAddress):
-    filter = Q(prefix=prefix)
-
-    if child_model == IPAddress:
-        filter |= Q(address__net_contained_or_equal=prefix.prefix, vrf=prefix.vrf)
-    elif child_model == IPRange:
-        filter |= Q(
-            start_address__net_contained_or_equal=prefix.prefix,
-            end_address__net_contained_or_equal=prefix.prefix,
-            vrf=prefix.vrf
-        )
-
-    addresses = child_model.objects.filter(filter)
-    for address in addresses:
-        # If addresses prefix is not set then this model is the only option
-        if not address.prefix:
-            address.prefix = prefix
-        # This address has a different VRF so the prefix cannot be the parent prefix
-        elif address.prefix != address.find_prefix(address):
-            address.prefix = address.find_prefix(address)
-        else:
-            pass
-
-    # Update the addresses
-    child_model.objects.bulk_update(addresses, ['prefix'], batch_size=100)
-
-
-def delete_object_prefix(prefix, child_model, child_objects):
-    if not prefix.parent or prefix.vrf != prefix.parent.vrf:
-        # Prefix will be Set Null
-        return
-
-    # Set prefix to prefix parent
-    for address in child_objects:
-        address.prefix = prefix.parent
-
-    # Run a bulk update
-    child_model.objects.bulk_update(child_objects, ['prefix'], batch_size=100)
-
-
-def update_ipaddress_prefix(prefix, delete=False):
-    if delete:
-        delete_object_prefix(prefix, IPAddress, prefix.ip_addresses.all())
-    else:
-        update_object_prefix(prefix, child_model=IPAddress)
-
-
-def update_iprange_prefix(prefix, delete=False):
-    if delete:
-        delete_object_prefix(prefix, IPRange, prefix.ip_ranges.all())
-    else:
-        update_object_prefix(prefix, child_model=IPRange)
-
-
-def update_prefix_parents(prefix, delete=False, created=False):
-    if delete:
-        # Set prefix to prefix parent
-        prefixes = prefix.children.all()
-        for address in prefixes:
-            address.parent = prefix.parent
-
-        # Run a bulk update
-        Prefix.objects.bulk_update(prefixes, ['parent'], batch_size=100)
-    else:
-        # Build filter to get prefixes that will be impacted by this change:
-        # * Parent prefix is this prefixes parent, and;
-        # * Prefix is contained by this prefix, and;
-        # * Prefix is either within this VRF or there is no VRF and this prefix is a container prefix
-        filter = Q(
-            parent=prefix.parent,
-            vrf=prefix.vrf,
-            prefix__net_contained=str(prefix.prefix)
-        )
-        is_container = False
-        if prefix.status == PrefixStatusChoices.STATUS_CONTAINER and prefix.vrf is None:
-            is_container = True
-            filter |= Q(
-                parent=prefix.parent,
-                vrf=None,
-                prefix__net_contained=str(prefix.prefix),
-            )
-
-        # Get all impacted prefixes.  Ensure we use distinct() to weed out duplicate prefixes from joins
-        prefixes = Prefix.objects.filter(filter)
-        # Include children
-        if not created:
-            prefixes |= prefix.children.all()
-
-        for pfx in prefixes.distinct():
-            # Update parent criteria:
-            # * This prefix contains the child prefix, has a parent that is the prefixes parent and is "In-VRF"
-            # * This prefix does not contain the child prefix
-            if pfx.vrf != prefix.vrf and not (prefix.vrf is None and is_container):
-                # Prefix is contained but not in-VRF
-                # print(f'{pfx} is no longer "in-VRF"')
-                pfx.parent = prefix.parent
-            elif pfx.prefix in prefix.prefix and pfx.parent != prefix and pfx.parent == prefix.parent:
-                # Prefix is in-scope
-                # print(f'{pfx} is in {prefix}')
-                pfx.parent = prefix
-            elif pfx.prefix not in prefix.prefix and pfx.parent == prefix:
-                # Prefix has fallen out of scope
-                # print(f'{pfx} is not in {prefix}')
-                pfx.parent = prefix.parent
-        rows = Prefix.objects.bulk_update(prefixes, ['parent'], batch_size=100)
-        print(rows)
-
-
 @receiver(post_save, sender=Prefix)
 def handle_prefix_saved(instance, created, **kwargs):
 
     # Prefix has changed (or new instance has been created)
     if created or instance.vrf_id != instance._vrf_id or instance.prefix != instance._prefix:
 
-        update_ipaddress_prefix(instance)
-        update_iprange_prefix(instance)
-        update_prefix_parents(instance, created=created)
         update_parents_children(instance)
         update_children_depth(instance)
 
@@ -156,13 +45,6 @@ def handle_prefix_saved(instance, created, **kwargs):
             update_children_depth(old_prefix)
 
 
-@receiver(pre_delete, sender=Prefix)
-def pre_handle_prefix_deleted(instance, **kwargs):
-    update_ipaddress_prefix(instance, delete=True)
-    update_iprange_prefix(instance, delete=True)
-    update_prefix_parents(instance, delete=True)
-
-
 @receiver(post_delete, sender=Prefix)
 def handle_prefix_deleted(instance, **kwargs):
 

netbox/ipam/tests/test_models.py

@@ -1,3 +1,5 @@
+from time import sleep
+
 from django.contrib.contenttypes.models import ContentType
 from django.core.exceptions import ValidationError
 from django.test import TestCase, override_settings
@@ -653,6 +655,120 @@ class TestPrefixHierarchy(TestCase):
         self.assertEqual(prefixes[3]._children, 0)
 
 
+class TestTriggers(TestCase):
+    """
+    Test the automatic updating of depth and child count in response to changes made within
+    the prefix hierarchy.
+    """
+
+    @classmethod
+    def setUpTestData(cls):
+
+        prefixes = (
+            # IPv4
+            Prefix(prefix='10.0.0.0/8'),
+            Prefix(prefix='10.0.0.0/16'),
+            Prefix(prefix='10.0.0.0/24'),
+            Prefix(prefix='192.168.0.0/16'),
+            # IPv6
+            Prefix(prefix='2001:db8::/32'),
+            Prefix(prefix='2001:db8::/40'),
+            Prefix(prefix='2001:db8::/48'),
+        )
+
+        for prefix in prefixes:
+            prefix.clean()
+            prefix.save()
+
+        vrfs = (
+            VRF(name='VRF A'),
+            VRF(name='VRF B'),
+        )
+
+        for prefix in prefixes:
+            prefix.clean()
+            prefix.save()
+
+        for vrf in vrfs:
+            vrf.clean()
+            vrf.save()
+
+    def test_current_hierarchy(self):
+        self.assertIsNone(Prefix.objects.get(prefix='10.0.0.0/8').parent)
+        self.assertIsNone(Prefix.objects.get(prefix='192.168.0.0/16').parent)
+        self.assertIsNone(Prefix.objects.get(prefix='2001:db8::/32').parent)
+
+        self.assertIsNotNone(Prefix.objects.get(prefix='10.0.0.0/16').parent)
+        self.assertIsNotNone(Prefix.objects.get(prefix='10.0.0.0/24').parent)
+
+        self.assertIsNotNone(Prefix.objects.get(prefix='2001:db8::/40').parent)
+        self.assertIsNotNone(Prefix.objects.get(prefix='2001:db8::/48').parent)
+
+    def test_basic_insert(self):
+        pfx = Prefix.objects.create(prefix='2001:db8::/44')
+        self.assertIsNotNone(Prefix.objects.get(prefix='2001:db8::/48').parent)
+        self.assertEqual(Prefix.objects.get(prefix='2001:db8::/48').parent, pfx)
+
+    def test_vrf_insert(self):
+        vrf = VRF.objects.get(name='VRF A')
+        pfx = Prefix.objects.create(prefix='2001:db8::/44', vrf=vrf)
+        parent = Prefix.objects.get(prefix='2001:db8::/40')
+        self.assertIsNotNone(Prefix.objects.get(prefix='2001:db8::/48').parent)
+        self.assertNotEqual(Prefix.objects.get(prefix='2001:db8::/48').parent, pfx)
+        self.assertEqual(Prefix.objects.get(prefix='2001:db8::/48').parent, parent)
+
+        prefixes = (
+            Prefix(prefix='10.2.0.0/16', vrf=vrf),
+            Prefix(prefix='10.2.0.0/24', vrf=vrf),
+        )
+
+        for prefix in prefixes:
+            prefix.clean()
+            prefix.save()
+
+        self.assertIsNone(Prefix.objects.get(pk=prefixes[0].pk).parent)
+        self.assertEqual(Prefix.objects.get(pk=prefixes[1].pk).parent, prefixes[0])
+
+        new_pfx = Prefix.objects.create(prefix='10.2.0.0/23', vrf=vrf)
+
+        self.assertIsNone(Prefix.objects.get(pk=prefixes[0].pk).parent)
+        self.assertEqual(new_pfx.parent, prefixes[0])
+        self.assertEqual(Prefix.objects.get(pk=prefixes[1].pk).parent, new_pfx)
+
+    def test_basic_delete(self):
+        prefixes = (
+            Prefix(prefix='10.2.0.0/16'),
+            Prefix(prefix='10.2.0.0/23'),
+            Prefix(prefix='10.2.0.0/24'),
+        )
+        for prefix in prefixes:
+            prefix.clean()
+            prefix.save()
+
+    def test_vrf_delete(self):
+        vrf = VRF.objects.get(name='VRF A')
+
+        prefixes = (
+            Prefix(prefix='10.2.0.0/16', vrf=vrf),
+            Prefix(prefix='10.2.0.0/23', vrf=vrf),
+            Prefix(prefix='10.2.0.0/24', vrf=vrf),
+        )
+
+        for prefix in prefixes:
+            prefix.clean()
+            prefix.save()
+
+        self.assertIsNone(prefixes[0].parent)
+        self.assertEqual(prefixes[1].parent, prefixes[0])
+        self.assertEqual(prefixes[2].parent, prefixes[1])
+
+        prefixes[1].delete()
+        prefixes[2].refresh_from_db()
+
+        self.assertIsNone(prefixes[0].parent)
+        self.assertEqual(prefixes[2].parent, prefixes[0])
+
+
 class TestIPAddress(TestCase):
     """
     Test the automatic updating of depth and child count in response to changes made within

netbox/ipam/triggers.py

@@ -0,0 +1,108 @@
+ipam_prefix_delete_adjust_prefix_parent = """
+-- Update Child Prefix's with Prefix's PARENT
+UPDATE ipam_prefix SET parent_id=OLD.parent_id WHERE parent_id=OLD.id;
+RETURN OLD;
+"""
+
+
+ipam_prefix_delete_adjust_ipaddress_prefix = """
+-- Update IP Address with prefix's PARENT
+UPDATE ipam_ipaddress SET prefix_id=OLD.parent_id WHERE prefix_id=OLD.id;
+RETURN OLD;
+"""
+
+
+ipam_prefix_delete_adjust_iprange_prefix = """
+-- Update IP Range with prefix's PARENT
+UPDATE ipam_iprange SET prefix_id=OLD.parent_id WHERE prefix_id=OLD.id;
+RETURN OLD;
+"""
+
+
+ipam_prefix_insert_adjust_prefix_parent = """
+UPDATE ipam_prefix
+SET parent_id=NEW.id 
+WHERE 
+    prefix << NEW.prefix
+    AND
+    (
+        (vrf_id = NEW.vrf_id OR (vrf_id IS NULL AND NEW.vrf_id IS NULL))
+        OR
+        (
+            NEW.vrf_id IS NULL
+            AND
+            NEW.status = 'container'
+            AND
+            NOT EXISTS(
+                SELECT 1 FROM ipam_prefix p WHERE p.prefix >> ipam_prefix.prefix AND p.vrf_id = ipam_prefix.vrf_id
+            )
+        )
+    )
+    AND id != NEW.id
+    AND NOT EXISTS (
+        SELECT 1 FROM ipam_prefix p
+        WHERE
+            p.prefix >> ipam_prefix.prefix
+            AND p.prefix << NEW.prefix
+            AND (
+                (p.vrf_id = ipam_prefix.vrf_id OR (p.vrf_id IS NULL AND ipam_prefix.vrf_id IS NULL))
+                OR
+                (p.vrf_id IS NULL AND p.status = 'container')
+            )
+            AND p.id != NEW.id
+    )
+;
+RETURN NEW;
+"""
+
+
+ipam_prefix_insert_adjust_ipaddress_prefix = """
+UPDATE ipam_prefix
+SET prefix_id=NEW.id
+WHERE
+    NEW.prefix >> ipaddress.address
+    AND
+    (
+        (NEW.vrf = ipaddress.vrf_id OR (NEW.vrf_id IS NULL and ipaddress.vrf_id IS NULL))
+        OR
+        (NEW.vrf_id IS NULL AND NEW.status = 'container')
+    )
+    AND (
+        ipaddress.prefix_id IS NULL
+        OR
+        EXISTS (
+            SELECT 1 from prefix p WHERE
+                p.id = ipaddress.prefix_id
+                AND NEW.prefix << p.prefix
+        )
+    )
+    AND 
+        -- Check to ensure current parent PREFIX is not in a VRF
+        NOT EXISTS (
+            SELECT 1 from prefix p WHERE (
+                p.id = ipaddress.prefix_id
+                AND
+                p.vrf_id IS NOT NULL
+                AND
+                ipaddress.vrf_id IS NOT NULL
+                AND
+                (
+                    NEW.vrf_id IS NULL AND NEW.status = 'container'
+                )
+            )
+        )
+    AND
+     NOT EXISTS (
+        SELECT 1 FROM prefix p
+        WHERE
+        p.prefix >> ipaddress.address
+        AND p.id != NEW.id
+        AND p.prefix << NEW.prefix
+        AND (
+            (p.vrf_id = ipaddress.vrf_id OR (p.vrf_id IS NULL AND ipaddress.vrf_id IS NULL))
+            OR
+            (p.vrf_id IS NULL AND NEW.status = 'container')
+        )
+     );
+RETURN NEW;
+"""

netbox/netbox/settings.py

@@ -425,6 +425,7 @@ INSTALLED_APPS = [
     'sorl.thumbnail',
     'taggit',
     'timezone_field',
+    'pgtrigger',
     'core',
     'account',
     'circuits',