1
0
mirror of https://github.com/uptrace/go-clickhouse.git synced 2025-06-08 23:26:11 +02:00

chore: automatically close block iterator

This commit is contained in:
Vladimir Mihailenco 2022-04-30 16:54:06 +03:00
parent 94da9c644c
commit c580ef9f5a
2 changed files with 50 additions and 5 deletions

View File

@ -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
}

View File

@ -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) {