Skip to content

Commit 8b0663d

Browse files
committed
feat(tree-builder): add utils for building and printing the tree
1 parent 0ab4973 commit 8b0663d

File tree

2 files changed

+82
-8
lines changed

2 files changed

+82
-8
lines changed

Diff for: internal/db/explainTreeBuilder.go

+54
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
package db
2+
3+
import (
4+
"fmt"
5+
)
6+
7+
type QueryPlanNode struct {
8+
ID string
9+
ParentID string
10+
NotUsed string
11+
Detail string
12+
Children []*QueryPlanNode
13+
}
14+
15+
func BuildQueryPlanTree(rows [][]string) (*QueryPlanNode, error) {
16+
var nodes []*QueryPlanNode
17+
nodeMap := make(map[string]*QueryPlanNode)
18+
19+
for _, row := range rows {
20+
id := row[0]
21+
parentId := row[1]
22+
notUsed := row[2]
23+
detail := row[3]
24+
25+
node := &QueryPlanNode{
26+
ID: id,
27+
ParentID: parentId,
28+
NotUsed: notUsed,
29+
Detail: detail,
30+
}
31+
32+
nodes = append(nodes, node)
33+
nodeMap[id] = node
34+
}
35+
36+
root := &QueryPlanNode{}
37+
for _, node := range nodes {
38+
if node.ParentID == "0" {
39+
root = node
40+
} else {
41+
parent := nodeMap[node.ParentID]
42+
parent.Children = append(parent.Children, node)
43+
}
44+
}
45+
46+
return root, nil
47+
}
48+
49+
func PrintQueryPlanTree(node *QueryPlanNode, indent string) {
50+
fmt.Printf("%s%s\n", indent, node.Detail)
51+
for _, child := range node.Children {
52+
PrintQueryPlanTree(child, indent+" ")
53+
}
54+
}

Diff for: internal/db/output.go

+28-8
Original file line numberDiff line numberDiff line change
@@ -14,15 +14,34 @@ type Printer interface {
1414
print(statementResult StatementResult, outF io.Writer) error
1515
}
1616

17+
type ExplainQueryPrinter struct{}
18+
19+
func (eqp ExplainQueryPrinter) print(statementResult StatementResult, outF io.Writer) error {
20+
data := [][]string{}
21+
22+
tableData, err := appendData(statementResult, data, TABLE)
23+
if err != nil {
24+
return err
25+
}
26+
27+
root, err := BuildQueryPlanTree(tableData)
28+
if err != nil {
29+
return err
30+
}
31+
println("QUERY PLAN")
32+
PrintQueryPlanTree(root, "")
33+
34+
return nil
35+
}
36+
1737
type TablePrinter struct {
1838
withoutHeader bool
1939
}
2040

2141
func (t TablePrinter) print(statementResult StatementResult, outF io.Writer) error {
2242
data := [][]string{}
2343
table := createTable(outF)
24-
showHeader := !IsResultComingFromExplainQueryPlan(statementResult) && !t.withoutHeader
25-
if showHeader {
44+
if !t.withoutHeader {
2645
table.SetHeader(statementResult.ColumnNames)
2746
}
2847

@@ -33,6 +52,7 @@ func (t TablePrinter) print(statementResult StatementResult, outF io.Writer) err
3352

3453
table.AppendBulk(tableData)
3554
table.Render()
55+
3656
return nil
3757
}
3858

@@ -102,14 +122,13 @@ func appendData(statementResult StatementResult, data [][]string, mode FormatTyp
102122
data = append(data, formattedRow)
103123
}
104124

105-
if IsResultComingFromExplainQueryPlan(statementResult) {
106-
data = [][]string{{"Tree goes here!"}}
107-
}
108-
109125
return data, nil
110126
}
111127

112-
func getPrinter(mode enums.PrintMode, withoutHeader bool) (Printer, error) {
128+
func getPrinter(mode enums.PrintMode, withoutHeader bool, isExplainQueryPlan bool) (Printer, error) {
129+
if isExplainQueryPlan {
130+
return &ExplainQueryPrinter{}, nil
131+
}
113132
switch mode {
114133
case enums.TABLE_MODE:
115134
return &TablePrinter{
@@ -149,7 +168,8 @@ func PrintStatementResult(statementResult StatementResult, outF io.Writer, witho
149168
return &UnableToPrintStatementResult{}
150169
}
151170

152-
printer, err := getPrinter(mode, withoutHeader)
171+
isExplainQueryPlan := IsResultComingFromExplainQueryPlan(statementResult)
172+
printer, err := getPrinter(mode, withoutHeader, isExplainQueryPlan)
153173
if err != nil {
154174
return err
155175
}

0 commit comments

Comments
 (0)