Skip to content

Commit 865024a

Browse files
asmyasnikovasmyasnikov
committed
replaced ensure... to explicit constructors of entities
1 parent af3be44 commit 865024a

File tree

2 files changed

+90
-119
lines changed

2 files changed

+90
-119
lines changed

‎internal/table/client.go‎

Lines changed: 21 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,6 @@ import (
55

66
"github.com/jonboulle/clockwork"
77
"github.com/ydb-platform/ydb-go-genproto/Ydb_Table_V1"
8-
"github.com/ydb-platform/ydb-go-genproto/protos/Ydb_Table"
98
"google.golang.org/grpc"
109

1110
"github.com/ydb-platform/ydb-go-sdk/v3/internal/allocator"
@@ -286,35 +285,44 @@ func (c *Client) BulkUpsert(
286285
a := allocator.New()
287286
defer a.Free()
288287

289-
config := c.retryOptions(opts...)
290-
config.RetryOptions = append(config.RetryOptions, retry.WithIdempotent(true))
288+
attempts, config := 0, c.retryOptions(opts...)
289+
config.RetryOptions = append(config.RetryOptions,
290+
retry.WithIdempotent(true),
291+
retry.WithTrace(&trace.Retry{
292+
OnRetry: func(info trace.RetryLoopStartInfo) func(trace.RetryLoopDoneInfo) {
293+
return func(info trace.RetryLoopDoneInfo) {
294+
attempts = info.Attempts
295+
}
296+
},
297+
}),
298+
)
291299

292-
attempts, onDone := 0, trace.TableOnBulkUpsert(config.Trace, &ctx,
300+
onDone := trace.TableOnBulkUpsert(config.Trace, &ctx,
293301
stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/v3/internal/table.(*Client).BulkUpsert"),
294302
)
295303
defer func() {
296304
onDone(finalErr, attempts)
297305
}()
298306

299-
request := Ydb_Table.BulkUpsertRequest{
300-
Table: tableName,
301-
}
302-
finalErr = data.ApplyBulkUpsertRequest(a, (*table.BulkUpsertRequest)(&request))
303-
if finalErr != nil {
304-
return finalErr
307+
request, err := data.ToYDB(a, tableName)
308+
if err != nil {
309+
return xerrors.WithStackTrace(err)
305310
}
306311

307-
finalErr = retry.Retry(ctx,
312+
err = retry.Retry(ctx,
308313
func(ctx context.Context) (err error) {
309314
attempts++
310-
_, err = c.client.BulkUpsert(ctx, &request)
315+
_, err = c.client.BulkUpsert(ctx, request)
311316

312317
return err
313318
},
314319
config.RetryOptions...,
315320
)
321+
if err != nil {
322+
return xerrors.WithStackTrace(err)
323+
}
316324

317-
return xerrors.WithStackTrace(finalErr)
325+
return nil
318326
}
319327

320328
func executeTxOperation(ctx context.Context, c *Client, op table.TxOperation, tx table.Transaction) (err error) {

‎table/table.go‎

Lines changed: 69 additions & 106 deletions
Original file line numberDiff line numberDiff line change
@@ -588,87 +588,71 @@ func WithTrace(t trace.Table) traceOption { //nolint:gocritic
588588
return traceOption{t: &t}
589589
}
590590

591-
type (
592-
BulkUpsertRequest Ydb_Table.BulkUpsertRequest
593-
)
594-
595591
type BulkUpsertData interface {
596-
ApplyBulkUpsertRequest(a *allocator.Allocator, req *BulkUpsertRequest) error
592+
ToYDB(a *allocator.Allocator, tableName string) (*Ydb_Table.BulkUpsertRequest, error)
597593
}
598594

599595
type bulkUpsertRows struct {
600-
Rows value.Value
596+
rows value.Value
601597
}
602598

603-
func (data bulkUpsertRows) ApplyBulkUpsertRequest(a *allocator.Allocator, req *BulkUpsertRequest) error {
604-
req.Rows = value.ToYDB(data.Rows, a)
605-
606-
return nil
599+
func (data bulkUpsertRows) ToYDB(a *allocator.Allocator, tableName string) (*Ydb_Table.BulkUpsertRequest, error) {
600+
return &Ydb_Table.BulkUpsertRequest{
601+
Table: tableName,
602+
Rows: value.ToYDB(data.rows, a),
603+
}, nil
607604
}
608605

609606
func BulkUpsertDataRows(rows value.Value) bulkUpsertRows {
610607
return bulkUpsertRows{
611-
Rows: rows,
608+
rows: rows,
612609
}
613610
}
614611

615612
type bulkUpsertCsv struct {
616-
Data []byte
617-
Options []csvFormatOption
613+
data []byte
614+
opts []csvFormatOption
618615
}
619616

620617
type csvFormatOption interface {
621-
ApplyCsvFormatOption(req *BulkUpsertRequest) (err error)
618+
applyCsvFormatOption(dataFormat *Ydb_Table.BulkUpsertRequest_CsvSettings) (err error)
622619
}
623620

624-
func (data bulkUpsertCsv) ApplyBulkUpsertRequest(a *allocator.Allocator, req *BulkUpsertRequest) error {
625-
req.Data = data.Data
621+
func (data bulkUpsertCsv) ToYDB(a *allocator.Allocator, tableName string) (*Ydb_Table.BulkUpsertRequest, error) {
622+
var (
623+
request = &Ydb_Table.BulkUpsertRequest{
624+
Table: tableName,
625+
Data: data.data,
626+
}
627+
dataFormat = &Ydb_Table.BulkUpsertRequest_CsvSettings{
628+
CsvSettings: &Ydb_Formats.CsvSettings{},
629+
}
630+
)
626631

627-
var err error
628-
for _, opt := range data.Options {
632+
for _, opt := range data.opts {
629633
if opt != nil {
630-
err = opt.ApplyCsvFormatOption(req)
631-
if err != nil {
632-
return err
634+
if err := opt.applyCsvFormatOption(dataFormat); err != nil {
635+
return nil, err
633636
}
634637
}
635638
}
636639

637-
return err
640+
request.DataFormat = dataFormat
641+
642+
return request, nil
638643
}
639644

640645
func BulkUpsertDataCsv(data []byte, opts ...csvFormatOption) bulkUpsertCsv {
641646
return bulkUpsertCsv{
642-
Data: data,
643-
Options: opts,
644-
}
645-
}
646-
647-
func ensureCsvDataFormatSettings(req *BulkUpsertRequest) (format *Ydb_Formats.CsvSettings) {
648-
if settings, ok := req.DataFormat.(*Ydb_Table.BulkUpsertRequest_CsvSettings); ok {
649-
if settings.CsvSettings == nil {
650-
settings.CsvSettings = &Ydb_Formats.CsvSettings{}
651-
}
652-
653-
return settings.CsvSettings
654-
}
655-
656-
req.DataFormat = &Ydb_Table.BulkUpsertRequest_CsvSettings{
657-
CsvSettings: &Ydb_Formats.CsvSettings{},
647+
data: data,
648+
opts: opts,
658649
}
659-
660-
settings, ok := req.DataFormat.(*Ydb_Table.BulkUpsertRequest_CsvSettings)
661-
if !ok {
662-
return nil
663-
}
664-
665-
return settings.CsvSettings
666650
}
667651

668652
type csvHeaderOption struct{}
669653

670-
func (opt *csvHeaderOption) ApplyCsvFormatOption(req *BulkUpsertRequest) error {
671-
ensureCsvDataFormatSettings(req).Header = true
654+
func (opt *csvHeaderOption) applyCsvFormatOption(dataFormat *Ydb_Table.BulkUpsertRequest_CsvSettings) error {
655+
dataFormat.CsvSettings.Header = true
672656

673657
return nil
674658
}
@@ -678,114 +662,93 @@ func WithCsvHeader() csvFormatOption {
678662
return &csvHeaderOption{}
679663
}
680664

681-
type csvNullValueOption struct {
682-
Value []byte
683-
}
665+
type csvNullValueOption []byte
684666

685-
func (opt *csvNullValueOption) ApplyCsvFormatOption(req *BulkUpsertRequest) error {
686-
ensureCsvDataFormatSettings(req).NullValue = opt.Value
667+
func (nullValue csvNullValueOption) applyCsvFormatOption(dataFormat *Ydb_Table.BulkUpsertRequest_CsvSettings) error {
668+
dataFormat.CsvSettings.NullValue = nullValue
687669

688670
return nil
689671
}
690672

691673
// String value that would be interpreted as NULL.
692674
func WithCsvNullValue(value []byte) csvFormatOption {
693-
return &csvNullValueOption{value}
675+
return csvNullValueOption(value)
694676
}
695677

696-
type csvDelimiterOption struct {
697-
Value []byte
698-
}
678+
type csvDelimiterOption []byte
699679

700-
func (opt *csvDelimiterOption) ApplyCsvFormatOption(req *BulkUpsertRequest) error {
701-
ensureCsvDataFormatSettings(req).Delimiter = opt.Value
680+
func (delimeter csvDelimiterOption) applyCsvFormatOption(dataFormat *Ydb_Table.BulkUpsertRequest_CsvSettings) error {
681+
dataFormat.CsvSettings.Delimiter = delimeter
702682

703683
return nil
704684
}
705685

706686
// Fields delimiter in CSV file. It's "," if not set.
707687
func WithCsvDelimiter(value []byte) csvFormatOption {
708-
return &csvDelimiterOption{value}
688+
return csvDelimiterOption(value)
709689
}
710690

711-
type csvSkipRowsOption struct {
712-
Count uint32
713-
}
691+
type csvSkipRowsOption uint32
714692

715-
func (opt *csvSkipRowsOption) ApplyCsvFormatOption(req *BulkUpsertRequest) error {
716-
ensureCsvDataFormatSettings(req).SkipRows = opt.Count
693+
func (skipRows csvSkipRowsOption) applyCsvFormatOption(dataFormat *Ydb_Table.BulkUpsertRequest_CsvSettings) error {
694+
dataFormat.CsvSettings.SkipRows = uint32(skipRows)
717695

718696
return nil
719697
}
720698

721699
// Number of rows to skip before CSV data. It should be present only in the first upsert of CSV file.
722-
func WithCsvSkipRows(count uint32) csvFormatOption {
723-
return &csvSkipRowsOption{count}
700+
func WithCsvSkipRows(skipRows uint32) csvFormatOption {
701+
return csvSkipRowsOption(skipRows)
724702
}
725703

726704
type bulkUpsertArrow struct {
727-
Data []byte
728-
Options []arrowFormatOption
705+
data []byte
706+
opts []arrowFormatOption
729707
}
730708

731709
type arrowFormatOption interface {
732-
ApplyArrowFormatOption(req *BulkUpsertRequest) (err error)
710+
applyArrowFormatOption(req *Ydb_Table.BulkUpsertRequest_ArrowBatchSettings) (err error)
733711
}
734712

735-
func (data bulkUpsertArrow) ApplyBulkUpsertRequest(a *allocator.Allocator, req *BulkUpsertRequest) error {
736-
req.Data = data.Data
713+
func (data bulkUpsertArrow) ToYDB(a *allocator.Allocator, tableName string) (*Ydb_Table.BulkUpsertRequest, error) {
714+
var (
715+
request = &Ydb_Table.BulkUpsertRequest{
716+
Table: tableName,
717+
Data: data.data,
718+
}
719+
dataFormat = &Ydb_Table.BulkUpsertRequest_ArrowBatchSettings{
720+
ArrowBatchSettings: &Ydb_Formats.ArrowBatchSettings{},
721+
}
722+
)
737723

738-
var err error
739-
for _, opt := range data.Options {
724+
for _, opt := range data.opts {
740725
if opt != nil {
741-
err = opt.ApplyArrowFormatOption(req)
742-
if err != nil {
743-
return err
726+
if err := opt.applyArrowFormatOption(dataFormat); err != nil {
727+
return nil, err
744728
}
745729
}
746730
}
747731

748-
return err
732+
request.DataFormat = dataFormat
733+
734+
return request, nil
749735
}
750736

751737
func BulkUpsertDataArrow(data []byte, opts ...arrowFormatOption) bulkUpsertArrow {
752738
return bulkUpsertArrow{
753-
Data: data,
754-
Options: opts,
739+
data: data,
740+
opts: opts,
755741
}
756742
}
757743

758-
func ensureArrowDataFormatSettings(req *BulkUpsertRequest) (format *Ydb_Formats.ArrowBatchSettings) {
759-
if settings, ok := req.DataFormat.(*Ydb_Table.BulkUpsertRequest_ArrowBatchSettings); ok {
760-
if settings.ArrowBatchSettings == nil {
761-
settings.ArrowBatchSettings = &Ydb_Formats.ArrowBatchSettings{}
762-
}
763-
764-
return settings.ArrowBatchSettings
765-
}
766-
767-
req.DataFormat = &Ydb_Table.BulkUpsertRequest_ArrowBatchSettings{
768-
ArrowBatchSettings: &Ydb_Formats.ArrowBatchSettings{},
769-
}
770-
771-
settings, ok := req.DataFormat.(*Ydb_Table.BulkUpsertRequest_ArrowBatchSettings)
772-
if !ok {
773-
return nil
774-
}
775-
776-
return settings.ArrowBatchSettings
777-
}
778-
779-
type arrowSchemaOption struct {
780-
Schema []byte
781-
}
744+
type arrowSchemaOption []byte
782745

783-
func (opt *arrowSchemaOption) ApplyArrowFormatOption(req *BulkUpsertRequest) error {
784-
ensureArrowDataFormatSettings(req).Schema = opt.Schema
746+
func (schema arrowSchemaOption) applyArrowFormatOption(dataFormat *Ydb_Table.BulkUpsertRequest_ArrowBatchSettings) error {
747+
dataFormat.ArrowBatchSettings.Schema = schema
785748

786749
return nil
787750
}
788751

789752
func WithArrowSchema(schema []byte) arrowFormatOption {
790-
return &arrowSchemaOption{schema}
753+
return arrowSchemaOption(schema)
791754
}

0 commit comments

Comments
 (0)