Merge pull request #1803 from ydb-platform/issue_1511

Split bulk for bulkupsert

Author: kprokopenko

CHANGELOG.md

@@ -1,3 +1,5 @@
+* Added the ability to send BulkRequest exceeding the GrpcMaxMessageSize
+
 ## v3.110.0
 * Added read partitions in parallel for topic listener.
 

driver.go

@@ -470,6 +470,8 @@ func (d *Driver) connect(ctx context.Context) (err error) {
 					// prepend common params from root config
 					[]tableConfig.Option{
 						tableConfig.With(d.config.Common),
+
+						tableConfig.WithMaxRequestMessageSize(d.config.GrpcMaxMessageSize()),
 					},
 					d.tableOptions...,
 				)...,

internal/table/client.go

@@ -2,12 +2,14 @@ package table
 
 import (
 	"context"
+	"fmt"
 
 	"github.com/jonboulle/clockwork"
 	"github.com/ydb-platform/ydb-go-genproto/Ydb_Table_V1"
 	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb"
 	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb_Table"
 	"google.golang.org/grpc"
+	"google.golang.org/protobuf/proto"
 
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/pool"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/stack"
@@ -299,7 +301,7 @@ func (c *Client) BulkUpsert(
 		retry.WithTrace(&trace.Retry{
 			OnRetry: func(info trace.RetryLoopStartInfo) func(trace.RetryLoopDoneInfo) {
 				return func(info trace.RetryLoopDoneInfo) {
-					attempts = info.Attempts
+					attempts += max(info.Attempts-1, 0) // `max` guarded against negative values
 				}
 			},
 		}),
@@ -309,32 +311,111 @@ func (c *Client) BulkUpsert(
 		stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/v3/internal/table.(*Client).BulkUpsert"),
 	)
 	defer func() {
-		onDone(finalErr, attempts)
+		onDone(finalErr, attempts+1)
 	}()
 
 	request, err := data.ToYDB(tableName)
 	if err != nil {
 		return xerrors.WithStackTrace(err)
 	}
 
+	chunks := make([]*Ydb_Table.BulkUpsertRequest, 0, 1)
+
+	// We must send requests in chunks to avoid exceeding the maximum message size
+	chunks, err = chunkBulkUpsertRequest(chunks, request, c.config.MaxRequestMessageSize())
+	if err != nil {
+		return err
+	}
+
+	return c.sendBulkUpsertRequest(ctx, chunks, config.RetryOptions...)
+}
+
+func (c *Client) sendBulkUpsertRequest(
+	ctx context.Context,
+	requests []*Ydb_Table.BulkUpsertRequest,
+	retryOpts ...retry.Option,
+) error {
 	client := Ydb_Table_V1.NewTableServiceClient(c.cc)
 
-	err = retry.Retry(ctx,
-		func(ctx context.Context) (err error) {
-			attempts++
-			_, err = client.BulkUpsert(ctx, request)
+	for _, request := range requests {
+		err := retry.Retry(ctx,
+			func(ctx context.Context) (err error) {
+				_, err = client.BulkUpsert(ctx, request)
 
-			return err
-		},
-		config.RetryOptions...,
-	)
-	if err != nil {
-		return xerrors.WithStackTrace(err)
+				return err
+			},
+			retryOpts...,
+		)
+		if err != nil {
+			return xerrors.WithStackTrace(err)
+		}
 	}
 
 	return nil
 }
 
+// chunkBulkUpsertRequest splits a bulk upsert request into smaller chunks if it exceeds the maximum message size.
+// It recursively divides the request into smaller parts, ensuring each chunk is within the size limit.
+// Returns a slice of chunked bulk upsert requests or an error if the request cannot be split.
+func chunkBulkUpsertRequest(
+	dst []*Ydb_Table.BulkUpsertRequest,
+	req *Ydb_Table.BulkUpsertRequest,
+	maxBytes int,
+) ([]*Ydb_Table.BulkUpsertRequest, error) {
+	reqSize := proto.Size(req)
+
+	// not exceed the maximum size -> ret original request
+	if reqSize <= maxBytes {
+		return append(dst, req), nil
+	}
+
+	// not a row bulk upsert request -> ret original request
+	if req.GetRows() == nil || req.GetRows().GetValue() == nil {
+		return nil, xerrors.WithStackTrace(
+			xerrors.Wrap(
+				fmt.Errorf("ydb: request size (%d bytes) exceeds maximum size (%d bytes) "+
+					" but cannot be chunked (only row-based bulk upserts support chunking)", reqSize, maxBytes)))
+	}
+
+	n := len(req.GetRows().GetValue().GetItems())
+	if n == 0 {
+		return dst, nil
+	}
+
+	// we cannot split one item and one item is too big
+	if n == 1 {
+		return nil, xerrors.WithStackTrace(
+			xerrors.Wrap(
+				fmt.Errorf("ydb: single row size (%d bytes) exceeds maximum request size (%d bytes) "+
+					"- row is too large to process", reqSize, maxBytes)))
+	}
+
+	left, right := splitBulkUpsertRequestAt(req, n/2)
+
+	dst, err := chunkBulkUpsertRequest(dst, left, maxBytes)
+	if err != nil {
+		return nil, err
+	}
+
+	return chunkBulkUpsertRequest(dst, right, maxBytes)
+}
+
+// splitBulkUpsertRequestAt splits a bulk upsert request into two parts at the specified position.
+// It divides the request's items into two separate requests, with the first request containing
+// items from the start up to the specified position, and the second request containing the remaining items.
+// Returns two modified bulk upsert requests with their respective item sets.
+func splitBulkUpsertRequestAt(req *Ydb_Table.BulkUpsertRequest, pos int) (_, _ *Ydb_Table.BulkUpsertRequest) {
+	items := req.GetRows().GetValue().GetItems() // save original items
+	req.Rows.Value.Items = nil
+
+	right := proto.Clone(req).(*Ydb_Table.BulkUpsertRequest) //nolint:forcetypeassert
+
+	req.Rows.Value.Items = items[:pos]
+	right.Rows.Value.Items = items[pos:]
+
+	return req, right
+}
+
 func makeReadRowsRequest(
 	sessionID string,
 	path string,

internal/table/client_test.go

@@ -8,6 +8,8 @@ import (
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb_Table"
 	"google.golang.org/grpc"
 	"google.golang.org/protobuf/proto"
@@ -19,6 +21,7 @@ import (
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsync"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xtest"
 	"github.com/ydb-platform/ydb-go-sdk/v3/table"
+	"github.com/ydb-platform/ydb-go-sdk/v3/table/types"
 	"github.com/ydb-platform/ydb-go-sdk/v3/testutil"
 )
 
@@ -80,6 +83,54 @@ func TestRaceWgClosed(t *testing.T) {
 	}, xtest.StopAfter(27*time.Second))
 }
 
+func TestChunkBulkUpsertRequest(t *testing.T) {
+	t.Run("empty request", func(t *testing.T) {
+		input := newTestBulkRequest(t, 0)
+		got, err := chunkBulkUpsertRequest(nil, input, 100)
+		require.NoError(t, err)
+		assert.Len(t, got, 1)
+		assert.Equal(t, input, got[0])
+	})
+
+	t.Run("one chunk greater than maxSize", func(t *testing.T) {
+		input := newTestBulkRequest(t, 1)
+		_, err := chunkBulkUpsertRequest(nil, input, 10)
+		assert.Error(t, err)
+	})
+
+	t.Run("one request", func(t *testing.T) {
+		input := newTestBulkRequest(t, 50)
+		got, err := chunkBulkUpsertRequest(nil, input, 100)
+		require.NoError(t, err)
+		assert.Len(t, got, 2)
+		assert.Less(t, proto.Size(got[0]), 100)
+		assert.Less(t, proto.Size(got[1]), 100)
+	})
+
+	t.Run("zero max size", func(t *testing.T) {
+		input := newTestBulkRequest(t, 50)
+		_, err := chunkBulkUpsertRequest(nil, input, 0)
+		assert.Error(t, err)
+	})
+}
+
+func newTestBulkRequest(t *testing.T, itemsLen int) *Ydb_Table.BulkUpsertRequest {
+	t.Helper()
+
+	rows := make([]types.Value, itemsLen)
+
+	for i := range itemsLen {
+		rows[i] = types.StructValue()
+	}
+
+	req, err := table.BulkUpsertDataRows(
+		types.ListValue(rows...),
+	).ToYDB("testTable")
+	require.NoError(t, err)
+
+	return req
+}
+
 var okHandler = func(interface{}) (proto.Message, error) {
 	return &emptypb.Empty{}, nil
 }

internal/table/config/config.go

@@ -162,6 +162,13 @@ func WithIgnoreTruncated() Option {
 	}
 }
 
+// WithMaxRequestMessageSize sets the maximum size of request message in bytes.
+func WithMaxRequestMessageSize(maxMessageSize int) Option {
+	return func(c *Config) {
+		c.maxRequestMessageSize = maxMessageSize
+	}
+}
+
 // ExecuteDataQueryOverQueryService overrides Execute handle with query service execute with materialized result
 func ExecuteDataQueryOverQueryService(b bool) Option {
 	return func(c *Config) {
@@ -202,6 +209,8 @@ type Config struct {
 	useQuerySession                  bool
 	executeDataQueryOverQueryService bool
 
+	maxRequestMessageSize int
+
 	trace *trace.Table
 
 	clock clockwork.Clock
@@ -304,6 +313,13 @@ func (c *Config) DeleteTimeout() time.Duration {
 	return c.deleteTimeout
 }
 
+// MaxRequestMessageSize returns the maximum size in bytes for a single request message.
+//
+// If the value is exceeded, the request will be split into several parts.
+func (c *Config) MaxRequestMessageSize() int {
+	return c.maxRequestMessageSize
+}
+
 func defaults() *Config {
 	return &Config{
 		poolLimit:            DefaultSessionPoolSizeLimit,

tests/integration/table_bulk_upsert_test.go

@@ -7,6 +7,7 @@ import (
 	"context"
 	"fmt"
 	"os"
+	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -15,6 +16,7 @@ import (
 	"github.com/ydb-platform/ydb-go-sdk/v3/table"
 	"github.com/ydb-platform/ydb-go-sdk/v3/table/result/named"
 	"github.com/ydb-platform/ydb-go-sdk/v3/table/types"
+	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
 )
 
 func TestTableBulkUpsertSession(t *testing.T) {
@@ -75,6 +77,43 @@ func TestTableBulkUpsert(t *testing.T) {
 	}
 }
 
+func TestTableBulkUpsertGrpcMaxMessageSize(t *testing.T) {
+	const MB = 1024 * 1024
+
+	var (
+		scope     = newScope(t)
+		driver    = scope.Driver(ydb.WithGrpcMaxMessageSize(10 * MB))
+		tablePath = scope.TablePath()
+	)
+
+	s := strings.Repeat("a", 5*MB)
+
+	var rows []types.Value
+	for i := int64(0); i < 10; i++ { // => 50MB the whole request
+		rows = append(rows, types.StructValue(
+			types.StructFieldValue("id", types.Int64Value(i)),
+			types.StructFieldValue("val", types.TextValue(s)),
+		))
+	}
+
+	attemptsTraceChecker := trace.Table{
+		OnBulkUpsert: func(_ trace.TableBulkUpsertStartInfo) func(trace.TableBulkUpsertDoneInfo) {
+			return func(info trace.TableBulkUpsertDoneInfo) {
+				scope.Require.Equal(1, info.Attempts, "expected 1 attempt for successful request")
+			}
+		},
+	}
+
+	err := driver.Table().BulkUpsert(scope.Ctx, tablePath, table.BulkUpsertDataRows(
+		types.ListValue(rows...),
+	), table.WithTrace(attemptsTraceChecker))
+	scope.Require.NoError(err)
+
+	for i := int64(0); i < 10; i++ {
+		assertIdValue(scope.Ctx, t, tablePath, i, s)
+	}
+}
+
 func TestTableCsvBulkUpsert(t *testing.T) {
 	var (
 		scope     = newScope(t)