From 541c6fe20fec9f46051ce2e00944d37255b52d39 Mon Sep 17 00:00:00 2001
From: Marcus Gartner <marcus@lob.com>
Date: Wed, 23 Oct 2019 15:25:20 -0700
Subject: [PATCH] Add query planning time

---
 pkg/flame/flame.go      | 17 +++++++++++++++--
 pkg/flame/flame_test.go | 17 ++++++++++++-----
 pkg/html/html.go        |  9 ++++++---
 pkg/plan/plan.go        |  4 ++--
 pkg/plan/plan_test.go   |  8 ++++----
 5 files changed, 39 insertions(+), 16 deletions(-)

diff --git a/pkg/flame/flame.go b/pkg/flame/flame.go
index 4803faa..0edc63a 100644
--- a/pkg/flame/flame.go
+++ b/pkg/flame/flame.go
@@ -11,14 +11,27 @@ type Flame struct {
 	Name     string  `json:"name"`
 	Value    float64 `json:"value"`
 	Detail   string  `json:"detail"`
+	Color    string  `json:"color"`
 	Children []Flame `json:"children"`
 }
 
 func New(p plan.Plan) Flame {
-	// TODO add planning time frame and total
 	// TODO handle CTE InitPlan
+	planningFlame := Flame{
+		Name:   "Query Planning",
+		Value:  p.PlanningTime,
+		Detail: "Time to generate the query plan",
+		Color:  "#00C05A",
+	}
+
+	executionFlame := convert(p.ExecutionTree)
 
-	return convert(p.Root)
+	return Flame{
+		Name:     "Total",
+		Value:    planningFlame.Value + executionFlame.Value,
+		Detail:   "This node includes planning and execution time",
+		Children: []Flame{planningFlame, executionFlame},
+	}
 }
 
 func convert(n plan.Node) Flame {
diff --git a/pkg/flame/flame_test.go b/pkg/flame/flame_test.go
index 34adad3..4c3043e 100644
--- a/pkg/flame/flame_test.go
+++ b/pkg/flame/flame_test.go
@@ -12,7 +12,8 @@ func TestNew(t *testing.T) {
 
 	t.Run("creates a new Flame from a Plan", func(t *testing.T) {
 		p := plan.Plan{
-			Root: plan.Node{
+			PlanningTime: 0.01,
+			ExecutionTree: plan.Node{
 				Method:    "Limit",
 				TotalTime: 0.123,
 				Children: []plan.Node{
@@ -27,11 +28,17 @@ func TestNew(t *testing.T) {
 
 		f := New(p)
 
-		assert.Equal(t, "Limit", f.Name)
-		assert.Equal(t, 0.123, f.Value)
+		assert.Equal(t, "Total", f.Name)
+		assert.Equal(t, 0.133, f.Value)
 
-		assert.Equal(t, "Seq Scan on bears", f.Children[0].Name)
-		assert.Equal(t, 0.022, f.Children[0].Value)
+		assert.Equal(t, "Query Planning", f.Children[0].Name)
+		assert.Equal(t, 0.01, f.Children[0].Value)
+
+		assert.Equal(t, "Limit", f.Children[1].Name)
+		assert.Equal(t, 0.123, f.Children[1].Value)
+
+		assert.Equal(t, "Seq Scan on bears", f.Children[1].Children[0].Name)
+		assert.Equal(t, 0.022, f.Children[1].Children[0].Value)
 	})
 
 }
diff --git a/pkg/html/html.go b/pkg/html/html.go
index 5d517a9..ce6c2de 100644
--- a/pkg/html/html.go
+++ b/pkg/html/html.go
@@ -90,7 +90,7 @@ const flameTemplate = `
 
     <script src="https://d3js.org/d3.v4.min.js" charset="utf-8"></script>
     <script type="text/javascript" src=https://cdnjs.cloudflare.com/ajax/libs/d3-tip/0.9.1/d3-tip.min.js></script>
-    <script type="text/javascript" src="https://cdn.jsdelivr.net/gh/spiermar/d3-flame-graph@2.0.3/dist/d3-flamegraph.min.js"></script>
+    <script type="text/javascript" src="https://cdn.jsdelivr.net/gh/spiermar/d3-flame-graph@2.1.3/dist/d3-flamegraph.min.js"></script>
 
     <script type="text/javascript">
     var flameGraph = d3.flamegraph()
@@ -99,10 +99,13 @@ const flameTemplate = `
       .transitionDuration(750)
       .minFrameSize(5)
       .transitionEase(d3.easeCubic)
-      .sort(true)
+      .sort(false)
       .title("")
       .differential(false)
-      .selfValue(false);
+      .selfValue(false)
+      .setColorMapper(function(d, originalColor) {
+        return d.data.color || originalColor;
+      });
 
     var tip = d3.tip()
       .direction("s")
diff --git a/pkg/plan/plan.go b/pkg/plan/plan.go
index ff2b312..98b30d6 100644
--- a/pkg/plan/plan.go
+++ b/pkg/plan/plan.go
@@ -7,8 +7,8 @@ import (
 )
 
 type Plan struct {
-	Root Node `json:"Plan"`
-	// TODO PlanningTime float64 `json:"Planning Time"`
+	PlanningTime  float64 `json:"Planning Time"`
+	ExecutionTree Node    `json:"Plan"`
 }
 
 type Node struct {
diff --git a/pkg/plan/plan_test.go b/pkg/plan/plan_test.go
index b4efb50..4eaf0e7 100644
--- a/pkg/plan/plan_test.go
+++ b/pkg/plan/plan_test.go
@@ -14,11 +14,11 @@ func TestNew(t *testing.T) {
 
 		_, p := New(input)
 
-		assert.Equal(t, "Nested Loop", p.Root.Method)
-		assert.Equal(t, "", p.Root.Table)
-		assert.Equal(t, 0.049, p.Root.TotalTime)
+		assert.Equal(t, "Nested Loop", p.ExecutionTree.Method)
+		assert.Equal(t, "", p.ExecutionTree.Table)
+		assert.Equal(t, 0.049, p.ExecutionTree.TotalTime)
 
-		child := p.Root.Children[0]
+		child := p.ExecutionTree.Children[0]
 
 		assert.Equal(t, "Hash Join", child.Method)
 		assert.Equal(t, "users", child.Table)