From a20ccfee7ef28adc820da50d3bb3301183c1c08c Mon Sep 17 00:00:00 2001
From: Jeremy Stretch <jstretch@netboxlabs.com>
Date: Mon, 13 May 2024 13:26:14 -0400
Subject: [PATCH] Update queryset resolution methods for compatibility with
 Django 5.0

---
 netbox/utilities/fields.py    | 37 +++++++++++++++++++++++------------
 netbox/utilities/querysets.py | 10 +++++-----
 2 files changed, 29 insertions(+), 18 deletions(-)

diff --git a/netbox/utilities/fields.py b/netbox/utilities/fields.py
index 3f7ce6f54..2640f6886 100644
--- a/netbox/utilities/fields.py
+++ b/netbox/utilities/fields.py
@@ -70,14 +70,24 @@ class RestrictedGenericForeignKey(GenericForeignKey):
     #  1. Capture restrict_params from RestrictedPrefetch (hack)
     #  2. If restrict_params is set, call restrict() on the queryset for
     #     the related model
-    def get_prefetch_queryset(self, instances, queryset=None):
+    def get_prefetch_querysets(self, instances, querysets=None):
         restrict_params = {}
+        custom_queryset_dict = {}
 
         # Compensate for the hack in RestrictedPrefetch
-        if type(queryset) is dict:
-            restrict_params = queryset
-        elif queryset is not None:
-            raise ValueError(_("Custom queryset can't be used for this lookup."))
+        if type(querysets) is dict:
+            restrict_params = querysets
+
+        elif querysets is not None:
+            for queryset in querysets:
+                ct_id = self.get_content_type(
+                    model=queryset.query.model, using=queryset.db
+                ).pk
+                if ct_id in custom_queryset_dict:
+                    raise ValueError(
+                        "Only one queryset is allowed for each content type."
+                    )
+                custom_queryset_dict[ct_id] = queryset
 
         # For efficiency, group the instances by content type and then do one
         # query per model
@@ -100,15 +110,16 @@ class RestrictedGenericForeignKey(GenericForeignKey):
 
         ret_val = []
         for ct_id, fkeys in fk_dict.items():
-            instance = instance_dict[ct_id]
-            ct = self.get_content_type(id=ct_id, using=instance._state.db)
-            if restrict_params:
-                # Override the default behavior to call restrict() on each model's queryset
-                qs = ct.model_class().objects.filter(pk__in=fkeys).restrict(**restrict_params)
-                ret_val.extend(qs)
+            if ct_id in custom_queryset_dict:
+                # Return values from the custom queryset, if provided.
+                ret_val.extend(custom_queryset_dict[ct_id].filter(pk__in=fkeys))
             else:
-                # Default behavior
-                ret_val.extend(ct.get_all_objects_for_this_type(pk__in=fkeys))
+                instance = instance_dict[ct_id]
+                ct = self.get_content_type(id=ct_id, using=instance._state.db)
+                qs = ct.model_class().objects.filter(pk__in=fkeys)
+                if restrict_params:
+                    qs = qs.restrict(**restrict_params)
+                ret_val.extend(qs)
 
         # For doing the join in Python, we have to match both the FK val and the
         # content type, so we use a callable that returns a (fk, class) pair.
diff --git a/netbox/utilities/querysets.py b/netbox/utilities/querysets.py
index 50917dd0f..0a9389b1f 100644
--- a/netbox/utilities/querysets.py
+++ b/netbox/utilities/querysets.py
@@ -20,14 +20,14 @@ class RestrictedPrefetch(Prefetch):
 
         super().__init__(lookup, queryset=queryset, to_attr=to_attr)
 
-    def get_current_queryset(self, level):
+    def get_current_querysets(self, level):
         params = {
             'user': self.restrict_user,
             'action': self.restrict_action,
         }
 
-        if qs := super().get_current_queryset(level):
-            return qs.restrict(**params)
+        if querysets := super().get_current_querysets(level):
+            return [qs.restrict(**params) for qs in querysets]
 
         # Bit of a hack. If no queryset is defined, pass through the dict of restrict()
         # kwargs to be handled by the field. This is necessary e.g. for GenericForeignKey
@@ -49,11 +49,11 @@ class RestrictedQuerySet(QuerySet):
         permission_required = get_permission_for_model(self.model, action)
 
         # Bypass restriction for superusers and exempt views
-        if user.is_superuser or permission_is_exempt(permission_required):
+        if user and user.is_superuser or permission_is_exempt(permission_required):
             qs = self
 
         # User is anonymous or has not been granted the requisite permission
-        elif not user.is_authenticated or permission_required not in user.get_all_permissions():
+        elif user is None or not user.is_authenticated or permission_required not in user.get_all_permissions():
             qs = self.none()
 
         # Filter the queryset to include only objects with allowed attributes