From 0d68d0c059f9bb162622861e7c0c482c8bc907ab Mon Sep 17 00:00:00 2001
From: Jeremy Stretch <stretch@packetlife.net>
Date: Wed, 5 Aug 2020 15:55:47 -0400
Subject: [PATCH] Replace annotate_depth() with annotate_tree()

---
 netbox/ipam/querysets.py | 44 ++++++++++++++--------------------------
 netbox/ipam/views.py     | 13 +++---------
 2 files changed, 18 insertions(+), 39 deletions(-)

diff --git a/netbox/ipam/querysets.py b/netbox/ipam/querysets.py
index 6d2dc6f33..102954f13 100644
--- a/netbox/ipam/querysets.py
+++ b/netbox/ipam/querysets.py
@@ -3,34 +3,20 @@ from utilities.querysets import RestrictedQuerySet
 
 class PrefixQuerySet(RestrictedQuerySet):
 
-    def annotate_depth(self, limit=None):
+    def annotate_tree(self):
         """
-        Iterate through a QuerySet of Prefixes and annotate the hierarchical level of each. While it would be preferable
-        to do this using .annotate() on the QuerySet to count the unique parents of each prefix, that approach introduces
-        performance issues at scale.
-
-        Because we're adding a non-field attribute to the model, annotation must be made *after* any QuerySet
-        modifications.
+        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).
         """
-        queryset = self
-        stack = []
-        for p in queryset:
-            try:
-                prev_p = stack[-1]
-            except IndexError:
-                prev_p = None
-            if prev_p is not None:
-                while (p.prefix not in prev_p.prefix) or p.prefix == prev_p.prefix:
-                    stack.pop()
-                    try:
-                        prev_p = stack[-1]
-                    except IndexError:
-                        prev_p = None
-                        break
-            if prev_p is not None:
-                prev_p.has_children = True
-            stack.append(p)
-            p.depth = len(stack) - 1
-        if limit is None:
-            return queryset
-        return list(filter(lambda p: p.depth <= limit, queryset))
+        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))',
+            }
+        )
diff --git a/netbox/ipam/views.py b/netbox/ipam/views.py
index ec067e705..52db68127 100644
--- a/netbox/ipam/views.py
+++ b/netbox/ipam/views.py
@@ -221,9 +221,7 @@ class AggregateView(ObjectView):
             'site', 'role'
         ).order_by(
             'prefix'
-        ).annotate_depth(
-            limit=0
-        )
+        ).annotate_tree()
 
         # Add available prefixes to the table if requested
         if request.GET.get('show_available', 'true') == 'true':
@@ -326,11 +324,6 @@ class PrefixListView(ObjectListView):
     table = tables.PrefixDetailTable
     template_name = 'ipam/prefix_list.html'
 
-    def alter_queryset(self, request):
-        # Show only top-level prefixes by default (unless searching)
-        limit = None if request.GET.get('expand') or request.GET.get('q') else 0
-        return self.queryset.annotate_depth(limit=limit)
-
 
 class PrefixView(ObjectView):
     queryset = Prefix.objects.prefetch_related('vrf', 'site__region', 'tenant__group', 'vlan__group', 'role')
@@ -353,7 +346,7 @@ class PrefixView(ObjectView):
             prefix__net_contains=str(prefix.prefix)
         ).prefetch_related(
             'site', 'role'
-        ).annotate_depth()
+        ).annotate_tree()
         parent_prefix_table = tables.PrefixTable(list(parent_prefixes), orderable=False)
         parent_prefix_table.exclude = ('vrf',)
 
@@ -386,7 +379,7 @@ class PrefixPrefixesView(ObjectView):
         # Child prefixes table
         child_prefixes = prefix.get_child_prefixes().restrict(request.user, 'view').prefetch_related(
             'site', 'vlan', 'role',
-        ).annotate_depth(limit=0)
+        ).annotate_tree()
 
         # Add available prefixes to the table if requested
         if child_prefixes and request.GET.get('show_available', 'true') == 'true':