Replace annotate_depth() with annotate_tree()

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))',
+            }
+        )