more tests, added doc comment

Author: SpencerTorres

conn_batch.go

@@ -320,6 +320,11 @@ 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

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)
+}