diff --git a/README.md b/README.md index d747a7446..4eade6853 100644 --- a/README.md +++ b/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` ``` diff --git a/auth.go b/auth.go index e758e6d00..f6b157a12 100644 --- a/auth.go +++ b/auth.go @@ -346,7 +346,7 @@ func (mc *mysqlConn) handleAuthResult(oldAuthData []byte, plugin string) error { case 1: switch authData[0] { case cachingSha2PasswordFastAuthSuccess: - if err = mc.readResultOK(); err == nil { + if err = mc.resultUnchanged().readResultOK(); err == nil { return nil // auth successful } @@ -397,7 +397,7 @@ func (mc *mysqlConn) handleAuthResult(oldAuthData []byte, plugin string) error { return err } } - return mc.readResultOK() + return mc.resultUnchanged().readResultOK() default: return ErrMalformPkt @@ -426,7 +426,7 @@ func (mc *mysqlConn) handleAuthResult(oldAuthData []byte, plugin string) error { if err != nil { return err } - return mc.readResultOK() + return mc.resultUnchanged().readResultOK() } default: diff --git a/connection.go b/connection.go index 14a972b40..631a1dc24 100644 --- a/connection.go +++ b/connection.go @@ -23,9 +23,8 @@ import ( type mysqlConn struct { buf buffer netConn net.Conn - rawConn net.Conn // underlying connection when netConn is TLS connection. - affectedRows uint64 - insertId uint64 + rawConn net.Conn // underlying connection when netConn is TLS connection. + result mysqlResult // managed by clearResult() and handleOkPacket(). cfg *Config connector *connector maxAllowedPacket int @@ -155,6 +154,7 @@ func (mc *mysqlConn) cleanup() { if err := 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) if err == nil { - return &mysqlResult{ - affectedRows: int64(mc.affectedRows), - insertId: int64(mc.insertId), - }, err + copied := mc.result + return &copied, err } return nil, mc.markBadConn(err) } // Internal function to execute commands func (mc *mysqlConn) exec(query string) error { + handleOk := mc.clearResult() // Send command if err := mc.writeCommandPacketStr(comQuery, query); err != nil { return mc.markBadConn(err) } // Read Result - resLen, err := mc.readResultSetHeaderPacket() + resLen, err := handleOk.readResultSetHeaderPacket() if err != nil { return err } @@ -354,7 +351,7 @@ func (mc *mysqlConn) exec(query string) error { } } - return mc.discardResults() + return handleOk.discardResults() } func (mc *mysqlConn) Query(query string, args []driver.Value) (driver.Rows, error) { @@ -362,6 +359,8 @@ func (mc *mysqlConn) Query(query string, args []driver.Value) (driver.Rows, erro } func (mc *mysqlConn) query(query string, args []driver.Value) (*textRows, error) { + handleOk := mc.clearResult() + if mc.closed.Load() { mc.cfg.Logger.Print(ErrInvalidConn) return nil, driver.ErrBadConn @@ -382,7 +381,7 @@ func (mc *mysqlConn) query(query string, args []driver.Value) (*textRows, error) if err == nil { // Read Result var resLen int - resLen, err = mc.readResultSetHeaderPacket() + resLen, err = handleOk.readResultSetHeaderPacket() if err == 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(name string) ([]byte, error) { // Send command + handleOk := mc.clearResult() if err := mc.writeCommandPacketStr(comQuery, "SELECT @@"+name); err != nil { return nil, err } // Read Result - resLen, err := mc.readResultSetHeaderPacket() + resLen, err := handleOk.readResultSetHeaderPacket() if err == nil { rows := new(textRows) rows.mc = mc @@ -466,11 +466,12 @@ func (mc *mysqlConn) Ping(ctx context.Context) (err error) { } defer mc.finish() + handleOk := mc.clearResult() if err = mc.writeCommandPacket(comPing); err != nil { return mc.markBadConn(err) } - return mc.readResultOK() + return handleOk.readResultOK() } // BeginTx implements driver.ConnBeginTx interface diff --git a/driver_test.go b/driver_test.go index abf91a486..cd94c434e 100644 --- a/driver_test.go +++ b/driver_test.go @@ -2154,11 +2154,51 @@ func TestRejectReadOnly(t *testing.T) { } func TestPing(t *testing.T) { + ctx := context.Background() runTests(t, dsn, func(dbt *DBTest) { if err := dbt.db.Ping(); err != nil { dbt.fail("Ping", "Ping", err) } }) + + runTests(t, dsn, func(dbt *DBTest) { + conn, err := dbt.db.Conn(ctx) + if err != nil { + dbt.fail("db", "Conn", err) + } + + // Check that affectedRows and insertIds are cleared after each call. + conn.Raw(func(conn interface{}) error { + c := conn.(*mysqlConn) + + // Issue a query that sets affectedRows and insertIds. + q, err := c.Query(`SELECT 1`, nil) + if err != nil { + dbt.fail("Conn", "Query", err) + } + if got, want := c.result.affectedRows, []int64{0}; !reflect.DeepEqual(got, want) { + dbt.Fatalf("bad affectedRows: got %v, want=%v", got, want) + } + if got, 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. + for i := 0; i < 2; i++ { + if err := c.Ping(ctx); err != nil { + dbt.fail("Pinger", "Ping", err) + } + if got, want := c.result.affectedRows, []int64(nil); !reflect.DeepEqual(got, want) { + t.Errorf("bad affectedRows: got %v, want=%v", got, want) + } + if got, want := c.result.insertIds, []int64(nil); !reflect.DeepEqual(got, want) { + t.Errorf("bad affectedRows: got %v, want=%v", got, want) + } + } + return nil + }) + }) } // See Issue #799 @@ -2378,6 +2418,42 @@ func TestMultiResultSetNoSelect(t *testing.T) { }) } +func TestExecMultipleResults(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) + if err != nil { + t.Fatalf("failed to connect: %v", err) + } + conn.Raw(func(conn interface{}) 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) + if err != nil { + t.Fatalf("insert statements failed: %v", err) + } + mres := res.(Result) + if got, 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. + if got, want := mres.AllLastInsertIds(), []int64{1, 3}; !reflect.DeepEqual(got, want) { + t.Errorf("bad AllLastInsertIds: got %v, want %v", got, want) + } + return nil + }) + }) +} + // tests if rows are set in a proper state if some results were ignored before // calling rows.NextResultSet. func TestSkipResults(t *testing.T) { @@ -2399,6 +2475,42 @@ func TestSkipResults(t *testing.T) { }) } +func TestQueryMultipleResults(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) + if err != nil { + t.Fatalf("failed to connect: %v", err) + } + conn.Raw(func(conn interface{}) error { + qr := conn.(driver.Queryer) + + c := conn.(*mysqlConn) + + // Demonstrate that repeated queries reset the affectedRows + for i := 0; i < 2; i++ { + _, err := qr.Query(` + INSERT INTO test (value) VALUES ('a'), ('b'); + INSERT INTO test (value) VALUES ('c'), ('d'), ('e'); + `, nil) + if err != nil { + t.Fatalf("insert statements failed: %v", err) + } + if got, want := c.result.affectedRows, []int64{2, 3}; !reflect.DeepEqual(got, want) { + t.Errorf("bad affectedRows: got %v, want=%v", got, want) + } + } + return nil + }) + }) +} + func TestPingContext(t *testing.T) { runTests(t, dsn, func(dbt *DBTest) { ctx, cancel := context.WithCancel(context.Background()) diff --git a/infile.go b/infile.go index 3279dcffd..cfd41914e 100644 --- a/infile.go +++ b/infile.go @@ -93,7 +93,7 @@ func deferredClose(err *error, closer io.Closer) { const defaultPacketSize = 16 * 1024 // 16KB is small enough for disk readahead and large enough for TCP -func (mc *mysqlConn) handleInFileRequest(name string) (err error) { +func (mc *okHandler) handleInFileRequest(name string) (err error) { var rdr io.Reader var data []byte packetSize := defaultPacketSize @@ -154,7 +154,7 @@ func (mc *mysqlConn) handleInFileRequest(name string) (err error) { for err == nil { n, err = rdr.Read(data[4:]) if n > 0 { - if ioErr := mc.writePacket(data[:4+n]); ioErr != nil { + if ioErr := mc.conn().writePacket(data[:4+n]); ioErr != nil { return ioErr } } @@ -168,7 +168,7 @@ func (mc *mysqlConn) handleInFileRequest(name string) (err error) { if data == nil { data = make([]byte, 4) } - if ioErr := mc.writePacket(data[:4]); ioErr != nil { + if ioErr := mc.conn().writePacket(data[:4]); ioErr != nil { return ioErr } @@ -177,6 +177,6 @@ func (mc *mysqlConn) handleInFileRequest(name string) (err error) { return mc.readResultOK() } - mc.readPacket() + mc.conn().readPacket() return err } diff --git a/packets.go b/packets.go index c10072c94..1a7f2c376 100644 --- a/packets.go +++ b/packets.go @@ -511,7 +511,9 @@ func (mc *mysqlConn) readAuthResult() ([]byte, string, error) { switch data[0] { case iOK: - return nil, "", mc.handleOkPacket(data) + // resultUnchanged, since auth happens before any queries or + // commands have been executed. + return nil, "", mc.resultUnchanged().handleOkPacket(data) case iAuthMoreData: return data[1:], "", err @@ -535,8 +537,8 @@ func (mc *mysqlConn) readAuthResult() ([]byte, string, error) { } // Returns error if Packet is not an 'Result OK'-Packet -func (mc *mysqlConn) readResultOK() error { - data, err := mc.readPacket() +func (mc *okHandler) readResultOK() error { + data, err := mc.conn().readPacket() if err != nil { return err } @@ -544,13 +546,17 @@ func (mc *mysqlConn) readResultOK() error { if data[0] == iOK { return mc.handleOkPacket(data) } - return mc.handleErrorPacket(data) + return mc.conn().handleErrorPacket(data) } // Result Set Header Packet // http://dev.mysql.com/doc/internals/en/com-query-response.html#packet-ProtocolText::Resultset -func (mc *mysqlConn) readResultSetHeaderPacket() (int, error) { - data, err := mc.readPacket() +func (mc *okHandler) readResultSetHeaderPacket() (int, error) { + // handleOkPacket replaces both values; other cases leave the values unchanged. + mc.result.affectedRows = append(mc.result.affectedRows, 0) + mc.result.insertIds = append(mc.result.insertIds, 0) + + data, err := mc.conn().readPacket() if err == nil { switch data[0] { @@ -558,7 +564,7 @@ func (mc *mysqlConn) readResultSetHeaderPacket() (int, error) { return 0, mc.handleOkPacket(data) case iERR: - return 0, mc.handleErrorPacket(data) + return 0, mc.conn().handleErrorPacket(data) case iLocalInFile: return 0, mc.handleInFileRequest(string(data[1:])) @@ -623,18 +629,61 @@ func readStatus(b []byte) statusFlag { return statusFlag(b[0]) | statusFlag(b[1])<<8 } +// Returns an instance of okHandler for codepaths where mysqlConn.result doesn't +// need to be cleared first (e.g. during authentication, or while additional +// resultsets are being fetched.) +func (mc *mysqlConn) resultUnchanged() *okHandler { + return (*okHandler)(mc) +} + +// okHandler represents the state of the connection when mysqlConn.result has +// been prepared for processing of OK packets. +// +// To correctly populate mysqlConn.result (updated by handleOkPacket()), all +// callpaths must either: +// +// 1. first clear it using clearResult(), or +// 2. confirm that they don't need to (by calling resultUnchanged()). +// +// Both return an instance of type *okHandler. +type okHandler mysqlConn + +// Exposees the underlying type's methods. +func (mc *okHandler) conn() *mysqlConn { + return (*mysqlConn)(mc) +} + +// clearResult clears the connection's stored affectedRows and insertIds +// fields. +// +// It returns a handler that can process OK responses. +func (mc *mysqlConn) clearResult() *okHandler { + mc.result = mysqlResult{} + return (*okHandler)(mc) +} + // Ok Packet // http://dev.mysql.com/doc/internals/en/generic-response-packets.html#packet-OK_Packet -func (mc *mysqlConn) handleOkPacket(data []byte) error { +func (mc *okHandler) handleOkPacket(data []byte) error { var n, m int + var affectedRows, insertId uint64 // 0x00 [1 byte] // Affected rows [Length Coded Binary] - mc.affectedRows, _, n = readLengthEncodedInteger(data[1:]) + affectedRows, _, n = readLengthEncodedInteger(data[1:]) // Insert id [Length Coded Binary] - mc.insertId, _, m = readLengthEncodedInteger(data[1+n:]) + insertId, _, m = readLengthEncodedInteger(data[1+n:]) + + // Update for the current statement result (only used by + // readResultSetHeaderPacket). + if len(mc.result.affectedRows) > 0 { + mc.result.affectedRows[len(mc.result.affectedRows)-1] = int64(affectedRows) + } + if len(mc.result.insertIds) > 0 { + mc.result.insertIds[len(mc.result.insertIds)-1] = int64(insertId) + } // server_status [2 bytes] mc.status = readStatus(data[1+n+m : 1+n+m+2]) @@ -1165,7 +1214,9 @@ func (stmt *mysqlStmt) writeExecutePacket(args []driver.Value) error { return mc.writePacket(data) } -func (mc *mysqlConn) discardResults() error { +// For each remaining resultset in the stream, discards its rows and updates +// mc.affectedRows and mc.insertIds. +func (mc *okHandler) discardResults() error { for mc.status&statusMoreResultsExists != 0 { resLen, err := mc.readResultSetHeaderPacket() if err != nil { @@ -1173,11 +1224,11 @@ func (mc *mysqlConn) discardResults() error { } if resLen > 0 { // columns - if err := mc.readUntilEOF(); err != nil { + if err := mc.conn().readUntilEOF(); err != nil { return err } // rows - if err := mc.readUntilEOF(); err != nil { + if err := mc.conn().readUntilEOF(); err != nil { return err } } diff --git a/result.go b/result.go index c6438d034..36a432e81 100644 --- a/result.go +++ b/result.go @@ -8,15 +8,44 @@ package mysql +import "database/sql/driver" + +// Result exposes data not available through *connection.Result. +// +// This is accessible by executing statements using sql.Conn.Raw() and +// downcasting the returned result: +// +// res, err := rawConn.Exec(...) +// res.(mysql.Result).AllRowsAffected() +// +type Result interface { + driver.Result + // AllRowsAffected returns a slice containing the affected rows for each + // executed statement. + AllRowsAffected() []int64 + // AllLastInsertIds returns a slice containing the last inserted ID for each + // executed statement. + AllLastInsertIds() []int64 +} + type mysqlResult struct { - affectedRows int64 - insertId int64 + // One entry in both slices is created for every executed statement result. + affectedRows []int64 + insertIds []int64 } func (res *mysqlResult) LastInsertId() (int64, error) { - return res.insertId, nil + return res.insertIds[len(res.insertIds)-1], nil } func (res *mysqlResult) RowsAffected() (int64, error) { - return res.affectedRows, nil + return res.affectedRows[len(res.affectedRows)-1], nil +} + +func (res *mysqlResult) AllLastInsertIds() []int64 { + return append([]int64{}, res.insertIds...) // defensive copy +} + +func (res *mysqlResult) AllRowsAffected() []int64 { + return append([]int64{}, res.affectedRows...) // defensive copy } diff --git a/rows.go b/rows.go index 888bdb5f0..63d0ed2d5 100644 --- a/rows.go +++ b/rows.go @@ -123,7 +123,8 @@ func (rows *mysqlRows) Close() (err error) { err = mc.readUntilEOF() } if err == nil { - if err = mc.discardResults(); err != nil { + handleOk := mc.clearResult() + if err = handleOk.discardResults(); err != nil { return err } } @@ -160,7 +161,9 @@ func (rows *mysqlRows) nextResultSet() (int, error) { return 0, io.EOF } rows.rs = resultSet{} - return rows.mc.readResultSetHeaderPacket() + // rows.mc.affectedRows and rows.mc.insertIds accumulate on each call to + // nextResultSet. + return rows.mc.resultUnchanged().readResultSetHeaderPacket() } func (rows *mysqlRows) nextNotEmptyResultSet() (int, error) { diff --git a/statement.go b/statement.go index d8b3975a5..31e7799c4 100644 --- a/statement.go +++ b/statement.go @@ -61,12 +61,10 @@ func (stmt *mysqlStmt) Exec(args []driver.Value) (driver.Result, error) { } mc := stmt.mc - - mc.affectedRows = 0 - mc.insertId = 0 + handleOk := stmt.mc.clearResult() // Read Result - resLen, err := mc.readResultSetHeaderPacket() + resLen, err := handleOk.readResultSetHeaderPacket() if err != nil { return nil, err } @@ -83,14 +81,12 @@ func (stmt *mysqlStmt) Exec(args []driver.Value) (driver.Result, error) { } } - if err := mc.discardResults(); err != nil { + if err := handleOk.discardResults(); err != nil { return nil, err } - return &mysqlResult{ - affectedRows: int64(mc.affectedRows), - insertId: int64(mc.insertId), - }, nil + copied := mc.result + return &copied, nil } func (stmt *mysqlStmt) Query(args []driver.Value) (driver.Rows, error) { @@ -111,7 +107,8 @@ func (stmt *mysqlStmt) query(args []driver.Value) (*binaryRows, error) { mc := stmt.mc // Read Result - resLen, err := mc.readResultSetHeaderPacket() + handleOk := stmt.mc.clearResult() + resLen, err := handleOk.readResultSetHeaderPacket() if err != nil { return nil, err }