Introduce rebuild_prefixes management command

netbox/ipam/utils.py

@@ -91,3 +91,63 @@ def add_available_vlans(vlan_group, vlans):
     vlans.sort(key=lambda v: v.vid if type(v) == VLAN else v['vid'])
 
     return vlans
+
+
+def rebuild_prefixes(vrf):
+    """
+    Rebuild the prefix hierarchy for all prefixes in the specified VRF (or global table).
+    """
+    def contains(parent, child):
+        return child in parent and child != parent
+
+    def push_to_stack(prefix):
+        # Increment child count on parent nodes
+        for n in stack:
+            n['children'] += 1
+        stack.append({
+            'pk': prefix['pk'],
+            'prefix': prefix['prefix'],
+            'children': 0,
+        })
+
+    stack = []
+    update_queue = []
+    prefixes = Prefix.objects.filter(vrf=vrf).values('pk', 'prefix')
+
+    # Iterate through all Prefixes in the VRF, growing and shrinking the stack as we go
+    for i, p in enumerate(prefixes):
+
+        # Grow the stack if this is a child of the most recent prefix
+        if not stack or contains(stack[-1]['prefix'], p['prefix']):
+            push_to_stack(p)
+
+        # Handle duplicate prefixes
+        elif stack[-1]['prefix'] == p['prefix']:
+            update_queue.append(
+                Prefix(pk=p['pk'], _depth=len(stack) - 1, _children=stack[-1]['children'])
+            )
+
+        # If this is a sibling or parent of the most recent prefix, pop nodes from the
+        # stack until we reach a parent prefix (or the root)
+        else:
+            while stack and not contains(stack[-1]['prefix'], p['prefix']):
+                node = stack.pop()
+                update_queue.append(
+                    Prefix(pk=node['pk'], _depth=len(stack), _children=node['children'])
+                )
+            push_to_stack(p)
+
+        # Flush the update queue once it reaches 100 Prefixes
+        if len(update_queue) >= 100:
+            Prefix.objects.bulk_update(update_queue, ['_depth', '_children'])
+            update_queue = []
+
+    # Clear out any prefixes remaining in the stack
+    while stack:
+        node = stack.pop()
+        update_queue.append(
+            Prefix(pk=node['pk'], _depth=len(stack), _children=node['children'])
+        )
+
+    # Final flush of any remaining Prefixes
+    Prefix.objects.bulk_update(update_queue, ['_depth', '_children'])