Merge pull request #6488 from netbox-community/6087-prefix-depth-chil…

…dren

Closes #6087: Cache prefix depth & children count

netbox/ipam/filtersets.py

@@ -209,6 +209,12 @@ class PrefixFilterSet(PrimaryModelFilterSet, TenancyFilterSet):
  method='search_contains',
  label='Prefixes which contain this prefix or IP',
  )
+ depth = MultiValueNumberFilter(
+ field_name='_depth'
+ )
+ children = MultiValueNumberFilter(
+ field_name='_children'
+ )
  mask_length = django_filters.NumberFilter(
  field_name='prefix',
  lookup_expr='net_mask_length'

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/0047_prefix_depth_children.py

@@ -0,0 +1,21 @@
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+ dependencies = [
+ ('ipam', '0046_set_vlangroup_scope_types'),
+ ]
+
+ operations = [
+ migrations.AddField(
+ model_name='prefix',
+ name='_children',
+ field=models.PositiveBigIntegerField(default=0, editable=False),
+ ),
+ migrations.AddField(
+ model_name='prefix',
+ name='_depth',
+ field=models.PositiveSmallIntegerField(default=0, editable=False),
+ ),
+ ]

netbox/ipam/migrations/0048_prefix_populate_depth_children.py

@@ -0,0 +1,46 @@
+from django.db import migrations
+
+from ipam.utils import rebuild_prefixes
+
+
+def push_to_stack(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,
+ })
+
+
+def populate_prefix_hierarchy(apps, schema_editor):
+ """
+ Populate _depth and _children attrs for all Prefixes.
+ """
+ Prefix = apps.get_model('ipam', 'Prefix')
+ VRF = apps.get_model('ipam', 'VRF')
+
+ total_count = Prefix.objects.count()
+ print(f'\nUpdating {total_count} prefixes...')
+
+ # Rebuild the global table
+ rebuild_prefixes(None)
+
+ # Iterate through all VRFs, rebuilding each
+ for vrf in VRF.objects.all():
+ rebuild_prefixes(vrf)
+
+
+class Migration(migrations.Migration):
+
+ dependencies = [
+ ('ipam', '0047_prefix_depth_children'),
+ ]
+
+ operations = [
+ migrations.RunPython(
+ code=populate_prefix_hierarchy,
+ reverse_code=migrations.RunPython.noop
+ ),
+ ]

netbox/ipam/models/ip.py

@@ -293,6 +293,16 @@ class Prefix(PrimaryModel):
  blank=True
  )
 
