fix Close() on rows when a timeout during processing(porting ClickHou…

…se/clickhouse-go#548)
timeplus-io · Dec 18, 2023 · 22c4d2d · 22c4d2d
1 parent c66adfd
commit 22c4d2d
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 6 deletions.
diff --git 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
@@ -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)
 	}()