From 87b14e876189d5023679b15a38f95d7e10114bb7 Mon Sep 17 00:00:00 2001
From: Jens Langhammer <jens.langhammer@beryju.org>
Date: Mon, 7 Sep 2020 11:27:02 +0200
Subject: [PATCH] flows/planner: optimise db queries during plan building

---
 passbook/flows/models.py  |  4 ++--
 passbook/flows/planner.py | 19 ++++++++-----------
 2 files changed, 10 insertions(+), 13 deletions(-)

diff --git a/passbook/flows/models.py b/passbook/flows/models.py
index 75bf63a48..87c60b193 100644
--- a/passbook/flows/models.py
+++ b/passbook/flows/models.py
@@ -11,7 +11,7 @@ from rest_framework.serializers import BaseSerializer
 from structlog import get_logger
 
 from passbook.core.types import UIUserSettings
-from passbook.lib.models import SerializerModel
+from passbook.lib.models import InheritanceForeignKey, SerializerModel
 from passbook.policies.models import PolicyBindingModel
 
 if TYPE_CHECKING:
@@ -150,7 +150,7 @@ class FlowStageBinding(SerializerModel, PolicyBindingModel):
     fsb_uuid = models.UUIDField(primary_key=True, editable=False, default=uuid4)
 
     target = models.ForeignKey("Flow", on_delete=models.CASCADE)
-    stage = models.ForeignKey(Stage, on_delete=models.CASCADE)
+    stage = InheritanceForeignKey(Stage, on_delete=models.CASCADE)
 
     re_evaluate_policies = models.BooleanField(
         default=False,
diff --git a/passbook/flows/planner.py b/passbook/flows/planner.py
index 8cd2e1be9..248686a15 100644
--- a/passbook/flows/planner.py
+++ b/passbook/flows/planner.py
@@ -156,25 +156,22 @@ class FlowPlanner:
             if default_context:
                 plan.context = default_context
             # Check Flow policies
-            for stage in (
-                self.flow.stages.order_by("flowstagebinding__order")
-                .select_subclasses()
-                .select_related()
-            ):
-                binding: FlowStageBinding = stage.flowstagebinding_set.get(
-                    target__pk=self.flow.pk
-                )
+            for binding in FlowStageBinding.objects.filter(
+                target__pk=self.flow.pk
+            ).order_by("order"):
                 engine = PolicyEngine(binding, user, request)
                 engine.request.context = plan.context
                 engine.build()
                 if engine.passing:
-                    LOGGER.debug("f(plan): Stage passing", stage=stage, flow=self.flow)
-                    plan.stages.append(stage)
+                    LOGGER.debug(
+                        "f(plan): Stage passing", stage=binding.stage, flow=self.flow
+                    )
+                    plan.stages.append(binding.stage)
                     marker = StageMarker()
                     if binding.re_evaluate_policies:
                         LOGGER.debug(
                             "f(plan): Stage has re-evaluate marker",
-                            stage=stage,
+                            stage=binding.stage,
                             flow=self.flow,
                         )
                         marker = ReevaluateMarker(binding=binding, user=user)