Add query graph visuals

Signed-off-by: Florent Poinsard <[email protected]>
planetscale · Nov 27, 2024 · 5ad4205 · 5ad4205
1 parent 367ffb5
commit 5ad4205
Show file tree

Hide file tree

Showing 6 changed files with 236 additions and 42 deletions.
diff --git a/go/cmd/summarize.go b/go/cmd/summarize.go
@@ -24,6 +24,7 @@ import (
 
 func summarizeCmd() *cobra.Command {
 	var hotMetric string
+	var showGraph bool
 
 	cmd := &cobra.Command{
 		Use:     "summarize old_file.json [new_file.json]",
@@ -32,11 +33,12 @@ func summarizeCmd() *cobra.Command {
 		Example: "vt summarize old.json new.json",
 		Args:    cobra.RangeArgs(1, 2),
 		Run: func(_ *cobra.Command, args []string) {
-			summarize.Run(args, hotMetric)
+			summarize.Run(args, hotMetric, showGraph)
 		},
 	}
 
 	cmd.Flags().StringVar(&hotMetric, "hot-metric", "total-time", "Metric to determine hot queries (options: usage-count, total-rows-examined, avg-rows-examined, avg-time, total-time)")
+	cmd.Flags().BoolVar(&showGraph, "graph", false, "Show the query graph in the browser")
 
 	return cmd
 }
diff --git a/go/summarize/force-graph.go b/go/summarize/force-graph.go
@@ -0,0 +1,190 @@
+/*
+Copyright 2024 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package summarize
+
+import (
+	"encoding/json"
+	"fmt"
+	"html/template"
+	"net/http"
+)
+
+type (
+	node struct {
+		ID string `json:"id"`
+	}
+
+	link struct {
+		Source    string `json:"source"`
+		SourceIdx int    `json:"source_idx"`
+		Target    string `json:"target"`
+		TargetIdx int    `json:"target_idx"`
+		Value     int    `json:"value"`
+	}
+
+	forceGraphData struct {
+		Nodes []node `json:"nodes"`
+		Links []link `json:"links"`
+	}
+)
+
+func createForceGraphData(s *Summary) forceGraphData {
+	var data forceGraphData
+
+	idxTableNode := make(map[string]int)
+	for _, table := range s.tables {
+		if len(table.JoinPredicates) > 0 {
+			data.Nodes = append(data.Nodes, node{ID: table.Table})
+			idxTableNode[table.Table] = len(data.Nodes) - 1
+		}
+	}
+	for _, join := range s.joins {
+		data.Links = append(data.Links, link{
+			Source:    join.Tbl1,
+			SourceIdx: idxTableNode[join.Tbl1],
+			Target:    join.Tbl2,
+			TargetIdx: idxTableNode[join.Tbl2],
+			Value:     join.Occurrences,
+		})
+	}
+	return data
+}
+
+func renderQueryGraph(s *Summary) {
+	data := createForceGraphData(s)
+
+	// Start the HTTP server
+	http.HandleFunc("/", func(w http.ResponseWriter, _ *http.Request) {
+		serveIndex(w, data)
+	})
+
+	port := "1010"
+	fmt.Printf("Server started at http://localhost:%s\n", port)
+	// nolint: gosec,nolintlint // this is all ran locally so no need to care about vulnerabilities around timeouts
+	if err := http.ListenAndServe(":"+port, nil); err != nil {
+		exit(err.Error())
+	}
+}
+
+// Function to dynamically generate and serve index.html
+func serveIndex(w http.ResponseWriter, data forceGraphData) {
+	dataBytes, err := json.Marshal(data)
+	if err != nil {
+		exit(err.Error())
+	}
+
+	tmpl, err := template.New("index").Parse(templateHTML)
+	if err != nil {
+		http.Error(w, "Failed to parse template", http.StatusInternalServerError)
+		return
+	}
+
+	// nolint: gosec,nolintlint // this is all ran locally so no need to care about vulnerabilities around escaping
+	if err := tmpl.Execute(w, template.JS(dataBytes)); err != nil {
+		http.Error(w, "Failed to execute template", http.StatusInternalServerError)
+		return
+	}
+}
+
+const templateHTML = `<head>
+    <style> body { margin: 0; } </style>
+    <script src="//unpkg.com/force-graph"></script>
+</head>
+<body>
+    <div id="graph"></div>
+    <script>
+		let data = {{.}};
+        data.links.forEach(link => {
+            const a = data.nodes[link.source_idx];
+            const b = data.nodes[link.target_idx];
+            !a.neighbors && (a.neighbors = []);
+            !b.neighbors && (b.neighbors = []);
+            a.neighbors.push(b);
+            b.neighbors.push(a);
+
+            !a.links && (a.links = []);
+            !b.links && (b.links = []);
+            a.links.push(link);
+            b.links.push(link);
+        });
+
+        const highlightNodes = new Set();
+        const highlightLinks = new Set();
+        let hoverNode = null;
+
+        const Graph = ForceGraph()
+        (document.getElementById('graph'))
+            .graphData(data)
+            .nodeId('id')
+            .nodeLabel('id')
+            .linkWidth('value')
+            .linkLabel('value')
+            .onLinkHover(link => {
+                highlightNodes.clear();
+                highlightLinks.clear();
+
+                if (link) {
+                    highlightLinks.add(link);
+                    highlightNodes.add(link.source);
+                    highlightNodes.add(link.target);
+                }
+            })
+            .autoPauseRedraw(false) // keep redrawing after engine has stopped
+            .linkWidth(link => highlightLinks.has(link) ? 5 : 1)
+            .linkDirectionalParticles(4)
+            .linkDirectionalParticleWidth(link => highlightLinks.has(link) ? 4 : 0)
+            .nodeCanvasObject((node, ctx, globalScale) => {
+                const label = node.id;
+                const fontSize = 12/globalScale;
+                ctx.font = fontSize+'px Sans-Serif';
+                const textWidth = ctx.measureText(label).width;
+                const bckgDimensions = [textWidth, fontSize].map(n => n + fontSize * 1); // some padding
+
+                ctx.fillStyle = 'rgb(0,14,71)';
+                ctx.fillRect(node.x - bckgDimensions[0] / 2, node.y - bckgDimensions[1] / 2, ...bckgDimensions);
+
+                ctx.textAlign = 'center';
+                ctx.textBaseline = 'middle';
+                ctx.fillStyle = 'rgb(255,255,255)';
+                if (highlightNodes.has(node)) {
+                    ctx.fillStyle = node === hoverNode ? 'red' : 'orange';
+                }
+                ctx.fillText(label, node.x, node.y);
+
+                node.__bckgDimensions = bckgDimensions;
+
+                if (highlightNodes.has(node)) {
+                    ctx.beginPath();
+                    ctx.fill();
+                }
+            })
+            .onNodeHover(node => {
+                highlightNodes.clear();
+                highlightLinks.clear();
+                if (node) {
+                    highlightNodes.add(node);
+                    node.neighbors.forEach(neighbor => highlightNodes.add(neighbor));
+                    node.links.forEach(link => highlightLinks.add(link));
+                }
+
+                hoverNode = node || null;
+            })
+            .d3Force('link').strength(link => {
+                return data.links[link.index].value * 0.2
+            });
+    </script>
+</body>`
diff --git a/go/summarize/markdown.go b/go/summarize/markdown.go
@@ -18,14 +18,12 @@ package summarize
 
 import (
 	"fmt"
-	"maps"
 	"slices"
 	"sort"
 	"strconv"
 	"strings"
 
 	humanize "github.com/dustin/go-humanize"
-	"vitess.io/vitess/go/vt/vtgate/planbuilder/operators"
 
 	"github.com/vitessio/vt/go/keys"
 	"github.com/vitessio/vt/go/markdown"
@@ -148,50 +146,16 @@ func renderColumnUsageTable(md *markdown.MarkDown, summary *TableSummary) {
 }
 
 func renderTablesJoined(md *markdown.MarkDown, summary *Summary) {
-	type joinDetails struct {
-		Tbl1, Tbl2  string
-		Occurrences int
-		predicates  []operators.JoinPredicate
-	}
-
-	var joins []joinDetails
-	for tables, predicates := range summary.queryGraph {
-		occurrences := 0
-		for _, count := range predicates {
-			occurrences += count
-		}
-		joinPredicates := slices.Collect(maps.Keys(predicates))
-		sort.Slice(joinPredicates, func(i, j int) bool {
-			return joinPredicates[i].String() < joinPredicates[j].String()
-		})
-		joins = append(joins, joinDetails{
-			Tbl1:        tables.Tbl1,
-			Tbl2:        tables.Tbl2,
-			Occurrences: occurrences,
-			predicates:  joinPredicates,
-		})
-	}
-
-	if len(joins) == 0 {
+	if len(summary.joins) == 0 {
 		return
 	}
 
 	if len(summary.queryGraph) > 0 {
 		md.PrintHeader("Tables Joined", 2)
 	}
 
-	sort.Slice(joins, func(i, j int) bool {
-		if joins[i].Occurrences != joins[j].Occurrences {
-			return joins[i].Occurrences > joins[j].Occurrences
-		}
-		if joins[i].Tbl1 != joins[j].Tbl1 {
-			return joins[i].Tbl1 < joins[j].Tbl1
-		}
-		return joins[i].Tbl2 < joins[j].Tbl2
-	})
-
 	md.Println("```")
-	for _, join := range joins {
+	for _, join := range summary.joins {
 		md.Printf("%s ↔ %s (Occurrences: %d)\n", join.Tbl1, join.Tbl2, join.Occurrences)
 		for i, pred := range join.predicates {
 			var s string

diff --git a/go/summarize/summarize-keys.go b/go/summarize/summarize-keys.go
@@ -19,6 +19,7 @@ package summarize
 import (
 	"fmt"
 	"iter"
+	"maps"
 	"slices"
 	"sort"
 
@@ -53,6 +54,12 @@ type (
 		Tbl1, Tbl2 string
 	}
 
+	joinDetails struct {
+		Tbl1, Tbl2  string
+		Occurrences int
+		predicates  []operators.JoinPredicate
+	}
+
 	queryGraph map[graphKey]map[operators.JoinPredicate]int
 )
 
@@ -236,6 +243,32 @@ func summarizeKeysQueries(summary *Summary, queries *keys.Output) {
 			summary.queryGraph.AddJoinPredicate(key, pred)
 		}
 	}
+
+	for tables, predicates := range summary.queryGraph {
+		occurrences := 0
+		for _, count := range predicates {
+			occurrences += count
+		}
+		joinPredicates := slices.Collect(maps.Keys(predicates))
+		sort.Slice(joinPredicates, func(i, j int) bool {
+			return joinPredicates[i].String() < joinPredicates[j].String()
+		})
+		summary.joins = append(summary.joins, joinDetails{
+			Tbl1:        tables.Tbl1,
+			Tbl2:        tables.Tbl2,
+			Occurrences: occurrences,
+			predicates:  joinPredicates,
+		})
+	}
+	sort.Slice(summary.joins, func(i, j int) bool {
+		if summary.joins[i].Occurrences != summary.joins[j].Occurrences {
+			return summary.joins[i].Occurrences > summary.joins[j].Occurrences
+		}
+		if summary.joins[i].Tbl1 != summary.joins[j].Tbl1 {
+			return summary.joins[i].Tbl1 < summary.joins[j].Tbl1
+		}
+		return summary.joins[i].Tbl2 < summary.joins[j].Tbl2
+	})
 }
 
 func checkQueryForHotness(hotQueries *[]keys.QueryAnalysisResult, query keys.QueryAnalysisResult, metricReader getMetric) {

diff --git a/go/summarize/summarize.go b/go/summarize/summarize.go
@@ -36,7 +36,7 @@ type (
 
 type summaryWorker = func(s *Summary) error
 
-func Run(files []string, hotMetric string) {
+func Run(files []string, hotMetric string, showGraph bool) {
 	var traces []traceSummary
 	var workers []summaryWorker
 
@@ -61,7 +61,10 @@ func Run(files []string, hotMetric string) {
 
 	traceCount := len(traces)
 	if traceCount <= 0 {
-		printSummary(hotMetric, workers)
+		s := printSummary(hotMetric, workers)
+		if showGraph {
+			renderQueryGraph(s)
+		}
 		return
 	}
 
@@ -74,7 +77,7 @@ func Run(files []string, hotMetric string) {
 	}
 }
 
-func printSummary(hotMetric string, workers []summaryWorker) {
+func printSummary(hotMetric string, workers []summaryWorker) *Summary {
 	s := NewSummary(hotMetric)
 	for _, worker := range workers {
 		err := worker(s)
@@ -83,6 +86,7 @@ func printSummary(hotMetric string, workers []summaryWorker) {
 		}
 	}
 	s.PrintMarkdown(os.Stdout, time.Now())
+	return s
 }
 
 func checkTraceConditions(traces []traceSummary, workers []summaryWorker, hotMetric string) {

diff --git a/go/summarize/summary.go b/go/summarize/summary.go
@@ -36,6 +36,7 @@ type (
 		hotQueryFn    getMetric
 		analyzedFiles []string
 		queryGraph    queryGraph
+		joins         []joinDetails
 		hasRowCount   bool
 	}