Merge pull request #67 from planetscale/foreach

refactor: add & use foreachSQLQuery

go/data/query.go

@@ -52,6 +52,37 @@ type (
 	}
 )
 
+// ForeachSQLQuery reads a query log file and calls the provided function for each normal SQL query in the log.
+// If the query log contains directives, they will be read and queries will be skipped as necessary.
+func ForeachSQLQuery(loader IteratorLoader, f func(Query) error) error {
+	skip := false
+	for {
+		query, kontinue := loader.Next()
+		if !kontinue {
+			break
+		}
+
+		switch query.Type {
+		case Skip, Error, VExplain:
+			skip = true
+		case Unknown:
+			return fmt.Errorf("unknown command type: %s", query.Type)
+		case Comment, CommentWithCommand, EmptyLine, WaitForAuthoritative, SkipIfBelowVersion:
+			// no-op for keys
+		case QueryT:
+			if skip {
+				skip = false
+				continue
+			}
+			if err := f(query); err != nil {
+				return err
+			}
+		}
+	}
+
+	return nil
+}
+
 // 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 {

go/keys/keys.go

@@ -86,28 +86,10 @@ func run(out io.Writer, cfg Config) error {
 	}
 
 	loader := cfg.Loader.Load(cfg.FileName)
-	skip := false
-	for {
-		query, kontinue := loader.Next()
-		if !kontinue {
-			break
-		}
-
-		switch query.Type {
-		case data.Skip, data.Error, data.VExplain:
-			skip = true
-		case data.Unknown:
-			return fmt.Errorf("unknown command type: %s", query.Type)
-		case data.Comment, data.CommentWithCommand, data.EmptyLine, data.WaitForAuthoritative, data.SkipIfBelowVersion:
-			// no-op for keys
-		case data.QueryT:
-			if skip {
-				skip = false
-				continue
-			}
-			process(query, si, ql)
-		}
-	}
+	_ = data.ForeachSQLQuery(loader, func(q data.Query) error {
+		process(q, si, ql)
+		return nil
+	})
 
 	closeErr := loader.Close()
 	jsonWriteErr := ql.writeJSONTo(out)