Merge branch 'master' into tls-in-driver

Author: atercattus

.github/workflows/ci.yml

@@ -5,7 +5,7 @@ jobs:
   test:
     strategy:
       matrix:
-        go: [ 1.16, 1.15 ]
+        go: [ 1.18, 1.17, 1.16, 1.15 ]
     name: Tests Go ${{ matrix.go }}
     runs-on: ubuntu-18.04
 

client/auth.go

@@ -139,6 +139,20 @@ func (c *Conn) genAuthResponse(authData []byte) ([]byte, bool, error) {
 	}
 }
 
+// generate connection attributes data
+func (c *Conn) genAttributes() []byte {
+	if len(c.attributes) == 0 {
+		return nil
+	}
+
+	attrData := make([]byte, 0)
+	for k, v := range c.attributes {
+		attrData = append(attrData, PutLengthEncodedString([]byte(k))...)
+		attrData = append(attrData, PutLengthEncodedString([]byte(v))...)
+	}
+	return append(PutLengthEncodedInt(uint64(len(attrData))), attrData...)
+}
+
 // See: http://dev.mysql.com/doc/internals/en/connection-phase-packets.html#packet-Protocol::HandshakeResponse
 func (c *Conn) writeAuthHandshake() error {
 	if !authPluginAllowed(c.authPluginName) {
@@ -195,6 +209,12 @@ func (c *Conn) writeAuthHandshake() error {
 		capability |= CLIENT_CONNECT_WITH_DB
 		length += len(c.db) + 1
 	}
+	// connection attributes
+	attrData := c.genAttributes()
+	if len(attrData) > 0 {
+		capability |= CLIENT_CONNECT_ATTRS
+		length += len(attrData)
+	}
 
 	data := make([]byte, length+4)
 
@@ -264,6 +284,12 @@ func (c *Conn) writeAuthHandshake() error {
 	// Assume native client during response
 	pos += copy(data[pos:], c.authPluginName)
 	data[pos] = 0x00
+	pos++
+
+	// connection attributes
+	if len(attrData) > 0 {
+		copy(data[pos:], attrData)
+	}
 
 	return c.WritePacket(data)
 }

client/auth_test.go

@@ -0,0 +1,42 @@
+package client
+
+import (
+	"bytes"
+	"testing"
+
+	"github.com/go-mysql-org/go-mysql/mysql"
+)
+
+func TestConnGenAttributes(t *testing.T) {
+	c := &Conn{
+		// example data from
+		// https://dev.mysql.com/doc/internals/en/connection-phase-packets.html#packet-Protocol::HandshakeResponse41
+		attributes: map[string]string{
+			"_os":             "debian6.0",
+			"_client_name":    "libmysql",
+			"_pid":            "22344",
+			"_client_version": "5.6.6-m9",
+			"_platform":       "x86_64",
+			"foo":             "bar",
+		},
+	}
+
+	data := c.genAttributes()
+
+	// the order of the attributes map cannot be guaranteed so to test the content
+	// of the attribute data we need to check its partial contents
+
+	if len(data) != 98 {
+		t.Fatalf("unexpected data length, got %d", len(data))
+	}
+	if data[0] != 0x61 {
+		t.Fatalf("unexpected length-encoded int, got %#x", data[0])
+	}
+
+	for k, v := range c.attributes {
+		fixt := append(mysql.PutLengthEncodedString([]byte(k)), mysql.PutLengthEncodedString([]byte(v))...)
+		if !bytes.Contains(data, fixt) {
+			t.Fatalf("%s attribute not found", k)
+		}
+	}
+}

client/conn.go

@@ -10,6 +10,7 @@ import (
 
 	. "github.com/go-mysql-org/go-mysql/mysql"
 	"github.com/go-mysql-org/go-mysql/packet"
+	"github.com/go-mysql-org/go-mysql/utils"
 	"github.com/pingcap/errors"
 )
 
@@ -27,6 +28,8 @@ type Conn struct {
 	// client-set capabilities only
 	ccaps uint32
 
+	attributes map[string]string
+
 	status uint16
 
 	charset string
@@ -43,6 +46,9 @@ type SelectPerRowCallback func(row []FieldValue) error
 // This function will be called once per result from ExecuteSelectStreaming
 type SelectPerResultCallback func(result *Result) error
 
+// This function will be called once per result from ExecuteMultiple
+type ExecPerResultCallback func(result *Result, err error)
+
 func getNetProto(addr string) string {
 	proto := "tcp"
 	if strings.Contains(addr, "/") {
@@ -198,6 +204,68 @@ func (c *Conn) Execute(command string, args ...interface{}) (*Result, error) {
 	}
 }
 
+// ExecuteMultiple will call perResultCallback for every result of the multiple queries
+// that are executed.
+//
+// When ExecuteMultiple is used, the connection should have the SERVER_MORE_RESULTS_EXISTS
+// flag set to signal the server multiple queries are executed. Handling the responses
+// is up to the implementation of perResultCallback.
+//
+// Example:
+//
+//      queries := "SELECT 1; SELECT NOW();"
+//      conn.ExecuteMultiple(queries, func(result *mysql.Result, err error) {
+//          // Use the result as you want
+//      })
+//
+func (c *Conn) ExecuteMultiple(query string, perResultCallback ExecPerResultCallback) (*Result, error) {
+	if err := c.writeCommandStr(COM_QUERY, query); err != nil {
+		return nil, errors.Trace(err)
+	}
+
+	var buf []byte
+	var err error
+	var result *Result
+	defer utils.ByteSlicePut(buf)
+
+	for {
+		buf, err = c.ReadPacketReuseMem(utils.ByteSliceGet(16)[:0])
+		if err != nil {
+			return nil, errors.Trace(err)
+		}
+
+		switch buf[0] {
+		case OK_HEADER:
+			result, err = c.handleOKPacket(buf)
+		case ERR_HEADER:
+			err = c.handleErrorPacket(append([]byte{}, buf...))
+			result = nil
+		case LocalInFile_HEADER:
+			err = ErrMalformPacket
+			result = nil
+		default:
+			result, err = c.readResultset(buf, false)
+		}
+
+		// call user-defined callback
+		perResultCallback(result, err)
+
+		// if there was an error of this was the last result, stop looping
+		if err != nil || result.Status&SERVER_MORE_RESULTS_EXISTS == 0 {
+			break
+		}
+	}
+
+	// return an empty result(set) signaling we're done streaming a multiple
+	// streaming session
+	// if this would end up in WriteValue, it would just be ignored as all
+	// responses should have been handled in perResultCallback
+	return &Result{Resultset: &Resultset{
+		Streaming:     StreamingMultiple,
+		StreamingDone: true,
+	}}, nil
+}
+
 // ExecuteSelectStreaming will call perRowCallback for every row in resultset
 //   WITHOUT saving any row data to Result.{Values/RawPkg/RowDatas} fields.
 // When given, perResultCallback will be called once per result
@@ -236,6 +304,10 @@ func (c *Conn) Rollback() error {
 	return errors.Trace(err)
 }
 
+func (c *Conn) SetAttributes(attributes map[string]string) {
+	c.attributes = attributes
+}
+
 func (c *Conn) SetCharset(charset string) error {
 	if c.charset == charset {
 		return nil

client/conn_test.go

@@ -2,6 +2,7 @@ package client
 
 import (
 	"fmt"
+	"strings"
 
 	. "github.com/pingcap/check"
 
@@ -16,7 +17,10 @@ type connTestSuite struct {
 func (s *connTestSuite) SetUpSuite(c *C) {
 	var err error
 	addr := fmt.Sprintf("%s:%s", *testHost, s.port)
-	s.c, err = Connect(addr, *testUser, *testPassword, "")
+	s.c, err = Connect(addr, *testUser, *testPassword, "", func(c *Conn) {
+		// required for the ExecuteMultiple test
+		c.SetCapability(mysql.CLIENT_MULTI_STATEMENTS)
+	})
 	if err != nil {
 		c.Fatal(err)
 	}
@@ -78,6 +82,46 @@ func (s *connTestSuite) testExecute_DropTable(c *C) {
 	c.Assert(err, IsNil)
 }
 
+func (s *connTestSuite) TestExecuteMultiple(c *C) {
+	queries := []string{
+		`INSERT INTO ` + testExecuteSelectStreamingTablename + ` (id, str) VALUES (999, "executemultiple")`,
+		`SELECT id FROM ` + testExecuteSelectStreamingTablename + ` LIMIT 2`,
+		`DELETE FROM ` + testExecuteSelectStreamingTablename + ` WHERE id=999`,
+		`THIS IS BOGUS()`,
+	}
+
+	count := 0
+	result, err := s.c.ExecuteMultiple(strings.Join(queries, "; "), func(result *mysql.Result, err error) {
+		switch count {
+		// the INSERT/DELETE query have no resultset, but should have set affectedrows
+		// the err should be nil
+		// also, since this is not the last query, the SERVER_MORE_RESULTS_EXISTS
+		// flag should be set
+		case 0, 2:
+			c.Assert(result.Status&mysql.SERVER_MORE_RESULTS_EXISTS, Not(Equals), 0)
+			c.Assert(result.Resultset, IsNil)
+			c.Assert(result.AffectedRows, Equals, uint64(1))
+			c.Assert(err, IsNil)
+		case 1:
+			// the SELECT query should have an resultset
+			// still not the last query, flag should be set
+			c.Assert(result.Status&mysql.SERVER_MORE_RESULTS_EXISTS, Not(Equals), 0)
+			c.Assert(result.Resultset, NotNil)
+			c.Assert(err, IsNil)
+		case 3:
+			// this query is obviously bogus so the error should be non-nil
+			c.Assert(result, IsNil)
+			c.Assert(err, NotNil)
+		}
+		count++
+	})
+
+	c.Assert(count, Equals, 4)
+	c.Assert(err, IsNil)
+	c.Assert(result.StreamingDone, Equals, true)
+	c.Assert(result.Streaming, Equals, mysql.StreamingMultiple)
+}
+
 func (s *connTestSuite) TestExecuteSelectStreaming(c *C) {
 	var (
 		expectedRowId                int64

client/resp.go

@@ -309,7 +309,7 @@ func (c *Conn) readResultsetStreaming(data []byte, binary bool, result *Result,
 	}
 
 	// this is a streaming resultset
-	result.Resultset.Streaming = true
+	result.Resultset.Streaming = StreamingSelect
 
 	if err := c.readResultColumns(result); err != nil {
 		return errors.Trace(err)

go.mod

@@ -14,6 +14,7 @@ require (
 	github.com/shopspring/decimal v0.0.0-20180709203117-cd690d0c9e24
 	github.com/siddontang/go v0.0.0-20180604090527-bdc77568d726
 	github.com/siddontang/go-log v0.0.0-20180807004314-8d05993dda07
+	github.com/stretchr/testify v1.7.0
 	go.uber.org/multierr v1.6.0 // indirect
 	go.uber.org/zap v1.16.0 // indirect
 	golang.org/x/text v0.3.6 // indirect

go.sum

@@ -51,10 +51,12 @@ github.com/siddontang/go v0.0.0-20180604090527-bdc77568d726 h1:xT+JlYxNGqyT+XcU8
 github.com/siddontang/go v0.0.0-20180604090527-bdc77568d726/go.mod h1:3yhqj7WBBfRhbBlzyOC3gUxftwsU0u8gqevxwIHQpMw=
 github.com/siddontang/go-log v0.0.0-20180807004314-8d05993dda07 h1:oI+RNwuC9jF2g2lP0u0cVEEZrc/AYBCuFdvwrLWM/6Q=
 github.com/siddontang/go-log v0.0.0-20180807004314-8d05993dda07/go.mod h1:yFdBgwXP24JziuRl2NMUahT7nGLNOKi1SIiFxMttVD4=
+github.com/stretchr/objx v0.1.0 h1:4G4v2dO3VZwixGIRoQ5Lfboy6nUhCyYzaqnIAPPhYs4=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
-github.com/stretchr/testify v1.4.0 h1:2E4SXV/wtOkTonXsotYi4li6zVWxYlZuYNCXe9XRJyk=
 github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81PSLYec5m4=
+github.com/stretchr/testify v1.7.0 h1:nwc3DEeHmmLAfoZucVR881uASk0Mfjw8xYJ99tb5CcY=
+github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
 go.uber.org/atomic v1.3.2/go.mod h1:gD2HeocX3+yG+ygLZcrzQJaqmWj9AIm7n08wl/qW/PE=
 go.uber.org/atomic v1.6.0/go.mod h1:sABNBOSYdrvTF6hTgEIbc7YasKWGhgEQZyfxyTvoXHQ=
@@ -112,5 +114,7 @@ gopkg.in/natefinch/lumberjack.v2 v2.0.0 h1:1Lc07Kr7qY4U2YPouBjpCLxpiyxIVoxqXgkXL
 gopkg.in/natefinch/lumberjack.v2 v2.0.0/go.mod h1:l0ndWWf7gzL7RNwBG7wST/UCcT4T24xpD6X8LsfU/+k=
 gopkg.in/yaml.v2 v2.2.2 h1:ZCJp+EgiOT7lHqUV2J862kp8Qj64Jo6az82+3Td9dZw=
 gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
+gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c h1:dUUwHk2QECo/6vqA44rthZ8ie2QXMNeKRTHCNY2nXvo=
+gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
 honnef.co/go/tools v0.0.1-2019.2.3 h1:3JgtbtFHMiCmsznwGVTUWbgGov+pVqnlf1dEJTNAXeM=
 honnef.co/go/tools v0.0.1-2019.2.3/go.mod h1:a3bituU0lyd329TUQxRnasdCoJDkEUEAqEt0JzvZhAg=