Filter VLANs and VLANGroups by site or cluster site for VM

netbox/ipam/querysets.py

@@ -81,31 +81,34 @@ class VLANQuerySet(RestrictedQuerySet):
 
         # Find all relevant VLANGroups
         q = Q()
+        site = vm.site or vm.cluster.site
         if vm.cluster:
-            if vm.cluster.site:
-                if vm.cluster.site.region:
-                    q |= Q(
-                        scope_type=ContentType.objects.get_by_natural_key('dcim', 'region'),
-                        scope_id__in=vm.cluster.site.region.get_ancestors(include_self=True)
-                    )
-                if vm.cluster.site.group:
-                    q |= Q(
-                        scope_type=ContentType.objects.get_by_natural_key('dcim', 'sitegroup'),
-                        scope_id__in=vm.cluster.site.group.get_ancestors(include_self=True)
-                    )
-                q |= Q(
-                    scope_type=ContentType.objects.get_by_natural_key('dcim', 'site'),
-                    scope_id=vm.cluster.site_id
-                )
+            # Add VLANGroups scoped to the assigned cluster (or its group)
+            q |= Q(
+                scope_type=ContentType.objects.get_by_natural_key('virtualization', 'cluster'),
+                scope_id=vm.cluster_id
+            )
             if vm.cluster.group:
                 q |= Q(
                     scope_type=ContentType.objects.get_by_natural_key('virtualization', 'clustergroup'),
                     scope_id=vm.cluster.group_id
                 )
+        if site:
+            # Add VLANGroups scoped to the assigned site (or its group or region)
             q |= Q(
-                scope_type=ContentType.objects.get_by_natural_key('virtualization', 'cluster'),
-                scope_id=vm.cluster_id
+                scope_type=ContentType.objects.get_by_natural_key('dcim', 'site'),
+                scope_id=site.pk
             )
+            if site.region:
+                q |= Q(
+                    scope_type=ContentType.objects.get_by_natural_key('dcim', 'region'),
+                    scope_id__in=site.region.get_ancestors(include_self=True)
+                )
+            if site.group:
+                q |= Q(
+                    scope_type=ContentType.objects.get_by_natural_key('dcim', 'sitegroup'),
+                    scope_id__in=site.group.get_ancestors(include_self=True)
+                )
         vlan_groups = VLANGroup.objects.filter(q)
 
         # Return all applicable VLANs
@@ -114,7 +117,7 @@ class VLANQuerySet(RestrictedQuerySet):
             Q(group__scope_id__isnull=True, site__isnull=True) |  # Global group VLANs
             Q(group__isnull=True, site__isnull=True)  # Global VLANs
         )
-        if vm.cluster and vm.cluster.site:
-            q |= Q(site=vm.cluster.site)
+        if site:
+            q |= Q(site=site)
 
         return self.filter(q)