policy(minor): improve loading of policy subclasses

passbook/policy/engine.py

@@ -8,8 +8,8 @@ from django.http import HttpRequest
 from structlog import get_logger
 
 from passbook.core.models import Policy, User
-from passbook.policy.struct import PolicyRequest, PolicyResult
-from passbook.policy.task import PolicyTask
+from passbook.policy.struct import PolicyRequest
+from passbook.policy.process import PolicyProcess
 
 LOGGER = get_logger()
 
@@ -23,7 +23,7 @@ class PolicyEngine:
     __request: HttpRequest
     __user: User
 
-    __proc_list: List[Tuple[Connection, PolicyTask]] = []
+    __proc_list: List[Tuple[Connection, PolicyProcess]] = []
 
     def __init__(self, policies, user: User = None, request: HttpRequest = None):
         self.policies = policies
@@ -40,6 +40,13 @@ class PolicyEngine:
         self.__request = request
         return self
 
+    def _select_subclasses(self) -> List[Policy]:
+        """Make sure all Policies are their respective classes"""
+        return Policy.objects \
+            .filter(pk__in=[x.pk for x in self.policies]) \
+            .select_subclasses() \
+            .order_by('order')
+
     def build(self) -> 'PolicyEngine':
         """Build task group"""
         if not self.__user:
@@ -47,18 +54,15 @@ class PolicyEngine:
         cached_policies = []
         request = PolicyRequest(self.__user)
         request.http_request = self.__request
-        for policy in self.policies:
+        for policy in self._select_subclasses():
             cached_policy = cache.get(_cache_key(policy, self.__user), None)
             if cached_policy:
                 LOGGER.debug("Taking result from cache", policy=policy.pk.hex)
                 cached_policies.append(cached_policy)
             else:
-                LOGGER.debug("Looking up real class of policy...")
-                # TODO: Rewrite this to lookup all policies at once
-                policy = Policy.objects.get_subclass(pk=policy.pk)
                 LOGGER.debug("Evaluating policy", policy=policy.pk.hex)
                 our_end, task_end = Pipe(False)
-                task = PolicyTask()
+                task = PolicyProcess()
                 task.ret = task_end
                 task.request = request
                 task.policy = policy

passbook/policy/process.py

@@ -8,13 +8,13 @@ from passbook.core.models import Policy
 from passbook.policy.exceptions import PolicyException
 from passbook.policy.struct import PolicyRequest, PolicyResult
 
-LOGGER = get_logger(__name__)
+LOGGER = get_logger()
 
 
 def _cache_key(policy, user):
     return "policy_%s#%s" % (policy.uuid, user.pk)
 
-class PolicyTask(Process):
+class PolicyProcess(Process):
     """Evaluate a single policy within a seprate process"""
 
     ret: Connection
@@ -23,8 +23,8 @@ class PolicyTask(Process):
 
     def run(self):
         """Task wrapper to run policy checking"""
-        LOGGER.debug("Running policy `%s`#%s for user %s...", self.policy.name,
-                     self.policy.pk.hex, self.request.user)
+        LOGGER.debug("Running policy", policy=self.policy,
+                     user=self.request.user, process="PolicyProcess")
         try:
             policy_result = self.policy.passes(self.request)
         except PolicyException as exc:
@@ -33,7 +33,8 @@ class PolicyTask(Process):
         # Invert result if policy.negate is set
         if self.policy.negate:
             policy_result = not policy_result
-        LOGGER.debug("Policy %r#%s got %s", self.policy.name, self.policy.pk.hex, policy_result)
+        LOGGER.debug("Got result", policy=self.policy, result=policy_result,
+                     process="PolicyProcess")
         # cache_key = _cache_key(self.policy, self.request.user)
         # cache.set(cache_key, (self.policy.action, policy_result, message))
         # LOGGER.debug("Cached entry as %s", cache_key)