Add support for CSV parsing

go/cmd/flag_config.go

@@ -0,0 +1,93 @@
+/*
+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 cmd
+
+import (
+	"errors"
+	"fmt"
+
+	"github.com/spf13/cobra"
+
+	"github.com/vitessio/vt/go/data"
+)
+
+type csvFlags struct {
+	header                                                                                                         bool
+	queryField, connectionIDField, queryTimeField, lockTimeField, rowsSentField, rowsExaminedField, timestampField int
+}
+
+const allowedInputTypes = "'sql', 'mysql-log', 'vtgate-log', 'csv'"
+
+func addInputTypeFlag(cmd *cobra.Command, s *string) {
+	*s = "sql"
+	cmd.Flags().StringVar(s, "input-type", "sql", fmt.Sprintf("Specifies the type of input file: %s", allowedInputTypes))
+}
+
+func addCSVConfigFlag(cmd *cobra.Command, c *csvFlags) {
+	cmd.Flags().BoolVar(&c.header, "csv-header", false, "Indicates that the CSV file has a header row")
+	cmd.Flags().IntVar(&c.queryField, "csv-query-field", 0, "Column index or name for the query field (required)")
+	cmd.Flags().IntVar(&c.connectionIDField, "csv-connection-id-field", 0, "Column index or name for the connection ID field")
+	cmd.Flags().IntVar(&c.queryTimeField, "csv-query-time-field", 0, "Column index or name for the query time field")
+	cmd.Flags().IntVar(&c.lockTimeField, "csv-lock-time-field", 0, "Column index or name for the lock time field")
+	cmd.Flags().IntVar(&c.rowsSentField, "csv-rows-sent-field", 0, "Column index or name for the rows sent field")
+	cmd.Flags().IntVar(&c.rowsExaminedField, "csv-rows-examined-field", 0, "Column index or name for the rows examined field")
+	cmd.Flags().IntVar(&c.timestampField, "csv-timestamp-field", 0, "Column index or name for the timestamp field")
+}
+
+func configureLoader(inputType string, needsBindVars bool, csvConfig data.CSVConfig) (data.Loader, error) {
+	switch inputType {
+	case "sql":
+		return data.SlowQueryLogLoader{}, nil
+	case "mysql-log":
+		return data.MySQLLogLoader{}, nil
+	case "vtgate-log":
+		return data.VtGateLogLoader{NeedsBindVars: needsBindVars}, nil
+	case "csv":
+		if csvConfig.QueryField == -1 {
+			return nil, errors.New("must specify query field for CSV loader")
+		}
+		return data.CSVLogLoader{Config: csvConfig}, nil
+	default:
+		return nil, fmt.Errorf("invalid input type: must be %s", allowedInputTypes)
+	}
+}
+
+func csvFlagsToConfig(cmd *cobra.Command, flags csvFlags) data.CSVConfig {
+	var c data.CSVConfig
+	if cmd.Flags().Changed("csv-query-field") {
+		c.QueryField = flags.queryField
+	}
+	if cmd.Flags().Changed("csv-connection-id-field") {
+		c.ConnectionIDField = &flags.connectionIDField
+	}
+	if cmd.Flags().Changed("csv-query-time-field") {
+		c.QueryTimeField = &flags.queryTimeField
+	}
+	if cmd.Flags().Changed("csv-lock-time-field") {
+		c.LockTimeField = &flags.lockTimeField
+	}
+	if cmd.Flags().Changed("csv-rows-sent-field") {
+		c.RowsSentField = &flags.rowsSentField
+	}
+	if cmd.Flags().Changed("csv-rows-examined-field") {
+		c.RowsExaminedField = &flags.rowsExaminedField
+	}
+	if cmd.Flags().Changed("csv-timestamp-field") {
+		c.TimestampField = &flags.timestampField
+	}
+	return c
+}

go/cmd/keys.go

@@ -17,8 +17,6 @@ limitations under the License.
 package cmd
 
 import (
-	"fmt"
-
 	"github.com/spf13/cobra"
 
 	"github.com/vitessio/vt/go/data"
@@ -27,18 +25,22 @@ import (
 
 func keysCmd() *cobra.Command {
 	var inputType string
-
+	flags := new(csvFlags)
+	var csvConfig data.CSVConfig
 	cmd := &cobra.Command{
 		Use:     "keys ",
 		Short:   "Runs vexplain keys on all queries of the test file",
 		Example: "vt keys file.test",
 		Args:    cobra.ExactArgs(1),
+		PreRun: func(cmd *cobra.Command, _ []string) {
+			csvConfig = csvFlagsToConfig(cmd, *flags)
+		},
 		RunE: func(_ *cobra.Command, args []string) error {
 			cfg := keys.Config{
 				FileName: args[0],
 			}
 
-			loader, err := configureLoader(inputType, false)
+			loader, err := configureLoader(inputType, false, csvConfig)
 			if err != nil {
 				return err
 			}
@@ -49,26 +51,7 @@ func keysCmd() *cobra.Command {
 	}
 
 	addInputTypeFlag(cmd, &inputType)
+	addCSVConfigFlag(cmd, flags)
 
 	return cmd
 }
-
-const allowedInputTypes = "'sql', 'mysql-log' or 'vtgate-log'"
-
-func addInputTypeFlag(cmd *cobra.Command, s *string) {
-	*s = "sql"
-	cmd.Flags().StringVar(s, "input-type", "sql", fmt.Sprintf("Specifies the type of input file: %s", allowedInputTypes))
-}
-
-func configureLoader(inputType string, needsBindVars bool) (data.Loader, error) {
-	switch inputType {
-	case "sql":
-		return data.SlowQueryLogLoader{}, nil
-	case "mysql-log":
-		return data.MySQLLogLoader{}, nil
-	case "vtgate-log":
-		return data.VtGateLogLoader{NeedsBindVars: needsBindVars}, nil
-	default:
-		return nil, fmt.Errorf("invalid input type: must be %s", allowedInputTypes)
-	}
-}

go/cmd/tester.go

@@ -21,23 +21,29 @@ import (
 
 	"github.com/spf13/cobra"
 
+	"github.com/vitessio/vt/go/data"
 	vttester "github.com/vitessio/vt/go/tester"
 )
 
 func testerCmd() *cobra.Command {
 	var cfg vttester.Config
 	var inputType string
+	flags := new(csvFlags)
+	var csvConfig data.CSVConfig
 
 	cmd := &cobra.Command{
 		Aliases: []string{"test"},
 		Use:     "tester ",
 		Short:   "Test the given workload against both Vitess and MySQL.",
 		Example: "vt tester ",
 		Args:    cobra.MinimumNArgs(1),
+		PreRun: func(cmd *cobra.Command, _ []string) {
+			csvConfig = csvFlagsToConfig(cmd, *flags)
+		},
 		RunE: func(cmd *cobra.Command, args []string) error {
 			cfg.Tests = args
 			cfg.Compare = true
-			loader, err := configureLoader(inputType, true)
+			loader, err := configureLoader(inputType, true, csvConfig)
 			if err != nil {
 				return err
 			}
@@ -52,25 +58,31 @@ func testerCmd() *cobra.Command {
 	cmd.Flags().BoolVar(&cfg.OLAP, "olap", false, "Use OLAP to run the queries.")
 	cmd.Flags().BoolVar(&cfg.XUnit, "xunit", false, "Get output in an xml file instead of errors directory")
 	addInputTypeFlag(cmd, &inputType)
+	addCSVConfigFlag(cmd, flags)
 
 	return cmd
 }
 
 func tracerCmd() *cobra.Command {
 	var cfg vttester.Config
 	var inputType string
+	flags := new(csvFlags)
+	var csvConfig data.CSVConfig
 
 	cmd := &cobra.Command{
 		Use:   "trace ",
 		Short: "Runs the given workload and does a `vexplain trace` on all queries.",
 		Args:  cobra.MinimumNArgs(1),
+		PreRun: func(cmd *cobra.Command, _ []string) {
+			csvConfig = csvFlagsToConfig(cmd, *flags)
+		},
 		RunE: func(cmd *cobra.Command, args []string) error {
 			if cfg.TraceFile == "" {
 				return errors.New("flag --trace-file is required when tracing")
 			}
 			cfg.Tests = args
 			cfg.Compare = false
-			loader, err := configureLoader(inputType, true)
+			loader, err := configureLoader(inputType, true, csvConfig)
 			if err != nil {
 				return err
 			}
@@ -82,6 +94,7 @@ func tracerCmd() *cobra.Command {
 
 	commonFlags(cmd, &cfg)
 	addInputTypeFlag(cmd, &inputType)
+	addCSVConfigFlag(cmd, flags)
 
 	return cmd
 }

go/cmd/transactions.go

@@ -19,24 +19,30 @@ package cmd
 import (
 	"github.com/spf13/cobra"
 
+	"github.com/vitessio/vt/go/data"
 	"github.com/vitessio/vt/go/transactions"
 )
 
 func transactionsCmd() *cobra.Command {
 	var inputType string
+	flags := new(csvFlags)
+	var csvConfig data.CSVConfig
 
 	cmd := &cobra.Command{
 		Use:     "transactions ",
 		Aliases: []string{"txs"},
 		Short:   "Analyze transactions on a query log",
 		Example: "vt transactions file.log",
 		Args:    cobra.ExactArgs(1),
+		PreRun: func(cmd *cobra.Command, _ []string) {
+			csvConfig = csvFlagsToConfig(cmd, *flags)
+		},
 		RunE: func(_ *cobra.Command, args []string) error {
 			cfg := transactions.Config{
 				FileName: args[0],
 			}
 
-			loader, err := configureLoader(inputType, false)
+			loader, err := configureLoader(inputType, false, csvConfig)
 			if err != nil {
 				return err
 			}
@@ -48,6 +54,7 @@ func transactionsCmd() *cobra.Command {
 	}
 
 	addInputTypeFlag(cmd, &inputType)
+	addCSVConfigFlag(cmd, flags)
 
 	return cmd
 }