basic pipelining PREPARE + EXECUTE when possible

Author: rusher

connection.go

@@ -208,6 +208,17 @@ func (mc *mysqlConn) Prepare(query string) (driver.Stmt, error) {
 	if mc.closed.Load() {
 		return nil, driver.ErrBadConn
 	}
+	if mc.clientExtCapabilities&clientStmtBulkOperations != 0 {
+		// can delay PREPARE
+		stmt := &mysqlStmt{
+			mc:           mc,
+			id:           0xffffffff,
+			paramCount:   -1,
+			initialQuery: query,
+		}
+		return stmt, nil
+	}
+
 	// Send command
 	err := mc.writeCommandPacketStr(comStmtPrepare, query)
 	if err != nil {
@@ -221,28 +232,32 @@ func (mc *mysqlConn) Prepare(query string) (driver.Stmt, error) {
 	}
 
 	// Read Result
+	err = mc.readPrepareResult(stmt)
+	return stmt, err
+}
+
+func (mc *mysqlConn) readPrepareResult(stmt *mysqlStmt) error {
 	columnCount, err := stmt.readPrepareResultPacket()
 	if err == nil {
 		if stmt.paramCount > 0 {
 			if err = mc.skipColumns(stmt.paramCount); err != nil {
-				return nil, err
+				return err
 			}
 		}
 
 		if columnCount > 0 {
 			if mc.clientExtCapabilities&clientCacheMetadata != 0 {
 				if stmt.columns, err = mc.readColumns(int(columnCount)); err != nil {
-					return nil, err
+					return err
 				}
 			} else {
 				if err = mc.skipColumns(int(columnCount)); err != nil {
-					return nil, err
+					return err
 				}
 			}
 		}
 	}
-
-	return stmt, err
+	return err
 }
 
 func (mc *mysqlConn) interpolateParams(query string, args []driver.Value) (string, error) {
@@ -256,7 +271,7 @@ func (mc *mysqlConn) interpolateParams(query string, args []driver.Value) (strin
 		// can not take the buffer. Something must be wrong with the connection
 		mc.cleanup()
 		// interpolateParams would be called before sending any query.
-		// So its safe to retry.
+		// So it's safe to retry.
 		return "", driver.ErrBadConn
 	}
 	buf = buf[:0]

packets.go

@@ -310,7 +310,7 @@ func (mc *mysqlConn) initClientCapabilities(serverCapabilities capabilityFlag, c
 func (mc *mysqlConn) writeHandshakeResponsePacket(authResp []byte, serverCapabilities capabilityFlag, serverExtendedCapabilities extendedCapabilityFlag, plugin string) error {
 	// Adjust client flags based on server support
 	mc.clientCapabilities = mc.initClientCapabilities(serverCapabilities, mc.cfg)
-	mc.clientExtCapabilities = clientCacheMetadata & serverExtendedCapabilities
+	mc.clientExtCapabilities = (clientCacheMetadata | clientStmtBulkOperations) & serverExtendedCapabilities
 
 	sendConnectAttrs := mc.clientCapabilities&clientConnectAttrs != 0
 
@@ -1016,7 +1016,7 @@ func (stmt *mysqlStmt) writeCommandLongData(paramID int, arg []byte) error {
 // Execute Prepared Statement
 // http://dev.mysql.com/doc/internals/en/com-stmt-execute.html
 func (stmt *mysqlStmt) writeExecutePacket(args []driver.Value) error {
-	if len(args) != stmt.paramCount {
+	if stmt.paramCount != -1 && len(args) != stmt.paramCount {
 		return fmt.Errorf(
 			"argument count mismatch (got: %d; has: %d)",
 			len(args),
@@ -1028,7 +1028,7 @@ func (stmt *mysqlStmt) writeExecutePacket(args []driver.Value) error {
 	mc := stmt.mc
 
 	// Determine threshold dynamically to avoid packet size shortage.
-	longDataSize := mc.maxAllowedPacket / (stmt.paramCount + 1)
+	longDataSize := mc.maxAllowedPacket / (len(args) + 1)
 	if longDataSize < 64 {
 		longDataSize = 64
 	}
@@ -1066,7 +1066,7 @@ func (stmt *mysqlStmt) writeExecutePacket(args []driver.Value) error {
 
 		var nullMask []byte
 		if maskLen, typesLen := (len(args)+7)/8, 1+2*len(args); pos+maskLen+typesLen >= cap(data) {
-			// buffer has to be extended but we don't know by how much so
+			// buffer has to be extended, but we don't know by how much so
 			// we depend on append after all data with known sizes fit.
 			// We stop at that because we deal with a lot of columns here
 			// which makes the required allocation size hard to guess.
@@ -1118,7 +1118,7 @@ func (stmt *mysqlStmt) writeExecutePacket(args []driver.Value) error {
 			case uint64:
 				paramTypes[i+i] = byte(fieldTypeLongLong)
 				paramTypes[i+i+1] = 0x80 // type is unsigned
-				paramValues = binary.LittleEndian.AppendUint64(paramValues, uint64(v))
+				paramValues = binary.LittleEndian.AppendUint64(paramValues, v)
 
 			case float64:
 				paramTypes[i+i] = byte(fieldTypeDouble)

statement.go

@@ -17,10 +17,11 @@ import (
 )
 
 type mysqlStmt struct {
-	mc         *mysqlConn
-	id         uint32
-	paramCount int
-	columns    []mysqlField
+	mc           *mysqlConn
+	id           uint32
+	paramCount   int
+	columns      []mysqlField
+	initialQuery string
 }
 
 func (stmt *mysqlStmt) Close() error {
@@ -52,16 +53,53 @@ func (stmt *mysqlStmt) CheckNamedValue(nv *driver.NamedValue) (err error) {
 }
 
 func (stmt *mysqlStmt) Exec(args []driver.Value) (driver.Result, error) {
-	if stmt.mc.closed.Load() {
+	mc := stmt.mc
+	if mc.closed.Load() {
 		return nil, driver.ErrBadConn
 	}
+
+	var prepareSequence uint8
+	var prepareCompressSequence uint8
+
+	if mc.clientExtCapabilities&clientStmtBulkOperations != 0 {
+		// Send command
+		err := mc.writeCommandPacketStr(comStmtPrepare, stmt.initialQuery)
+		if err != nil {
+			mc.log(err)
+			return nil, driver.ErrBadConn
+		}
+		prepareSequence = mc.sequence
+		prepareCompressSequence = mc.compressSequence
+	}
+
 	// Send command
 	err := stmt.writeExecutePacket(args)
 	if err != nil {
 		return nil, stmt.mc.markBadConn(err)
 	}
 
-	mc := stmt.mc
+	if stmt.mc.clientExtCapabilities&clientStmtBulkOperations != 0 {
+		// Read Prepare Result
+		var executeSequence uint8
+		var executeCompressSequence uint8
+		executeSequence = mc.sequence
+		executeCompressSequence = mc.compressSequence
+
+		mc.sequence = prepareSequence
+		mc.compressSequence = prepareCompressSequence
+
+		err = mc.readPrepareResult(stmt)
+
+		mc.sequence = executeSequence
+		mc.compressSequence = executeCompressSequence
+		if err != nil {
+			// skip executeResult (will return an error)
+			handleOk := stmt.mc.clearResult()
+			_, _, _ = handleOk.readResultSetHeaderPacket()
+			return nil, err
+		}
+	}
+
 	handleOk := stmt.mc.clearResult()
 
 	// Read Result
@@ -98,13 +136,48 @@ func (stmt *mysqlStmt) query(args []driver.Value) (*binaryRows, error) {
 	if stmt.mc.closed.Load() {
 		return nil, driver.ErrBadConn
 	}
+
+	mc := stmt.mc
+	var prepareSequence uint8
+	var prepareCompressSequence uint8
+
+	if mc.clientExtCapabilities&clientStmtBulkOperations != 0 {
+		// Send command
+		err := mc.writeCommandPacketStr(comStmtPrepare, stmt.initialQuery)
+		if err != nil {
+			mc.log(err)
+			return nil, driver.ErrBadConn
+		}
+		prepareSequence = mc.sequence
+		prepareCompressSequence = mc.compressSequence
+	}
 	// Send command
 	err := stmt.writeExecutePacket(args)
 	if err != nil {
 		return nil, stmt.mc.markBadConn(err)
 	}
 
-	mc := stmt.mc
+	if stmt.mc.clientExtCapabilities&clientStmtBulkOperations != 0 {
+		// Read Prepare Result
+		var executeSequence uint8
+		var executeCompressSequence uint8
+		executeSequence = mc.sequence
+		executeCompressSequence = mc.compressSequence
+
+		mc.sequence = prepareSequence
+		mc.compressSequence = prepareCompressSequence
+
+		err = mc.readPrepareResult(stmt)
+
+		mc.sequence = executeSequence
+		mc.compressSequence = executeCompressSequence
+		if err != nil {
+			// skip executeResult (will return an error)
+			handleOk := stmt.mc.clearResult()
+			_, _, _ = handleOk.readResultSetHeaderPacket()
+			return nil, err
+		}
+	}
 
 	// Read Result
 	handleOk := stmt.mc.clearResult()