Merge pull request #1491 from UgnineSirdis/bulk-upsert-from-csv

Support bulk upsert from scv, arrow and ydb internal formats

Author: asmyasnikov

CHANGELOG.md

@@ -1,8 +1,10 @@
+* Support bulk upsert from scv, arrow and ydb internal formats in table client
+
 ## v3.82.0
-* Fixed error on experimental TopicListener.Close
+* Fixed error on experimental `TopicListener.Close`
 * Disabled reporting of `ydb_go_sdk_query_session_count` when metrics are disabled
 * Disabled reporting of `ydb_go_sdk_ydb_query_session_create_latency` histogram metrics when metrics are disabled
-* Allowed skip column for `ScanStruct` by tag `-` 
+* Allowed skip column for `ScanStruct` by tag `-`
 
 ## v3.81.4
 * Returned `topicwriter.ErrQueueLimitExceed`, accidental removed at `v3.81.0`
@@ -14,15 +16,15 @@
 * Removed `experimantal` comment for query service client
 
 ## v3.81.1
-* Fixed nil pointer dereference panic on failed `ydb.Open` 
+* Fixed nil pointer dereference panic on failed `ydb.Open`
 * Added ip discovery. Server can show own ip address and target hostname in the ListEndpoint message. These fields are used to bypass DNS resolving.
 
 ## v3.81.0
 * Added error ErrMessagesPutToInternalQueueBeforeError to topic writer
 * Added write to topics within transactions
 
 ## v3.80.10
-* Added `ydb.WithSessionPoolSessionUsageLimit()` option for limitation max count of session usage  
+* Added `ydb.WithSessionPoolSessionUsageLimit()` option for limitation max count of session usage
 * Refactored experimental topic iterators in `topicsugar` package
 
 ## v3.80.9

internal/table/client.go

