Bugfix/issue 62 fix Close() on rows when a timeout during processing (#…

…64) porting: ClickHouse/clickhouse-go#548
timeplus-io · Dec 18, 2023 · d47a4ac · d47a4ac
1 parent 86e77b8
commit d47a4ac
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)
  }()