+ # Cached depth & child counts
+ _depth = models.PositiveSmallIntegerField(
+ default=0,
+ editable=False
+ )
+ _children = models.PositiveBigIntegerField(
+ default=0,
+ editable=False
+ )
+
  objects = PrefixQuerySet.as_manager()
 
  csv_headers = [
@@ -306,6 +316,13 @@ class Meta:
  ordering = (F('vrf').asc(nulls_first=True), 'prefix', 'pk') # (vrf, prefix) may be non-unique
  verbose_name_plural = 'prefixes'
 
+ def __init__(self, *args, **kwargs):
+ super().__init__(*args, **kwargs)
+
+ # Cache the original prefix and VRF so we can check if they have changed on post_save
+ self._prefix = self.prefix
+ self._vrf = self.vrf
+
  def __str__(self):
  return str(self.prefix)
 
@@ -373,6 +390,14 @@ def family(self):
  return self.prefix.version
  return None
 
+ @property
+ def depth(self):
+ return self._depth
+
+ @property
+ def children(self):
+ return self._children
+
  def _set_prefix_length(self, value):
  """
  Expose the IPNetwork object's prefixlen attribute on the parent model so that it can be manipulated directly,
@@ -385,6 +410,26 @@ def _set_prefix_length(self, value):
  def get_status_class(self):
  return PrefixStatusChoices.CSS_CLASSES.get(self.status)
 
+ def get_parents(self, include_self=False):
+ """
+ Return all containing Prefixes in the hierarchy.
+ """
+ lookup = 'net_contains_or_equals' if include_self else 'net_contains'
+ return Prefix.objects.filter(**{
+ 'vrf': self.vrf,
+ f'prefix__{lookup}': self.prefix
+ })
+
+ def get_children(self, include_self=False):
+ """
+ Return all covered Prefixes in the hierarchy.
+ """
+ lookup = 'net_contained_or_equal' if include_self else 'net_contained'
+ return Prefix.objects.filter(**{
+ 'vrf': self.vrf,
+ f'prefix__{lookup}': self.prefix
+ })
+
  def get_duplicates(self):
  return Prefix.objects.filter(vrf=self.vrf, prefix=str(self.prefix)).exclude(pk=self.pk)
 

netbox/ipam/querysets.py

@@ -1,27 +1,32 @@
 from django.contrib.contenttypes.models import ContentType
 from django.db.models import Q
+from django.db.models.expressions import RawSQL
 
 from utilities.querysets import RestrictedQuerySet
 
 
 class PrefixQuerySet(RestrictedQuerySet):
 
- def annotate_tree(self):
+ def annotate_hierarchy(self):
  """
- Annotate the number of parent and child prefixes for each Prefix. Raw SQL is needed for these subqueries
- because we need to cast NULL VRF values to integers for comparison. (NULL != NULL).
+ Annotate the depth and number of child prefixes for each Prefix. Cast null VRF values to zero for
+ comparison. (NULL != NULL).
  """
- return self.extra(
- select={
- 'parents': 'SELECT COUNT(U0."prefix") AS "c" '
- 'FROM "ipam_prefix" U0 '
- 'WHERE (U0."prefix" >> "ipam_prefix"."prefix" '
- 'AND COALESCE(U0."vrf_id", 0) = COALESCE("ipam_prefix"."vrf_id", 0))',
- 'children': 'SELECT COUNT(U1."prefix") AS "c" '
- 'FROM "ipam_prefix" U1 '
- 'WHERE (U1."prefix" << "ipam_prefix"."prefix" '
- 'AND COALESCE(U1."vrf_id", 0) = COALESCE("ipam_prefix"."vrf_id", 0))',
- }
+ return self.annotate(
+ hierarchy_depth=RawSQL(
+ 'SELECT COUNT(DISTINCT U0."prefix") AS "c" '
+ 'FROM "ipam_prefix" U0 '
+ 'WHERE (U0."prefix" >> "ipam_prefix"."prefix" '
+ 'AND COALESCE(U0."vrf_id", 0) = COALESCE("ipam_prefix"."vrf_id", 0))',
+ ()
+ ),
+ hierarchy_children=RawSQL(
+ 'SELECT COUNT(U1."prefix") AS "c" '
+ 'FROM "ipam_prefix" U1 '
+ 'WHERE (U1."prefix" << "ipam_prefix"."prefix" '
+ 'AND COALESCE(U1."vrf_id", 0) = COALESCE("ipam_prefix"."vrf_id", 0))',
+ ()
+ )
  )
 
 

netbox/ipam/signals.py

@@ -1,9 +1,52 @@
-from django.db.models.signals import pre_delete
+from django.db.models.signals import post_delete, post_save, pre_delete
 from django.dispatch import receiver
 
 from dcim.models import Device
 from virtualization.models import VirtualMachine
-from .models import IPAddress
+from .models import IPAddress, Prefix
+
+
+def update_parents_children(prefix):
+ """
+ Update depth on prefix & containing prefixes
+ """
+ parents = prefix.get_parents(include_self=True).annotate_hierarchy()
+ for parent in parents:
+ parent._children = parent.hierarchy_children
+ Prefix.objects.bulk_update(parents, ['_children'], batch_size=100)
+
+
+def update_children_depth(prefix):
+ """
+ Update children count on prefix & contained prefixes
+ """
+ children = prefix.get_children(include_self=True).annotate_hierarchy()
+ for child in children:
+ child._depth = child.hierarchy_depth
+ Prefix.objects.bulk_update(children, ['_depth'], batch_size=100)
+
+
+@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 != instance._vrf or instance.prefix != instance._prefix:
+
+ update_parents_children(instance)
+ update_children_depth(instance)
+
+ # If this is not a new prefix, clean up parent/children of previous prefix
+ if not created:
+ old_prefix = Prefix(vrf=instance._vrf, prefix=instance._prefix)
+ update_parents_children(old_prefix)
+ update_children_depth(old_prefix)
+
+
+@receiver(post_delete, sender=Prefix)
+def handle_prefix_deleted(instance, **kwargs):
+
+ update_parents_children(instance)
+ update_children_depth(instance)
 
 
 @receiver(pre_delete, sender=IPAddress)

netbox/ipam/tables.py

@@ -15,7 +15,7 @@
 
 PREFIX_LINK = """
 {% load helpers %}
-{% for i in record.parents|as_range %}
+{% for i in record.depth|as_range %}
  <i class="mdi mdi-circle-small"></i>
 {% endfor %}
 <a href="{% if record.pk %}{% url 'ipam:prefix' pk=record.pk %}{% else %}{% url 'ipam:prefix_add' %}?prefix={{ record }}{% if object.vrf %}&vrf={{ object.vrf.pk }}{% endif %}{% if object.site %}&site={{ object.site.pk }}{% endif %}{% if object.tenant %}&tenant_group={{ object.tenant.group.pk }}&tenant={{ object.tenant.pk }}{% endif %}{% endif %}">{{ record.prefix }}</a>
@@ -262,6 +262,19 @@ class PrefixTable(BaseTable):
  template_code=PREFIX_LINK,
  attrs={'td': {'class': 'text-nowrap'}}
  )
