Feature: Add support for multiple result sets through the driver.RowsNextResultSet interface

README.md

@@ -27,7 +27,7 @@ cd dbs
 dolt clone <REMOTE URL>
 ```
 
-Finally you can create the dbs directory as shown above and then create the database in code using a SQL `CREATE TABLE` statement
+Finally, you can create the dbs directory as shown above and then create the database in code using a SQL `CREATE TABLE` statement
 
 ### Connecting to the Database
 
@@ -61,3 +61,31 @@ clientfoundrows - If set to true, returns the number of matching rows instead of
 #### Example DSN
 
 `file:///path/to/dbs?commitname=Your%20Name&[email protected]&database=databasename`
+
+### Multi-Statement Support
+
+If you pass the `multistatements=true` parameter in the DSN, you can execute multiple statements in one query. The returned 
+rows allow you to iterate over the returned result sets by using the `NextResultSet` method, just like you can with the
+MySQL driver. 
+
+```go
+rows, err := db.Query("SELECT * from someTable; SELECT * from anotherTable;")
+// If an error is returned, it means it came from the first statement
+if err != nil {
+	panic(err)
+}
+
+for rows.Next() {
+	// process the first result set
+}
+
+if rows.NextResultSet() {
+    for rows.Next() {
+        // process the second result set
+    }
+} else {
+	// If NextResultSet returns false when there were more statements, it means there was an error,
+	// which you can access through rows.Err()
+	panic(rows.Err())
+}
+```

conn.go

@@ -5,12 +5,11 @@ import (
 	"database/sql"
 	"database/sql/driver"
 	"fmt"
-	"io"
 	"time"
 
 	"github.com/dolthub/dolt/go/cmd/dolt/commands/engine"
-
 	gms "github.com/dolthub/go-mysql-server/sql"
+	"github.com/dolthub/vitess/go/vt/sqlparser"
 )
 
 var _ driver.Conn = (*DoltConn)(nil)
@@ -22,71 +21,54 @@ type DoltConn struct {
 	DataSource *DoltDataSource
 }
 
