Skip to content

Commit

Permalink
include Variant and Dynamic array support
Browse files Browse the repository at this point in the history
  • Loading branch information
SpencerTorres committed Feb 16, 2025
1 parent 3bd05b7 commit 8374597
Show file tree
Hide file tree
Showing 4 changed files with 165 additions and 13 deletions.
15 changes: 12 additions & 3 deletions lib/column/dynamic.go
Original file line number Diff line number Diff line change
Expand Up @@ -313,8 +313,13 @@ func (c *Dynamic) encodeData(buffer *proto.Buffer) {
c.variant.encodeData(buffer)
}

func (c *Dynamic) Encode(buffer *proto.Buffer) {
func (c *Dynamic) WriteStatePrefix(buffer *proto.Buffer) error {
c.encodeHeader(buffer)

return nil
}

func (c *Dynamic) Encode(buffer *proto.Buffer) {
c.encodeData(buffer)
}

Expand Down Expand Up @@ -393,13 +398,17 @@ func (c *Dynamic) decodeData(reader *proto.Reader, rows int) error {
return nil
}

func (c *Dynamic) Decode(reader *proto.Reader, rows int) error {
func (c *Dynamic) ReadStatePrefix(reader *proto.Reader) error {
err := c.decodeHeader(reader)
if err != nil {
return fmt.Errorf("failed to decode dynamic header: %w", err)
}

err = c.decodeData(reader, rows)
return nil
}

func (c *Dynamic) Decode(reader *proto.Reader, rows int) error {
err := c.decodeData(reader, rows)
if err != nil {
return fmt.Errorf("failed to decode dynamic data: %w", err)
}
Expand Down
15 changes: 12 additions & 3 deletions lib/column/variant.go
Original file line number Diff line number Diff line change
Expand Up @@ -275,8 +275,13 @@ func (c *Variant) encodeData(buffer *proto.Buffer) {
}
}

func (c *Variant) Encode(buffer *proto.Buffer) {
func (c *Variant) WriteStatePrefix(buffer *proto.Buffer) error {
c.encodeHeader(buffer)

return nil
}

func (c *Variant) Encode(buffer *proto.Buffer) {
c.encodeData(buffer)
}

Expand Down Expand Up @@ -336,13 +341,17 @@ func (c *Variant) decodeData(reader *proto.Reader, rows int) error {
return nil
}

func (c *Variant) Decode(reader *proto.Reader, rows int) error {
func (c *Variant) ReadStatePrefix(reader *proto.Reader) error {
err := c.decodeHeader(reader)
if err != nil {
return fmt.Errorf("failed to decode variant header: %w", err)
}

err = c.decodeData(reader, rows)
return nil
}

func (c *Variant) Decode(reader *proto.Reader, rows int) error {
err := c.decodeData(reader, rows)
if err != nil {
return fmt.Errorf("failed to decode variant data: %w", err)
}
Expand Down
73 changes: 70 additions & 3 deletions tests/dynamic_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,6 @@ import (
"context"
"fmt"
"github.com/ClickHouse/clickhouse-go/v2"
"github.com/ClickHouse/clickhouse-go/v2/lib/chcol"
"github.com/ClickHouse/clickhouse-go/v2/lib/driver"
"github.com/stretchr/testify/require"
"testing"
Expand Down Expand Up @@ -87,7 +86,7 @@ func TestDynamic(t *testing.T) {
rows, err := conn.Query(ctx, "SELECT c FROM test_dynamic")
require.NoError(t, err)

var row chcol.Dynamic
var row clickhouse.Dynamic

require.True(t, rows.Next())
err = rows.Scan(&row)
Expand Down Expand Up @@ -135,6 +134,74 @@ func TestDynamic(t *testing.T) {
require.Equal(t, colMapStringInt64, row.Any())
}

func TestDynamicArray(t *testing.T) {
ctx := context.Background()
conn := setupDynamicTest(t)

const ddl = `
CREATE TABLE IF NOT EXISTS test_dynamic (
c Array(Dynamic)
) Engine = MergeTree() ORDER BY tuple()
`
require.NoError(t, conn.Exec(ctx, ddl))
defer func() {
require.NoError(t, conn.Exec(ctx, "DROP TABLE IF EXISTS test_dynamic"))
}()

batch, err := conn.PrepareBatch(ctx, "INSERT INTO test_dynamic (c)")
require.NoError(t, err)

batch.Append([]clickhouse.Dynamic{
clickhouse.NewDynamicWithType(int64(42), "Int64"),
clickhouse.NewDynamicWithType(true, "Bool"),
})
require.NoError(t, batch.Send())

rows, err := conn.Query(ctx, "SELECT c FROM test_dynamic")
require.NoError(t, err)

var arrRow []clickhouse.Dynamic

require.True(t, rows.Next())
err = rows.Scan(&arrRow)
require.NoError(t, err)
require.Len(t, arrRow, 2)

require.Equal(t, int64(42), arrRow[0].Any())
require.Equal(t, true, arrRow[1].Any())
}

func TestDynamicEmptyArray(t *testing.T) {
ctx := context.Background()
conn := setupDynamicTest(t)

const ddl = `
CREATE TABLE IF NOT EXISTS test_dynamic (
c Array(Dynamic)
) Engine = MergeTree() ORDER BY tuple()
`
require.NoError(t, conn.Exec(ctx, ddl))
defer func() {
require.NoError(t, conn.Exec(ctx, "DROP TABLE IF EXISTS test_dynamic"))
}()

batch, err := conn.PrepareBatch(ctx, "INSERT INTO test_dynamic (c)")
require.NoError(t, err)

batch.Append([]clickhouse.Dynamic{})
require.NoError(t, batch.Send())

rows, err := conn.Query(ctx, "SELECT c FROM test_dynamic")
require.NoError(t, err)

var arrRow []clickhouse.Dynamic

require.True(t, rows.Next())
err = rows.Scan(&arrRow)
require.NoError(t, err)
require.Len(t, arrRow, 0)
}

func TestDynamic_ScanWithType(t *testing.T) {
ctx := context.Background()
conn := setupDynamicTest(t)
Expand All @@ -160,7 +227,7 @@ func TestDynamic_ScanWithType(t *testing.T) {
rows, err := conn.Query(ctx, "SELECT c FROM test_dynamic")
require.NoError(t, err)

var row chcol.Dynamic
var row clickhouse.Dynamic

require.True(t, rows.Next())
err = rows.Scan(&row)
Expand Down
75 changes: 71 additions & 4 deletions tests/variant_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,6 @@ import (
"time"

"github.com/ClickHouse/clickhouse-go/v2"
"github.com/ClickHouse/clickhouse-go/v2/lib/chcol"
"github.com/ClickHouse/clickhouse-go/v2/lib/driver"
"github.com/stretchr/testify/require"
)
Expand Down Expand Up @@ -101,7 +100,7 @@ func TestVariant(t *testing.T) {
rows, err := conn.Query(ctx, "SELECT c FROM test_variant")
require.NoError(t, err)

var row chcol.Variant
var row clickhouse.Variant

require.True(t, rows.Next())
err = rows.Scan(&row)
Expand Down Expand Up @@ -154,6 +153,74 @@ func TestVariant(t *testing.T) {
require.Equal(t, colMapStringInt64, row.Any())
}

func TestVariantArray(t *testing.T) {
ctx := context.Background()
conn := setupVariantTest(t)

const ddl = `
CREATE TABLE IF NOT EXISTS test_variant (
c Array(Variant(Int64))
) Engine = MergeTree() ORDER BY tuple()
`
require.NoError(t, conn.Exec(ctx, ddl))
defer func() {
require.NoError(t, conn.Exec(ctx, "DROP TABLE IF EXISTS test_variant"))
}()

batch, err := conn.PrepareBatch(ctx, "INSERT INTO test_variant (c)")
require.NoError(t, err)

batch.Append([]clickhouse.Variant{
clickhouse.NewVariantWithType(int64(42), "Int64"),
clickhouse.NewVariantWithType(int64(84), "Int64"),
})
require.NoError(t, batch.Send())

rows, err := conn.Query(ctx, "SELECT c FROM test_variant")
require.NoError(t, err)

var arrRow []clickhouse.Variant

require.True(t, rows.Next())
err = rows.Scan(&arrRow)
require.NoError(t, err)
require.Len(t, arrRow, 2)

require.Equal(t, int64(42), arrRow[0].Any())
require.Equal(t, int64(84), arrRow[1].Any())
}

func TestVariantEmptyArray(t *testing.T) {
ctx := context.Background()
conn := setupVariantTest(t)

const ddl = `
CREATE TABLE IF NOT EXISTS test_variant (
c Array(Variant(Int64))
) Engine = MergeTree() ORDER BY tuple()
`
require.NoError(t, conn.Exec(ctx, ddl))
defer func() {
require.NoError(t, conn.Exec(ctx, "DROP TABLE IF EXISTS test_variant"))
}()

batch, err := conn.PrepareBatch(ctx, "INSERT INTO test_variant (c)")
require.NoError(t, err)

batch.Append([]clickhouse.Variant{})
require.NoError(t, batch.Send())

rows, err := conn.Query(ctx, "SELECT c FROM test_variant")
require.NoError(t, err)

var arrRow []clickhouse.Variant

require.True(t, rows.Next())
err = rows.Scan(&arrRow)
require.NoError(t, err)
require.Len(t, arrRow, 0)
}

func TestVariant_ScanWithType(t *testing.T) {
ctx := context.Background()
conn := setupVariantTest(t)
Expand All @@ -179,7 +246,7 @@ func TestVariant_ScanWithType(t *testing.T) {
rows, err := conn.Query(ctx, "SELECT c FROM test_variant")
require.NoError(t, err)

var row chcol.Variant
var row clickhouse.Variant

require.True(t, rows.Next())
err = rows.Scan(&row)
Expand Down Expand Up @@ -235,7 +302,7 @@ func TestVariant_BatchFlush(t *testing.T) {

i := 0
for rows.Next() {
var row chcol.Variant
var row clickhouse.Variant
err = rows.Scan(&row)

if i%2 == 0 {
Expand Down

0 comments on commit 8374597

Please sign in to comment.