diff --git a/netbox/utilities/fields.py b/netbox/utilities/fields.py index 87eb8f2fe..b2bc4d2cd 100644 --- a/netbox/utilities/fields.py +++ b/netbox/utilities/fields.py @@ -78,10 +78,18 @@ class NaturalOrderingField(models.CharField): class RestrictedGenericForeignKey(GenericForeignKey): - # Replicated from GenericForeignKey + # Replicated largely from GenericForeignKey. Changes include: + # 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): + restrict_params = {} + # Compensate for the hack in RestrictedPrefetch - restrict_params = queryset if type(queryset) is dict else {} + if type(queryset) is dict: + restrict_params = queryset + elif queryset is not None: + raise ValueError("Custom queryset can't be used for this lookup.") # For efficiency, group the instances by content type and then do one # query per model diff --git a/netbox/utilities/querysets.py b/netbox/utilities/querysets.py index a7f65d46c..0e5f1cd5c 100644 --- a/netbox/utilities/querysets.py +++ b/netbox/utilities/querysets.py @@ -21,13 +21,12 @@ class RestrictedPrefetch(Prefetch): 'action': self.restrict_action, } - qs = super().get_current_queryset(level) - if qs: - return qs.filter(**params) + if qs := super().get_current_queryset(level): + return qs.restrict(**params) # Bit of a hack. If no queryset is defined, pass through the dict of restrict() - # kwargs to be called later. This is necessary e.g. for GenericForeignKey fields, - # which do not permit setting a queryset on a Prefetch object. + # kwargs to be handled by the field. This is necessary e.g. for GenericForeignKey + # fields, which do not permit setting a queryset on a Prefetch object. return params