Cover ReadSplits method

.github/workflows/build.yml

@@ -37,10 +37,12 @@ jobs:
       - name: Run integration tests
         run: |
           go test -c -o fq-connector-go-tests -coverpkg=./... -covermode=atomic ./tests
-          ./fq-connector-go-tests -projectPath=$(pwd) -test.coverprofile=coverage_integratoin_tests.out 
+          ./fq-connector-go-tests -projectPath=$(pwd) -test.coverprofile=coverage_integration_tests.out 
       - name: Union coverage
         run: |
-          cat coverage* | grep -v pb.go | grep -v mock.go > coverage.out
+          cat coverage_unit_tests.out | grep -v 'pb.go\|mock.go\|library' > coverage.out
+          cat coverage_integration_tests.out | grep -v 'atomic\|pb.go\|mock.go\|library' >> coverage.out
+          go tool cover -func=coverage.out
       - name: Upload coverage to Codecov
         uses: codecov/codecov-action@v4-beta
         with:

Makefile

@@ -1,3 +1,5 @@
+PROJECT_PATH = $(shell pwd)
+
 build:
 	go build -o fq-connector-go ./app
 
@@ -10,14 +12,21 @@ lint:
 unit_test:
 	go test -v ./app/...
 
-PROJECT_PATH = $(shell pwd)
-integration_test: integration_test_clean
-	docker-compose -f ./tests/infra/datasource/docker-compose.yaml up -d
+integration_test: integration_test_env
 	go test -c -o fq-connector-go-tests ./tests
 	./fq-connector-go-tests -projectPath=$(PROJECT_PATH)
 
-integration_test_clean:
+integration_test_env:
 	docker-compose -f ./tests/infra/datasource/docker-compose.yaml rm -f -v 
+	docker-compose -f ./tests/infra/datasource/docker-compose.yaml up -d
 
+test_coverage: integration_test_env
+	go test -coverpkg=./... -coverprofile=coverage_unit_tests.out -covermode=atomic ./app/...
+	go test -c -o fq-connector-go-tests -coverpkg=./... -covermode=atomic ./tests
+	./fq-connector-go-tests -projectPath=$(PROJECT_PATH) -test.coverprofile=coverage_integration_tests.out 
+	cat coverage_unit_tests.out | grep -v 'pb.go\|mock.go\|library' > coverage.out
+	cat coverage_integration_tests.out | grep -v 'atomic\|pb.go\|mock.go\|library' >> coverage.out
+	go tool cover -func=coverage.out
+
 build_image_base: 
 	docker build -t ghcr.io/ydb-platform/fq-connector-go:base -f ./Dockerfile.base .

app/server/data_source_collection.go

