YQ Connector: YQ-2745: integration tests with Arrow data checks for C…

…lickHouse

Makefile

@@ -17,6 +17,7 @@ integration_test: integration_test_env
 	./fq-connector-go-tests -projectPath=$(PROJECT_PATH)
 
 integration_test_env:
+	docker-compose -f ./tests/infra/datasource/docker-compose.yaml stop
 	docker-compose -f ./tests/infra/datasource/docker-compose.yaml rm -f -v 
 	docker-compose -f ./tests/infra/datasource/docker-compose.yaml up -d
 

app/client/app.go

@@ -1,13 +1,12 @@
 package client
 
 import (
-	"bytes"
 	"context"
 	"fmt"
 	"io/ioutil"
 	"os"
 
-	"github.com/apache/arrow/go/v13/arrow/ipc"
+	"github.com/apache/arrow/go/v13/arrow"
 	"github.com/spf13/cobra"
 	"go.uber.org/zap"
 	"google.golang.org/protobuf/encoding/prototext"
@@ -67,7 +66,11 @@ func callServer(logger *zap.Logger, cfg *config.TClientConfig) error {
 
 	switch cfg.DataSourceInstance.Kind {
 	case api_common.EDataSourceKind_CLICKHOUSE, api_common.EDataSourceKind_POSTGRESQL, api_common.EDataSourceKind_YDB:
-		splits, err = prepareSplits(cl, tableName, cfg.DataSourceInstance)
+		typeMappingSettings := &api_service_protos.TTypeMappingSettings{
+			DateTimeFormat: api_service_protos.EDateTimeFormat_YQL_FORMAT,
+		}
+
+		splits, err = prepareSplits(cl, cfg.DataSourceInstance, typeMappingSettings, tableName)
 
 		if err != nil {
 			return fmt.Errorf("prepare splits: %w", err)
@@ -86,11 +89,12 @@ func callServer(logger *zap.Logger, cfg *config.TClientConfig) error {
 
 func prepareSplits(
 	cl Client,
-	tableName string,
 	dsi *api_common.TDataSourceInstance,
+	typeMappingSettings *api_service_protos.TTypeMappingSettings,
+	tableName string,
 ) ([]*api_service_protos.TSplit, error) {
 	// DescribeTable
-	describeTableResponse, err := cl.DescribeTable(context.TODO(), dsi, tableName)
+	describeTableResponse, err := cl.DescribeTable(context.TODO(), dsi, typeMappingSettings, tableName)
 	if err != nil {
 		return nil, fmt.Errorf("describe table: %w", err)
 	}
@@ -126,38 +130,25 @@ func readSplits(
 		return fmt.Errorf("read splits: %w", err)
 	}
 
-	if err := dumpReadResponses(logger, readSplitsResponses); err != nil {
-		return fmt.Errorf("dump read responses: %w", err)
+	records, err := common.ReadResponsesToArrowRecords(readSplitsResponses)
+	if err != nil {
+		return fmt.Errorf("read responses to arrow records: %w", err)
 	}
 
+	dumpReadResponses(logger, records)
+
 	return nil
 }
 
 func dumpReadResponses(
 	logger *zap.Logger,
-	responses []*api_service_protos.TReadSplitsResponse,
-) error {
-	for _, resp := range responses {
-		buf := bytes.NewBuffer(resp.GetArrowIpcStreaming())
-
-		reader, err := ipc.NewReader(buf)
-		if err != nil {
-			return fmt.Errorf("new reader: %w", err)
-		}
-
-		for reader.Next() {
-			record := reader.Record()
-			logger.Debug("schema", zap.String("schema", record.Schema().String()))
-
-			for i, column := range record.Columns() {
-				logger.Debug("column", zap.Int("id", i), zap.String("data", column.String()))
-			}
+	records []arrow.Record,
+) {
+	for _, record := range records {
+		for i, column := range record.Columns() {
+			logger.Debug("column", zap.Int("id", i), zap.String("data", column.String()))
 		}
-
-		reader.Release()
 	}
-
-	return nil
 }
 
 var Cmd = &cobra.Command{

app/client/client.go

@@ -24,6 +24,7 @@ type Client interface {
 	DescribeTable(
 		ctx context.Context,
 		dsi *api_common.TDataSourceInstance,
+		typeMappingSettings *api_service_protos.TTypeMappingSettings,
 		tableName string,
 	) (*api_service_protos.TDescribeTableResponse, error)
 
@@ -49,11 +50,13 @@ type clientImpl struct {
 func (c *clientImpl) DescribeTable(
 	ctx context.Context,
 	dsi *api_common.TDataSourceInstance,
+	typeMappingSettings *api_service_protos.TTypeMappingSettings,
 	tableName string,
 ) (*api_service_protos.TDescribeTableResponse, error) {
 	request := &api_service_protos.TDescribeTableRequest{
-		DataSourceInstance: dsi,
-		Table:              tableName,
+		DataSourceInstance:  dsi,
+		Table:               tableName,
+		TypeMappingSettings: typeMappingSettings,
 	}
 
 	return c.client.DescribeTable(ctx, request)

app/common/api_helpers.go

@@ -1,6 +1,12 @@
 package common
 
 import (
+	"bytes"
+	"fmt"
+
+	"github.com/apache/arrow/go/v13/arrow"
+	"github.com/apache/arrow/go/v13/arrow/ipc"
+
 	api_service_protos "github.com/ydb-platform/fq-connector-go/api/service/protos"
 )
 
@@ -42,3 +48,27 @@ func SchemaToSelectWhatItems(
 
 	return out
 }
+
+func ReadResponsesToArrowRecords(responses []*api_service_protos.TReadSplitsResponse) ([]arrow.Record, error) {
+	var out []arrow.Record
+
+	for _, resp := range responses {
+		buf := bytes.NewBuffer(resp.GetArrowIpcStreaming())
+
+		reader, err := ipc.NewReader(buf)
+		if err != nil {
+			return nil, fmt.Errorf("new reader: %w", err)
+		}
+
+		for reader.Next() {
+			record := reader.Record()
+
+			record.Retain()
+			out = append(out, record)
+		}
+
+		reader.Release()
+	}
+
+	return out, nil
+}

app/server/utils/arrow_helpers.go → app/common/arrow_helpers.go

@@ -1,17 +1,38 @@
-package utils
+package common
 
 import (
 	"fmt"
+	"time"
 
 	"github.com/apache/arrow/go/v13/arrow"
 	"github.com/apache/arrow/go/v13/arrow/array"
 	"github.com/apache/arrow/go/v13/arrow/memory"
 	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb"
 
 	api_service_protos "github.com/ydb-platform/fq-connector-go/api/service/protos"
-	"github.com/ydb-platform/fq-connector-go/app/common"
 )
 
+type ValueType interface {
+	bool |
+		int8 | int16 | int32 | int64 |
+		uint8 | uint16 | uint32 | uint64 |
+		float32 | float64 |
+		string | []byte |
+		time.Time
+}
+
+type ArrowArrayType[VT ValueType] interface {
+	*array.Boolean |
+		*array.Int8 | *array.Int16 | *array.Int32 | *array.Int64 |
+		*array.Uint8 | *array.Uint16 | *array.Uint32 | *array.Uint64 |
+		*array.Float32 | *array.Float64 |
+		*array.String | *array.Binary
+
+	Len() int
+	Value(int) VT
+	IsNull(int) bool
+}
+
 type ArrowBuilder[VT ValueType] interface {
 	AppendNull()
 	Append(value VT)
@@ -49,7 +70,7 @@ func SelectWhatToArrowSchema(selectWhat *api_service_protos.TSelect_TWhat) (*arr
 		default:
 			err := fmt.Errorf(
 				"only primitive and optional types are supported, got '%T' instead: %w",
-				t, common.ErrDataTypeNotSupported,
+				t, ErrDataTypeNotSupported,
 			)
 
 			return nil, err
@@ -86,7 +107,7 @@ func YdbTypesToArrowBuilders(ydbTypes []*Ydb.Type, arrowAllocator memory.Allocat
 		default:
 			err := fmt.Errorf(
 				"only primitive and optional types are supported, got '%T' instead: %w",
-				t, common.ErrDataTypeNotSupported,
+				t, ErrDataTypeNotSupported,
 			)
 
 			return nil, err
@@ -140,7 +161,7 @@ func ydbTypeToArrowBuilder(typeID Ydb.Type_PrimitiveTypeId, arrowAllocator memor
 	case Ydb.Type_TIMESTAMP:
 		builder = array.NewUint64Builder(arrowAllocator)
 	default:
-		return nil, fmt.Errorf("register type '%v': %w", typeID, common.ErrDataTypeNotSupported)
+		return nil, fmt.Errorf("register type '%v': %w", typeID, ErrDataTypeNotSupported)
 	}
 
 	return builder, nil
@@ -188,7 +209,7 @@ func ydbTypeToArrowField(typeID Ydb.Type_PrimitiveTypeId, column *Ydb.Column) (a
 	case Ydb.Type_TIMESTAMP:
 		field = arrow.Field{Name: column.Name, Type: arrow.PrimitiveTypes.Uint64}
 	default:
-		return arrow.Field{}, fmt.Errorf("register type '%v': %w", typeID, common.ErrDataTypeNotSupported)
+		return arrow.Field{}, fmt.Errorf("register type '%v': %w", typeID, ErrDataTypeNotSupported)
 	}
 
 	return field, nil

app/server/utils/time.go → app/common/time.go

@@ -1,10 +1,8 @@
-package utils
+package common
 
 import (
 	"fmt"
 	"time"
-
-	"github.com/ydb-platform/fq-connector-go/app/common"
 )
 
 var (
@@ -15,7 +13,7 @@ var (
 
 func TimeToYDBDate(t *time.Time) (uint16, error) {
 	if t.Before(minYDBTime) || t.After(maxYDBTime) {
-		return 0, fmt.Errorf("convert '%v' to YDB Date: %w", t, common.ErrValueOutOfTypeBounds)
+		return 0, fmt.Errorf("convert '%v' to YDB Date: %w", t, ErrValueOutOfTypeBounds)
 	}
 
 	days := t.Sub(minYDBTime).Hours() / 24
@@ -25,7 +23,7 @@ func TimeToYDBDate(t *time.Time) (uint16, error) {
 
 func TimeToYDBDatetime(t *time.Time) (uint32, error) {
 	if t.Before(minYDBTime) || t.After(maxYDBTime) {
-		return 0, fmt.Errorf("convert '%v' to YDB Date: %w", t, common.ErrValueOutOfTypeBounds)
+		return 0, fmt.Errorf("convert '%v' to YDB Date: %w", t, ErrValueOutOfTypeBounds)
 	}
 
 	seconds := t.Unix()
@@ -35,10 +33,23 @@ func TimeToYDBDatetime(t *time.Time) (uint32, error) {
 
 func TimeToYDBTimestamp(t *time.Time) (uint64, error) {
 	if t.Before(minYDBTime) || t.After(maxYDBTime) {
-		return 0, fmt.Errorf("convert '%v' to YDB Date: %w", t, common.ErrValueOutOfTypeBounds)
+		return 0, fmt.Errorf("convert '%v' to YDB Date: %w", t, ErrValueOutOfTypeBounds)
 	}
 
 	seconds := t.UnixMicro()
 
 	return uint64(seconds), nil
 }
+
+type ydbTime interface {
+	uint16 | uint32 | uint64
+}
+
+func MustTimeToYDBType[OUT ydbTime](f func(t *time.Time) (OUT, error), t time.Time) OUT {
+	res, err := f(&t)
+	if err != nil {
+		panic(err)
+	}
+
+	return res
+}

app/server/utils/time_test.go → app/common/time_test.go

@@ -1,13 +1,11 @@
-package utils
+package common
 
 import (
 	"errors"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/require"
-
-	"github.com/ydb-platform/fq-connector-go/app/common"
 )
 
 func TestTimeToYDBDate(t *testing.T) {
@@ -31,12 +29,12 @@ func TestTimeToYDBDate(t *testing.T) {
 		{
 			input:  time.Date(1969, 12, 31, 23, 59, 00, 00, time.UTC),
 			output: 0,
-			err:    common.ErrValueOutOfTypeBounds,
+			err:    ErrValueOutOfTypeBounds,
 		},
 		{
 			input:  time.Date(9999, 01, 01, 00, 00, 00, 00, time.UTC),
 			output: 0,
-			err:    common.ErrValueOutOfTypeBounds,
+			err:    ErrValueOutOfTypeBounds,
 		},
 	}
 
@@ -48,7 +46,7 @@ func TestTimeToYDBDate(t *testing.T) {
 			require.Equal(t, tc.output, output)
 
 			if tc.err != nil {
-				require.True(t, errors.Is(tc.err, common.ErrValueOutOfTypeBounds))
+				require.True(t, errors.Is(tc.err, ErrValueOutOfTypeBounds))
 			} else {
 				require.NoError(t, err)
 			}
@@ -77,12 +75,12 @@ func TestTimeToYDBDatetime(t *testing.T) {
 		{
 			input:  time.Date(1969, 12, 31, 23, 59, 00, 00, time.UTC),
 			output: 0,
-			err:    common.ErrValueOutOfTypeBounds,
+			err:    ErrValueOutOfTypeBounds,
 		},
 		{
 			input:  time.Date(9999, 01, 01, 00, 00, 00, 00, time.UTC),
 			output: 0,
-			err:    common.ErrValueOutOfTypeBounds,
+			err:    ErrValueOutOfTypeBounds,
 		},
 	}
 
@@ -94,7 +92,7 @@ func TestTimeToYDBDatetime(t *testing.T) {
 			require.Equal(t, tc.output, output)
 
 			if tc.err != nil {
-				require.True(t, errors.Is(tc.err, common.ErrValueOutOfTypeBounds))
+				require.True(t, errors.Is(tc.err, ErrValueOutOfTypeBounds))
 			} else {
 				require.NoError(t, err)
 			}
@@ -123,12 +121,12 @@ func TestTimeToYDBTimestamp(t *testing.T) {
 		{
 			input:  time.Date(1969, 12, 31, 23, 59, 00, 00, time.UTC),
 			output: 0,
-			err:    common.ErrValueOutOfTypeBounds,
+			err:    ErrValueOutOfTypeBounds,
 		},
 		{
 			input:  time.Date(29427, 01, 01, 00, 00, 00, 00, time.UTC),
 			output: 0,
-			err:    common.ErrValueOutOfTypeBounds,
+			err:    ErrValueOutOfTypeBounds,
 		},
 	}
 
@@ -140,7 +138,7 @@ func TestTimeToYDBTimestamp(t *testing.T) {
 			require.Equal(t, tc.output, output)
 
 			if tc.err != nil {
-				require.True(t, errors.Is(tc.err, common.ErrValueOutOfTypeBounds))
+				require.True(t, errors.Is(tc.err, ErrValueOutOfTypeBounds))
 			} else {
 				require.NoError(t, err)
 			}