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:
parent
94da9c644c
commit
c580ef9f5a
@ -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
|
||||
}
|
||||
|
18
ch/proto.go
18
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) {
|
||||
|
Loading…
x
Reference in New Issue
Block a user