diff --git a/netbox/ipam/management/__init__.py b/netbox/ipam/management/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/netbox/ipam/management/commands/__init__.py b/netbox/ipam/management/commands/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/netbox/ipam/management/commands/rebuild_prefixes.py b/netbox/ipam/management/commands/rebuild_prefixes.py new file mode 100644 index 000000000..5d614b834 --- /dev/null +++ b/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.')) diff --git a/netbox/ipam/migrations/0048_prefix_populate_depth_children.py b/netbox/ipam/migrations/0048_prefix_populate_depth_children.py index b265e7f6f..dc5070656 100644 --- a/netbox/ipam/migrations/0048_prefix_populate_depth_children.py +++ b/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: + # Rebuild the global table + rebuild_prefixes(None) - 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'] - ) - ) - - # 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): diff --git a/netbox/ipam/utils.py b/netbox/ipam/utils.py index 0414a01e0..953a4a468 100644 --- a/netbox/ipam/utils.py +++ b/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'])