Exec() now provides access to status of multiple statements. (#1309)

It now reports the last inserted ID and affected row count for all statements, not just the last one. This is useful to execute batches of statements such as UPDATE with minimal roundtrips. Co-authored-by: Inada Naoki <[email protected]>

Author: mherr-google

README.md

@@ -305,6 +305,22 @@ Allow multiple statements in one query. This can be used to bach multiple querie
 
 When `multiStatements` is used, `?` parameters must only be used in the first statement. [interpolateParams](#interpolateparams) can be used to avoid this limitation unless prepared statement is used explicitly.
 
+It's possible to access the last inserted ID and number of affected rows for multiple statements by using `sql.Conn.Raw()` and the `mysql.Result`. For example:
+
+```go
+conn, _:= db.Conn(ctx)
+conn.Raw(func(conn interface{}) error{
+ex:= conn.(driver.Execer)
+res, err:= ex.Exec(`
+ UPDATE point SET x = 1 WHERE y = 2;
+ UPDATE point SET x = 2 WHERE y = 3;
+`, nil)
+// Both slices have 2 elements.
+ log.Print(res.(mysql.Result).AllRowsAffected())
+ log.Print(res.(mysql.Result).AllLastInsertIds())
+})
+```
+
 ##### `parseTime`
 
 ```

auth.go

@@ -346,7 +346,7 @@ func (mc *mysqlConn) handleAuthResult(oldAuthData []byte, plugin string) error{
 case1:
 switchauthData[0]{
 casecachingSha2PasswordFastAuthSuccess:
-iferr=mc.readResultOK(); err==nil{
+iferr=mc.resultUnchanged().readResultOK(); err==nil{
 returnnil// auth successful
  }
 
@@ -397,7 +397,7 @@ func (mc *mysqlConn) handleAuthResult(oldAuthData []byte, plugin string) error{
 returnerr
  }
  }
-returnmc.readResultOK()
+returnmc.resultUnchanged().readResultOK()
 
 default:
 returnErrMalformPkt
@@ -426,7 +426,7 @@ func (mc *mysqlConn) handleAuthResult(oldAuthData []byte, plugin string) error{
 iferr!=nil{
 returnerr
  }
-returnmc.readResultOK()
+returnmc.resultUnchanged().readResultOK()
  }
 
 default:

connection.go

@@ -23,9 +23,8 @@ import (
 typemysqlConnstruct{
 bufbuffer
 netConn net.Conn
-rawConn net.Conn// underlying connection when netConn is TLS connection.
-affectedRowsuint64
-insertIduint64
+rawConn net.Conn// underlying connection when netConn is TLS connection.
+resultmysqlResult// managed by clearResult() and handleOkPacket().
 cfg*Config
 connector*connector
 maxAllowedPacketint
@@ -155,6 +154,7 @@ func (mc *mysqlConn) cleanup(){
 iferr:=mc.netConn.Close(); err!=nil{
 mc.cfg.Logger.Print(err)
  }
+mc.clearResult()
 }
 
 func (mc*mysqlConn) error() error{
@@ -316,28 +316,25 @@ func (mc *mysqlConn) Exec(query string, args []driver.Value) (driver.Result, err
  }
 query=prepared
  }
-mc.affectedRows=0
-mc.insertId=0
 
 err:=mc.exec(query)
 iferr==nil{
-return&mysqlResult{
-affectedRows: int64(mc.affectedRows),
-insertId: int64(mc.insertId),
- }, err
+copied:=mc.result
+return&copied, err
  }
 returnnil, mc.markBadConn(err)
 }
 
 // Internal function to execute commands
 func (mc*mysqlConn) exec(querystring) error{
+handleOk:=mc.clearResult()
 // Send command
 iferr:=mc.writeCommandPacketStr(comQuery, query); err!=nil{
 returnmc.markBadConn(err)
  }
 
 // Read Result
-resLen, err:=mc.readResultSetHeaderPacket()
+resLen, err:=handleOk.readResultSetHeaderPacket()
 iferr!=nil{
 returnerr
  }
@@ -354,14 +351,16 @@ func (mc *mysqlConn) exec(query string) error{
  }
  }
 
-returnmc.discardResults()
+returnhandleOk.discardResults()
 }
 
 func (mc*mysqlConn) Query(querystring, args []driver.Value) (driver.Rows, error){
 returnmc.query(query, args)
 }
 
 func (mc*mysqlConn) query(querystring, args []driver.Value) (*textRows, error){
+handleOk:=mc.clearResult()
+
 ifmc.closed.Load(){
 mc.cfg.Logger.Print(ErrInvalidConn)
 returnnil, driver.ErrBadConn
@@ -382,7 +381,7 @@ func (mc *mysqlConn) query(query string, args []driver.Value) (*textRows, error)
 iferr==nil{
 // Read Result
 varresLenint
-resLen, err=mc.readResultSetHeaderPacket()
+resLen, err=handleOk.readResultSetHeaderPacket()
 iferr==nil{
 rows:=new(textRows)
 rows.mc=mc
@@ -410,12 +409,13 @@ func (mc *mysqlConn) query(query string, args []driver.Value) (*textRows, error)
 // The returned byte slice is only valid until the next read
 func (mc*mysqlConn) getSystemVar(namestring) ([]byte, error){
 // Send command
+handleOk:=mc.clearResult()
 iferr:=mc.writeCommandPacketStr(comQuery, "SELECT @@"+name); err!=nil{
 returnnil, err
  }
 
 // Read Result
-resLen, err:=mc.readResultSetHeaderPacket()
+resLen, err:=handleOk.readResultSetHeaderPacket()
 iferr==nil{
 rows:=new(textRows)
 rows.mc=mc
@@ -466,11 +466,12 @@ func (mc *mysqlConn) Ping(ctx context.Context) (err error){
  }
 defermc.finish()
 
+handleOk:=mc.clearResult()
 iferr=mc.writeCommandPacket(comPing); err!=nil{
 returnmc.markBadConn(err)
  }
 
-returnmc.readResultOK()
+returnhandleOk.readResultOK()
 }
 
 // BeginTx implements driver.ConnBeginTx interface

driver_test.go

@@ -2154,11 +2154,51 @@ func TestRejectReadOnly(t *testing.T){
 }
 
 funcTestPing(t*testing.T){
+ctx:=context.Background()
 runTests(t, dsn, func(dbt*DBTest){
 iferr:=dbt.db.Ping(); err!=nil{
 dbt.fail("Ping", "Ping", err)
  }
  })
+
+runTests(t, dsn, func(dbt*DBTest){
+conn, err:=dbt.db.Conn(ctx)
+iferr!=nil{
+dbt.fail("db", "Conn", err)
+ }
+
+// Check that affectedRows and insertIds are cleared after each call.
+conn.Raw(func(conninterface{}) error{
+c:=conn.(*mysqlConn)
+
+// Issue a query that sets affectedRows and insertIds.
+q, err:=c.Query(`SELECT 1`, nil)
+iferr!=nil{
+dbt.fail("Conn", "Query", err)
+ }
+ifgot, want:=c.result.affectedRows, []int64{0}; !reflect.DeepEqual(got, want){
+dbt.Fatalf("bad affectedRows: got %v, want=%v", got, want)
+ }
+ifgot, want:=c.result.insertIds, []int64{0}; !reflect.DeepEqual(got, want){
+dbt.Fatalf("bad insertIds: got %v, want=%v", got, want)
+ }
+q.Close()
+
+// Verify that Ping() clears both fields.
+fori:=0; i<2; i++{
+iferr:=c.Ping(ctx); err!=nil{
+dbt.fail("Pinger", "Ping", err)
+ }
+ifgot, want:=c.result.affectedRows, []int64(nil); !reflect.DeepEqual(got, want){
+t.Errorf("bad affectedRows: got %v, want=%v", got, want)
+ }
+ifgot, want:=c.result.insertIds, []int64(nil); !reflect.DeepEqual(got, want){
+t.Errorf("bad affectedRows: got %v, want=%v", got, want)
+ }
+ }
+returnnil
+ })
+ })
 }
 
 // See Issue #799
@@ -2378,6 +2418,42 @@ func TestMultiResultSetNoSelect(t *testing.T){
  })
 }
 
+funcTestExecMultipleResults(t*testing.T){
+ctx:=context.Background()
+runTestsWithMultiStatement(t, dsn, func(dbt*DBTest){
+dbt.mustExec(`
+ CREATE TABLE test (
+ id INT NOT NULL AUTO_INCREMENT,
+ value VARCHAR(255),
+ PRIMARY KEY (id)
+ )`)
+conn, err:=dbt.db.Conn(ctx)
+iferr!=nil{
+t.Fatalf("failed to connect: %v", err)
+ }
+conn.Raw(func(conninterface{}) error{
+ex:=conn.(driver.Execer)
+res, err:=ex.Exec(`
+ INSERT INTO test (value) VALUES ('a'), ('b');
+ INSERT INTO test (value) VALUES ('c'), ('d'), ('e');
+ `, nil)
+iferr!=nil{
+t.Fatalf("insert statements failed: %v", err)
+ }
+mres:=res.(Result)
+ifgot, want:=mres.AllRowsAffected(), []int64{2, 3}; !reflect.DeepEqual(got, want){
+t.Errorf("bad AllRowsAffected: got %v, want=%v", got, want)
+ }
+// For INSERTs containing multiple rows, LAST_INSERT_ID() returns the
+// first inserted ID, not the last.
+ifgot, want:=mres.AllLastInsertIds(), []int64{1, 3}; !reflect.DeepEqual(got, want){
+t.Errorf("bad AllLastInsertIds: got %v, want %v", got, want)
+ }
+returnnil
+ })
+ })
+}
+
 // tests if rows are set in a proper state if some results were ignored before
 // calling rows.NextResultSet.
 funcTestSkipResults(t*testing.T){
@@ -2399,6 +2475,42 @@ func TestSkipResults(t *testing.T){
  })
 }
 
+funcTestQueryMultipleResults(t*testing.T){
+ctx:=context.Background()
+runTestsWithMultiStatement(t, dsn, func(dbt*DBTest){
+dbt.mustExec(`
+ CREATE TABLE test (
+ id INT NOT NULL AUTO_INCREMENT,
+ value VARCHAR(255),
+ PRIMARY KEY (id)
+ )`)
+conn, err:=dbt.db.Conn(ctx)
+iferr!=nil{
+t.Fatalf("failed to connect: %v", err)
+ }
+conn.Raw(func(conninterface{}) error{
+qr:=conn.(driver.Queryer)
+
+c:=conn.(*mysqlConn)
+
+// Demonstrate that repeated queries reset the affectedRows
+fori:=0; i<2; i++{
+_, err:=qr.Query(`
+ INSERT INTO test (value) VALUES ('a'), ('b');
+ INSERT INTO test (value) VALUES ('c'), ('d'), ('e');
+ `, nil)
+iferr!=nil{
+t.Fatalf("insert statements failed: %v", err)
+ }
+ifgot, want:=c.result.affectedRows, []int64{2, 3}; !reflect.DeepEqual(got, want){
+t.Errorf("bad affectedRows: got %v, want=%v", got, want)
+ }
+ }
+returnnil
+ })
+ })
+}
+
 funcTestPingContext(t*testing.T){
 runTests(t, dsn, func(dbt*DBTest){
 ctx, cancel:=context.WithCancel(context.Background())

infile.go

@@ -93,7 +93,7 @@ func deferredClose(err *error, closer io.Closer){
 
 constdefaultPacketSize=16*1024// 16KB is small enough for disk readahead and large enough for TCP
 
-func (mc*mysqlConn) handleInFileRequest(namestring) (errerror){
+func (mc*okHandler) handleInFileRequest(namestring) (errerror){
 varrdr io.Reader
 vardata []byte
 packetSize:=defaultPacketSize
@@ -154,7 +154,7 @@ func (mc *mysqlConn) handleInFileRequest(name string) (err error){
 forerr==nil{
 n, err=rdr.Read(data[4:])
 ifn>0{
-ifioErr:=mc.writePacket(data[:4+n]); ioErr!=nil{
+ifioErr:=mc.conn().writePacket(data[:4+n]); ioErr!=nil{
 returnioErr
  }
  }
@@ -168,7 +168,7 @@ func (mc *mysqlConn) handleInFileRequest(name string) (err error){
 ifdata==nil{
 data=make([]byte, 4)
  }
-ifioErr:=mc.writePacket(data[:4]); ioErr!=nil{
+ifioErr:=mc.conn().writePacket(data[:4]); ioErr!=nil{
 returnioErr
  }
 
@@ -177,6 +177,6 @@ func (mc *mysqlConn) handleInFileRequest(name string) (err error){
 returnmc.readResultOK()
  }
 
-mc.readPacket()
+mc.conn().readPacket()
 returnerr
 }