From ae966e2aebf2e40a3b6d6e9607d5acc3ed07a653 Mon Sep 17 00:00:00 2001
From: tangenta <tangenta@126.com>
Date: Sat, 20 Dec 2025 12:37:04 +0800
Subject: [PATCH] support recording data from multiple result set

Signed-off-by: tangenta <tangenta@126.com>
---
 r/multi_query.result | 11 +++++++++++
 src/main.go          | 44 +++++++++++++++++++++++++++-----------------
 t/multi_query.test   | 12 ++++++++++++
 3 files changed, 50 insertions(+), 17 deletions(-)
 create mode 100644 r/multi_query.result
 create mode 100644 t/multi_query.test

diff --git a/r/multi_query.result b/r/multi_query.result
new file mode 100644
index 0000000..eef9183
--- /dev/null
+++ b/r/multi_query.result
@@ -0,0 +1,11 @@
+set global tidb_enable_procedure = ON;
+CREATE PROCEDURE multi_select()
+BEGIN
+SELECT 1;
+SELECT 2, 3;
+END//
+CALL multi_select();
+1
+1
+2	3
+2	3
diff --git a/src/main.go b/src/main.go
index 8f328ca..a0e0836 100644
--- a/src/main.go
+++ b/src/main.go
@@ -953,15 +953,17 @@ func (rows *byteRows) Swap(i, j int) {
 	rows.data[i], rows.data[j] = rows.data[j], rows.data[i]
 }
 
-func dumpToByteRows(rows *sql.Rows) (*byteRows, error) {
-	cols, err := rows.Columns()
-	if err != nil {
-		return nil, errors.Trace(err)
-	}
+func dumpToByteRows(rows *sql.Rows) ([]*byteRows, error) {
+	var result []*byteRows
 
-	data := make([]byteRow, 0, 8)
-	args := make([]interface{}, len(cols))
 	for {
+		cols, err := rows.Columns()
+		if err != nil {
+			return nil, errors.Trace(err)
+		}
+
+		data := make([]byteRow, 0, 8)
+		args := make([]interface{}, len(cols))
 		for rows.Next() {
 			tmp := make([][]byte, len(cols))
 			for i := 0; i < len(args); i++ {
@@ -974,16 +976,19 @@ func dumpToByteRows(rows *sql.Rows) (*byteRows, error) {
 
 			data = append(data, byteRow{tmp})
 		}
+
+		result = append(result, &byteRows{cols: cols, data: data})
+
 		if !rows.NextResultSet() {
 			break
 		}
 	}
-	err = rows.Err()
+	err := rows.Err()
 	if err != nil {
 		return nil, errors.Trace(err)
 	}
 
-	return &byteRows{cols: cols, data: data}, nil
+	return result, nil
 }
 
 func (t *tester) executeStmt(query string) error {
@@ -992,15 +997,20 @@ func (t *tester) executeStmt(query string) error {
 	if err != nil {
 		return errors.Trace(err)
 	}
+	defer raw.Close()
 
-	rows, err := dumpToByteRows(raw)
+	allRows, err := dumpToByteRows(raw)
 	if err != nil {
 		return errors.Trace(err)
 	}
 
-	if t.enableResultLog && (len(rows.cols) > 0 || len(rows.data) > 0) {
-		if err = t.writeQueryResult(rows); err != nil {
-			return errors.Trace(err)
+	if t.enableResultLog {
+		for _, rows := range allRows {
+			if len(rows.cols) > 0 || len(rows.data) > 0 {
+				if err = t.writeQueryResult(rows); err != nil {
+					return errors.Trace(err)
+				}
+			}
 		}
 	}
 
@@ -1023,14 +1033,14 @@ func (t *tester) executeStmt(query string) error {
 			return errors.Trace(err)
 		}
 
-		rows, err := dumpToByteRows(raw)
+		allRows, err := dumpToByteRows(raw)
 		if err != nil {
 			return errors.Trace(err)
 		}
 
-		if len(rows.data) > 0 {
-			sort.Sort(rows)
-			return t.writeQueryResult(rows)
+		if len(allRows) > 0 && len(allRows[0].data) > 0 {
+			sort.Sort(allRows[0])
+			return t.writeQueryResult(allRows[0])
 		}
 	}
 	return nil
diff --git a/t/multi_query.test b/t/multi_query.test
new file mode 100644
index 0000000..c7ce2fa
--- /dev/null
+++ b/t/multi_query.test
@@ -0,0 +1,12 @@
+set global tidb_enable_procedure = ON;
+DELIMITER //
+
+CREATE PROCEDURE multi_select()
+BEGIN
+    SELECT 1;
+    SELECT 2, 3;
+END//
+
+DELIMITER ;
+
+CALL multi_select();