support recording data from multiple result set

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

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

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();