Move schema info definition from flag to parameters, add FileType to schema info and deduce that type from the json input itself

Signed-off-by: Rohit Nayak <rohit@planetscale.com>

Author: rohit-nayak-ps

go/cmd/summarize.go

@@ -24,7 +24,6 @@ import (
 
 func summarizeCmd() *cobra.Command {
 	var hotMetric string
-	var schemaInfoPath string
 
 	cmd := &cobra.Command{
 		Use:     "summarize old_file.json [new_file.json]",
@@ -33,12 +32,11 @@ 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, schemaInfoPath)
+			summarize.Run(args, hotMetric)
 		},
 	}
 
 	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().StringVar(&schemaInfoPath, "schema-info-path", "", "Path to output of 'vt schema' command (optional)")
 
 	return cmd
 }

go/summarize/summarize.go

@@ -39,10 +39,25 @@ type (
 	}
 )
 
-func Run(files []string, hotMetric, schemaInfoPath string) {
-	traces := make([]readingSummary, len(files))
+func Run(files []string, hotMetric string) {
+	var traceFiles []string
+	var dbInfoPath string
+	// todo: add file types for other json types. Right now just checks for dbinfo files, else defaults
+	for _, file := range files {
+		typ, _ := getFileType(file)
+		switch typ {
+		case dbInfoFile:
+			fmt.Printf("dbinfo file: %s\n", file)
+			dbInfoPath = file
+		default:
+			fmt.Printf("trace file: %s\n", file)
+			traceFiles = append(traceFiles, file)
+		}
+	}
+
+	traces := make([]readingSummary, len(traceFiles))
 	var err error
-	for i, arg := range files {
+	for i, arg := range traceFiles {
 		traces[i], err = readTraceFile(arg)
 		if err != nil {
 			exit(err.Error())
@@ -62,7 +77,7 @@ func Run(files []string, hotMetric, schemaInfoPath string) {
 	if firstTrace.AnalysedQueries == nil {
 		printTraceSummary(os.Stdout, terminalWidth(), highlightQuery, firstTrace)
 	} else {
-		printKeysSummary(os.Stdout, firstTrace, time.Now(), hotMetric, schemaInfoPath)
+		printKeysSummary(os.Stdout, firstTrace, time.Now(), hotMetric, dbInfoPath)
 	}
 }
 

go/summarize/utils.go

@@ -0,0 +1,79 @@
+package summarize
+
+import (
+	"encoding/json"
+	"errors"
+	"fmt"
+	"os"
+)
+
+type fileType int
+
+const (
+	unknownFile fileType = iota
+	traceFile
+	keysFile
+	dbInfoFile
+)
+
+var fileTypeMap = map[string]fileType{
+	"trace":  traceFile,
+	"keys":   keysFile,
+	"dbinfo": dbInfoFile,
+}
+
+// getFileType reads the first key-value pair from a JSON file and returns the type of the file
+// Note:
+func getFileType(filename string) (fileType, error) {
+	// read json file
+	file, err := os.Open(filename)
+	if err != nil {
+		return unknownFile, errors.New(fmt.Sprintf("error opening file: %v", err))
+	}
+	defer file.Close()
+
+	decoder := json.NewDecoder(file)
+
+	token, err := decoder.Token()
+	if err != nil {
+		return unknownFile, errors.New(fmt.Sprintf("Error reading token: %v", err))
+	}
+
+	// Ensure the token is the start of an object
+	if delim, ok := token.(json.Delim); !ok || delim != '{' {
+		return unknownFile, errors.New(fmt.Sprintf("Expected start of object '{'"))
+	}
+
+	// Read the key-value pairs
+	for decoder.More() {
+		// Read the key
+		keyToken, err := decoder.Token()
+		if err != nil {
+			return unknownFile, errors.New(fmt.Sprintf("Error reading key token: %v", err))
+		}
+
+		key, ok := keyToken.(string)
+		if !ok {
+			return unknownFile, errors.New(fmt.Sprintf("Expected key to be a string: %s", keyToken))
+		}
+
+		// Read the value
+		valueToken, err := decoder.Token()
+		if err != nil {
+			return unknownFile, errors.New(fmt.Sprintf("Error reading value token: %v", err))
+		}
+
+		// Check if the key is "FileType"
+		if key == "FileType" {
+			if fileType, ok := fileTypeMap[valueToken.(string)]; ok {
+				return fileType, nil
+			} else {
+				return unknownFile, errors.New(fmt.Sprintf("Unknown FileType: %s", valueToken))
+			}
+		} else {
+			// Currently we expect the first key to be FileType, for optimization reasons
+			return unknownFile, nil
+		}
+	}
+	return unknownFile, nil
+}

go/summarize/utils_test.go

@@ -0,0 +1,45 @@
+package summarize
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+// Test utils getFileType function
+func TestGetFileType(t *testing.T) {
+	type testCase struct {
+		filename      string
+		expectedType  fileType
+		expectedError string
+	}
+	testCases := []testCase{
+		{
+			filename:     "../testdata/keys-log.json",
+			expectedType: unknownFile,
+		},
+		{
+			filename:     "../testdata/sakila-schema-info.json",
+			expectedType: dbInfoFile,
+		},
+		{
+			filename:      "../testdata/mysql.query.log",
+			expectedType:  unknownFile,
+			expectedError: "Error reading token",
+		},
+	}
+	for _, tc := range testCases {
+		t.Run(tc.filename, func(t *testing.T) {
+			ft, err := getFileType(tc.filename)
+			if tc.expectedError != "" {
+				require.Error(t, err)
+			}
+			if err != nil {
+				require.Contains(t, err.Error(), tc.expectedError)
+			}
+			if ft != tc.expectedType {
+				t.Errorf("Expected type: %v, got: %v", tc.expectedType, ft)
+			}
+		})
+	}
+}

go/testdata/sakila-schema-info.json

@@ -1,4 +1,5 @@
 {
+  "FileType": "dbinfo",
   "Tables": [
     {
       "Name": "actor",