1
0
mirror of https://github.com/netbox-community/netbox.git synced 2024-05-10 07:54:54 +00:00

Introduce unrestricted() method on RestrictedQuerySet

This commit is contained in:
Jeremy Stretch
2020-06-16 12:20:21 -04:00
parent ce5fd7955f
commit ffb43a8534

View File

@ -21,20 +21,22 @@ class RestrictedQuerySet(QuerySet):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
# Initialize the is_restricted flag to False. This indicates that the QuerySet has not yet been restricted.
self.is_restricted = False
# Initialize the allow_evaluation flag to False. This indicates that the QuerySet has not yet been restricted.
self.allow_evaluation = False
def _check_restriction(self):
# Raise a warning if the QuerySet is evaluated without first calling restrict().
if not getattr(self, 'is_restricted', False):
# Raise a warning if the QuerySet is evaluated without first calling restrict() or unrestricted().
if not getattr(self, 'allow_evaluation', False):
logger = logging.getLogger('netbox.RestrictedQuerySet')
logger.warning(f'Evaluation of RestrictedQuerySet prior to calling restrict(): {self.model}')
logger.warning(
f'Evaluation of RestrictedQuerySet prior to calling restrict() or unrestricted(): {self.model}'
)
def _clone(self):
# Persist the is_restricted flag when cloning the QuerySet.
# Persist the allow_evaluation flag when cloning the QuerySet.
c = super()._clone()
c.is_restricted = self.is_restricted
c.allow_evaluation = self.allow_evaluation
return c
@ -46,6 +48,14 @@ class RestrictedQuerySet(QuerySet):
self._check_restriction()
return super().count()
def unrestricted(self):
"""
Bypass restriction for the QuerySet. This is necessary in cases where we are not interacting with the objects
directly (e.g. when filtering by related object).
"""
self.allow_evaluation = True
return self
def restrict(self, user, action):
"""
Filter the QuerySet to return only objects on which the specified user has been granted the specified
@ -75,7 +85,7 @@ class RestrictedQuerySet(QuerySet):
attrs |= Q(**perm_attrs)
qs = self.filter(attrs)
# Mark the QuerySet as having been restricted
qs.is_restricted = True
# Allow QuerySet evaluation
qs.allow_evaluation = True
return qs