Skip to content

Commit cd37406

Browse files
hongkercolincchenhongkerjkaflik
authored
Optional flag to close query with flush (#1276)
* close query when use batch.Flush * modify message by author * Update 1271_test.go change name * Update options.go change WithCloseQuery to WithCloseOnFlush * modify property name * change property name * modify property name --------- Co-authored-by: colincchen <[email protected]> Co-authored-by: hongker <[email protected]> Co-authored-by: Kuba Kaflik <[email protected]>
1 parent 1ae716e commit cd37406

File tree

3 files changed

+159
-18
lines changed

3 files changed

+159
-18
lines changed

conn_batch.go

Lines changed: 23 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -83,14 +83,15 @@ func (c *connect) prepareBatch(ctx context.Context, query string, opts driver.Pr
8383
}
8484

8585
b := &batch{
86-
ctx: ctx,
87-
query: query,
88-
conn: c,
89-
block: block,
90-
released: false,
91-
connRelease: release,
92-
connAcquire: acquire,
93-
onProcess: onProcess,
86+
ctx: ctx,
87+
query: query,
88+
conn: c,
89+
block: block,
90+
released: false,
91+
connRelease: release,
92+
connAcquire: acquire,
93+
onProcess: onProcess,
94+
closeOnFlush: opts.CloseOnFlush,
9495
}
9596

9697
if opts.ReleaseConnection {
@@ -101,16 +102,17 @@ func (c *connect) prepareBatch(ctx context.Context, query string, opts driver.Pr
101102
}
102103

103104
type batch struct {
104-
err error
105-
ctx context.Context
106-
query string
107-
conn *connect
108-
sent bool // sent signalize that batch is send to ClickHouse.
109-
released bool // released signalize that conn was returned to pool and can't be used.
110-
block *proto.Block
111-
connRelease func(*connect, error)
112-
connAcquire func(context.Context) (*connect, error)
113-
onProcess *onProcess
105+
err error
106+
ctx context.Context
107+
query string
108+
conn *connect
109+
sent bool // sent signalize that batch is send to ClickHouse.
110+
released bool // released signalize that conn was returned to pool and can't be used.
111+
closeOnFlush bool // closeOnFlush signalize that batch should close query and release conn when use Flush
112+
block *proto.Block
113+
connRelease func(*connect, error)
114+
connAcquire func(context.Context) (*connect, error)
115+
onProcess *onProcess
114116
}
115117

116118
func (b *batch) release(err error) {
@@ -305,6 +307,9 @@ func (b *batch) Flush() error {
305307
if err := b.conn.sendData(b.block, ""); err != nil {
306308
return err
307309
}
310+
if b.closeOnFlush {
311+
b.release(b.closeQuery())
312+
}
308313
}
309314
b.block.Reset()
310315
return nil

lib/driver/options.go

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package driver
22

33
type PrepareBatchOptions struct {
44
ReleaseConnection bool
5+
CloseOnFlush bool
56
}
67

78
type PrepareBatchOption func(options *PrepareBatchOptions)
@@ -11,3 +12,10 @@ func WithReleaseConnection() PrepareBatchOption {
1112
options.ReleaseConnection = true
1213
}
1314
}
15+
16+
// WithCloseOnFlush closes batch INSERT query when Flush is executed
17+
func WithCloseOnFlush() PrepareBatchOption {
18+
return func(options *PrepareBatchOptions) {
19+
options.CloseOnFlush = true
20+
}
21+
}

tests/issues/1271_test.go

Lines changed: 128 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,128 @@
1+
package issues
2+
3+
import (
4+
"context"
5+
"fmt"
6+
"github.com/ClickHouse/clickhouse-go/v2/lib/driver"
7+
"github.com/google/uuid"
8+
"math/rand"
9+
"runtime"
10+
"strings"
11+
"testing"
12+
"time"
13+
14+
"github.com/ClickHouse/clickhouse-go/v2"
15+
clickhouse_tests "github.com/ClickHouse/clickhouse-go/v2/tests"
16+
"github.com/stretchr/testify/require"
17+
)
18+
19+
// test for https://github.com/ClickHouse/clickhouse-go/issues/1271
20+
func Test1271(t *testing.T) {
21+
var (
22+
conn, err = clickhouse_tests.GetConnection("issues", clickhouse.Settings{
23+
"max_execution_time": 60,
24+
"allow_experimental_object_type": true,
25+
}, nil, &clickhouse.Compression{
26+
Method: clickhouse.CompressionLZ4,
27+
})
28+
)
29+
defer func() {
30+
conn.Exec(context.Background(), "DROP TABLE flush_with_close_query_example")
31+
conn.Close()
32+
}()
33+
conn.Exec(context.Background(), "DROP TABLE IF EXISTS flush_with_close_query_example")
34+
ctx := context.Background()
35+
err = conn.Exec(ctx, `
36+
CREATE TABLE IF NOT EXISTS flush_with_close_query_example (
37+
Col1 UInt64
38+
, Col2 String
39+
, Col3 FixedString(3)
40+
, Col4 UUID
41+
, Col5 Map(String, UInt64)
42+
, Col6 Array(String)
43+
, Col7 Tuple(String, UInt64, Array(Map(String, UInt64)))
44+
, Col8 DateTime
45+
) Engine = MergeTree() ORDER BY tuple()
46+
`)
47+
require.NoError(t, err)
48+
49+
batch, err := conn.PrepareBatch(ctx, "INSERT INTO flush_with_close_query_example", driver.WithCloseOnFlush())
50+
require.NoError(t, err)
51+
// 1 million rows should only take < 1s on most desktops
52+
for i := 0; i < 100_000; i++ {
53+
require.NoError(t, batch.Append(
54+
uint64(i),
55+
RandAsciiString(5),
56+
RandAsciiString(3),
57+
uuid.New(),
58+
map[string]uint64{"key": uint64(i)}, // Map(String, UInt64)
59+
[]string{RandAsciiString(1), RandAsciiString(1), RandAsciiString(1), RandAsciiString(1), RandAsciiString(1), RandAsciiString(1)}, // Array(String)
60+
[]any{ // Tuple(String, UInt64, Array(Map(String, UInt64)))
61+
RandAsciiString(10), uint64(i), []map[string]uint64{
62+
{"key": uint64(i)},
63+
{"key": uint64(i + 1)},
64+
{"key": uint64(i) + 2},
65+
},
66+
},
67+
time.Now().Add(time.Duration(i)*time.Second),
68+
))
69+
if i > 0 && i%10000 == 0 {
70+
require.NoError(t, batch.Flush())
71+
PrintMemUsage()
72+
}
73+
}
74+
require.NoError(t, batch.Flush())
75+
// confirm we have the right count
76+
var col1 uint64
77+
require.NoError(t, conn.QueryRow(ctx, "SELECT count() FROM flush_with_close_query_example").Scan(&col1))
78+
require.Equal(t, uint64(100_000), col1)
79+
}
80+
81+
const letterBytes = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
82+
const (
83+
letterIdxBits = 6 // 6 bits to represent a letter index
84+
letterIdxMask = 1<<letterIdxBits - 1 // All 1-bits, as many as letterIdxBits
85+
letterIdxMax = 63 / letterIdxBits // # of letter indices fitting in 63 bits
86+
)
87+
88+
func RandAsciiString(n int) string {
89+
return randString(n, letterBytes)
90+
}
91+
92+
var src = rand.NewSource(time.Now().UnixNano())
93+
94+
func randString(n int, bytes string) string {
95+
sb := strings.Builder{}
96+
sb.Grow(n)
97+
// A src.Int63() generates 63 random bits, enough for letterIdxMax characters!
98+
for i, cache, remain := n-1, src.Int63(), letterIdxMax; i >= 0; {
99+
if remain == 0 {
100+
cache, remain = src.Int63(), letterIdxMax
101+
}
102+
if idx := int(cache & letterIdxMask); idx < len(bytes) {
103+
sb.WriteByte(bytes[idx])
104+
i--
105+
}
106+
cache >>= letterIdxBits
107+
remain--
108+
}
109+
110+
return sb.String()
111+
}
112+
113+
// PrintMemUsage outputs the current, total and OS memory being used. As well as the number
114+
// of garage collection cycles completed.
115+
// thanks to https://golangcode.com/print-the-current-memory-usage/
116+
func PrintMemUsage() {
117+
var m runtime.MemStats
118+
runtime.ReadMemStats(&m)
119+
// For info on each, see: https://golang.org/pkg/runtime/#MemStats
120+
fmt.Printf("Alloc = %v MiB", bToMb(m.Alloc))
121+
fmt.Printf("\tTotalAlloc = %v MiB", bToMb(m.TotalAlloc))
122+
fmt.Printf("\tSys = %v MiB", bToMb(m.Sys))
123+
fmt.Printf("\tNumGC = %v\n", m.NumGC)
124+
}
125+
126+
func bToMb(b uint64) uint64 {
127+
return b / 1024 / 1024
128+
}

0 commit comments

Comments
 (0)