Merge pull request #356 from ydb-platform/truncated

* Changed behavior on `result.Err()` on truncated result (return non-…

Author: asmyasnikov

CHANGELOG.md

@@ -1,3 +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/client.go

@@ -111,87 +111,96 @@ func (c *Client) createSession(ctx context.Context, opts ...createSessionOption)
 	if c.isClosed() {
 		return nil, errClosedClient
 	}
+
+	options := createSessionOptions{}
+	for _, o := range opts {
+		o(&options)
+	}
+
 	defer func() {
 		if s == nil {
 			return
 		}
-		options := createSessionOptions{}
-		for _, o := range opts {
-			o(&options)
-		}
 		for _, onCreate := range options.onCreate {
 			onCreate(s)
 		}
 		s.onClose = append(s.onClose, options.onClose...)
 	}()
-	type result struct {
-		s   *session
-		err error
-	}
-
-	ch := make(chan result)
 
-	c.spawnedGoroutines.Add(1)
-	go func() {
-		defer c.spawnedGoroutines.Done()
+	select {
+	case <-c.done:
+		return nil, xerrors.WithStackTrace(errClosedClient)
 
-		var (
+	default:
+		type result struct {
 			s   *session
 			err error
-		)
-
-		createSessionCtx := xcontext.WithoutDeadline(ctx)
-
-		if timeout := c.config.CreateSessionTimeout(); timeout > 0 {
-			var cancel context.CancelFunc
-			createSessionCtx, cancel = context.WithTimeout(createSessionCtx, timeout)
-			defer cancel()
 		}
 
-		s, err = c.build(createSessionCtx)
+		ch := make(chan result)
 
-		closeSession := func(s *session) {
-			if s == nil {
-				return
-			}
+		c.spawnedGoroutines.Add(1)
+		go func() {
+			defer c.spawnedGoroutines.Done()
+
+			var (
+				s   *session
+				err error
+			)
 
-			closeSessionCtx := xcontext.WithoutDeadline(ctx)
+			createSessionCtx := xcontext.WithoutDeadline(ctx)
 
-			if timeout := c.config.DeleteTimeout(); timeout > 0 {
+			if timeout := c.config.CreateSessionTimeout(); timeout > 0 {
 				var cancel context.CancelFunc
-				createSessionCtx, cancel = context.WithTimeout(closeSessionCtx, timeout)
+				createSessionCtx, cancel = context.WithTimeout(createSessionCtx, timeout)
 				defer cancel()
 			}
 
-			_ = s.Close(ctx)
-		}
+			s, err = c.build(createSessionCtx)
 
-		select {
-		case ch <- result{
-			s:   s,
-			err: err,
-		}: // nop
+			closeSession := func(s *session) {
+				if s == nil {
+					return
+				}
+
+				closeSessionCtx := xcontext.WithoutDeadline(ctx)
+
+				if timeout := c.config.DeleteTimeout(); timeout > 0 {
+					var cancel context.CancelFunc
+					createSessionCtx, cancel = context.WithTimeout(closeSessionCtx, timeout)
+					defer cancel()
+				}
+
+				_ = s.Close(ctx)
+			}
 
+			select {
+			case ch <- result{
+				s:   s,
+				err: err,
+			}: // nop
+
+			case <-c.done:
+				closeSession(s)
+
+			case <-ctx.Done():
+				closeSession(s)
+			}
+		}()
+
+		select {
 		case <-c.done:
-			closeSession(s)
+			return nil, xerrors.WithStackTrace(errClosedClient)
 
 		case <-ctx.Done():
-			closeSession(s)
-		}
-	}()
+			return nil, xerrors.WithStackTrace(ctx.Err())
 
-	select {
-	case r := <-ch:
-		if r.err != nil {
-			return nil, xerrors.WithStackTrace(r.err)
+		case r := <-ch:
+			if r.err != nil {
+				return nil, xerrors.WithStackTrace(r.err)
+			}
+			return r.s, nil
 		}
-		return r.s, nil
-
-	case <-c.done:
-		return nil, xerrors.WithStackTrace(errClosedClient)
-
-	case <-ctx.Done():
-		return nil, xerrors.WithStackTrace(ctx.Err())
 	}
 }
 

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
 }
 

internal/table/scanner/scanner.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"database/sql"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io"
 	"math"
@@ -22,13 +23,16 @@ import (
 	"github.com/ydb-platform/ydb-go-sdk/v3/table/types"
 )
 
+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
 
@@ -206,11 +210,25 @@ func (s *scanner) Truncated() bool {
 	return s.set.Truncated
 }
 
+// Truncated returns true if current result set has been truncated by server
+func (s *scanner) truncated() bool {
+	if s.set == nil {
+		return false
+	}
+	return s.set.Truncated
+}
+
 // Err returns error caused Scanner to be broken.
 func (s *scanner) Err() error {
 	s.errMtx.RLock()
 	defer s.errMtx.RUnlock()
-	return s.err
+	if s.err != nil {
+		return s.err
+	}
+	if !s.ignoreTruncated && s.truncated() {
+		return xerrors.WithStackTrace(errTruncated)
+	}
+	return nil
 }
 
 // Must not be exported.

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)

table/table_e2e_test.go

@@ -1510,7 +1510,7 @@ func TestLongStream(t *testing.T) {
 					time.Sleep(discoveryInterval)
 				}
 				if err = res.Err(); err != nil {
-					return fmt.Errorf("received error: %w (duration: %v)", err, time.Since(start))
+					return fmt.Errorf("received error (duration: %v): %w", time.Since(start), err)
 				}
 				if rowsCount != upsertRowsCount {
 					return fmt.Errorf("wrong rows count: %v, expected: %d (duration: %v)",
@@ -1551,7 +1551,7 @@ func TestLongStream(t *testing.T) {
 					time.Sleep(discoveryInterval)
 				}
 				if err = res.Err(); err != nil {
-					return fmt.Errorf("received error: %w (duration: %v)", err, time.Since(start))
+					return fmt.Errorf("received error (duration: %v): %w", time.Since(start), err)
 				}
 				if rowsCount != upsertRowsCount {
 					return fmt.Errorf("wrong rows count: %v, expected: %d (duration: %v)",