Merge pull request #1566 from ClickHouse/fix_cancel_query

Add `Close` function to batch interface

Author: SpencerTorres

clickhouse.go

@@ -147,6 +147,8 @@ func (ch *clickhouse) Exec(ctx context.Context, query string, args ...any) error
 	if err != nil {
 		return err
 	}
+	conn.debugf("[acquired] connection [%d]", conn.id)
+
 	if err := conn.exec(ctx, query, args...); err != nil {
 		ch.release(conn, err)
 		return err
@@ -345,6 +347,8 @@ func (ch *clickhouse) release(conn *connect, err error) {
 		return
 	}
 	conn.released = true
+	conn.debugf("[released] connection [%d]", conn.id)
+
 	select {
 	case <-ch.open:
 	default:

conn_batch.go

@@ -320,6 +320,26 @@ func (b *batch) closeQuery() error {
 	return nil
 }
 
+// Close will end the current INSERT without sending the currently buffered rows, and release the connection.
+// This may result in zero row inserts if no rows were appended.
+// If a batch was already sent this does nothing.
+// This should be called via defer after a batch is opened to prevent
+// batches from falling out of scope and timing out.
+func (b *batch) Close() error {
+	if b.sent || b.released {
+		return nil
+	}
+
+	if err := b.closeQuery(); err != nil {
+		return err
+	}
+	b.sent = true
+
+	b.release(nil)
+
+	return nil
+}
+
 type batchColumn struct {
 	err     error
 	batch   driver.Batch

conn_http_batch.go

@@ -109,6 +109,11 @@ func (b *httpBatch) Flush() error {
 	return nil
 }
 
+func (b *httpBatch) Close() error {
+	b.sent = true
+	return nil
+}
+
 func (b *httpBatch) Abort() error {
 	defer func() {
 		b.sent = true

examples/clickhouse_api/batch.go

@@ -53,6 +53,8 @@ func BatchInsert() error {
 	if err != nil {
 		return err
 	}
+	defer batch.Close()
+
 	for i := 0; i < 1000; i++ {
 		err := batch.Append(
 			uint8(42),

lib/driver/driver.go

@@ -87,6 +87,7 @@ type (
 		IsSent() bool
 		Rows() int
 		Columns() []column.Interface
+		Close() error
 	}
 	BatchColumn interface {
 		Append(any) error

tests/abort_test.go

@@ -56,3 +56,25 @@ func TestAbort(t *testing.T) {
 		}
 	}
 }
+
+func TestBatchClose(t *testing.T) {
+	conn, err := GetNativeConnection(nil, nil, &clickhouse.Compression{
+		Method: clickhouse.CompressionLZ4,
+	})
+	require.NoError(t, err)
+	ctx := context.Background()
+
+	batch, err := conn.PrepareBatch(ctx, "INSERT INTO function null('x UInt64') VALUES (1)")
+	require.NoError(t, err)
+	require.NoError(t, batch.Close())
+	require.NoError(t, batch.Close()) // No error on multiple calls
+
+	batch, err = conn.PrepareBatch(ctx, "INSERT INTO function null('x UInt64') VALUES (1)")
+	require.NoError(t, err)
+	if assert.NoError(t, batch.Append(uint8(1))) && assert.NoError(t, batch.Send()) {
+		var col1 uint8
+		if err := conn.QueryRow(ctx, "SELECT 1").Scan(&col1); assert.NoError(t, err) {
+			assert.Equal(t, uint8(1), col1)
+		}
+	}
+}

tests/batch_test.go

@@ -19,6 +19,7 @@ package tests
 
 import (
 	"testing"
+	"time"
 
 	"github.com/ClickHouse/clickhouse-go/v2"
 	"github.com/stretchr/testify/require"
@@ -51,3 +52,73 @@ func TestBatchContextCancellation(t *testing.T) {
 	// assert if connection is properly released after context cancellation
 	require.NoError(t, conn.Exec(context.Background(), "SELECT 1"))
 }
+
+func TestBatchCloseConnectionReleased(t *testing.T) {
+	te, err := GetTestEnvironment(testSet)
+	require.NoError(t, err)
+	opts := ClientOptionsFromEnv(te, clickhouse.Settings{}, false)
+	opts.MaxOpenConns = 1
+	conn, err := GetConnectionWithOptions(&opts)
+	require.NoError(t, err)
+
+	b, err := conn.PrepareBatch(context.Background(), "INSERT INTO function null('x UInt64')")
+	require.NoError(t, err)
+	for i := 0; i < 100; i++ {
+		require.NoError(t, b.Append(i))
+	}
+
+	err = b.Close()
+	require.NoError(t, err)
+
+	// assert if connection is properly released after close called
+	require.NoError(t, conn.Exec(context.Background(), "SELECT 1"))
+}
+
+func TestBatchSendConnectionReleased(t *testing.T) {
+	te, err := GetTestEnvironment(testSet)
+	require.NoError(t, err)
+	opts := ClientOptionsFromEnv(te, clickhouse.Settings{}, false)
+	opts.MaxOpenConns = 1
+	conn, err := GetConnectionWithOptions(&opts)
+	require.NoError(t, err)
+
+	b, err := conn.PrepareBatch(context.Background(), "INSERT INTO function null('x UInt64')")
+	require.NoError(t, err)
+	for i := 0; i < 100; i++ {
+		require.NoError(t, b.Append(i))
+	}
+
+	err = b.Send()
+	require.NoError(t, err)
+
+	// Close should be deferred after the batch is opened
+	// Validate that it can be called after Send
+	err = b.Close()
+	require.NoError(t, err)
+
+	// assert if connection is properly released after Send called
+	require.NoError(t, conn.Exec(context.Background(), "SELECT 1"))
+}
+
+// This test validates that connections are blocked if a batch is not properly
+// cleaned up. This isn't required behavior, but this test confirms it happens.
+func TestBatchCloseConnectionHold(t *testing.T) {
+	te, err := GetTestEnvironment(testSet)
+	require.NoError(t, err)
+	opts := ClientOptionsFromEnv(te, clickhouse.Settings{}, false)
+	opts.MaxOpenConns = 1
+	opts.DialTimeout = 2 * time.Second // Lower timeout for faster acquire error
+	conn, err := GetConnectionWithOptions(&opts)
+	require.NoError(t, err)
+
+	b, err := conn.PrepareBatch(context.Background(), "INSERT INTO function null('x UInt64')")
+	require.NoError(t, err)
+	for i := 0; i < 100; i++ {
+		require.NoError(t, b.Append(i))
+	}
+
+	// batch.Close() should be called here
+
+	// assert if connection is blocked if close is not called.
+	require.ErrorIs(t, conn.Exec(context.Background(), "SELECT 1"), clickhouse.ErrAcquireConnTimeout)
+}