feat: Add ability to have the formatted text plan be populated and inserted for presto_query_plans MySQL table from the loadjson command

cmd/loadjson/main.go

@@ -9,6 +9,7 @@ import (
 	"os/signal"
 	"path/filepath"
 	"pbench/log"
+	"pbench/prestoapi"
 	"pbench/stage"
 	"pbench/utils"
 	"reflect"
@@ -17,7 +18,6 @@ import (
 	"time"
 
 	"github.com/prestodb/presto-go-client/v2/queryjson"
-
 	"github.com/spf13/cobra"
 )
 
@@ -205,6 +205,18 @@ func processFile(ctx context.Context, path string) {
 			log.Error().Err(err).Str("path", path).Msg("failed to pre-process query info JSON")
 			return
 		}
+
+		// Generate the text plan from the JSON plan
+		if queryInfo.AssembledQueryPlanJson != "" {
+			textPlan, err := prestoapi.FormatQueryPlanAsText(queryInfo.AssembledQueryPlanJson)
+			if err != nil {
+				// Log the error but don't fail the entire operation
+				// The json_plan will still be available
+				log.Error().Err(err).Str("path", path).Msg("failed to generate text plan")
+			} else {
+				queryInfo.TextPlan = textPlan
+			}
+		}
 	}
 	if ExtractPlanJson {
 		planJsonFilePath := filepath.Join(OutputPath, fileName[0:len(fileName)-len(filepath.Ext(fileName))]+".plan.json")

go.mod

@@ -8,7 +8,7 @@ require (
 	github.com/go-sql-driver/mysql v1.9.3
 	github.com/influxdata/influxdb-client-go/v2 v2.14.0
 	github.com/pkg/errors v0.9.1
-	github.com/prestodb/presto-go-client/v2 v2.1.1
+	github.com/prestodb/presto-go-client/v2 v2.1.2
 	github.com/rs/zerolog v1.34.0
 	github.com/spf13/cobra v1.10.2
 	github.com/stretchr/testify v1.11.1

go.sum

@@ -61,8 +61,8 @@ github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
 github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
-github.com/prestodb/presto-go-client/v2 v2.1.1 h1:3mIaMitux4x/kV5CCpT0XQEcyPuPxi7eIfdXtB9VGJg=
-github.com/prestodb/presto-go-client/v2 v2.1.1/go.mod h1:xXanlGM7ptbiRDX9rn4G7GXlITtR+pWX5MoR5G7t8cI=
+github.com/prestodb/presto-go-client/v2 v2.1.2 h1:UQK1ipQAHoQl2uShbBjhTXjo5HCkoMHOY8E2dIzlNy0=
+github.com/prestodb/presto-go-client/v2 v2.1.2/go.mod h1:xXanlGM7ptbiRDX9rn4G7GXlITtR+pWX5MoR5G7t8cI=
 github.com/rs/xid v1.6.0/go.mod h1:7XoLgs4eV+QndskICGsho+ADou8ySMSjJKDIan90Nz0=
 github.com/rs/zerolog v1.34.0 h1:k43nTLIwcTVQAncfCw4KZ2VY6ukYoZaBPNOE8txlOeY=
 github.com/rs/zerolog v1.34.0/go.mod h1:bJsvje4Z08ROH4Nhs5iH600c3IkWhwp44iRc54W6wYQ=

prestoapi/plan_formatter.go

@@ -0,0 +1,85 @@
+package prestoapi
+
+import (
+	"encoding/json"
+	"fmt"
+	"strings"
+)
+
+// PlanNode represents a node in the query execution plan
+type PlanNode struct {
+	ID            string                   `json:"id"`
+	Name          string                   `json:"name"`
+	Identifier    string                   `json:"identifier"`
+	Details       string                   `json:"details"`
+	Children      []PlanNode               `json:"children"`
+	RemoteSources []interface{}            `json:"remoteSources"`
+	Estimates     []map[string]interface{} `json:"estimates"`
+}
+
+// StagePlanWrapper wraps the plan for a stage
+type StagePlanWrapper struct {
+	Plan PlanNode `json:"plan"`
+}
+
+// FormatQueryPlanAsText converts a JSON query plan to human-readable text format
+// The input should be the AssembledQueryPlanJson which is a map of stage IDs to plan wrappers
+func FormatQueryPlanAsText(jsonPlan string) (string, error) {
+	if jsonPlan == "" {
+		return "", nil
+	}
+
+	// Parse the JSON plan - it's a map of stage IDs to plan wrappers
+	var stages map[string]StagePlanWrapper
+	if err := json.Unmarshal([]byte(jsonPlan), &stages); err != nil {
+		return "", fmt.Errorf("failed to parse query plan JSON: %w", err)
+	}
+
+	var result strings.Builder
+
+	// Process each stage (fragment) in order
+	// Note: stages are typically numbered 0, 1, 2, etc.
+	for stageID, stageWrapper := range stages {
+		result.WriteString(fmt.Sprintf("Fragment %s\n", stageID))
+		result.WriteString(formatPlanNode(&stageWrapper.Plan, 0))
+		result.WriteString("\n")
+	}
+
+	return result.String(), nil
+}
+
+// formatPlanNode recursively formats a plan node and its children
+func formatPlanNode(node *PlanNode, depth int) string {
+	if node == nil {
+		return ""
+	}
+
+	var result strings.Builder
+	indent := strings.Repeat("    ", depth)
+
+	// Format the node header with identifier if present
+	if node.Identifier != "" {
+		result.WriteString(fmt.Sprintf("%s- %s[PlanNodeId %s]%s\n",
+			indent, node.Name, node.ID, node.Identifier))
+	} else {
+		result.WriteString(fmt.Sprintf("%s- %s[PlanNodeId %s]\n",
+			indent, node.Name, node.ID))
+	}
+
+	// Add details if present (already formatted with proper indentation in the JSON)
+	if node.Details != "" {
+		detailLines := strings.Split(strings.TrimSpace(node.Details), "\n")
+		for _, line := range detailLines {
+			if line != "" {
+				result.WriteString(fmt.Sprintf("%s        %s\n", indent, line))
+			}
+		}
+	}
+
+	// Recursively format children with increased indentation
+	for i := range node.Children {
+		result.WriteString(formatPlanNode(&node.Children[i], depth+1))
+	}
+
+	return result.String()
+}

prestoapi/plan_formatter_test.go

@@ -0,0 +1,52 @@
+package prestoapi
+
+import (
+	"strings"
+	"testing"
+)
+
+func TestFormatQueryPlanAsText(t *testing.T) {
+	// Sample JSON plan from the user
+	jsonPlan := `{"0":{"plan":{"id":"22","name":"Output","identifier":"[n_name, revenue]","details":"revenue := sum (4:5)\n","children":[{"id":"1311","name":"TopN","identifier":"[1 by (sum DESC_NULLS_LAST)]","details":"","children":[{"id":"1310","name":"TopNPartial","identifier":"[1 by (sum DESC_NULLS_LAST)]","details":"","children":[{"id":"14","name":"Aggregate(FINAL)[n_name]","identifier":"","details":"sum := \"presto.default.sum\"((sum_13)) (4:5)\n","children":[{"id":"1705","name":"LocalExchange","identifier":"[SINGLE] ()","details":"","children":[{"id":"1703","name":"Aggregate(PARTIAL)[n_name]","identifier":"","details":"sum_13 := \"presto.default.sum\"((expr)) (4:5)\n","children":[{"id":"340","name":"Project","identifier":"[projectLocality = LOCAL]","details":"expr := (l_extendedprice) * ((DOUBLE'1.0') - (l_discount)) (8:6)\n","children":[{"id":"1410","name":"InnerJoin","identifier":"[(\"l_suppkey\" = \"s_suppkey\") AND (\"c_nationkey\" = \"s_nationkey\")]","details":"Distribution: REPLICATED\n","children":[{"id":"1409","name":"InnerJoin","identifier":"[(\"l_orderkey\" = \"o_orderkey\")]","details":"Distribution: REPLICATED\n","children":[{"id":"3","name":"TableScan","identifier":"[TableHandle {connectorId='tpch', connectorHandle='lineitem:sf1.0', layout='Optional[lineitem:sf1.0]'}]","details":"l_orderkey := tpch:l_orderkey (8:5)\nl_extendedprice := tpch:l_extendedprice (8:5)\nl_suppkey := tpch:l_suppkey (8:5)\nl_discount := tpch:l_discount (8:5)\n","children":[],"remoteSources":[],"estimates":[]}],"remoteSources":[],"estimates":[]}],"remoteSources":[],"estimates":[]}],"remoteSources":[],"estimates":[]}],"remoteSources":[],"estimates":[]}],"remoteSources":[],"estimates":[]}],"remoteSources":[],"estimates":[]}],"remoteSources":[],"estimates":[]}],"remoteSources":[],"estimates":[]}],"remoteSources":[],"estimates":[]}}}`
+
+	result, err := FormatQueryPlanAsText(jsonPlan)
+	if err != nil {
+		t.Fatalf("FormatQueryPlanAsText failed: %v", err)
+	}
+
+	// Verify the output contains expected elements
+	expectedStrings := []string{
+		"Fragment 0",
+		"- Output[PlanNodeId 22][n_name, revenue]",
+		"revenue := sum (4:5)",
+		"- TopN[PlanNodeId 1311][1 by (sum DESC_NULLS_LAST)]",
+		"- TableScan[PlanNodeId 3]",
+		"l_orderkey := tpch:l_orderkey",
+	}
+
+	for _, expected := range expectedStrings {
+		if !strings.Contains(result, expected) {
+			t.Errorf("Expected output to contain %q, but it didn't.\nGot:\n%s", expected, result)
+		}
+	}
+
+	// Print the result for manual inspection
+	t.Logf("Formatted plan:\n%s", result)
+}
+
+func TestFormatQueryPlanAsText_EmptyInput(t *testing.T) {
+	result, err := FormatQueryPlanAsText("")
+	if err != nil {
+		t.Fatalf("FormatQueryPlanAsText with empty input failed: %v", err)
+	}
+	if result != "" {
+		t.Errorf("Expected empty result for empty input, got: %q", result)
+	}
+}
+
+func TestFormatQueryPlanAsText_InvalidJSON(t *testing.T) {
+	_, err := FormatQueryPlanAsText("invalid json")
+	if err == nil {
+		t.Error("Expected error for invalid JSON, got nil")
+	}
+}