From 47de8c5963e67cf580581e0265a44a273730b543 Mon Sep 17 00:00:00 2001
From: AlanCoding <arominge@redhat.com>
Date: Fri, 28 Oct 2016 11:20:25 -0400
Subject: [PATCH] keep queries for building WF DAG to a minimum

---
 awx/main/scheduler/dag_workflow.py            |  3 +-
 .../tests/functional/models/test_workflow.py  | 28 +++++++++++++++++++
 2 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/awx/main/scheduler/dag_workflow.py b/awx/main/scheduler/dag_workflow.py
index 6d2b349365..0ea002bd1b 100644
--- a/awx/main/scheduler/dag_workflow.py
+++ b/awx/main/scheduler/dag_workflow.py
@@ -9,7 +9,8 @@ class WorkflowDAG(SimpleDAG):
             self._init_graph(workflow_job)
 
     def _init_graph(self, workflow_job):
-        workflow_nodes = workflow_job.workflow_job_nodes.all()
+        node_qs = workflow_job.workflow_job_nodes
+        workflow_nodes = node_qs.prefetch_related('success_nodes', 'failure_nodes', 'always_nodes').all()
         for workflow_node in workflow_nodes:
             self.add_node(workflow_node)
 
diff --git a/awx/main/tests/functional/models/test_workflow.py b/awx/main/tests/functional/models/test_workflow.py
index cc61c34ed9..c2131c1f90 100644
--- a/awx/main/tests/functional/models/test_workflow.py
+++ b/awx/main/tests/functional/models/test_workflow.py
@@ -6,6 +6,34 @@ import pytest
 from awx.main.models.workflow import WorkflowJob, WorkflowJobNode, WorkflowJobTemplateNode
 from awx.main.models.jobs import Job
 from awx.main.models.projects import ProjectUpdate
+from awx.main.scheduler.dag_workflow import WorkflowDAG
+
+# Django
+from django.test import TransactionTestCase
+
+
+@pytest.mark.django_db
+class TestWorkflowDAGFunctional(TransactionTestCase):
+
+    def workflow_job(self):
+        wfj = WorkflowJob.objects.create()
+        nodes = [WorkflowJobNode.objects.create(workflow_job=wfj) for i in range(0, 5)]
+        nodes[0].success_nodes.add(nodes[1])
+        nodes[1].success_nodes.add(nodes[2])
+        nodes[0].failure_nodes.add(nodes[3])
+        nodes[3].failure_nodes.add(nodes[4])
+        return wfj
+
+    def test_build_WFJT_dag(self):
+        '''
+        Test that building the graph uses 4 queries
+         1 to get the nodes
+         3 to get the related success, failure, and always connections
+        '''
+        dag = WorkflowDAG()
+        wfj = self.workflow_job()
+        with self.assertNumQueries(4):
+            dag._init_graph(wfj)
 
 @pytest.mark.django_db
 class TestWorkflowJob: