Introduce rebuild_prefixes management command

netbox/ipam/management/commands/rebuild_prefixes.py

@@ -0,0 +1,27 @@
+from django.core.management.base import BaseCommand
+
+from ipam.models import Prefix, VRF
+from ipam.utils import rebuild_prefixes
+
+
+class Command(BaseCommand):
+    help = "Rebuild the prefix hierarchy (depth and children counts)"
+
+    def handle(self, *model_names, **options):
+        self.stdout.write(f'Rebuilding {Prefix.objects.count()} prefixes...')
+
+        # Reset existing counts
+        Prefix.objects.update(_depth=0, _children=0)
+
+        # Rebuild the global table
+        global_count = Prefix.objects.filter(vrf__isnull=True).count()
+        self.stdout.write(f'Global: {global_count} prefixes...')
+        rebuild_prefixes(None)
+
+        # Rebuild each VRF
+        for vrf in VRF.objects.all():
+            vrf_count = Prefix.objects.filter(vrf=vrf).count()
+            self.stdout.write(f'VRF {vrf}: {vrf_count} prefixes...')
+            rebuild_prefixes(vrf)
+
+        self.stdout.write(self.style.SUCCESS('Finished.'))

netbox/ipam/migrations/0048_prefix_populate_depth_children.py

@@ -1,5 +1,7 @@
 from django.db import migrations
 
+from ipam.utils import rebuild_prefixes
+
 
 def push_to_stack(stack, prefix):
     # Increment child count on parent nodes
@@ -22,54 +24,12 @@ def populate_prefix_hierarchy(apps, schema_editor):
     total_count = Prefix.objects.count()
     print(f'\nUpdating {total_count} prefixes...')
 
-    # Iterate through all VRFs and the global table
-    vrfs = [None] + list(VRF.objects.values_list('pk', flat=True))
-    for vrf in vrfs:
-
-        stack = []
-        update_queue = []
-
-        # Iterate through all Prefixes in the VRF, growing and shrinking the stack as we go
-        prefixes = Prefix.objects.filter(vrf=vrf).values('pk', 'prefix')
-        for i, p in enumerate(prefixes):
-
-            # Grow the stack if this is a child of the most recent prefix
-            if not stack or p['prefix'] in stack[-1]['prefix']:
-                push_to_stack(stack, p)
-
-            # 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 p['prefix'] not in stack[-1]['prefix'] and p['prefix'] != stack[-1]['prefix']:
-                    node = stack.pop()
-                    update_queue.append(
-                        Prefix(
-                            pk=node['pk'],
-                            _depth=len(stack),
-                            _children=node['children']
-                        )
-                    )
-                push_to_stack(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 = []
-                    print(f'  [{i}/{total_count}]')
-
-        # 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']
-                )
-            )
+    # Rebuild the global table
+    rebuild_prefixes(None)
 
-        # Final flush of any remaining Prefixes
-        Prefix.objects.bulk_update(update_queue, ['_depth', '_children'])
+    # Iterate through all VRFs, rebuilding each
+    for vrf in VRF.objects.all():
+        rebuild_prefixes(vrf)
 
 
 class Migration(migrations.Migration):

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'])