From 22c4d2dcb09913e988f02169a66ca7a7ae99ba78 Mon Sep 17 00:00:00 2001
From: Leo Cai <leo.cai@timeplus.io>
Date: Mon, 18 Dec 2023 13:59:21 +0800
Subject: [PATCH] fix Close() on rows when a timeout during processing(porting
 ClickHouse/clickhouse-go#548)

---
 clickhouse_rows.go | 25 ++++++++++++++++++++-----
 conn_query.go      |  2 +-
 2 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/clickhouse_rows.go b/clickhouse_rows.go
index b4c730f..6c76037 100644
--- a/clickhouse_rows.go
+++ b/clickhouse_rows.go
@@ -95,11 +95,26 @@ func (r *rows) Columns() []string {
 }
 
 func (r *rows) Close() error {
-	for range r.stream {
-	}
-	for err := range r.errors {
-		if err != nil {
-			r.err = err
+	active := 2
+	for {
+		select {
+		case _, ok := <-r.stream:
+			if !ok {
+				active--
+				if active == 0 {
+					return r.err
+				}
+			}
+		case err, ok := <-r.errors:
+			if err != nil {
+				r.err = err
+			}
+			if !ok {
+				active--
+				if active == 0 {
+					return r.err
+				}
+			}
 		}
 	}
 	return r.err
diff --git a/conn_query.go b/conn_query.go
index 7a0c618..42895dc 100644
--- a/conn_query.go
+++ b/conn_query.go
@@ -66,8 +66,8 @@ func (c *connect) query(ctx context.Context, release func(*connect, error), quer
 		if err != nil {
 			errors <- err
 		}
-		close(errors)
 		close(stream)
+		close(errors)
 		release(c, err)
 	}()