From 538984c6d2a806f2de4e1af4e5d5f510840c6929 Mon Sep 17 00:00:00 2001
From: Johannes Erwerle <jo@swagspace.org>
Date: Mon, 21 Feb 2022 12:26:12 +0100
Subject: [PATCH] Fixes 8553: Fix Provider network, ASN, and contact options
 missings from global search selector

---
 netbox/netbox/constants.py | 374 +++++++++++++++++++++----------------
 netbox/netbox/forms.py     |  49 ++---
 2 files changed, 226 insertions(+), 197 deletions(-)

diff --git a/netbox/netbox/constants.py b/netbox/netbox/constants.py
index e29da6617..45de4d5b2 100644
--- a/netbox/netbox/constants.py
+++ b/netbox/netbox/constants.py
@@ -1,4 +1,5 @@
 from collections import OrderedDict
+from typing import Dict
 
 from circuits.filtersets import CircuitFilterSet, ProviderFilterSet, ProviderNetworkFilterSet
 from circuits.models import Circuit, ProviderNetwork, Provider
@@ -26,169 +27,212 @@ from virtualization.models import Cluster, VirtualMachine
 from virtualization.tables import ClusterTable, VirtualMachineTable
 
 SEARCH_MAX_RESULTS = 15
-SEARCH_TYPES = OrderedDict((
-    # Circuits
-    ('provider', {
-        'queryset': Provider.objects.annotate(
-            count_circuits=count_related(Circuit, 'provider')
-        ),
-        'filterset': ProviderFilterSet,
-        'table': ProviderTable,
-        'url': 'circuits:provider_list',
-    }),
-    ('circuit', {
-        'queryset': Circuit.objects.prefetch_related(
-            'type', 'provider', 'tenant', 'terminations__site'
-        ),
-        'filterset': CircuitFilterSet,
-        'table': CircuitTable,
-        'url': 'circuits:circuit_list',
-    }),
-    ('providernetwork', {
-        'queryset': ProviderNetwork.objects.prefetch_related('provider'),
-        'filterset': ProviderNetworkFilterSet,
-        'table': ProviderNetworkTable,
-        'url': 'circuits:providernetwork_list',
-    }),
-    # DCIM
-    ('site', {
-        'queryset': Site.objects.prefetch_related('region', 'tenant'),
-        'filterset': SiteFilterSet,
-        'table': SiteTable,
-        'url': 'dcim:site_list',
-    }),
-    ('rack', {
-        'queryset': Rack.objects.prefetch_related('site', 'location', 'tenant', 'role'),
-        'filterset': RackFilterSet,
-        'table': RackTable,
-        'url': 'dcim:rack_list',
-    }),
-    ('rackreservation', {
-        'queryset': RackReservation.objects.prefetch_related('site', 'rack', 'user'),
-        'filterset': RackReservationFilterSet,
-        'table': RackReservationTable,
-        'url': 'dcim:rackreservation_list',
-    }),
-    ('location', {
-        'queryset': Location.objects.add_related_count(
-            Location.objects.add_related_count(
-                Location.objects.all(),
-                Device,
-                'location',
-                'device_count',
-                cumulative=True
+
+CIRCUIT_TYPES = OrderedDict(
+    (
+        ('provider', {
+            'queryset': Provider.objects.annotate(
+                count_circuits=count_related(Circuit, 'provider')
             ),
-            Rack,
-            'location',
-            'rack_count',
-            cumulative=True
-        ).prefetch_related('site'),
-        'filterset': LocationFilterSet,
-        'table': LocationTable,
-        'url': 'dcim:location_list',
-    }),
-    ('devicetype', {
-        'queryset': DeviceType.objects.prefetch_related('manufacturer').annotate(
-            instance_count=count_related(Device, 'device_type')
-        ),
-        'filterset': DeviceTypeFilterSet,
-        'table': DeviceTypeTable,
-        'url': 'dcim:devicetype_list',
-    }),
-    ('device', {
-        'queryset': Device.objects.prefetch_related(
-            'device_type__manufacturer', 'device_role', 'tenant', 'site', 'rack', 'primary_ip4', 'primary_ip6',
-        ),
-        'filterset': DeviceFilterSet,
-        'table': DeviceTable,
-        'url': 'dcim:device_list',
-    }),
-    ('virtualchassis', {
-        'queryset': VirtualChassis.objects.prefetch_related('master').annotate(
-            member_count=count_related(Device, 'virtual_chassis')
-        ),
-        'filterset': VirtualChassisFilterSet,
-        'table': VirtualChassisTable,
-        'url': 'dcim:virtualchassis_list',
-    }),
-    ('cable', {
-        'queryset': Cable.objects.all(),
-        'filterset': CableFilterSet,
-        'table': CableTable,
-        'url': 'dcim:cable_list',
-    }),
-    ('powerfeed', {
-        'queryset': PowerFeed.objects.all(),
-        'filterset': PowerFeedFilterSet,
-        'table': PowerFeedTable,
-        'url': 'dcim:powerfeed_list',
-    }),
-    # Virtualization
-    ('cluster', {
-        'queryset': Cluster.objects.prefetch_related('type', 'group').annotate(
-            device_count=count_related(Device, 'cluster'),
-            vm_count=count_related(VirtualMachine, 'cluster')
-        ),
-        'filterset': ClusterFilterSet,
-        'table': ClusterTable,
-        'url': 'virtualization:cluster_list',
-    }),
-    ('virtualmachine', {
-        'queryset': VirtualMachine.objects.prefetch_related(
-            'cluster', 'tenant', 'platform', 'primary_ip4', 'primary_ip6',
-        ),
-        'filterset': VirtualMachineFilterSet,
-        'table': VirtualMachineTable,
-        'url': 'virtualization:virtualmachine_list',
-    }),
-    # IPAM
-    ('vrf', {
-        'queryset': VRF.objects.prefetch_related('tenant'),
-        'filterset': VRFFilterSet,
-        'table': VRFTable,
-        'url': 'ipam:vrf_list',
-    }),
-    ('aggregate', {
-        'queryset': Aggregate.objects.prefetch_related('rir'),
-        'filterset': AggregateFilterSet,
-        'table': AggregateTable,
-        'url': 'ipam:aggregate_list',
-    }),
-    ('prefix', {
-        'queryset': Prefix.objects.prefetch_related('site', 'vrf__tenant', 'tenant', 'vlan', 'role'),
-        'filterset': PrefixFilterSet,
-        'table': PrefixTable,
-        'url': 'ipam:prefix_list',
-    }),
-    ('ipaddress', {
-        'queryset': IPAddress.objects.prefetch_related('vrf__tenant', 'tenant'),
-        'filterset': IPAddressFilterSet,
-        'table': IPAddressTable,
-        'url': 'ipam:ipaddress_list',
-    }),
-    ('vlan', {
-        'queryset': VLAN.objects.prefetch_related('site', 'group', 'tenant', 'role'),
-        'filterset': VLANFilterSet,
-        'table': VLANTable,
-        'url': 'ipam:vlan_list',
-    }),
-    ('asn', {
-        'queryset': ASN.objects.prefetch_related('rir', 'tenant'),
-        'filterset': ASNFilterSet,
-        'table': ASNTable,
-        'url': 'ipam:asn_list',
-    }),
-    # Tenancy
-    ('tenant', {
-        'queryset': Tenant.objects.prefetch_related('group'),
-        'filterset': TenantFilterSet,
-        'table': TenantTable,
-        'url': 'tenancy:tenant_list',
-    }),
-    ('contact', {
-        'queryset': Contact.objects.prefetch_related('group', 'assignments').annotate(assignment_count=count_related(ContactAssignment, 'contact')),
-        'filterset': ContactFilterSet,
-        'table': ContactTable,
-        'url': 'tenancy:contact_list',
-    }),
-))
+            'filterset': ProviderFilterSet,
+            'table': ProviderTable,
+            'url': 'circuits:provider_list',
+        }),
+        ('circuit', {
+            'queryset': Circuit.objects.prefetch_related(
+                'type', 'provider', 'tenant', 'terminations__site'
+            ),
+            'filterset': CircuitFilterSet,
+            'table': CircuitTable,
+            'url': 'circuits:circuit_list',
+        }),
+        ('providernetwork', {
+            'queryset': ProviderNetwork.objects.prefetch_related('provider'),
+            'filterset': ProviderNetworkFilterSet,
+            'table': ProviderNetworkTable,
+            'url': 'circuits:providernetwork_list',
+        }),
+    )
+)
+
+
+DCIM_TYPES = OrderedDict(
+    (
+        ('site', {
+            'queryset': Site.objects.prefetch_related('region', 'tenant'),
+            'filterset': SiteFilterSet,
+            'table': SiteTable,
+            'url': 'dcim:site_list',
+        }),
+        ('rack', {
+            'queryset': Rack.objects.prefetch_related('site', 'location', 'tenant', 'role'),
+            'filterset': RackFilterSet,
+            'table': RackTable,
+            'url': 'dcim:rack_list',
+        }),
+        ('rackreservation', {
+            'queryset': RackReservation.objects.prefetch_related('site', 'rack', 'user'),
+            'filterset': RackReservationFilterSet,
+            'table': RackReservationTable,
+            'url': 'dcim:rackreservation_list',
+        }),
+        ('location', {
+            'queryset': Location.objects.add_related_count(
+                Location.objects.add_related_count(
+                    Location.objects.all(),
+                    Device,
+                    'location',
+                    'device_count',
+                    cumulative=True
+                ),
+                Rack,
+                'location',
+                'rack_count',
+                cumulative=True
+            ).prefetch_related('site'),
+            'filterset': LocationFilterSet,
+            'table': LocationTable,
+            'url': 'dcim:location_list',
+        }),
+        ('devicetype', {
+            'queryset': DeviceType.objects.prefetch_related('manufacturer').annotate(
+                instance_count=count_related(Device, 'device_type')
+            ),
+            'filterset': DeviceTypeFilterSet,
+            'table': DeviceTypeTable,
+            'url': 'dcim:devicetype_list',
+        }),
+        ('device', {
+            'queryset': Device.objects.prefetch_related(
+                'device_type__manufacturer', 'device_role', 'tenant', 'site', 'rack', 'primary_ip4', 'primary_ip6',
+            ),
+            'filterset': DeviceFilterSet,
+            'table': DeviceTable,
+            'url': 'dcim:device_list',
+        }),
+        ('virtualchassis', {
+            'queryset': VirtualChassis.objects.prefetch_related('master').annotate(
+                member_count=count_related(Device, 'virtual_chassis')
+            ),
+            'filterset': VirtualChassisFilterSet,
+            'table': VirtualChassisTable,
+            'url': 'dcim:virtualchassis_list',
+        }),
+        ('cable', {
+            'queryset': Cable.objects.all(),
+            'filterset': CableFilterSet,
+            'table': CableTable,
+            'url': 'dcim:cable_list',
+        }),
+        ('powerfeed', {
+            'queryset': PowerFeed.objects.all(),
+            'filterset': PowerFeedFilterSet,
+            'table': PowerFeedTable,
+            'url': 'dcim:powerfeed_list',
+        }),
+    )
+)
+
+IPAM_TYPES = OrderedDict(
+    (
+        ('vrf', {
+            'queryset': VRF.objects.prefetch_related('tenant'),
+            'filterset': VRFFilterSet,
+            'table': VRFTable,
+            'url': 'ipam:vrf_list',
+        }),
+        ('aggregate', {
+            'queryset': Aggregate.objects.prefetch_related('rir'),
+            'filterset': AggregateFilterSet,
+            'table': AggregateTable,
+            'url': 'ipam:aggregate_list',
+        }),
+        ('prefix', {
+            'queryset': Prefix.objects.prefetch_related('site', 'vrf__tenant', 'tenant', 'vlan', 'role'),
+            'filterset': PrefixFilterSet,
+            'table': PrefixTable,
+            'url': 'ipam:prefix_list',
+        }),
+        ('ipaddress', {
+            'queryset': IPAddress.objects.prefetch_related('vrf__tenant', 'tenant'),
+            'filterset': IPAddressFilterSet,
+            'table': IPAddressTable,
+            'url': 'ipam:ipaddress_list',
+        }),
+        ('vlan', {
+            'queryset': VLAN.objects.prefetch_related('site', 'group', 'tenant', 'role'),
+            'filterset': VLANFilterSet,
+            'table': VLANTable,
+            'url': 'ipam:vlan_list',
+        }),
+        ('asn', {
+            'queryset': ASN.objects.prefetch_related('rir', 'tenant'),
+            'filterset': ASNFilterSet,
+            'table': ASNTable,
+            'url': 'ipam:asn_list',
+        }),
+    )
+)
+
+TENANCY_TYPES = OrderedDict(
+    (
+        ('tenant', {
+            'queryset': Tenant.objects.prefetch_related('group'),
+            'filterset': TenantFilterSet,
+            'table': TenantTable,
+            'url': 'tenancy:tenant_list',
+        }),
+        ('contact', {
+            'queryset': Contact.objects.prefetch_related('group', 'assignments').annotate(
+                assignment_count=count_related(ContactAssignment, 'contact')),
+            'filterset': ContactFilterSet,
+            'table': ContactTable,
+            'url': 'tenancy:contact_list',
+        }),
+    )
+)
+
+VIRTUALIZATION_TYPES = OrderedDict(
+    (
+        ('cluster', {
+            'queryset': Cluster.objects.prefetch_related('type', 'group').annotate(
+                device_count=count_related(Device, 'cluster'),
+                vm_count=count_related(VirtualMachine, 'cluster')
+            ),
+            'filterset': ClusterFilterSet,
+            'table': ClusterTable,
+            'url': 'virtualization:cluster_list',
+        }),
+        ('virtualmachine', {
+            'queryset': VirtualMachine.objects.prefetch_related(
+                'cluster', 'tenant', 'platform', 'primary_ip4', 'primary_ip6',
+            ),
+            'filterset': VirtualMachineFilterSet,
+            'table': VirtualMachineTable,
+            'url': 'virtualization:virtualmachine_list',
+        }),
+    )
+)
+
+SEARCH_TYPE_HIERARCHY = OrderedDict(
+    (
+        ("Circuits", CIRCUIT_TYPES),
+        ("DCIM", DCIM_TYPES),
+        ("IPAM", IPAM_TYPES),
+        ("Tenancy", TENANCY_TYPES),
+        ("Virtualization", VIRTUALIZATION_TYPES),
+    )
+)
+
+
+def build_search_types() -> Dict[str, Dict]:
+    result = dict()
+
+    for app_types in SEARCH_TYPE_HIERARCHY.values():
+        for name, items in app_types.items():
+            result[name] = items
+
+    return result
+
+
+SEARCH_TYPES = build_search_types()
diff --git a/netbox/netbox/forms.py b/netbox/netbox/forms.py
index b5d68c1fc..d220527fa 100644
--- a/netbox/netbox/forms.py
+++ b/netbox/netbox/forms.py
@@ -1,39 +1,24 @@
 from django import forms
 
 from utilities.forms import BootstrapMixin
+from netbox.constants import SEARCH_TYPE_HIERARCHY
 
-OBJ_TYPE_CHOICES = (
-    ('', 'All Objects'),
-    ('Circuits', (
-        ('provider', 'Providers'),
-        ('circuit', 'Circuits'),
-    )),
-    ('DCIM', (
-        ('site', 'Sites'),
-        ('rack', 'Racks'),
-        ('rackreservation', 'Rack reservations'),
-        ('location', 'Locations'),
-        ('devicetype', 'Device Types'),
-        ('device', 'Devices'),
-        ('virtualchassis', 'Virtual chassis'),
-        ('cable', 'Cables'),
-        ('powerfeed', 'Power feeds'),
-    )),
-    ('IPAM', (
-        ('vrf', 'VRFs'),
-        ('aggregate', 'Aggregates'),
-        ('prefix', 'Prefixes'),
-        ('ipaddress', 'IP Addresses'),
-        ('vlan', 'VLANs'),
-    )),
-    ('Tenancy', (
-        ('tenant', 'Tenants'),
-    )),
-    ('Virtualization', (
-        ('cluster', 'Clusters'),
-        ('virtualmachine', 'Virtual Machines'),
-    )),
-)
+
+def build_search_choices():
+    result = list()
+    result.append(('', 'All Objects'))
+    for category, items in SEARCH_TYPE_HIERARCHY.items():
+        subcategories = list()
+        for slug, obj in items.items():
+            name = obj['queryset'].model._meta.verbose_name_plural
+            name = name[0].upper() + name[1:]
+            subcategories.append((slug, name))
+        result.append((category, tuple(subcategories)))
+
+    return tuple(result)
+
+
+OBJ_TYPE_CHOICES = build_search_choices()
 
 
 def build_options():