Use IteratorLoader to process queries

Signed-off-by: Florent Poinsard <[email protected]>

Author: frouioui

go/data/query.go

@@ -24,7 +24,7 @@ import (
 
 type (
 	Loader interface {
-		Load(url string) ([]Query, error)
+		Load(filename string) IteratorLoader
 	}
 
 	IteratorLoader interface {

go/data/query_log_parse.go

@@ -44,12 +44,6 @@ type (
 	}
 )
 
-func (MySQLLogLoader) Load(fileName string) ([]Query, error) {
-	loader := MySQLLogLoader{}.Loadit(fileName)
-
-	return makeSlice(loader)
-}
-
 func makeSlice(loader IteratorLoader) ([]Query, error) {
 	var queries []Query
 	for {
@@ -151,7 +145,7 @@ func (s *logReaderState) NextLine() (string, bool) {
 	return s.scanner.Text(), true
 }
 
-func (MySQLLogLoader) Loadit(fileName string) IteratorLoader {
+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)

go/data/query_log_parse_test.go

@@ -23,13 +23,15 @@ 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.Equal(t, 1517, len(gotQueries), "expected 1517 queries") //nolint:testifylint // too many elements for the output to be readable
 }
 
 func TestSmallSnippet(t *testing.T) {
-	gotQueries, err := 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{
 		{

go/data/sql_script_loader.go

@@ -106,7 +106,7 @@ func readData(url string) ([]byte, error) {
 	return io.ReadAll(res.Body)
 }
 
-func (SQLScriptLoader) Loadit(filename string) IteratorLoader {
+func (SQLScriptLoader) Load(filename string) IteratorLoader {
 	var scanner *bufio.Scanner
 	var fd *os.File
 
@@ -133,11 +133,6 @@ func (SQLScriptLoader) Loadit(filename string) IteratorLoader {
 	}
 }
 
-func (SQLScriptLoader) Load(url string) ([]Query, error) {
-	loader := SQLScriptLoader{}.Loadit(url)
-	return makeSlice(loader)
-}
-
 // Helper function to parse individual queries
 func parseQuery(rs Query) (*Query, error) {
 	realS := rs.Query

go/data/vtgate_log_parse.go

@@ -45,7 +45,7 @@ type (
 	}
 )
 
-func (vll VtGateLogLoader) Loadit(fileName string) IteratorLoader {
+func (vll VtGateLogLoader) Load(fileName string) IteratorLoader {
 	reg := regexp.MustCompile(`\t"([^"]+)"\t(\{(?:[^{}]|(?:\{[^{}]*\}))*\}|"[^"]+")`)
 	fd, err := os.OpenFile(fileName, os.O_RDONLY, 0)
 	if err != nil {
@@ -129,11 +129,6 @@ func (s *vtgateLogReaderState) Next() (Query, bool) {
 	return Query{}, false
 }
 
-func (vll VtGateLogLoader) Load(fileName string) (queries []Query, err error) {
-	loader := vll.Loadit(fileName)
-	return makeSlice(loader)
-}
-
 func addBindVarsToQuery(query string, bvs map[string]*querypb.BindVariable) (string, error) {
 	parser, err := sqlparser.New(sqlparser.Options{})
 	if err != nil {

go/data/vtgate_log_parse_test.go

@@ -25,12 +25,14 @@ import (
 )
 
 func TestParseVtGateQueryLogRedacted(t *testing.T) {
-	_, err := VtGateLogLoader{NeedsBindVars: true}.Load("./testdata/vtgate.query.log.redacted")
+	loader := VtGateLogLoader{NeedsBindVars: true}.Load("./testdata/vtgate.query.log.redacted")
+	_, err := makeSlice(loader)
 	require.EqualError(t, err, "line 1: query has redacted bind variables, cannot parse them")
 }
 
 func TestParseVtGateQueryLog(t *testing.T) {
-	gotQueries, err := VtGateLogLoader{NeedsBindVars: true}.Load("./testdata/vtgate.query.log")
+	loader := VtGateLogLoader{NeedsBindVars: true}.Load("./testdata/vtgate.query.log")
+	gotQueries, err := makeSlice(loader)
 	require.NoError(t, err)
 
 	require.Len(t, gotQueries, 25)
@@ -47,7 +49,8 @@ func TestParseVtGateQueryLog(t *testing.T) {
 }
 
 func TestParseVtGateQueryLogNoBindVars(t *testing.T) {
-	gotQueries, err := VtGateLogLoader{NeedsBindVars: false}.Load("./testdata/vtgate.query.log")
+	loader := VtGateLogLoader{NeedsBindVars: false}.Load("./testdata/vtgate.query.log")
+	gotQueries, err := makeSlice(loader)
 	require.NoError(t, err)
 	require.Len(t, gotQueries, 25)
 

go/keys/keys.go

@@ -49,13 +49,15 @@ func run(out io.Writer, cfg Config) error {
 		queries: make(map[string]*QueryAnalysisResult),
 		failed:  make(map[string]*QueryFailedResult),
 	}
-	queries, err := cfg.Loader.Load(cfg.FileName)
-	if err != nil {
-		return err
-	}
 
+	loader := cfg.Loader.Load(cfg.FileName)
 	skip := false
-	for _, query := range queries {
+	for {
+		query, kontinue := loader.Next()
+		if !kontinue {
+			break
+		}
+
 		switch query.Type {
 		case data.Skip, data.Error, data.VExplain:
 			skip = true
@@ -72,6 +74,10 @@ func run(out io.Writer, cfg Config) error {
 		}
 	}
 
+	if err := loader.Close(); err != nil {
+		return err
+	}
+
 	return ql.writeJSONTo(out)
 }
 

go/tester/tester.go

@@ -223,15 +223,22 @@ func (t *Tester) Run() (err error) {
 			err = errors.Join(err, postErr)
 		}()
 	}
-	queries, err := t.loader.Load(t.name)
+
+	loader := t.loader.Load(t.name)
+	for {
+		query, kontinue := loader.Next()
+		if !kontinue {
+			break
+		}
+		t.handleQuery(query)
+	}
+
+	err = loader.Close()
 	if err != nil {
 		t.reporter.AddFailure(err)
 		return err
 	}
 
-	for _, q := range queries {
-		t.handleQuery(q)
-	}
 	fmt.Printf("%s\n", t.reporter.Report())
 
 	return nil