Add a new implementation for handling multiple queries without needing to split them

go/cmd/vtgateclienttest/services/callerid.go

@@ -29,6 +29,7 @@ import (
 	querypb "vitess.io/vitess/go/vt/proto/query"
 	vtgatepb "vitess.io/vitess/go/vt/proto/vtgate"
 	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
+	"vitess.io/vitess/go/vt/sqlparser"
 	"vitess.io/vitess/go/vt/vtgate/vtgateservice"
 )
 
@@ -104,3 +105,40 @@ func (c *callerIDClient) StreamExecute(ctx context.Context, mysqlCtx vtgateservi
 	}
 	return c.fallbackClient.StreamExecute(ctx, mysqlCtx, session, sql, bindVariables, callback)
 }
+
+// ExecuteMulti is part of the VTGateService interface
+func (c *callerIDClient) ExecuteMulti(ctx context.Context, mysqlCtx vtgateservice.MySQLConnection, session *vtgatepb.Session, sqlString string) (newSession *vtgatepb.Session, qrs []*sqltypes.Result, err error) {
+	queries, err := sqlparser.NewTestParser().SplitStatementToPieces(sqlString)
+	if err != nil {
+		return session, nil, err
+	}
+	var result *sqltypes.Result
+	for _, query := range queries {
+		session, result, err = c.Execute(ctx, mysqlCtx, session, query, nil, false)
+		if err != nil {
+			return session, qrs, err
+		}
+		qrs = append(qrs, result)
+	}
+	return session, qrs, nil
+}
+
+// StreamExecuteMulti is part of the VTGateService interface
+func (c *callerIDClient) StreamExecuteMulti(ctx context.Context, mysqlCtx vtgateservice.MySQLConnection, session *vtgatepb.Session, sqlString string, callback func(qr sqltypes.QueryResponse, more bool, firstPacket bool) error) (*vtgatepb.Session, error) {
+	queries, err := sqlparser.NewTestParser().SplitStatementToPieces(sqlString)
+	if err != nil {
+		return session, err
+	}
+	for idx, query := range queries {
+		firstPacket := true
+		session, err = c.StreamExecute(ctx, mysqlCtx, session, query, nil, func(result *sqltypes.Result) error {
+			err = callback(sqltypes.QueryResponse{QueryResult: result}, idx < len(queries)-1, firstPacket)
+			firstPacket = false
+			return err
+		})
+		if err != nil {
+			return session, err
+		}
+	}
+	return session, nil
+}

go/cmd/vtgateclienttest/services/echo.go

