feat(tree-builder): add utils for building and printing the tree

Author: douglasmdev

internal/db/explainTreeBuilder.go

@@ -0,0 +1,54 @@
+package db
+
+import (
+	"fmt"
+)
+
+type QueryPlanNode struct {
+	ID       string
+	ParentID string
+	NotUsed  string
+	Detail   string
+	Children []*QueryPlanNode
+}
+
+func BuildQueryPlanTree(rows [][]string) (*QueryPlanNode, error) {
+	var nodes []*QueryPlanNode
+	nodeMap := make(map[string]*QueryPlanNode)
+
+	for _, row := range rows {
+		id := row[0]
+		parentId := row[1]
+		notUsed := row[2]
+		detail := row[3]
+
+		node := &QueryPlanNode{
+			ID:       id,
+			ParentID: parentId,
+			NotUsed:  notUsed,
+			Detail:   detail,
+		}
+
+		nodes = append(nodes, node)
+		nodeMap[id] = node
+	}
+
+	root := &QueryPlanNode{}
+	for _, node := range nodes {
+		if node.ParentID == "0" {
+			root = node
+		} else {
+			parent := nodeMap[node.ParentID]
+			parent.Children = append(parent.Children, node)
+		}
+	}
+
+	return root, nil
+}
+
+func PrintQueryPlanTree(node *QueryPlanNode, indent string) {
+	fmt.Printf("%s%s\n", indent, node.Detail)
+	for _, child := range node.Children {
+		PrintQueryPlanTree(child, indent+"  ")
+	}
+}

internal/db/output.go

@@ -14,15 +14,34 @@ type Printer interface {
 	print(statementResult StatementResult, outF io.Writer) error
 }
 
+type ExplainQueryPrinter struct{}
+
+func (eqp ExplainQueryPrinter) print(statementResult StatementResult, outF io.Writer) error {
+	data := [][]string{}
+
+	tableData, err := appendData(statementResult, data, TABLE)
+	if err != nil {
+		return err
+	}
+
+	root, err := BuildQueryPlanTree(tableData)
+	if err != nil {
+		return err
+	}
+	println("QUERY PLAN")
+	PrintQueryPlanTree(root, "")
+
+	return nil
+}
+
 type TablePrinter struct {
 	withoutHeader bool
 }
 
 func (t TablePrinter) print(statementResult StatementResult, outF io.Writer) error {
 	data := [][]string{}
 	table := createTable(outF)
-	showHeader := !IsResultComingFromExplainQueryPlan(statementResult) && !t.withoutHeader
-	if showHeader {
+	if !t.withoutHeader {
 		table.SetHeader(statementResult.ColumnNames)
 	}
 
@@ -33,6 +52,7 @@ func (t TablePrinter) print(statementResult StatementResult, outF io.Writer) err
 
 	table.AppendBulk(tableData)
 	table.Render()
+
 	return nil
 }
 
@@ -102,14 +122,13 @@ func appendData(statementResult StatementResult, data [][]string, mode FormatTyp
 		data = append(data, formattedRow)
 	}
 
-	if IsResultComingFromExplainQueryPlan(statementResult) {
-		data = [][]string{{"Tree goes here!"}}
-	}
-
 	return data, nil
 }
 
-func getPrinter(mode enums.PrintMode, withoutHeader bool) (Printer, error) {
+func getPrinter(mode enums.PrintMode, withoutHeader bool, isExplainQueryPlan bool) (Printer, error) {
+	if isExplainQueryPlan {
+		return &ExplainQueryPrinter{}, nil
+	}
 	switch mode {
 	case enums.TABLE_MODE:
 		return &TablePrinter{
@@ -149,7 +168,8 @@ func PrintStatementResult(statementResult StatementResult, outF io.Writer, witho
 		return &UnableToPrintStatementResult{}
 	}
 
-	printer, err := getPrinter(mode, withoutHeader)
+	isExplainQueryPlan := IsResultComingFromExplainQueryPlan(statementResult)
+	printer, err := getPrinter(mode, withoutHeader, isExplainQueryPlan)
 	if err != nil {
 		return err
 	}