From d47a4accd64ac69f8045568e2dfc3b63ee494621 Mon Sep 17 00:00:00 2001
From: Leo Cai <143980630+leo-cai-timeplus@users.noreply.github.com>
Date: Mon, 18 Dec 2023 14:05:17 +0800
Subject: [PATCH] Bugfix/issue 62 fix Close() on rows when a timeout during
 processing (#64)

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