@@ -54,18 +54,18 @@ func (dsc *DataSourceCollection) DoReadSplit(
 	request *api_service_protos.TReadSplitsRequest,
 	split *api_service_protos.TSplit,
 ) error {
-	switch kind := request.GetDataSourceInstance().GetKind(); kind {
+	switch kind := split.GetSelect().GetDataSourceInstance().GetKind(); kind {
 	case api_common.EDataSourceKind_CLICKHOUSE, api_common.EDataSourceKind_POSTGRESQL, api_common.EDataSourceKind_YDB:
 		ds, err := dsc.rdbms.Make(logger, kind)
 		if err != nil {
 			return err
 		}
 
-		return readSplit[any](logger, stream, request, split, ds, dsc.memoryAllocator, dsc.readLimiterFactory, dsc.cfg)
+		return readSplit[any](logger, stream, request.GetFormat(), split, ds, dsc.memoryAllocator, dsc.readLimiterFactory, dsc.cfg)
 	case api_common.EDataSourceKind_S3:
 		ds := s3.NewDataSource()
 
-		return readSplit[string](logger, stream, request, split, ds, dsc.memoryAllocator, dsc.readLimiterFactory, dsc.cfg)
+		return readSplit[string](logger, stream, request.GetFormat(), split, ds, dsc.memoryAllocator, dsc.readLimiterFactory, dsc.cfg)
 	default:
 		return fmt.Errorf("unsupported data source type '%v': %w", kind, utils.ErrDataSourceNotSupported)
 	}
@@ -74,7 +74,7 @@ func (dsc *DataSourceCollection) DoReadSplit(
 func readSplit[T utils.Acceptor](
 	logger *zap.Logger,
 	stream api_service.Connector_ReadSplitsServer,
-	request *api_service_protos.TReadSplitsRequest,
+	format api_service_protos.TReadSplitsRequest_EFormat,
 	split *api_service_protos.TSplit,
 	dataSource datasource.DataSource[T],
 	memoryAllocator memory.Allocator,
@@ -86,7 +86,7 @@ func readSplit[T utils.Acceptor](
 	columnarBufferFactory, err := paging.NewColumnarBufferFactory[T](
 		logger,
 		memoryAllocator,
-		request.Format,
+		format,
 		split.Select.What)
 	if err != nil {
 		return fmt.Errorf("new columnar buffer factory: %w", err)
@@ -109,7 +109,6 @@ func readSplit[T utils.Acceptor](
 	streamer := streaming.NewStreamer(
 		logger,
 		stream,
-		request,
 		split,
 		sink,
 		dataSource,

app/server/service_connector.go

@@ -38,7 +38,7 @@ func (s *serviceConnector) DescribeTable(
 	ctx context.Context,
 	request *api_service_protos.TDescribeTableRequest,
 ) (*api_service_protos.TDescribeTableResponse, error) {
-	logger := utils.AnnotateLogger(s.logger, "DescribeTable", request.DataSourceInstance)
+	logger := utils.AnnotateLoggerForUnaryCall(s.logger, "DescribeTable", request.DataSourceInstance)
 	logger.Info("request handling started", zap.String("table", request.GetTable()))
 
 	if err := ValidateDescribeTableRequest(logger, request); err != nil {
@@ -65,7 +65,7 @@ func (s *serviceConnector) DescribeTable(
 }
 
 func (s *serviceConnector) ListSplits(request *api_service_protos.TListSplitsRequest, stream api_service.Connector_ListSplitsServer) error {
-	logger := utils.AnnotateLogger(s.logger, "ListSplits", nil)
+	logger := utils.AnnotateLoggerWithMethod(s.logger, "ListSplits")
 	logger.Info("request handling started", zap.Int("total selects", len(request.Selects)))
 
 	if err := ValidateListSplitsRequest(logger, request); err != nil {
@@ -77,7 +77,7 @@ func (s *serviceConnector) ListSplits(request *api_service_protos.TListSplitsReq
 	totalSplits := 0
 
 	for _, slct := range request.Selects {
-		if err := s.doListSplitsHandleSelect(stream, slct, &totalSplits); err != nil {
+		if err := s.doListSplitsHandleSelect(logger, stream, slct, &totalSplits); err != nil {
 			logger.Error("request handling failed", zap.Error(err))
 
 			return err
@@ -90,11 +90,12 @@ func (s *serviceConnector) ListSplits(request *api_service_protos.TListSplitsReq
 }
 
 func (s *serviceConnector) doListSplitsHandleSelect(
+	logger *zap.Logger,
 	stream api_service.Connector_ListSplitsServer,
 	slct *api_service_protos.TSelect,
 	totalSplits *int,
 ) error {
-	logger := utils.AnnotateLogger(s.logger, "ListSplits", slct.DataSourceInstance)
+	logger = utils.AnnotateLoggerWithDataSourceInstance(logger, slct.DataSourceInstance)
 
 	args := []zap.Field{
 		zap.Int("split_id", *totalSplits),
@@ -146,8 +147,9 @@ func (*serviceConnector) doListSplitsResponse(
 
 func (s *serviceConnector) ReadSplits(
 	request *api_service_protos.TReadSplitsRequest,
-	stream api_service.Connector_ReadSplitsServer) error {
-	logger := utils.AnnotateLogger(s.logger, "ReadSplits", request.DataSourceInstance)
+	stream api_service.Connector_ReadSplitsServer,
+) error {
+	logger := utils.AnnotateLoggerWithMethod(s.logger, "ReadSplits")
 	logger.Info("request handling started", zap.Int("total_splits", len(request.Splits)))
 
 	err := s.doReadSplits(logger, request, stream)
@@ -176,7 +178,9 @@ func (s *serviceConnector) doReadSplits(
 	}
 
 	for i, split := range request.Splits {
-		splitLogger := logger.With(zap.Int("split_id", i))
+		splitLogger := utils.
+			AnnotateLoggerWithDataSourceInstance(logger, split.Select.DataSourceInstance).
+			With(zap.Int("split_id", i))
 
 		err := s.dataSourceCollection.DoReadSplit(
 			splitLogger,

app/server/streaming/streamer.go

@@ -16,7 +16,6 @@ import (
 
 type Streamer[T utils.Acceptor] struct {
 	stream     api_service.Connector_ReadSplitsServer
-	request    *api_service_protos.TReadSplitsRequest
 	dataSource datasource.DataSource[T]
 	split      *api_service_protos.TSplit
 	sink       paging.Sink[T]
@@ -98,7 +97,6 @@ func (s *Streamer[T]) Run() error {
 func NewStreamer[T utils.Acceptor](
 	logger *zap.Logger,
 	stream api_service.Connector_ReadSplitsServer,
-	request *api_service_protos.TReadSplitsRequest,
 	split *api_service_protos.TSplit,
 	sink paging.Sink[T],
 	dataSource datasource.DataSource[T],
@@ -109,7 +107,6 @@ func NewStreamer[T utils.Acceptor](
 		logger:     logger,
 		stream:     stream,
 		split:      split,
-		request:    request,
 		dataSource: dataSource,
 		sink:       sink,
 		ctx:        ctx,

app/server/streaming/streamer_test.go

@@ -143,7 +143,6 @@ func (tc testCaseStreaming) messageParams() (sentMessages, rowsInLastMessage int
 
 func (tc testCaseStreaming) execute(t *testing.T) {
 	logger := utils.NewTestLogger(t)
-	request := &api_service_protos.TReadSplitsRequest{}
 	split := utils.MakeTestSplit()
 
 	ctx, cancel := context.WithCancel(context.Background())
@@ -258,7 +257,7 @@ func (tc testCaseStreaming) execute(t *testing.T) {
 	)
 	require.NoError(t, err)
 
-	streamer := NewStreamer(logger, stream, request, split, sink, dataSource)
+	streamer := NewStreamer(logger, stream, split, sink, dataSource)
 
 	err = streamer.Run()
 

app/server/utils/logger.go

@@ -16,22 +16,27 @@ import (
 
 // TODO: it's better to do this in GRPC middleware
 
-func AnnotateLogger(l *zap.Logger, method string, dsi *api_common.TDataSourceInstance) *zap.Logger {
-	logger := l.With(zap.String("method", method))
-
-	if dsi != nil {
-		logger = logger.With(
-			zap.String("data_source_kind", api_common.EDataSourceKind_name[int32(dsi.GetKind())]),
-			zap.String("host", dsi.GetEndpoint().GetHost()),
-			zap.Uint32("port", dsi.GetEndpoint().GetPort()),
-			zap.String("database", dsi.GetDatabase()),
-			zap.Bool("use_tls", dsi.GetUseTls()),
-			zap.String("protocol", dsi.GetProtocol().String()),
-			// TODO: can we print just a login without a password?
-		)
-	}
+func AnnotateLoggerWithMethod(l *zap.Logger, method string) *zap.Logger {
+	return l.With(zap.String("method", method))
+}
+
+func AnnotateLoggerWithDataSourceInstance(l *zap.Logger, dsi *api_common.TDataSourceInstance) *zap.Logger {
+	return l.With(
+		zap.String("data_source_kind", api_common.EDataSourceKind_name[int32(dsi.GetKind())]),
+		zap.String("host", dsi.GetEndpoint().GetHost()),
+		zap.Uint32("port", dsi.GetEndpoint().GetPort()),
+		zap.String("database", dsi.GetDatabase()),
+		zap.Bool("use_tls", dsi.GetUseTls()),
+		zap.String("protocol", dsi.GetProtocol().String()),
+		// TODO: can we print just a login without a password?
+	)
+}
+
+func AnnotateLoggerForUnaryCall(l *zap.Logger, method string, dsi *api_common.TDataSourceInstance) *zap.Logger {
+	l = AnnotateLoggerWithMethod(l, method)
+	l = AnnotateLoggerWithDataSourceInstance(l, dsi)
 
-	return logger
+	return l
 }
 
 func LogCloserError(logger *zap.Logger, closer io.Closer, msg string) {

app/server/validate.go

@@ -37,14 +37,24 @@ func ValidateListSplitsRequest(logger *zap.Logger, request *api_service_protos.T
 }
 
 func ValidateReadSplitsRequest(logger *zap.Logger, request *api_service_protos.TReadSplitsRequest) error {
-	if err := validateDataSourceInstance(logger, request.GetDataSourceInstance()); err != nil {
-		return fmt.Errorf("validate data source instance: %w", err)
-	}
-
 	if len(request.Splits) == 0 {
 		return fmt.Errorf("splits are empty: %w", utils.ErrInvalidRequest)
 	}
 
+	for i, split := range request.Splits {
+		if err := validateSplit(logger, split); err != nil {
+			return fmt.Errorf("validate split #%d: %w", i, err)
+		}
+	}
+
+	return nil
+}
+
+func validateSplit(logger *zap.Logger, split *api_service_protos.TSplit) error {
+	if err := validateSelect(logger, split.Select); err != nil {
+		return fmt.Errorf("validate select: %w", err)
+	}
+
 	return nil
 }
 
@@ -61,6 +71,10 @@ func validateSelect(logger *zap.Logger, slct *api_service_protos.TSelect) error
 }
 
 func validateDataSourceInstance(logger *zap.Logger, dsi *api_common.TDataSourceInstance) error {
+	if dsi == nil {
+		return fmt.Errorf("empty data source instance: %w", utils.ErrInvalidRequest)
+	}
+
 	if dsi.GetKind() == api_common.EDataSourceKind_DATA_SOURCE_KIND_UNSPECIFIED {
 		return fmt.Errorf("empty kind: %w", utils.ErrInvalidRequest)
 	}