@@ -27,6 +27,7 @@ import (
 	"vitess.io/vitess/go/mysql/collations"
 	"vitess.io/vitess/go/sqltypes"
 	"vitess.io/vitess/go/vt/callerid"
+	"vitess.io/vitess/go/vt/sqlparser"
 	"vitess.io/vitess/go/vt/vtgate/vtgateservice"
 
 	binlogdatapb "vitess.io/vitess/go/vt/proto/binlogdata"
@@ -130,6 +131,43 @@ func (c *echoClient) StreamExecute(ctx context.Context, mysqlCtx vtgateservice.M
 	return c.fallbackClient.StreamExecute(ctx, mysqlCtx, session, sql, bindVariables, callback)
 }
 
+// ExecuteMulti is part of the VTGateService interface
+func (c *echoClient) ExecuteMulti(ctx context.Context, mysqlCtx vtgateservice.MySQLConnection, session *vtgatepb.Session, sqlString string) (newSession *vtgatepb.Session, qrs []*sqltypes.Result, err error) {
+	queries, err := sqlparser.NewTestParser().SplitStatementToPieces(sqlString)
+	if err != nil {
+		return session, nil, err
+	}
+	var result *sqltypes.Result
+	for _, query := range queries {
+		session, result, err = c.Execute(ctx, mysqlCtx, session, query, nil, false)
+		if err != nil {
+			return session, qrs, err
+		}
+		qrs = append(qrs, result)
+	}
+	return session, qrs, nil
+}
+
+// StreamExecuteMulti is part of the VTGateService interface
+func (c *echoClient) StreamExecuteMulti(ctx context.Context, mysqlCtx vtgateservice.MySQLConnection, session *vtgatepb.Session, sqlString string, callback func(qr sqltypes.QueryResponse, more bool, firstPacket bool) error) (*vtgatepb.Session, error) {
+	queries, err := sqlparser.NewTestParser().SplitStatementToPieces(sqlString)
+	if err != nil {
+		return session, err
+	}
+	for idx, query := range queries {
+		firstPacket := true
+		session, err = c.StreamExecute(ctx, mysqlCtx, session, query, nil, func(result *sqltypes.Result) error {
+			err = callback(sqltypes.QueryResponse{QueryResult: result}, idx < len(queries)-1, firstPacket)
+			firstPacket = false
+			return err
+		})
+		if err != nil {
+			return session, err
+		}
+	}
+	return session, nil
+}
+
 func (c *echoClient) ExecuteBatch(ctx context.Context, session *vtgatepb.Session, sqlList []string, bindVariablesList []map[string]*querypb.BindVariable) (*vtgatepb.Session, []sqltypes.QueryResponse, error) {
 	if len(sqlList) > 0 && strings.HasPrefix(sqlList[0], EchoPrefix) {
 		var queryResponse []sqltypes.QueryResponse

go/cmd/vtgateclienttest/services/errors.go

@@ -21,6 +21,7 @@ import (
 	"strings"
 
 	"vitess.io/vitess/go/sqltypes"
+	"vitess.io/vitess/go/vt/sqlparser"
 	"vitess.io/vitess/go/vt/vterrors"
 	"vitess.io/vitess/go/vt/vtgate/vtgateservice"
 
@@ -146,6 +147,43 @@ func (c *errorClient) StreamExecute(ctx context.Context, mysqlCtx vtgateservice.
 	return c.fallbackClient.StreamExecute(ctx, mysqlCtx, session, sql, bindVariables, callback)
 }
 
+// ExecuteMulti is part of the VTGateService interface
+func (c *errorClient) ExecuteMulti(ctx context.Context, mysqlCtx vtgateservice.MySQLConnection, session *vtgatepb.Session, sqlString string) (newSession *vtgatepb.Session, qrs []*sqltypes.Result, err error) {
+	queries, err := sqlparser.NewTestParser().SplitStatementToPieces(sqlString)
+	if err != nil {
+		return session, nil, err
+	}
+	var result *sqltypes.Result
+	for _, query := range queries {
+		session, result, err = c.Execute(ctx, mysqlCtx, session, query, nil, false)
+		if err != nil {
+			return session, qrs, err
+		}
+		qrs = append(qrs, result)
+	}
+	return session, qrs, nil
+}
+
+// StreamExecuteMulti is part of the VTGateService interface
+func (c *errorClient) StreamExecuteMulti(ctx context.Context, mysqlCtx vtgateservice.MySQLConnection, session *vtgatepb.Session, sqlString string, callback func(qr sqltypes.QueryResponse, more bool, firstPacket bool) error) (*vtgatepb.Session, error) {
+	queries, err := sqlparser.NewTestParser().SplitStatementToPieces(sqlString)
+	if err != nil {
+		return session, err
+	}
+	for idx, query := range queries {
+		firstPacket := true
+		session, err = c.StreamExecute(ctx, mysqlCtx, session, query, nil, func(result *sqltypes.Result) error {
+			err = callback(sqltypes.QueryResponse{QueryResult: result}, idx < len(queries)-1, firstPacket)
+			firstPacket = false
+			return err
+		})
+		if err != nil {
+			return session, err
+		}
+	}
+	return session, nil
+}
+
 func (c *errorClient) Prepare(ctx context.Context, session *vtgatepb.Session, sql string) (*vtgatepb.Session, []*querypb.Field, uint16, error) {
 	if err := requestToPartialError(sql, session); err != nil {
 		return session, nil, 0, err

go/cmd/vtgateclienttest/services/fallback.go

@@ -59,6 +59,14 @@ func (c fallbackClient) StreamExecute(ctx context.Context, mysqlCtx vtgateservic
 	return c.fallback.StreamExecute(ctx, mysqlCtx, session, sql, bindVariables, callback)
 }
 
+func (c fallbackClient) ExecuteMulti(ctx context.Context, mysqlCtx vtgateservice.MySQLConnection, session *vtgatepb.Session, sqlString string) (newSession *vtgatepb.Session, qrs []*sqltypes.Result, err error) {
+	return c.fallback.ExecuteMulti(ctx, mysqlCtx, session, sqlString)
+}
+
+func (c fallbackClient) StreamExecuteMulti(ctx context.Context, mysqlCtx vtgateservice.MySQLConnection, session *vtgatepb.Session, sqlString string, callback func(qr sqltypes.QueryResponse, more bool, firstPacket bool) error) (*vtgatepb.Session, error) {
+	return c.fallback.StreamExecuteMulti(ctx, mysqlCtx, session, sqlString, callback)
+}
+
 func (c fallbackClient) Prepare(ctx context.Context, session *vtgatepb.Session, sql string) (*vtgatepb.Session, []*querypb.Field, uint16, error) {
 	return c.fallback.Prepare(ctx, session, sql)
 }

go/cmd/vtgateclienttest/services/terminal.go

@@ -74,6 +74,14 @@ func (c *terminalClient) Prepare(ctx context.Context, session *vtgatepb.Session,
 	return session, nil, 0, errTerminal
 }
 
+func (c *terminalClient) ExecuteMulti(ctx context.Context, mysqlCtx vtgateservice.MySQLConnection, session *vtgatepb.Session, sqlString string) (newSession *vtgatepb.Session, qrs []*sqltypes.Result, err error) {
+	return session, nil, errTerminal
+}
+
+func (c *terminalClient) StreamExecuteMulti(ctx context.Context, mysqlCtx vtgateservice.MySQLConnection, session *vtgatepb.Session, sqlString string, callback func(qr sqltypes.QueryResponse, more bool, firstPacket bool) error) (*vtgatepb.Session, error) {
+	return session, errTerminal
+}
+
 func (c *terminalClient) CloseSession(ctx context.Context, session *vtgatepb.Session) error {
 	return errTerminal
 }

go/flags/endtoend/vtcombo.txt

@@ -230,6 +230,7 @@ Flags:
       --mysql_allow_clear_text_without_tls                               If set, the server will allow the use of a clear text password over non-SSL connections.
       --mysql_auth_server_impl string                                    Which auth server implementation to use. Options: none, ldap, clientcert, static, vault. (default "static")
       --mysql_default_workload string                                    Default session workload (OLTP, OLAP, DBA) (default "OLTP")
+      --mysql_multi_query_protocol                                       If set, the server will use the new implementation of handling queries where-in multiple queries are sent together.
       --mysql_port int                                                   mysql port (default 3306)
       --mysql_server_bind_address string                                 Binds on this address when listening to MySQL binary protocol. Useful to restrict listening to 'localhost' only for instance.
       --mysql_server_flush_delay duration                                Delay after which buffered response will be flushed to the client. (default 100ms)

go/flags/endtoend/vtgate.txt

@@ -143,6 +143,7 @@ Flags:
       --mysql_ldap_auth_config_file string                               JSON File from which to read LDAP server config.
       --mysql_ldap_auth_config_string string                             JSON representation of LDAP server config.
       --mysql_ldap_auth_method string                                    client-side authentication method to use. Supported values: mysql_clear_password, dialog. (default "mysql_clear_password")
+      --mysql_multi_query_protocol                                       If set, the server will use the new implementation of handling queries where-in multiple queries are sent together.
       --mysql_server_bind_address string                                 Binds on this address when listening to MySQL binary protocol. Useful to restrict listening to 'localhost' only for instance.
       --mysql_server_flush_delay duration                                Delay after which buffered response will be flushed to the client. (default 100ms)
       --mysql_server_port int                                            If set, also listen for MySQL binary protocol connections on this port. (default -1)

go/mysql/conn.go

@@ -30,6 +30,8 @@ import (
 	"sync/atomic"
 	"time"
 
+	"github.com/spf13/pflag"
+
 	"vitess.io/vitess/go/bucketpool"
 	"vitess.io/vitess/go/mysql/collations"
 	"vitess.io/vitess/go/mysql/sqlerror"
@@ -38,6 +40,7 @@ import (
 	"vitess.io/vitess/go/vt/log"
 	querypb "vitess.io/vitess/go/vt/proto/query"
 	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
+	"vitess.io/vitess/go/vt/servenv"
 	"vitess.io/vitess/go/vt/vterrors"
 )
 
@@ -67,6 +70,19 @@ const (
 	ephemeralRead
 )
 
+var (
+	mysqlMultiQuery = false
+)
+
+func registerConnFlags(fs *pflag.FlagSet) {
+	fs.BoolVar(&mysqlMultiQuery, "mysql_multi_query_protocol", mysqlMultiQuery, "If set, the server will use the new implementation of handling queries where-in multiple queries are sent together.")
+}
+
+func init() {
+	servenv.OnParseFor("vtgate", registerConnFlags)
+	servenv.OnParseFor("vtcombo", registerConnFlags)
+}
+
 // A Getter has a Get()
 type Getter interface {
 	Get() *querypb.VTGateCallerID
@@ -914,6 +930,9 @@ func (c *Conn) handleNextCommand(handler Handler) bool {
 		res := c.execQuery("use "+sqlescape.EscapeID(db), handler, false)
 		return res != connErr
 	case ComQuery:
+		if mysqlMultiQuery {
+			return c.handleComQueryMulti(handler, data)
+		}
 		return c.handleComQuery(handler, data)
 	case ComPing:
 		return c.handleComPing()
@@ -1279,6 +1298,124 @@ func (c *Conn) handleComPing() bool {
 	return true
 }
 
+// handleComQueryMulti is a newer version of handleComQuery that uses
+// the StreamExecuteMulti and ExecuteMulti RPC calls to push the splitting of statements
+// down to Vtgate.
+func (c *Conn) handleComQueryMulti(handler Handler, data []byte) (kontinue bool) {
+	c.startWriterBuffering()
+	defer func() {
+		if err := c.endWriterBuffering(); err != nil {
+			log.Errorf("conn %v: flush() failed: %v", c.ID(), err)
+			kontinue = false
+		}
+	}()
+
+	queryStart := time.Now()
+	query := c.parseComQuery(data)
+	c.recycleReadPacket()
+
+	res := c.execQueryMulti(query, handler)
+	if res != execSuccess {
+		return res != connErr
+	}
+
+	timings.Record(queryTimingKey, queryStart)
+	return true
+}
+
+// execQueryMulti is a newer version of execQuery that uses
+// the StreamExecuteMulti and ExecuteMulti RPC calls to push the splitting of statements
+// down to Vtgate.
+func (c *Conn) execQueryMulti(query string, handler Handler) execResult {
+	// needsEndPacket signifies whether we have need to send the last packet to the client
+	// for a given query. This is used to determine whether we should send an
+	// end packet after the query is done or not. Initially we don't need to send an end packet
+	// so we initialize this value to false.
+	needsEndPacket := false
+	var res = execSuccess
+
+	err := handler.ComQueryMulti(c, query, func(qr sqltypes.QueryResponse, more bool, firstPacket bool) error {
+		flag := c.StatusFlags
+		if more {
+			flag |= ServerMoreResultsExists
+		}
+
+		// firstPacket tells us that this is the start of a new query result.
+		// If we haven't sent a last packet yet, we should send the end result packet.
+		if firstPacket && needsEndPacket {
+			if err := c.writeEndResult(true, 0, 0, handler.WarningCount(c)); err != nil {
+				log.Errorf("Error writing result to %s: %v", c, err)
+				return err
+			}
+		}
+
+		// We receive execution errors in a query as part of the QueryResponse.
+		// We check for those errors and send a error packet. If we are unable
+		// to send the error packet, then there is a connection error too.
+		if qr.QueryError != nil {
+			res = execErr
+			if !c.writeErrorPacketFromErrorAndLog(qr.QueryError) {
+				res = connErr
+			}
+			return nil
+		}
+
+		if firstPacket {
+			// The first packet signifies the start of a new query result.
+			// So we reset the needsEndPacket variable to signify we haven't sent the last
+			// packet for this query.
+			needsEndPacket = true
+			if len(qr.QueryResult.Fields) == 0 {
+
+				// A successful callback with no fields means that this was a
+				// DML or other write-only operation.
+				//
+				// We should not send any more packets after this, but make sure
+				// to extract the affected rows and last insert id from the result
+				// struct here since clients expect it.
+				ok := PacketOK{
+					affectedRows:     qr.QueryResult.RowsAffected,
+					lastInsertID:     qr.QueryResult.InsertID,
+					statusFlags:      flag,
+					warnings:         handler.WarningCount(c),
+					info:             "",
+					sessionStateData: qr.QueryResult.SessionStateChanges,
+				}
+				needsEndPacket = false
+				return c.writeOKPacket(&ok)
+			}
+
+			if err := c.writeFields(qr.QueryResult); err != nil {
+				return err
+			}
+		}
+
+		return c.writeRows(qr.QueryResult)
+	})
+
+	if res != execSuccess {
+		// We failed during the stream itself.
+		return res
+	}
+
+	if err != nil {
+		// We can't send an error in the middle of a stream.
+		// All we can do is abort the send, which will cause a 2013.
+		log.Errorf("Error in the middle of a stream to %s: %v", c, err)
+		return connErr
+	}
+
+	// If we haven't sent the final packet for the last query, we should send that too.
+	if needsEndPacket {
+		if err := c.writeEndResult(false, 0, 0, handler.WarningCount(c)); err != nil {
+			log.Errorf("Error writing result to %s: %v", c, err)
+			return connErr
+		}
+	}
+
+	return execSuccess
+}
+
 var errEmptyStatement = sqlerror.NewSQLError(sqlerror.EREmptyQuery, sqlerror.SSClientError, "Query was empty")
 
 func (c *Conn) handleComQuery(handler Handler, data []byte) (kontinue bool) {

go/mysql/conn_test.go

@@ -1127,6 +1127,25 @@ func (t testRun) ComQuery(c *Conn, query string, callback func(*sqltypes.Result)
 	return nil
 }
 
+func (t testRun) ComQueryMulti(c *Conn, sql string, callback func(qr sqltypes.QueryResponse, more bool, firstPacket bool) error) error {
+	queries, err := t.Env().Parser().SplitStatementToPieces(sql)
+	if err != nil {
+		return err
+	}
+	for i, query := range queries {
+		firstPacket := true
+		err = t.ComQuery(c, query, func(result *sqltypes.Result) error {
+			err = callback(sqltypes.QueryResponse{QueryResult: result}, i < len(queries)-1, firstPacket)
+			firstPacket = false
+			return err
+		})
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
 func (t testRun) ComPrepare(c *Conn, query string) ([]*querypb.Field, uint16, error) {
 	return nil, t.paramCounts, nil
 }