@@ -4,8 +4,11 @@ import (
 	"context"
 
 	"github.com/jonboulle/clockwork"
+	"github.com/ydb-platform/ydb-go-genproto/Ydb_Table_V1"
+	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb_Table"
 	"google.golang.org/grpc"
 
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/allocator"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/pool"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/stack"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/table/config"
@@ -27,6 +30,7 @@ func New(ctx context.Context, cc grpc.ClientConnInterface, config *config.Config
 	return &Client{
 		clock:  config.Clock(),
 		config: config,
+		client: Ydb_Table_V1.NewTableServiceClient(cc),
 		cc:     cc,
 		build: func(ctx context.Context) (s *session, err error) {
 			return newSession(ctx, cc, config)
@@ -86,6 +90,7 @@ func New(ctx context.Context, cc grpc.ClientConnInterface, config *config.Config
 type Client struct {
 	// read-only fields
 	config *config.Config
+	client Ydb_Table_V1.TableServiceClient
 	build  sessionBuilder
 	cc     grpc.ClientConnInterface
 	clock  clockwork.Clock
@@ -264,6 +269,54 @@ func (c *Client) DoTx(ctx context.Context, op table.TxOperation, opts ...table.O
 	}, config.RetryOptions...)
 }
 
+func (c *Client) BulkUpsert(
+	ctx context.Context,
+	tableName string,
+	data table.BulkUpsertData,
+	opts ...table.Option,
+) (finalErr error) {
+	if c == nil {
+		return xerrors.WithStackTrace(errNilClient)
+	}
+
+	if c.isClosed() {
+		return xerrors.WithStackTrace(errClosedClient)
+	}
+
+	a := allocator.New()
+	defer a.Free()
+
+	config := c.retryOptions(opts...)
+	config.RetryOptions = append(config.RetryOptions, retry.WithIdempotent(true))
+
+	attempts, onDone := 0, trace.TableOnBulkUpsert(config.Trace, &ctx,
+		stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/v3/internal/table.(*Client).BulkUpsert"),
+	)
+	defer func() {
+		onDone(finalErr, attempts)
+	}()
+
+	request := Ydb_Table.BulkUpsertRequest{
+		Table: tableName,
+	}
+	finalErr = data.ApplyBulkUpsertRequest(a, (*table.BulkUpsertRequest)(&request))
+	if finalErr != nil {
+		return finalErr
+	}
+
+	finalErr = retry.Retry(ctx,
+		func(ctx context.Context) (err error) {
+			attempts++
+			_, err = c.client.BulkUpsert(ctx, &request)
+
+			return err
+		},
+		config.RetryOptions...,
+	)
+
+	return xerrors.WithStackTrace(finalErr)
+}
+
 func executeTxOperation(ctx context.Context, c *Client, op table.TxOperation, tx table.Transaction) (err error) {
 	if panicCallback := c.config.PanicCallback(); panicCallback != nil {
 		defer func() {

internal/table/session.go

@@ -1261,12 +1261,11 @@ func (s *session) BulkUpsert(ctx context.Context, table string, rows value.Value
 		onDone      = trace.TableOnSessionBulkUpsert(
 			s.config.Trace(), &ctx,
 			stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/v3/internal/table.(*session).BulkUpsert"),
-			s,
 		)
 	)
 	defer func() {
 		defer a.Free()
-		onDone(err)
+		onDone(err, 1)
 	}()
 
 	for _, opt := range opts {

table/table.go

@@ -4,8 +4,10 @@ import (
 	"context"
 	"time"
 
+	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb_Formats"
 	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb_Table"
 
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/allocator"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/closer"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/params"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/tx"
@@ -68,6 +70,12 @@ type Client interface {
 	// If op TxOperation return non nil - transaction will be rollback
 	// Warning: if context without deadline or cancellation func than DoTx can run indefinitely
 	DoTx(ctx context.Context, op TxOperation, opts ...Option) error
+
+	// BulkUpsert upserts a batch of rows non-transactionally.
+	//
+	// Returns success only when all rows were successfully upserted. In case of an error some rows might
+	// be upserted and some might not.
+	BulkUpsert(ctx context.Context, table string, data BulkUpsertData, opts ...Option) error
 }
 
 type SessionStatus = string
@@ -179,6 +187,7 @@ type Session interface {
 		opts ...options.ExecuteScanQueryOption,
 	) (_ result.StreamResult, err error)
 
+	// Deprecated: use Client instance instead.
 	BulkUpsert(
 		ctx context.Context,
 		table string,
@@ -578,3 +587,205 @@ func (opt traceOption) ApplyTableOption(opts *Options) {
 func WithTrace(t trace.Table) traceOption { //nolint:gocritic
 	return traceOption{t: &t}
 }
+
+type (
+	BulkUpsertRequest Ydb_Table.BulkUpsertRequest
+)
+
+type BulkUpsertData interface {
+	ApplyBulkUpsertRequest(a *allocator.Allocator, req *BulkUpsertRequest) error
+}
+
+type bulkUpsertRows struct {
+	Rows value.Value
+}
+
+func (data bulkUpsertRows) ApplyBulkUpsertRequest(a *allocator.Allocator, req *BulkUpsertRequest) error {
+	req.Rows = value.ToYDB(data.Rows, a)
+
+	return nil
+}
+
+func NewBulkUpsertRows(rows value.Value) bulkUpsertRows {
+	return bulkUpsertRows{
+		Rows: rows,
+	}
+}
+
+type bulkUpsertCsv struct {
+	Data    []byte
+	Options []CsvFormatOption
+}
+
+type CsvFormatOption interface {
+	ApplyCsvFormatOption(req *BulkUpsertRequest) (err error)
+}
+
+func (data bulkUpsertCsv) ApplyBulkUpsertRequest(a *allocator.Allocator, req *BulkUpsertRequest) error {
+	req.Data = data.Data
+
+	var err error
+	for _, opt := range data.Options {
+		if opt != nil {
+			err = opt.ApplyCsvFormatOption(req)
+			if err != nil {
+				return err
+			}
+		}
+	}
+
+	return err
+}
+
+func NewBulkUpsertCsv(data []byte, opts ...CsvFormatOption) bulkUpsertCsv {
+	return bulkUpsertCsv{
+		Data:    data,
+		Options: opts,
+	}
+}
+
+func ensureCsvDataFormatSettings(req *BulkUpsertRequest) (format *Ydb_Formats.CsvSettings) {
+	if settings, ok := req.DataFormat.(*Ydb_Table.BulkUpsertRequest_CsvSettings); ok {
+		if settings.CsvSettings == nil {
+			settings.CsvSettings = &Ydb_Formats.CsvSettings{}
+		}
+
+		return settings.CsvSettings
+	}
+
+	req.DataFormat = &Ydb_Table.BulkUpsertRequest_CsvSettings{
+		CsvSettings: &Ydb_Formats.CsvSettings{},
+	}
+
+	settings, ok := req.DataFormat.(*Ydb_Table.BulkUpsertRequest_CsvSettings)
+	if !ok {
+		return nil
+	}
+
+	return settings.CsvSettings
+}
+
+type csvHeaderOption struct{}
+
+func (opt *csvHeaderOption) ApplyCsvFormatOption(req *BulkUpsertRequest) error {
+	ensureCsvDataFormatSettings(req).Header = true
+
+	return nil
+}
+
+// First not skipped line is a CSV header (list of column names).
+func WithCsvHeader() CsvFormatOption {
+	return &csvHeaderOption{}
+}
+
+type csvNullValueOption struct {
+	Value []byte
+}
+
+func (opt *csvNullValueOption) ApplyCsvFormatOption(req *BulkUpsertRequest) error {
+	ensureCsvDataFormatSettings(req).NullValue = opt.Value
+
+	return nil
+}
+
+// String value that would be interpreted as NULL.
+func WithCsvNullValue(value []byte) CsvFormatOption {
+	return &csvNullValueOption{value}
+}
+
+type csvDelimiterOption struct {
+	Value []byte
+}
+
+func (opt *csvDelimiterOption) ApplyCsvFormatOption(req *BulkUpsertRequest) error {
+	ensureCsvDataFormatSettings(req).Delimiter = opt.Value
+
+	return nil
+}
+
+// Fields delimiter in CSV file. It's "," if not set.
+func WithCsvDelimiter(value []byte) CsvFormatOption {
+	return &csvDelimiterOption{value}
+}
+
+type csvSkipRowsOption struct {
+	Count uint32
+}
+
+func (opt *csvSkipRowsOption) ApplyCsvFormatOption(req *BulkUpsertRequest) error {
+	ensureCsvDataFormatSettings(req).SkipRows = opt.Count
+
+	return nil
+}
+
+// Number of rows to skip before CSV data. It should be present only in the first upsert of CSV file.
+func WithCsvSkipRows(count uint32) CsvFormatOption {
+	return &csvSkipRowsOption{count}
+}
+
+type bulkUpsertArrow struct {
+	Data    []byte
+	Options []ArrowFormatOption
+}
+
+type ArrowFormatOption interface {
+	ApplyArrowFormatOption(req *BulkUpsertRequest) (err error)
+}
+
+func (data bulkUpsertArrow) ApplyBulkUpsertRequest(a *allocator.Allocator, req *BulkUpsertRequest) error {
+	req.Data = data.Data
+
+	var err error
+	for _, opt := range data.Options {
+		if opt != nil {
+			err = opt.ApplyArrowFormatOption(req)
+			if err != nil {
+				return err
+			}
+		}
+	}
+
+	return err
+}
+
+func NewBulkUpsertArrow(data []byte, opts ...ArrowFormatOption) bulkUpsertArrow {
+	return bulkUpsertArrow{
+		Data:    data,
+		Options: opts,
+	}
+}
+
+func ensureArrowDataFormatSettings(req *BulkUpsertRequest) (format *Ydb_Formats.ArrowBatchSettings) {
+	if settings, ok := req.DataFormat.(*Ydb_Table.BulkUpsertRequest_ArrowBatchSettings); ok {
+		if settings.ArrowBatchSettings == nil {
+			settings.ArrowBatchSettings = &Ydb_Formats.ArrowBatchSettings{}
+		}
+
+		return settings.ArrowBatchSettings
+	}
+
+	req.DataFormat = &Ydb_Table.BulkUpsertRequest_ArrowBatchSettings{
+		ArrowBatchSettings: &Ydb_Formats.ArrowBatchSettings{},
+	}
+
+	settings, ok := req.DataFormat.(*Ydb_Table.BulkUpsertRequest_ArrowBatchSettings)
+	if !ok {
+		return nil
+	}
+
+	return settings.ArrowBatchSettings
+}
+
+type arrowSchemaOption struct {
+	Schema []byte
+}
+
+func (opt *arrowSchemaOption) ApplyArrowFormatOption(req *BulkUpsertRequest) error {
+	ensureArrowDataFormatSettings(req).Schema = opt.Schema
+
+	return nil
+}
+
+func WithArrowSchema(schema []byte) ArrowFormatOption {
+	return &arrowSchemaOption{schema}
+}

tests/integration/helpers_test.go

@@ -180,7 +180,7 @@ func (scope *scopeT) Folder() string {
 				scope.Require.NoError(sugar.RemoveRecursive(scope.Ctx, driver, folderPath))
 			}
 		}
-		scope.Logf("Createing folder done: %v", folderPath)
+		scope.Logf("Creating folder done: %v", folderPath)
 		return fixenv.NewGenericResultWithCleanup(folderPath, clean), nil
 	}
 	return fixenv.CacheResult(scope.Env, f)
@@ -310,7 +310,7 @@ func (scope *scopeT) TableName(opts ...func(t *tableNameParams)) string {
 		createTableQueryTemplate: `
 			PRAGMA TablePathPrefix("{{.TablePathPrefix}}");
 			CREATE TABLE {{.TableName}} (
-				id Int64 NOT NULL, 
+				id Int64 NOT NULL,
 				val Text,
 				PRIMARY KEY (id)
 			)