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) # 1. Capture restrict_params from RestrictedPrefetch (hack)
# 2. If restrict_params is set, call restrict() on the queryset for # 2. If restrict_params is set, call restrict() on the queryset for
# the related model # the related model
def get_prefetch_queryset(self, instances, queryset=None): def get_prefetch_querysets(self, instances, querysets=None):
restrict_params = {} restrict_params = {}
custom_queryset_dict = {}
# Compensate for the hack in RestrictedPrefetch # Compensate for the hack in RestrictedPrefetch
if type(queryset) is dict: if type(querysets) is dict:
restrict_params = queryset restrict_params = querysets
elif queryset is not None:
raise ValueError(_("Custom queryset can't be used for this lookup.")) 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 # For efficiency, group the instances by content type and then do one
# query per model # query per model
@ -100,15 +110,16 @@ class RestrictedGenericForeignKey(GenericForeignKey):
ret_val = [] ret_val = []
for ct_id, fkeys in fk_dict.items(): for ct_id, fkeys in fk_dict.items():
instance = instance_dict[ct_id] if ct_id in custom_queryset_dict:
ct = self.get_content_type(id=ct_id, using=instance._state.db) # Return values from the custom queryset, if provided.
if restrict_params: ret_val.extend(custom_queryset_dict[ct_id].filter(pk__in=fkeys))
# 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)
else: else:
# Default behavior instance = instance_dict[ct_id]
ret_val.extend(ct.get_all_objects_for_this_type(pk__in=fkeys)) 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 # 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. # 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) super().__init__(lookup, queryset=queryset, to_attr=to_attr)
def get_current_queryset(self, level): def get_current_querysets(self, level):
params = { params = {
'user': self.restrict_user, 'user': self.restrict_user,
'action': self.restrict_action, 'action': self.restrict_action,
} }
if qs := super().get_current_queryset(level): if querysets := super().get_current_querysets(level):
return qs.restrict(**params) return [qs.restrict(**params) for qs in querysets]
# Bit of a hack. If no queryset is defined, pass through the dict of restrict() # 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 # 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) permission_required = get_permission_for_model(self.model, action)
# Bypass restriction for superusers and exempt views # 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 qs = self
# User is anonymous or has not been granted the requisite permission # 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() qs = self.none()
# Filter the queryset to include only objects with allowed attributes # Filter the queryset to include only objects with allowed attributes