+ depth = tables.Column(
+ accessor=Accessor('_depth'),
+ verbose_name='Depth'
+ )
+ children = LinkedCountColumn(
+ accessor=Accessor('_children'),
+ viewname='ipam:prefix_list',
+ url_params={
+ 'vrf_id': 'vrf_id',
+ 'within': 'prefix',
+ },
+ verbose_name='Children'
+ )
  status = ChoiceFieldColumn(
  default=AVAILABLE_LABEL
  )
@@ -287,7 +300,8 @@ class PrefixTable(BaseTable):
  class Meta(BaseTable.Meta):
  model = Prefix
  fields = (
- 'pk', 'prefix', 'status', 'children', 'vrf', 'tenant', 'site', 'vlan', 'role', 'is_pool', 'description',
+ 'pk', 'prefix', 'status', 'depth', 'children', 'vrf', 'tenant', 'site', 'vlan', 'role', 'is_pool',
+ 'description',
  )
  default_columns = ('pk', 'prefix', 'status', 'vrf', 'tenant', 'site', 'vlan', 'role', 'description')
  row_attrs = {

netbox/ipam/tests/test_filtersets.py

@@ -400,7 +400,8 @@ def setUpTestData(cls):
  Prefix(prefix='10.0.0.0/16'),
  Prefix(prefix='2001:db8::/32'),
  )
- Prefix.objects.bulk_create(prefixes)
+ for prefix in prefixes:
+ prefix.save()
 
  def test_family(self):
  params = {'family': '6'}
@@ -431,6 +432,18 @@ def test_contains(self):
  params = {'contains': '2001:db8:0:1::/64'}
  self.assertEqual(self.filterset(params, self.queryset).qs.count(), 2)
 
+ def test_depth(self):
+ params = {'depth': '0'}
+ self.assertEqual(self.filterset(params, self.queryset).qs.count(), 8)
+ params = {'depth__gt': '0'}
+ self.assertEqual(self.filterset(params, self.queryset).qs.count(), 2)
+
+ def test_children(self):
+ params = {'children': '0'}
+ self.assertEqual(self.filterset(params, self.queryset).qs.count(), 8)
+ params = {'children__gt': '0'}
+ self.assertEqual(self.filterset(params, self.queryset).qs.count(), 2)
+
  def test_mask_length(self):
  params = {'mask_length': '24'}
  self.assertEqual(self.filterset(params, self.queryset).qs.count(), 4)