Update queryset resolution methods for compatibility with Django 5.0

This commit is contained in:
Jeremy Stretch 2024-05-13 13:26:14 -04:00
parent c7850b586b
commit a20ccfee7e
2 changed files with 29 additions and 18 deletions

View File

@ -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.

View File

@ -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