From c580ef9f5ac15cddfcdea397838d7c7abecb1539 Mon Sep 17 00:00:00 2001
From: Vladimir Mihailenco <vladimir.webdev@gmail.com>
Date: Sat, 30 Apr 2022 16:54:06 +0300
Subject: [PATCH] chore: automatically close block iterator

---
 ch/db_test.go | 37 ++++++++++++++++++++++++++++++++++++-
 ch/proto.go   | 18 ++++++++++++++----
 2 files changed, 50 insertions(+), 5 deletions(-)

diff --git a/ch/db_test.go b/ch/db_test.go
index 7989776..6ca2562 100644
--- a/ch/db_test.go
+++ b/ch/db_test.go
@@ -5,6 +5,9 @@ import (
 	"database/sql"
 	"fmt"
 	"os"
+	"reflect"
+	"runtime"
+	"strings"
 	"testing"
 	"time"
 
@@ -305,12 +308,13 @@ func TestORM(t *testing.T) {
 		testORMSlice,
 		testORMColumnarStruct,
 		testORMInvalidEnumValue,
+		testORMInsertSelect,
 	}
 	for _, fn := range tests {
 		_, err := db.NewTruncateTable().Model((*Event)(nil)).Exec(ctx)
 		require.NoError(t, err)
 
-		t.Run("", func(t *testing.T) {
+		t.Run(funcName(fn), func(t *testing.T) {
 			fn(t, db)
 		})
 	}
@@ -493,6 +497,37 @@ func testORMInvalidEnumValue(t *testing.T, db *ch.DB) {
 	require.Equal(t, "invalid", dest.Kind)
 }
 
+func testORMInsertSelect(t *testing.T, db *ch.DB) {
+	ctx := context.Background()
+
+	for i := 0; i < 100; i++ {
+		src := &Event{
+			ID:        1,
+			Name:      "hello",
+			Count:     42,
+			Keys:      []string{"foo", "bar"},
+			Values:    [][]string{{}, {"hello", "world"}},
+			Kind:      "hello",
+			CreatedAt: time.Now(),
+		}
+		_, err := db.NewInsert().Model(src).Exec(ctx)
+		require.NoError(t, err)
+	}
+
+	var dest []Event
+	err := db.NewSelect().Model(&dest).Scan(ctx)
+	require.NoError(t, err)
+	require.Equal(t, 100, len(dest))
+}
+
+func funcName(x interface{}) string {
+	s := runtime.FuncForPC(reflect.ValueOf(x).Pointer()).Name()
+	if i := strings.LastIndexByte(s, '.'); i >= 0 {
+		return s[i+1:]
+	}
+	return s
+}
+
 func strptr(s string) *string {
 	return &s
 }
diff --git a/ch/proto.go b/ch/proto.go
index b4930b2..cd7f926 100644
--- a/ch/proto.go
+++ b/ch/proto.go
@@ -36,27 +36,37 @@ func newBlockIter(db *DB, cn *chpool.Conn) *blockIter {
 
 func (it *blockIter) Close() error {
 	if it.cn != nil {
-		it.db.releaseConn(it.cn, it.stickyErr)
-		it.cn = nil
+		it.close()
 	}
 	return nil
 }
 
+func (it *blockIter) close() {
+	it.db.releaseConn(it.cn, it.stickyErr)
+	it.cn = nil
+}
+
 func (it *blockIter) Err() error {
 	return it.stickyErr
 }
 
 func (it *blockIter) Next(ctx context.Context, block *chschema.Block) bool {
-	if it.stickyErr != nil {
+	if it.cn == nil {
 		return false
 	}
 
 	ok, err := it.read(ctx, block)
 	if err != nil {
 		it.stickyErr = err
+		it.close()
 		return false
 	}
-	return ok
+
+	if !ok {
+		it.close()
+		return false
+	}
+	return true
 }
 
 func (it *blockIter) read(ctx context.Context, block *chschema.Block) (bool, error) {