From 76bd962af0dd3a12a139ff1571050e39cf10b64b Mon Sep 17 00:00:00 2001
From: Stepan Rabotkin <srabotkin@ozon.ru>
Date: Wed, 2 Aug 2023 13:19:59 +0300
Subject: [PATCH] Fix panic and releasing in batch column

---
 conn_batch.go             | 27 ++++++++++++++-------------
 tests/issues/1053_test.go | 38 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 52 insertions(+), 13 deletions(-)
 create mode 100644 tests/issues/1053_test.go

diff --git a/conn_batch.go b/conn_batch.go
index 5bb72d63b1..c4ddb820f9 100644
--- a/conn_batch.go
+++ b/conn_batch.go
@@ -147,12 +147,15 @@ func (b *batch) IsSent() bool {
 
 func (b *batch) Column(idx int) driver.BatchColumn {
 	if len(b.block.Columns) <= idx {
-		b.release(nil)
+		err := &OpError{
+			Op:  "batch.Column",
+			Err: fmt.Errorf("invalid column index %d", idx),
+		}
+
+		b.release(err)
+
 		return &batchColumn{
-			err: &OpError{
-				Op:  "batch.Column",
-				Err: fmt.Errorf("invalid column index %d", idx),
-			},
+			err: err,
 		}
 	}
 	return &batchColumn{
@@ -229,13 +232,12 @@ type batchColumn struct {
 }
 
 func (b *batchColumn) Append(v any) (err error) {
-	if b.batch.IsSent() {
-		return ErrBatchAlreadySent
-	}
 	if b.err != nil {
-		b.release(b.err)
 		return b.err
 	}
+	if b.batch.IsSent() {
+		return ErrBatchAlreadySent
+	}
 	if _, err = b.column.Append(v); err != nil {
 		b.release(err)
 		return err
@@ -244,13 +246,12 @@ func (b *batchColumn) Append(v any) (err error) {
 }
 
 func (b *batchColumn) AppendRow(v any) (err error) {
-	if b.batch.IsSent() {
-		return ErrBatchAlreadySent
-	}
 	if b.err != nil {
-		b.release(b.err)
 		return b.err
 	}
+	if b.batch.IsSent() {
+		return ErrBatchAlreadySent
+	}
 	if err = b.column.AppendRow(v); err != nil {
 		b.release(err)
 		return err
diff --git a/tests/issues/1053_test.go b/tests/issues/1053_test.go
new file mode 100644
index 0000000000..0b5a538b46
--- /dev/null
+++ b/tests/issues/1053_test.go
@@ -0,0 +1,38 @@
+package issues
+
+import (
+	"context"
+	"github.com/ClickHouse/clickhouse-go/v2"
+	clickhouse_tests "github.com/ClickHouse/clickhouse-go/v2/tests"
+	"github.com/stretchr/testify/require"
+	"testing"
+)
+
+func Test1053(t *testing.T) {
+	var (
+		conn, err = clickhouse_tests.GetConnection("issues", clickhouse.Settings{
+			"max_execution_time": 60,
+		}, nil, &clickhouse.Compression{
+			Method: clickhouse.CompressionLZ4,
+		})
+	)
+	ctx := context.Background()
+	require.NoError(t, err)
+	const ddl = `
+		CREATE TABLE test_1053 (
+			Col1 UInt64
+		) Engine MergeTree() ORDER BY tuple()
+		`
+	defer func() {
+		conn.Exec(ctx, "DROP TABLE IF EXISTS test_1053")
+	}()
+	require.NoError(t, conn.Exec(ctx, ddl))
+
+	batch, err := conn.PrepareBatch(ctx, `INSERT INTO test_1053`)
+	require.NoError(t, err)
+
+	column := batch.Column(1000) // doesn't exist column
+
+	require.Error(t, column.Append(uint64(1)))
+	require.Error(t, column.AppendRow(uint64(1)))
+}