ability to send multiple resultsets from handleQuery

GoWebProd · GoWebProd · commit 33cb7b8e60fb · 2020-11-10T16:48:20.000+03:00
diff --git a/server/caching_sha2_cache_test.go b/server/caching_sha2_cache_test.go
@@ -105,7 +105,7 @@ func (s *cacheTestSuite) onAccept(c *C) {
 }
 
 func (s *cacheTestSuite) onConn(conn net.Conn, c *C) {
-	//co, err := NewConn(conn, *testUser, *testPassword, &testHandler{s})
+	// co, err := NewConn(conn, *testUser, *testPassword, &testHandler{s})
 	co, err := NewCustomizedConn(conn, s.server, s.credProvider, &testCacheHandler{s})
 	c.Assert(err, IsNil)
 	for {
@@ -137,7 +137,7 @@ func (s *cacheTestSuite) TestCache(c *C) {
 	t2 := time.Now()
 
 	d1 := int(t2.Sub(t1).Nanoseconds() / 1e6)
-	//log.Debugf("first connection took %d milliseconds", d1)
+	// log.Debugf("first connection took %d milliseconds", d1)
 
 	c.Assert(d1, GreaterEqual, delay)
 
@@ -154,7 +154,7 @@ func (s *cacheTestSuite) TestCache(c *C) {
 	t4 := time.Now()
 
 	d2 := int(t4.Sub(t3).Nanoseconds() / 1e6)
-	//log.Debugf("second connection took %d milliseconds", d2)
+	// log.Debugf("second connection took %d milliseconds", d2)
 
 	c.Assert(d2, Less, delay)
 	if s.db != nil {
@@ -178,7 +178,7 @@ func (h *testCacheHandler) handleQuery(query string, binary bool) (*mysql.Result
 	case "select":
 		var r *mysql.Resultset
 		var err error
-		//for handle go mysql driver select @@max_allowed_packet
+		// for handle go mysql driver select @@max_allowed_packet
 		if strings.Contains(strings.ToLower(query), "max_allowed_packet") {
 			r, err = mysql.BuildSimpleResultset([]string{"@@max_allowed_packet"}, [][]interface{}{
 				{mysql.MaxPayloadLen},
@@ -209,8 +209,13 @@ func (h *testCacheHandler) handleQuery(query string, binary bool) (*mysql.Result
 	return nil, nil
 }
 
-func (h *testCacheHandler) HandleQuery(query string) (*mysql.Result, error) {
-	return h.handleQuery(query, false)
+func (h *testCacheHandler) HandleQuery(query string) ([]*mysql.Result, error) {
+	res, err := h.handleQuery(query, false)
+	if err != nil {
+		return nil, err
+	}
+
+	return []*mysql.Result{res}, nil
 }
 
 func (h *testCacheHandler) HandleFieldList(table string, fieldWildcard string) ([]*mysql.Field, error) {
diff --git a/server/command.go b/server/command.go
@@ -9,24 +9,24 @@ import (
 )
 
 type Handler interface {
-	//handle COM_INIT_DB command, you can check whether the dbName is valid, or other.
+	// handle COM_INIT_DB command, you can check whether the dbName is valid, or other.
 	UseDB(dbName string) error
-	//handle COM_QUERY command, like SELECT, INSERT, UPDATE, etc...
-	//If Result has a Resultset (SELECT, SHOW, etc...), we will send this as the response, otherwise, we will send Result
-	HandleQuery(query string) (*Result, error)
-	//handle COM_FILED_LIST command
+	// handle COM_QUERY command, like SELECT, INSERT, UPDATE, etc...
+	// If Result has a Resultset (SELECT, SHOW, etc...), we will send this as the response, otherwise, we will send Result
+	HandleQuery(query string) ([]*Result, error)
+	// handle COM_FILED_LIST command
 	HandleFieldList(table string, fieldWildcard string) ([]*Field, error)
-	//handle COM_STMT_PREPARE, params is the param number for this statement, columns is the column number
-	//context will be used later for statement execute
+	// handle COM_STMT_PREPARE, params is the param number for this statement, columns is the column number
+	// context will be used later for statement execute
 	HandleStmtPrepare(query string) (params int, columns int, context interface{}, err error)
-	//handle COM_STMT_EXECUTE, context is the previous one set in prepare
-	//query is the statement prepare query, and args is the params for this statement
+	// handle COM_STMT_EXECUTE, context is the previous one set in prepare
+	// query is the statement prepare query, and args is the params for this statement
 	HandleStmtExecute(context interface{}, query string, args []interface{}) (*Result, error)
-	//handle COM_STMT_CLOSE, context is the previous one set in prepare
-	//this handler has no response
+	// handle COM_STMT_CLOSE, context is the previous one set in prepare
+	// this handler has no response
 	HandleStmtClose(context interface{}) error
-	//handle any other command that is not currently handled by the library,
-	//default implementation for this method will return an ER_UNKNOWN_ERROR
+	// handle any other command that is not currently handled by the library,
+	// default implementation for this method will return an ER_UNKNOWN_ERROR
 	HandleOtherCommand(cmd byte, data []byte) error
 }
 
@@ -134,7 +134,7 @@ type EmptyHandler struct {
 func (h EmptyHandler) UseDB(dbName string) error {
 	return nil
 }
-func (h EmptyHandler) HandleQuery(query string) (*Result, error) {
+func (h EmptyHandler) HandleQuery(query string) ([]*Result, error) {
 	return nil, fmt.Errorf("not supported now")
 }
 
diff --git a/server/resp.go b/server/resp.go
@@ -185,6 +185,26 @@ func (c *Conn) writeValue(value interface{}) error {
 		} else {
 			return c.writeOK(v)
 		}
+	case []*Result:
+		if len(v) == 0 {
+			return c.writeValue(nil)
+		}
+
+		c.status |= SERVER_MORE_RESULTS_EXISTS
+
+		for i, res := range v {
+			if i == len(v)-1 {
+				c.status &= ^SERVER_MORE_RESULTS_EXISTS
+			}
+
+			if err := c.writeValue(res); err != nil {
+				c.status &= ^SERVER_MORE_RESULTS_EXISTS
+
+				return err
+			}
+		}
+
+		return nil
 	case []*Field:
 		return c.writeFieldList(v)
 	case *Stmt:
diff --git a/server/server_test.go b/server/server_test.go
@@ -63,7 +63,7 @@ func Test(t *testing.T) {
 	inMemProvider.AddUser(*testUser, *testPassword)
 
 	servers := prepareServerConf()
-	//no TLS
+	// no TLS
 	for _, svr := range servers {
 		Suite(&serverTestSuite{
 			server:       svr,
@@ -135,7 +135,7 @@ func (s *serverTestSuite) onAccept(c *C) {
 }
 
 func (s *serverTestSuite) onConn(conn net.Conn, c *C) {
-	//co, err := NewConn(conn, *testUser, *testPassword, &testHandler{s})
+	// co, err := NewConn(conn, *testUser, *testPassword, &testHandler{s})
 	co, err := NewCustomizedConn(conn, s.server, s.credProvider, &testHandler{s})
 	c.Assert(err, IsNil)
 	// set SSL if defined
@@ -225,7 +225,7 @@ func (h *testHandler) handleQuery(query string, binary bool) (*mysql.Result, err
 	case "select":
 		var r *mysql.Resultset
 		var err error
-		//for handle go mysql driver select @@max_allowed_packet
+		// for handle go mysql driver select @@max_allowed_packet
 		if strings.Contains(strings.ToLower(query), "max_allowed_packet") {
 			r, err = mysql.BuildSimpleResultset([]string{"@@max_allowed_packet"}, [][]interface{}{
 				{mysql.MaxPayloadLen},
@@ -256,8 +256,14 @@ func (h *testHandler) handleQuery(query string, binary bool) (*mysql.Result, err
 	return nil, nil
 }
 
-func (h *testHandler) HandleQuery(query string) (*mysql.Result, error) {
-	return h.handleQuery(query, false)
+func (h *testHandler) HandleQuery(query string) ([]*mysql.Result, error) {
+	res, err := h.handleQuery(query, false)
+	if err != nil {
+		return nil, err
+
+	}
+
+	return []*mysql.Result{res}, nil
 }
 
 func (h *testHandler) HandleFieldList(table string, fieldWildcard string) ([]*mysql.Field, error) {
