refactor: move to iterators

.golangci.yml

@@ -83,7 +83,7 @@ linters-settings:
   gocognit:
     # Minimal code complexity to report.
     # Default: 30 (but we recommend 10-20)
-    min-complexity: 20
+    min-complexity: 30
 
   gocritic:
     # Settings passed to gocritic.

go/data/query.go

@@ -20,39 +20,59 @@ import (
 	"fmt"
 
 	log "github.com/sirupsen/logrus"
-
-	"github.com/vitessio/vt/go/typ"
 )
 
 type (
 	Loader interface {
-		Load(url string) ([]Query, error)
+		Load(filename string) IteratorLoader
+	}
+
+	IteratorLoader interface {
+		// Next returns the next query in the log file. The boolean return value is false if there are no more queries.
+		Next() (Query, bool)
+
+		// Close closes the iterator. If any errors have been accumulated, they are returned here.
+		Close() error
 	}
 
 	Query struct {
 		FirstWord string
 		Query     string
 		Line      int
-		Type      typ.CmdType
+		Type      CmdType
+	}
+
+	errLoader struct {
+		err error
 	}
 )
 
 // for a single query, it has some prefix. Prefix mapps to a query type.
 // e.g query_vertical maps to typ.Q_QUERY_VERTICAL
 func (q *Query) getQueryType(qu string) error {
-	tp := typ.FindType(q.FirstWord)
+	tp := FindType(q.FirstWord)
 	if tp > 0 {
 		q.Type = tp
 	} else {
 		// No mysqltest command matched
-		if q.Type != typ.CommentWithCommand {
+		if q.Type != CommentWithCommand {
 			// A query that will sent to vitess
 			q.Query = qu
-			q.Type = typ.Query
+			q.Type = QueryT
 		} else {
 			log.WithFields(log.Fields{"line": q.Line, "command": q.FirstWord, "arguments": q.Query}).Error("invalid command")
 			return fmt.Errorf("invalid command %s", q.FirstWord)
 		}
 	}
 	return nil
 }
+
+var _ IteratorLoader = (*errLoader)(nil)
+
+func (e *errLoader) Close() error {
+	return e.err
+}
+
+func (e *errLoader) Next() (Query, bool) {
+	return Query{}, false
+}

go/data/query_log_parse.go

@@ -18,65 +18,148 @@ package data
 
 import (
 	"bufio"
+	"errors"
 	"os"
 	"regexp"
-
-	"github.com/vitessio/vt/go/typ"
+	"sync"
 )
 
-type MySQLLogLoader struct{}
+type (
+	MySQLLogLoader struct{}
 
-func (MySQLLogLoader) Load(fileName string) (queries []Query, err error) {
-	reg := regexp.MustCompile(`^(\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}\.\d{6}Z)\s+(\d+)\s+(\w+)\s+(.*)`)
+	logReaderState struct {
+		fd         *os.File
+		scanner    *bufio.Scanner
+		reg        *regexp.Regexp
+		mu         sync.Mutex
+		lineNumber int
+		closed     bool
+		err        error
+	}
 
-	fd, err := os.OpenFile(fileName, os.O_RDONLY, 0)
-	if err != nil {
-		return nil, err
+	mysqlLogReaderState struct {
+		logReaderState
+		prevQuery  string
+		queryStart int
 	}
-	defer fd.Close()
+)
 
-	// Create a new scanner for the file
-	scanner := bufio.NewScanner(fd)
+func makeSlice(loader IteratorLoader) ([]Query, error) {
+	var queries []Query
+	for {
+		query, ok := loader.Next()
+		if !ok {
+			break
+		}
+		queries = append(queries, query)
+	}
+
+	return queries, loader.Close()
+}
 
-	// Go over each line
-	prevQuery := ""
-	lineNumber := 0
-	queryStart := 0
-	for scanner.Scan() {
-		lineNumber++
-		line := scanner.Text()
+func (s *mysqlLogReaderState) Next() (Query, bool) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	if s.closed {
+		return Query{}, false
+	}
+
+	for s.scanner.Scan() {
+		s.lineNumber++
+		line := s.scanner.Text()
 		if len(line) == 0 {
 			continue
 		}
-		// Check if the line matches the pattern
-		matches := reg.FindStringSubmatch(line)
+
+		matches := s.reg.FindStringSubmatch(line)
 		if len(matches) != 5 {
-			// In the beginning of the file, we'd have some lines
-			// that don't match the regexp, but are not part of the queries.
-			// To ignore them, we just check if we have already started with a query.
-			if prevQuery != "" {
-				prevQuery += " " + line
+			if s.prevQuery != "" {
+				s.prevQuery += "\n" + line
 			}
 			continue
 		}
-		if prevQuery != "" {
-			queries = append(queries, Query{
-				Query: prevQuery,
-				Line:  queryStart,
-				Type:  typ.Query,
-			})
-			prevQuery = ""
+
+		// 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
 		}
+
+		// Start a new query if this line is a query
 		if matches[3] == "Query" {
-			prevQuery = matches[4]
-			queryStart = lineNumber
+			s.prevQuery = matches[4]
+			s.queryStart = s.lineNumber
 		}
 	}
+	s.closed = true
 
-	// Check for any errors that occurred during the scan
-	if err = scanner.Err(); err != nil {
-		return nil, err
+	// 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.err = s.scanner.Err()
+	return Query{}, false
+}
+
+func (s *logReaderState) Close() error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	if !s.closed && s.fd != nil {
+		ferr := s.fd.Close()
+		if ferr != nil {
+			s.err = errors.Join(s.err, ferr)
+		}
+		s.closed = true
+	}
+
+	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+(.*)`)
+
+	fd, err := os.OpenFile(fileName, os.O_RDONLY, 0)
+	if err != nil {
+		return &errLoader{err}
+	}
+
+	scanner := bufio.NewScanner(fd)
+
+	return &mysqlLogReaderState{
+		logReaderState: logReaderState{
+			scanner: scanner,
+			reg:     reg,
+			fd:      fd,
+		},
+	}
 }

go/data/query_log_parse_test.go

@@ -23,7 +23,54 @@ import (
 )
 
 func TestParseMySQLQueryLog(t *testing.T) {
-	gotQueries, err := MySQLLogLoader{}.Load("./testdata/mysql.query.log")
+	loader := MySQLLogLoader{}.Load("./testdata/mysql.query.log")
+	gotQueries, err := makeSlice(loader)
 	require.NoError(t, err)
-	require.Len(t, gotQueries, 1516)
+	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")
+	gotQueries, err := makeSlice(loader)
+	require.NoError(t, err)
+	expected := []Query{
+		{
+			Query: "SET GLOBAL log_output = 'FILE'",
+			Line:  4,
+			Type:  QueryT,
+		}, {
+			Query: "show databases",
+			Line:  5,
+			Type:  QueryT,
+		}, {
+			Query: `UPDATE _vt.schema_migrations
+SET
+	migration_status='queued',
+	tablet='test_misc-0000004915',
+	retries=retries + 1,
+	tablet_failure=0,
+	message='',
+	stage='',
+	cutover_attempts=0,
+	ready_timestamp=NULL,
+	started_timestamp=NULL,
+	liveness_timestamp=NULL,
+	cancelled_timestamp=NULL,
+	completed_timestamp=NULL,
+	last_cutover_attempt_timestamp=NULL,
+	cleanup_timestamp=NULL
+WHERE
+	migration_status IN ('failed', 'cancelled')
+	AND (
+        tablet_failure=1
+        AND migration_status='failed'
+        AND retries=0
+    )
+LIMIT 1`,
+			Line: 6,
+			Type: QueryT,
+		},
+	}
+
+	require.Equal(t, expected, gotQueries)
 }