Просмотр исходного кода

Introduce unrestricted() method on RestrictedQuerySet

Jeremy Stretch 5 лет назад
Родитель
Сommit
ffb43a8534
1 измененных файлов с 19 добавлено и 9 удалено
  1. 19 9
      netbox/utilities/querysets.py

+ 19 - 9
netbox/utilities/querysets.py

@@ -21,20 +21,22 @@ class RestrictedQuerySet(QuerySet):
     def __init__(self, *args, **kwargs):
     def __init__(self, *args, **kwargs):
         super().__init__(*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):
     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 = 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):
     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 = super()._clone()
-        c.is_restricted = self.is_restricted
+        c.allow_evaluation = self.allow_evaluation
 
 
         return c
         return c
 
 
@@ -46,6 +48,14 @@ class RestrictedQuerySet(QuerySet):
         self._check_restriction()
         self._check_restriction()
         return super().count()
         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):
     def restrict(self, user, action):
         """
         """
         Filter the QuerySet to return only objects on which the specified user has been granted the specified
         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)
                     attrs |= Q(**perm_attrs)
             qs = self.filter(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
         return qs