[commands] faster sqldump query scanner (#8291)

* [commands] faster sqldump query scanner * imports seem to work * remove logs * more perf, don't parse twice * [ga-format-pr] Run go/utils/repofmt/format_repo.sh and go/Godeps/update.sh * fix delimiters * delete old code * more comments * simplify a bit * add line numbers back * delimiter line num * picky line nums * more comments * zach comment * more comments --------- Co-authored-by: max-hoffman <[email protected]>
dolthub · Aug 26, 2024 · f592150 · f592150
1 parent 7af1cc6
commit f592150
Show file tree

Hide file tree

Showing 7 changed files with 278 additions and 99 deletions.
diff --git a/go/cmd/dolt/cli/command.go b/go/cmd/dolt/cli/command.go
@@ -23,6 +23,8 @@ import (
 	"syscall"
 
 	"github.com/dolthub/go-mysql-server/sql"
+	querypb "github.com/dolthub/vitess/go/vt/proto/query"
+	"github.com/dolthub/vitess/go/vt/sqlparser"
 	"github.com/fatih/color"
 
 	eventsapi "github.com/dolthub/dolt/go/gen/proto/dolt/services/eventsapi/v1alpha1"
@@ -86,6 +88,7 @@ type SignalCommand interface {
 // SQL. The Queryist can be obtained from the CliContext passed into the Exec method by calling the QueryEngine method.
 type Queryist interface {
 	Query(ctx *sql.Context, query string) (sql.Schema, sql.RowIter, *sql.QueryFlags, error)
+	QueryWithBindings(ctx *sql.Context, query string, parsed sqlparser.Statement, bindings map[string]*querypb.BindVariable, qFlags *sql.QueryFlags) (sql.Schema, sql.RowIter, *sql.QueryFlags, error)
 }
 
 // This type is to store the content of a documented command, elsewhere we can transform this struct into

diff --git a/go/cmd/dolt/commands/engine/sqlengine.go b/go/cmd/dolt/commands/engine/sqlengine.go
@@ -30,6 +30,8 @@ import (
 	"github.com/dolthub/go-mysql-server/sql/mysql_db"
 	"github.com/dolthub/go-mysql-server/sql/rowexec"
 	_ "github.com/dolthub/go-mysql-server/sql/variables"
+	querypb "github.com/dolthub/vitess/go/vt/proto/query"
+	"github.com/dolthub/vitess/go/vt/sqlparser"
 	"github.com/sirupsen/logrus"
 
 	"github.com/dolthub/dolt/go/cmd/dolt/cli"
@@ -310,6 +312,10 @@ func (se *SqlEngine) Query(ctx *sql.Context, query string) (sql.Schema, sql.RowI
 	return se.engine.Query(ctx, query)
 }
 
+func (se *SqlEngine) QueryWithBindings(ctx *sql.Context, query string, parsed sqlparser.Statement, bindings map[string]*querypb.BindVariable, qFlags *sql.QueryFlags) (sql.Schema, sql.RowIter, *sql.QueryFlags, error) {
+	return se.engine.QueryWithBindings(ctx, query, parsed, bindings, qFlags)
+}
+
 // Analyze analyzes a node.
 func (se *SqlEngine) Analyze(ctx *sql.Context, n sql.Node, qFlags *sql.QueryFlags) (sql.Node, error) {
 	return se.engine.Analyzer.Analyze(ctx, n, nil, qFlags)

diff --git a/go/cmd/dolt/commands/filter-branch.go b/go/cmd/dolt/commands/filter-branch.go
@@ -282,7 +282,7 @@ func processFilterQuery(ctx context.Context, dEnv *env.DoltEnv, root doltdb.Root
 		return nil, err
 	}
 
-	scanner := NewSqlStatementScanner(strings.NewReader(query))
+	scanner := newStreamScanner(strings.NewReader(query))
 	if err != nil {
 		return nil, err
 	}

diff --git a/go/cmd/dolt/commands/sql.go b/go/cmd/dolt/commands/sql.go
@@ -612,7 +612,7 @@ func saveQuery(ctx *sql.Context, root doltdb.RootValue, query string, name strin
 
 // execBatchMode runs all the queries in the input reader
 func execBatchMode(ctx *sql.Context, qryist cli.Queryist, input io.Reader, continueOnErr bool, format engine.PrintResultFormat) error {
-	scanner := NewSqlStatementScanner(input)
+	scanner := newStreamScanner(input)
 	var query string
 	for scanner.Scan() {
 		if fileReadProg != nil {
@@ -630,7 +630,7 @@ func execBatchMode(ctx *sql.Context, qryist cli.Queryist, input io.Reader, conti
 		if err == sqlparser.ErrEmpty {
 			continue
 		} else if err != nil {
-			err = buildBatchSqlErr(scanner.statementStartLine, query, err)
+			err = buildBatchSqlErr(scanner.state.statementStartLine, query, err)
 			if !continueOnErr {
 				return err
 			} else {
@@ -642,7 +642,7 @@ func execBatchMode(ctx *sql.Context, qryist cli.Queryist, input io.Reader, conti
 		ctx.SetQueryTime(time.Now())
 		sqlSch, rowIter, _, err := processParsedQuery(ctx, query, qryist, sqlStatement)
 		if err != nil {
-			err = buildBatchSqlErr(scanner.statementStartLine, query, err)
+			err = buildBatchSqlErr(scanner.state.statementStartLine, query, err)
 			if !continueOnErr {
 				return err
 			} else {
@@ -661,7 +661,7 @@ func execBatchMode(ctx *sql.Context, qryist cli.Queryist, input io.Reader, conti
 			}
 			err = engine.PrettyPrintResults(ctx, format, sqlSch, rowIter)
 			if err != nil {
-				err = buildBatchSqlErr(scanner.statementStartLine, query, err)
+				err = buildBatchSqlErr(scanner.state.statementStartLine, query, err)
 				if !continueOnErr {
 					return err
 				} else {
@@ -673,7 +673,7 @@ func execBatchMode(ctx *sql.Context, qryist cli.Queryist, input io.Reader, conti
 	}
 
 	if err := scanner.Err(); err != nil {
-		return buildBatchSqlErr(scanner.statementStartLine, query, err)
+		return buildBatchSqlErr(scanner.state.statementStartLine, query, err)
 	}
 
 	return nil
@@ -1122,7 +1122,7 @@ func processParsedQuery(ctx *sql.Context, query string, qryist cli.Queryist, sql
 		}
 		return qryist.Query(ctx, query)
 	default:
-		return qryist.Query(ctx, query)
+		return qryist.QueryWithBindings(ctx, query, sqlStatement, nil, nil)
 	}
 }