From 6aa0acc0fca79327fe6a7b1c47b4133e86c7d2e8 Mon Sep 17 00:00:00 2001
From: Alexandr Mayorskiy <a.mayorskiy@corp.mail.ru>
Date: Tue, 17 Nov 2020 09:54:15 +0300
Subject: [PATCH] ability to send multiple resultsets from handleQuery

---
 mysql/result.go    | 14 ++++++++++++++
 mysql/resultset.go |  3 +++
 server/resp.go     | 15 ++++++++++++++-
 3 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/mysql/result.go b/mysql/result.go
index 797a4af75..f3c287743 100644
--- a/mysql/result.go
+++ b/mysql/result.go
@@ -19,3 +19,17 @@ func (r *Result) Close() {
 		r.Resultset = nil
 	}
 }
+
+func (r *Result) ChainResultSet(rs *Resultset) {
+	if r.Resultset == nil {
+		r.Resultset = rs
+		return
+	}
+
+	var lastRS *Resultset
+
+	for lastRS = r.Resultset; lastRS.Next != nil; lastRS = lastRS.Next {
+	}
+
+	lastRS.Next = rs
+}
diff --git a/mysql/resultset.go b/mysql/resultset.go
index 2da6e3b5f..f1d9fbedb 100644
--- a/mysql/resultset.go
+++ b/mysql/resultset.go
@@ -17,6 +17,9 @@ type Resultset struct {
 	RawPkg []byte
 
 	RowDatas []RowData
+
+	// In the case of multiple queries, we will have there a chaining list of separated Resultset
+	Next *Resultset
 }
 
 var (
diff --git a/server/resp.go b/server/resp.go
index db8632394..46ca6f9fa 100644
--- a/server/resp.go
+++ b/server/resp.go
@@ -181,7 +181,20 @@ func (c *Conn) writeValue(value interface{}) error {
 		return c.writeOK(nil)
 	case *Result:
 		if v != nil && v.Resultset != nil {
-			return c.writeResultset(v.Resultset)
+			for rs := v.Resultset; rs != nil; rs = rs.Next {
+				if rs.Next != nil {
+					c.status |= SERVER_MORE_RESULTS_EXISTS
+				}
+
+				err := c.writeResultset(rs)
+				c.status &= ^SERVER_MORE_RESULTS_EXISTS
+
+				if err != nil {
+					return err
+				}
+			}
+
+			return nil
 		} else {
 			return c.writeOK(v)
 		}