Signal optimizations

netbox/ipam/signals.py

@@ -29,119 +29,112 @@ def update_children_depth(prefix):
     Prefix.objects.bulk_update(children, ['_depth'], batch_size=100)
 
 
-def update_object_prefix(prefix, delete=False, parent_model=Prefix, child_model=IPAddress):
-    if delete:
-        # Get all possible addresses
-        addresses = child_model.objects.filter(prefix=prefix)
-        prefix = parent_model.objects.filter(
-            prefix__net_contains_or_equals=prefix.prefix,
+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
-        ).exclude(pk=prefix.pk).last()
+        )
 
-        for address in addresses:
-            # Set contained addresses to the containing prefix if it exists
+    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
-    else:
-        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
+        # 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):
-    update_object_prefix(prefix, delete, child_model=IPAddress)
+    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):
-    update_object_prefix(prefix, delete, child_model=IPRange)
+    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):
+def update_prefix_parents(prefix, delete=False, created=False):
     if delete:
-        # Get all possible addresses
+        # Set prefix to prefix parent
         prefixes = prefix.children.all()
+        for address in prefixes:
+            address.parent = prefix.parent
 
-        for pfx in prefixes:
-            # Set contained addresses to the containing prefix if it exists
-            pfx.parent = prefix.parent
+        # Run a bulk update
+        Prefix.objects.bulk_update(prefixes, ['parent'], batch_size=100)
     else:
-        # Get all possible addresses
-        prefixes = prefix.children.all() | Prefix.objects.filter(
-            Q(
-                parent=prefix.parent,
-                vrf=prefix.vrf,
-                prefix__net_contained=str(prefix.prefix)
-            ) | Q(
+        # 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,
-                status=PrefixStatusChoices.STATUS_CONTAINER,
                 prefix__net_contained=str(prefix.prefix),
             )
-        )
 
-        if isinstance(prefix.prefix, str):
-            prefix.prefix = IPNetwork(prefix.prefix)
-        for pfx in prefixes:
-            if isinstance(pfx.prefix, str):
-                pfx.prefix = IPNetwork(pfx.prefix)
-
-            if pfx.parent == prefix and pfx.prefix.ip not in prefix.prefix:
-                # Find new parents for orphaned prefixes
-                parent = Prefix.objects.exclude(pk=pfx.pk).filter(
-                    Q(
-                        vrf=pfx.vrf,
-                        prefix__net_contains=str(pfx.prefix)
-                    ) | Q(
-                        vrf=None,
-                        status=PrefixStatusChoices.STATUS_CONTAINER,
-                        prefix__net_contains=str(pfx.prefix),
-                    )
-                ).last()
-                # Set contained addresses to the containing prefix if it exists
-                pfx.parent = parent
-            elif pfx.parent == prefix and pfx.vrf != prefix.vrf:
-                # Find new parents for orphaned prefixes
-                parent = Prefix.objects.exclude(pk=pfx.pk).filter(
-                    Q(
-                        vrf=pfx.vrf,
-                        prefix__net_contains=str(pfx.prefix)
-                    ) | Q(
-                        vrf=None,
-                        status=PrefixStatusChoices.STATUS_CONTAINER,
-                        prefix__net_contains=str(pfx.prefix),
-                    )
-                ).last()
-                # Set contained addresses to the containing prefix if it exists
-                pfx.parent = parent
-            elif pfx.parent != prefix and pfx.vrf == prefix.vrf and pfx.prefix in prefix.prefix:
-                # Set the parent to the 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
-            else:
-                # No-OP as the prefix does not require modification
-                pass
-
-    # Update the prefixes
-    Prefix.objects.bulk_update(prefixes, ['parent'], batch_size=100)
+            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)
@@ -152,7 +145,7 @@ def handle_prefix_saved(instance, created, **kwargs):
 
         update_ipaddress_prefix(instance)
         update_iprange_prefix(instance)
-        update_prefix_parents(instance)
+        update_prefix_parents(instance, created=created)
         update_parents_children(instance)
         update_children_depth(instance)
 
@@ -165,8 +158,8 @@ def handle_prefix_saved(instance, created, **kwargs):
 
 @receiver(pre_delete, sender=Prefix)
 def pre_handle_prefix_deleted(instance, **kwargs):
-    update_ipaddress_prefix(instance, True)
-    update_iprange_prefix(instance, True)
+    update_ipaddress_prefix(instance, delete=True)
+    update_iprange_prefix(instance, delete=True)
     update_prefix_parents(instance, delete=True)
 
 
@@ -175,9 +168,6 @@ def handle_prefix_deleted(instance, **kwargs):
 
     update_parents_children(instance)
     update_children_depth(instance)
-    update_ipaddress_prefix(instance, delete=True)
-    update_iprange_prefix(instance, delete=True)
-    update_prefix_parents(instance, delete=True)
 
 
 @receiver(pre_delete, sender=IPAddress)