Merge pull request #59 from planetscale/support-slow-query-log-header

Incorporate Slow Query Log Metadata into Hot Query Summarization

go/cmd/keys.go

@@ -61,7 +61,7 @@ func addInputTypeFlag(cmd *cobra.Command, s *string) {
 func configureLoader(inputType string, needsBindVars bool) (data.Loader, error) {
 	switch inputType {
 	case "sql":
-		return data.SQLScriptLoader{}, nil
+		return data.SlowQueryLogLoader{}, nil
 	case "mysql-log":
 		return data.MySQLLogLoader{}, nil
 	case "vtgate-log":

go/cmd/summarize.go

@@ -23,14 +23,20 @@ import (
 )
 
 func summarizeCmd() *cobra.Command {
-	return &cobra.Command{
+	var hotMetric string
+
+	cmd := &cobra.Command{
 		Use:     "summarize old_file.json [new_file.json]",
 		Aliases: []string{"benchstat"},
 		Short:   "Compares and analyses a trace output",
 		Example: "vt summarize old.json new.json",
 		Args:    cobra.RangeArgs(1, 2),
 		Run: func(_ *cobra.Command, args []string) {
-			summarize.Run(args)
+			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)")
+
+	return cmd
 }

go/data/query.go

@@ -40,6 +40,11 @@ type (
 		Query     string
 		Line      int
 		Type      CmdType
+
+		// These fields are only set if the log file is a slow query log
+		QueryTime, LockTime    float64
+		RowsSent, RowsExamined int
+		Timestamp              int64
 	}
 
 	errLoader struct {

go/data/query_log_parse.go

@@ -19,6 +19,8 @@ package data
 import (
 	"bufio"
 	"errors"
+	"fmt"
+	"io"
 	"os"
 	"regexp"
 	"sync"
@@ -29,7 +31,7 @@ type (
 
 	logReaderState struct {
 		fd         *os.File
-		scanner    *bufio.Scanner
+		reader     *bufio.Reader
 		reg        *regexp.Regexp
 		mu         sync.Mutex
 		lineNumber int
@@ -65,9 +67,16 @@ func (s *mysqlLogReaderState) Next() (Query, bool) {
 		return Query{}, false
 	}
 
-	for s.scanner.Scan() {
-		s.lineNumber++
-		line := s.scanner.Text()
+	for {
+		line, done, err := s.readLine()
+		if err != nil {
+			s.err = fmt.Errorf("error reading file: %w", err)
+			return Query{}, false
+		}
+		if done {
+			break
+		}
+
 		if len(line) == 0 {
 			continue
 		}
@@ -82,20 +91,7 @@ func (s *mysqlLogReaderState) Next() (Query, bool) {
 
 		// If we have a previous query, return it before processing the new line
 		if s.prevQuery != "" {
-			query := Query{
-				Query: s.prevQuery,
-				Line:  s.queryStart,
-				Type:  QueryT,
-			}
-			s.prevQuery = ""
-
-			// If the new line is a query, store it for next iteration
-			if matches[3] == "Query" {
-				s.prevQuery = matches[4]
-				s.queryStart = s.lineNumber
-			}
-
-			return query, true
+			return s.processQuery(matches), true
 		}
 
 		// Start a new query if this line is a query
@@ -108,19 +104,63 @@ func (s *mysqlLogReaderState) Next() (Query, bool) {
 
 	// Return the last query if we have one
 	if s.prevQuery != "" {
-		query := Query{
-			Query: s.prevQuery,
-			Line:  s.queryStart,
-			Type:  QueryT,
-		}
-		s.prevQuery = ""
-		return query, true
+		return s.finalizeQuery(), true
 	}
 
-	s.err = s.scanner.Err()
 	return Query{}, false
 }
 
+func (s *logReaderState) readLine() (string, bool, error) {
+	s.lineNumber++
+	line, isPrefix, err := s.reader.ReadLine()
+	if err == io.EOF {
+		return "", true, nil
+	}
+	if err != nil {
+		return "", false, err
+	}
+
+	// Handle lines longer than the buffer size
+	totalLine := append([]byte{}, line...)
+	for isPrefix {
+		line, isPrefix, err = s.reader.ReadLine()
+		if err == io.EOF {
+			break
+		}
+		if err != nil {
+			return "", false, err
+		}
+		totalLine = append(totalLine, line...)
+	}
+	return string(totalLine), false, nil
+}
+
+func (s *mysqlLogReaderState) finalizeQuery() Query {
+	query := Query{
+		Query: s.prevQuery,
+		Line:  s.queryStart,
+		Type:  QueryT,
+	}
+	s.prevQuery = ""
+	return query
+}
+
+func (s *mysqlLogReaderState) processQuery(matches []string) Query {
+	query := Query{
+		Query: s.prevQuery,
+		Line:  s.queryStart,
+		Type:  QueryT,
+	}
+	s.prevQuery = ""
+
+	// If the new line is a query, store it for next iteration
+	if matches[3] == "Query" {
+		s.prevQuery = matches[4]
+		s.queryStart = s.lineNumber
+	}
+	return query
+}
+
 func (s *logReaderState) Close() error {
 	s.mu.Lock()
 	defer s.mu.Unlock()
@@ -136,15 +176,6 @@ func (s *logReaderState) Close() error {
 	return s.err
 }
 
-func (s *logReaderState) NextLine() (string, bool) {
-	more := s.scanner.Scan()
-	if !more {
-		return "", false
-	}
-
-	return s.scanner.Text(), true
-}
-
 func (MySQLLogLoader) Load(fileName string) IteratorLoader {
 	reg := regexp.MustCompile(`^(\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}\.\d{6}Z)\s+(\d+)\s+(\w+)\s+(.*)`)
 
@@ -153,13 +184,11 @@ func (MySQLLogLoader) Load(fileName string) IteratorLoader {
 		return &errLoader{err}
 	}
 
-	scanner := bufio.NewScanner(fd)
-
 	return &mysqlLogReaderState{
 		logReaderState: logReaderState{
-			scanner: scanner,
-			reg:     reg,
-			fd:      fd,
+			reader: bufio.NewReader(fd),
+			reg:    reg,
+			fd:     fd,
 		},
 	}
 }

go/data/query_log_parse_test.go

@@ -23,14 +23,14 @@ import (
 )
 
 func TestParseMySQLQueryLog(t *testing.T) {
-	loader := MySQLLogLoader{}.Load("./testdata/mysql.query.log")
+	loader := MySQLLogLoader{}.Load("../testdata/mysql.query.log")
 	gotQueries, err := makeSlice(loader)
 	require.NoError(t, err)
 	require.Equal(t, 1517, len(gotQueries), "expected 1517 queries") //nolint:testifylint // too many elements for the output to be readable
 }
 
 func TestSmallSnippet(t *testing.T) {
-	loader := MySQLLogLoader{}.Load("./testdata/mysql.small-query.log")
+	loader := MySQLLogLoader{}.Load("../testdata/mysql.small-query.log")
 	gotQueries, err := makeSlice(loader)
 	require.NoError(t, err)
 	expected := []Query{