-// Prepare returns a prepared statement, bound to this connection.
+// Prepare packages up |query| as a *doltStmt so it can be executed. If multistatements mode
+// has been enabled, then a *doltMultiStmt will be returned, capable of executing multiple statements.
 func (d *DoltConn) Prepare(query string) (driver.Stmt, error) {
-	multiStatements := d.DataSource.ParamIsTrue(MultiStatementsParam)
+	// Reuse the same ctx instance, but update the QueryTime to the current time.
+	// Statements are executed serially on a connection, so it's safe to reuse
+	// the same ctx instance and update the time.
+	d.gmsCtx.SetQueryTime(time.Now())
 
-	if multiStatements {
-		scanner := gms.NewMysqlParser()
-		parsed, prequery, remainder, err := scanner.Parse(d.gmsCtx, query, true)
-		if err != nil {
-			return nil, translateError(err)
-		}
+	if d.DataSource.ParamIsTrue(MultiStatementsParam) {
+		return d.prepareMultiStatement(query)
+	} else {
+		return d.prepareSingleStatement(query)
+	}
+}
 
-		for {
-			if len(remainder) == 0 {
-				query = prequery
-				break
-			}
-
-			err = func() error {
-				var rowIter gms.RowIter
-				_, rowIter, err = d.se.GetUnderlyingEngine().QueryWithBindings(d.gmsCtx, prequery, parsed, nil)
-				if err != nil {
-					return translateError(err)
-				}
-				defer rowIter.Close(d.gmsCtx)
-
-				for {
-					_, err := rowIter.Next(d.gmsCtx)
-					if err == io.EOF {
-						break
-					} else if err != nil {
-						return translateError(err)
-					}
-				}
-
-				return nil
-			}()
-			if err != nil {
-				return nil, err
-			}
-
-			parsed, prequery, remainder, err = scanner.Parse(d.gmsCtx, remainder, true)
-			if err != nil {
-				return nil, translateError(err)
-			}
-		}
-		if prequery != "" {
-			query = prequery
+// prepareSingleStatement creates a doltStmt from |query|.
+func (d *DoltConn) prepareSingleStatement(query string) (*doltStmt, error) {
+	return &doltStmt{
+		query:  query,
+		se:     d.se,
+		gmsCtx: d.gmsCtx,
+	}, nil
+}
+
+// prepareMultiStatement creates a doltStmt from each individual statement in |query|.
+func (d *DoltConn) prepareMultiStatement(query string) (*doltMultiStmt, error) {
+	var doltMultiStmt doltMultiStmt
+	scanner := gms.NewMysqlParser()
+
+	remainder := query
+	var err error
+	for remainder != "" {
+		_, query, remainder, err = scanner.Parse(d.gmsCtx, remainder, true)
+		if err == sqlparser.ErrEmpty {
+			// Skip over any empty statements
+			continue
+		} else if err != nil {
+			return nil, translateError(err)
 		}
-	}
 
-	if len(query) > 0 {
-		_, err := d.se.GetUnderlyingEngine().PrepareQuery(d.gmsCtx, query)
+		doltStmt, err := d.prepareSingleStatement(query)
 		if err != nil {
 			return nil, translateError(err)
 		}
+		doltMultiStmt.stmts = append(doltMultiStmt.stmts, doltStmt)
 	}
 
-	// Reuse the same ctx instance, but update the QueryTime to the current time. Since statements are
-	// executed serially on a connection, it's safe to reuse the same ctx instance and update the time.
-	d.gmsCtx.SetQueryTime(time.Now())
-	return &doltStmt{
-		query:  query,
-		se:     d.se,
-		gmsCtx: d.gmsCtx,
-	}, nil
+	return &doltMultiStmt, nil
 }
 
 // Close releases the resources held by the DoltConn instance

driver.go

@@ -42,7 +42,6 @@ type doltDriver struct {
 //
 // The path needs to point to a directory whose subdirectories are dolt databases.  If a "Create Database" command is
 // run a new subdirectory will be created in this path.
-// The supported parameters are
 func (d *doltDriver) Open(dataSource string) (driver.Conn, error) {
 	ctx := context.Background()
 	var fs filesys.Filesys = filesys.LocalFS
@@ -89,7 +88,7 @@ func (d *doltDriver) Open(dataSource string) (driver.Conn, error) {
 		ServerUser: "root",
 		Autocommit: true,
 	}
-	
+
 	se, err := engine.NewSqlEngine(ctx, mrEnv, seCfg)
 	if err != nil {
 		return nil, err
@@ -122,16 +121,16 @@ func (d *doltDriver) Open(dataSource string) (driver.Conn, error) {
 // with initialized environments for each of those subfolder data repositories. subfolders whose name starts with '.' are
 // skipped.
 func LoadMultiEnvFromDir(
-		ctx context.Context,
-		cfg config.ReadWriteConfig,
-		fs filesys.Filesys,
-		path, version string,
+	ctx context.Context,
+	cfg config.ReadWriteConfig,
+	fs filesys.Filesys,
+	path, version string,
 ) (*env.MultiRepoEnv, error) {
 
 	multiDbDirFs, err := fs.WithWorkingDir(path)
 	if err != nil {
 		return nil, errhand.VerboseErrorFromError(err)
 	}
-	
+
 	return env.MultiEnvForDirectory(ctx, cfg, multiDbDirFs, version, nil)
 }

example/main.go

@@ -39,16 +39,16 @@ func main() {
 	db, err := sql.Open("dolt", dataSource)
 	errExit("failed to open database using the dolt driver: %w", err)
 
-	err = printQuery(ctx, db, "CREATE DATABASE IF NOT EXISTS testdb;USE testdb;")
+	err = printQuery(ctx, db, "CREATE DATABASE IF NOT EXISTS testdb; USE testdb;")
 	errExit("", err)
 
 	err = printQuery(ctx, db, "USE testdb;")
 	errExit("", err)
 
-	printQuery(ctx, db, `CREATE TABLE IF NOT EXISTS t2(
-    pk int primary key auto_increment,
-    c1 varchar(32)
-)`)
+	err = printQuery(ctx, db, `CREATE TABLE IF NOT EXISTS t2(
+		pk int primary key auto_increment,
+		c1 varchar(32)
+	)`)
 	errExit("", err)
 
 	printQuery(ctx, db, "SHOW TABLES;")
@@ -63,21 +63,22 @@ func main() {
 	fmt.Println(result.LastInsertId())
 
 	err = printQuery(ctx, db, `CREATE TABLE IF NOT EXISTS t1 (
-	pk int PRIMARY KEY,
-	c1 varchar(512),
-	c2 float,
-	c3 bool,
-	c4 datetime
-);`)
+		pk int PRIMARY KEY,
+		c1 varchar(512),
+		c2 float,
+		c3 bool,
+		c4 datetime
+	);`)
+	errExit("", err)
 
 	err = printQuery(ctx, db, "SELECT * FROM t1;")
 	errExit("", err)
 
 	err = printQuery(ctx, db, `REPLACE INTO t1 VALUES 
-(1, 'this is a test', 0, 0, '1998-01-23 12:45:56'),
-(2, 'it is only a test', 1.0, 1, '2010-12-31 01:15:00'),
-(3, NULL, 3.335, 0, NULL),
-(4, 'something something', 3.5, 1, '2015-04-03 14:00:45');`)
+		(1, 'this is a test', 0, 0, '1998-01-23 12:45:56'),
+		(2, 'it is only a test', 1.0, 1, '2010-12-31 01:15:00'),
+		(3, NULL, 3.335, 0, NULL),
+		(4, 'something something', 3.5, 1, '2015-04-03 14:00:45');`)
 	errExit("", err)
 
 	err = printQuery(ctx, db, "SELECT * FROM t1;")

go.mod

@@ -5,9 +5,9 @@ go 1.22.2
 toolchain go1.22.3
 
 require (
-	github.com/dolthub/dolt/go v0.40.5-0.20240604165632-02f450318cb3
-	github.com/dolthub/go-mysql-server v0.18.2-0.20240604161217-d1dca79a32b8
-	github.com/dolthub/vitess v0.0.0-20240603172811-467efd832e48
+	github.com/dolthub/dolt/go v0.40.5-0.20240702155756-bcf4dd5f5cc1
+	github.com/dolthub/go-mysql-server v0.18.2-0.20240702022058-d7eb602c04ee
+	github.com/dolthub/vitess v0.0.0-20240709194214-7926ea9d425d
 	github.com/go-sql-driver/mysql v1.7.2-0.20231213112541-0004702b931d
 	github.com/stretchr/testify v1.8.4
 	gorm.io/driver/mysql v1.5.6
@@ -71,6 +71,7 @@ require (
 	github.com/mattn/go-colorable v0.1.13 // indirect
 	github.com/mattn/go-isatty v0.0.20 // indirect
 	github.com/mattn/go-runewidth v0.0.15 // indirect
+	github.com/mohae/uvarint v0.0.0-20160208145430-c3f9e62bf2b0 // indirect
 	github.com/oracle/oci-go-sdk/v65 v65.55.0 // indirect
 	github.com/pierrec/lz4/v4 v4.1.21 // indirect
 	github.com/pkg/errors v0.9.1 // indirect
@@ -94,17 +95,17 @@ require (
 	go.opentelemetry.io/otel/trace v1.23.1 // indirect
 	go.uber.org/multierr v1.11.0 // indirect
 	go.uber.org/zap v1.26.0 // indirect
-	golang.org/x/crypto v0.21.0 // indirect
+	golang.org/x/crypto v0.23.0 // indirect
 	golang.org/x/exp v0.0.0-20240205201215-2c58cdc269a3 // indirect
-	golang.org/x/mod v0.15.0 // indirect
-	golang.org/x/net v0.23.0 // indirect
+	golang.org/x/mod v0.17.0 // indirect
+	golang.org/x/net v0.25.0 // indirect
 	golang.org/x/oauth2 v0.17.0 // indirect
-	golang.org/x/sync v0.6.0 // indirect
-	golang.org/x/sys v0.18.0 // indirect
-	golang.org/x/term v0.18.0 // indirect
-	golang.org/x/text v0.14.0 // indirect
+	golang.org/x/sync v0.7.0 // indirect
+	golang.org/x/sys v0.20.0 // indirect
+	golang.org/x/term v0.20.0 // indirect
+	golang.org/x/text v0.16.0 // indirect
 	golang.org/x/time v0.5.0 // indirect
-	golang.org/x/tools v0.18.0 // indirect
+	golang.org/x/tools v0.21.1-0.20240508182429-e35e4ccd0d2d // indirect
 	golang.org/x/xerrors v0.0.0-20231012003039-104605ab7028 // indirect
 	google.golang.org/api v0.164.0 // indirect
 	google.golang.org/appengine v1.6.8 // indirect