add config option for disabling truncated flag

Author: asmyasnikov

CHANGELOG.md

@@ -1,4 +1,5 @@
 * Changed behavior on `result.Err()` on truncated unary result (returns non-retryable error now)
+* Added `ydb.WithIgnoreTruncated` option for disabling errors on truncated flag
 * Added simple transaction control constructors `table.OnlineReadOnlyTxControl()` and `table.StaleReadOnlyTxControl()`
 * Added transaction control specifier with context `ydb.WithTxControl`
 * Added value constructors `types.BytesValue`, `types.BytesValueFromString`, `types.TextValue`

internal/table/config/config.go

@@ -124,6 +124,13 @@ func WithTrace(trace trace.Table, opts ...trace.TableComposeOption) Option {
 	}
 }
 
+// WithIgnoreTruncated disables errors on truncated flag
+func WithIgnoreTruncated() Option {
+	return func(c *Config) {
+		c.ignoreTruncated = true
+	}
+}
+
 // Config is a configuration of table client
 type Config struct {
 	config.Common
@@ -134,6 +141,8 @@ type Config struct {
 	deleteTimeout        time.Duration
 	idleThreshold        time.Duration
 
+	ignoreTruncated bool
+
 	trace trace.Table
 }
 
@@ -159,6 +168,11 @@ func (c Config) KeepAliveMinSize() int {
 	return DefaultKeepAliveMinSize
 }
 
+// IgnoreTruncated specifies behavior on truncated flag
+func (c Config) IgnoreTruncated() bool {
+	return c.ignoreTruncated
+}
+
 // IdleKeepAliveThreshold is a number of keepAlive messages to call before the
 // session is removed if it is an excess session (see KeepAliveMinSize)
 // This means that session will be deleted after the expiration of lifetime = IdleThreshold * IdleKeepAliveThreshold

internal/table/scanner/result.go

@@ -70,24 +70,39 @@ type StreamResult interface {
 	resultWithError
 }
 
+type option func(r *baseResult)
+
+func WithIgnoreTruncated(ignoreTruncated bool) option {
+	return func(r *baseResult) {
+		r.scanner.ignoreTruncated = ignoreTruncated
+	}
+}
+
 func NewStream(
 	recv func(ctx context.Context) (*Ydb.ResultSet, *Ydb_TableStats.QueryStats, error),
 	onClose func(error) error,
+	opts ...option,
 ) StreamResult {
 	r := &streamResult{
 		recv:  recv,
 		close: onClose,
 	}
+	for _, o := range opts {
+		o(&r.baseResult)
+	}
 	return r
 }
 
-func NewUnary(sets []*Ydb.ResultSet, stats *Ydb_TableStats.QueryStats) UnaryResult {
+func NewUnary(sets []*Ydb.ResultSet, stats *Ydb_TableStats.QueryStats, opts ...option) UnaryResult {
 	r := &unaryResult{
 		baseResult: baseResult{
 			stats: stats,
 		},
 		sets: sets,
 	}
+	for _, o := range opts {
+		o(&r.baseResult)
+	}
 	return r
 }
 
@@ -98,16 +113,6 @@ func (r *baseResult) Reset(set *Ydb.ResultSet, columnNames ...string) {
 	}
 }
 
-func (r *unaryResult) Err() (err error) {
-	if err = r.baseResult.scanner.Err(); err != nil {
-		return err
-	}
-	if r.truncated() {
-		return xerrors.WithStackTrace(errTruncated)
-	}
-	return nil
-}
-
 func (r *unaryResult) NextResultSetErr(ctx context.Context, columns ...string) (err error) {
 	if r.isClosed() {
 		return xerrors.WithStackTrace(errAlreadyClosed)
@@ -124,10 +129,6 @@ func (r *unaryResult) NextResultSet(ctx context.Context, columns ...string) bool
 	return r.NextResultSetErr(ctx, columns...) == nil
 }
 
-func (r *streamResult) Err() (err error) {
-	return r.baseResult.scanner.Err()
-}
-
 func (r *streamResult) NextResultSetErr(ctx context.Context, columns ...string) (err error) {
 	if r.isClosed() {
 		return xerrors.WithStackTrace(errAlreadyClosed)

internal/table/scanner/scanner.go

@@ -26,12 +26,13 @@ import (
 var errTruncated = xerrors.Wrap(errors.New("truncated result"))
 
 type scanner struct {
-	set       *Ydb.ResultSet
-	row       *Ydb.Value
-	converter *rawConverter
-	stack     scanStack
-	nextRow   int
-	nextItem  int
+	set             *Ydb.ResultSet
+	row             *Ydb.Value
+	converter       *rawConverter
+	stack           scanStack
+	nextRow         int
+	nextItem        int
+	ignoreTruncated bool
 
 	columnIndexes []int
 
@@ -224,6 +225,9 @@ func (s *scanner) Err() error {
 	if s.err != nil {
 		return s.err
 	}
+	if !s.ignoreTruncated && s.truncated() {
+		return xerrors.WithStackTrace(errTruncated)
+	}
 	return nil
 }
 

internal/table/session.go

@@ -694,6 +694,7 @@ func (s *session) executeQueryResult(res *Ydb_Table.ExecuteQueryResult) (
 	r := scanner.NewUnary(
 		res.GetResultSets(),
 		res.GetQueryStats(),
+		scanner.WithIgnoreTruncated(s.config.IgnoreTruncated()),
 	)
 	return t, r, nil
 }
@@ -993,6 +994,7 @@ func (s *session) StreamReadTable(
 			onIntermediate(xerrors.HideEOF(err))(xerrors.HideEOF(err))
 			return err
 		},
+		scanner.WithIgnoreTruncated(true), // stream read table always returns truncated flag on last result set
 	), nil
 }
 
@@ -1075,6 +1077,7 @@ func (s *session) StreamExecuteScanQuery(
 			onIntermediate(xerrors.HideEOF(err))(xerrors.HideEOF(err))
 			return err
 		},
+		scanner.WithIgnoreTruncated(s.config.IgnoreTruncated()),
 	), nil
 }
 

internal/table/transaction.go

@@ -161,6 +161,7 @@ func (tx *transaction) CommitTx(
 	return scanner.NewUnary(
 		nil,
 		result.GetQueryStats(),
+		scanner.WithIgnoreTruncated(tx.s.config.IgnoreTruncated()),
 	), nil
 }
 

options.go

@@ -386,6 +386,14 @@ func WithSessionPoolDeleteTimeout(deleteTimeout time.Duration) Option {
 	}
 }
 
+// WithIgnoreTruncated disables errors on truncated flag
+func WithIgnoreTruncated() Option {
+	return func(ctx context.Context, c *connection) error {
+		c.tableOptions = append(c.tableOptions, tableConfig.WithIgnoreTruncated())
+		return nil
+	}
+}
+
 // WithPanicCallback specified behavior on panic
 // Warning: WithPanicCallback must be defined on start of all options
 // (before `WithTrace{Driver,Table,Scheme,Scripting,Coordination,Ratelimiter}` and other options)