Batch Columns() []column.Interface method (#1277)

* impl Batch.Block() method * replace Batch.Block() => Batch.Columns() * add TestBatchColumns
ClickHouse · Jun 6, 2024 · 8ad6ec6 · 8ad6ec6
1 parent 3a49215
commit 8ad6ec6
Show file tree

Hide file tree

Showing 4 changed files with 42 additions and 1 deletion.
diff --git a/conn_batch.go b/conn_batch.go
@@ -22,6 +22,7 @@ import (
 	"fmt"
 	"os"
 	"regexp"
+	"slices"
 	"strings"
 	"time"
 
@@ -319,6 +320,10 @@ func (b *batch) Rows() int {
 	return b.block.Rows()
 }
 
+func (b *batch) Columns() []column.Interface {
+	return slices.Clone(b.block.Columns)
+}
+
 func (b *batch) closeQuery() error {
 	if err := b.conn.sendData(&proto.Block{}, ""); err != nil {
 		return err

diff --git a/conn_http_batch.go b/conn_http_batch.go
@@ -23,6 +23,7 @@ import (
 	"fmt"
 	"io"
 	"regexp"
+	"slices"
 	"strings"
 
 	"github.com/ClickHouse/clickhouse-go/v2/lib/column"
@@ -240,4 +241,8 @@ func (b *httpBatch) Rows() int {
 	return b.block.Rows()
 }
 
+func (b *httpBatch) Columns() []column.Interface {
+	return slices.Clone(b.block.Columns)
+}
+
 var _ driver.Batch = (*httpBatch)(nil)
diff --git a/lib/driver/driver.go b/lib/driver/driver.go
@@ -22,6 +22,7 @@ import (
 	"reflect"
 	"time"
 
+	"github.com/ClickHouse/clickhouse-go/v2/lib/column"
 	"github.com/ClickHouse/clickhouse-go/v2/lib/proto"
 )
 
@@ -85,6 +86,7 @@ type (
 		Send() error
 		IsSent() bool
 		Rows() int
+		Columns() []column.Interface
 	}
 	BatchColumn interface {
 		Append(any) error

diff --git a/tests/batch_block_test.go b/tests/batch_block_test.go
@@ -20,10 +20,12 @@ package tests
 import (
 	"testing"
 
-	"github.com/ClickHouse/clickhouse-go/v2"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"golang.org/x/net/context"
+
+	"github.com/ClickHouse/clickhouse-go/v2"
+	"github.com/ClickHouse/clickhouse-go/v2/lib/column"
 )
 
 // TestBatchAppendRows tests experimental batch rows blocks append feature.
@@ -64,3 +66,30 @@ func TestBatchAppendRows(t *testing.T) {
 	require.NoError(t, row.Scan(&count))
 	assert.Equal(t, 1000000, int(count))
 }
+
+// TestBatchColumns tests Batch.Columns() method functionality
+func TestBatchColumns(t *testing.T) {
+	ctx := context.Background()
+	conn, err := GetNativeConnection(nil, nil, nil)
+	require.NoError(t, err)
+	// Prepare test table
+	require.NoError(t, conn.Exec(ctx, `
+		CREATE TABLE test_table (
+		    Col1 Int,
+			Col2 String
+		) Engine MergeTree() ORDER BY tuple()
+	`))
+	defer func() {
+		conn.Exec(ctx, "DROP TABLE IF EXISTS test_table")
+	}()
+
+	batch, err := conn.PrepareBatch(ctx, "INSERT INTO test_table")
+	require.NoError(t, err)
+	columns := batch.Columns()
+	if assert.Len(t, columns, 2) {
+		assert.IsType(t, new(column.Int32), columns[0])
+		assert.Equal(t, "Col1", columns[0].Name())
+		assert.IsType(t, new(column.String), columns[1])
+		assert.Equal(t, "Col2", columns[1].Name())
+	}
+}