From 11cfae5b1e3eab15f56476ff6ca23c7b6f2d3772 Mon Sep 17 00:00:00 2001
From: asmyasnikov <79263256394@ya.ru>
Date: Wed, 7 Sep 2022 13:47:06 +0300
Subject: [PATCH] * Changed type of truncated result error from
 `StreamExecuteScanQuery` to retryable error

---
 CHANGELOG.md                      |  2 +-
 MIGRATION_v2_v3.md                |  6 +++---
 internal/table/scanner/result.go  |  6 ++++++
 internal/table/scanner/scanner.go | 20 ++++++++++++--------
 internal/table/session.go         |  1 +
 5 files changed, 23 insertions(+), 12 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 337c56d9b..c747aed73 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,4 +1,4 @@
-* Marked the truncated result as retryable error
+* Changed type of truncated result error from `StreamExecuteScanQuery` to retryable error
 * Added closing sessions if node removed from discovery results
 * Moved session status type from `table/options` package to `table`
 * Changed session status source type from `uint32` to `string` alias 
diff --git a/MIGRATION_v2_v3.md b/MIGRATION_v2_v3.md
index 13958bfa8..758964837 100644
--- a/MIGRATION_v2_v3.md
+++ b/MIGRATION_v2_v3.md
@@ -146,8 +146,8 @@
 
 ## About truncated result <a name="truncated"></a>
 
-Call of `session.Execute` may returns a result with a flag `Truncated` because `YDB` have a default limit of rows is a 1000.
-In this case query must be changed for supporting pagination. Trucated flag in result must be checks explicitly.
+Call of `session.Execute` may return a result with a flag `Truncated` because `YDB` have a default limit of rows is a 1000.
+In this case query must be changed for supporting pagination. Truncated flag in result must be checks explicitly.
 - in `v2`:
   ```go
   var res *table.Result
@@ -178,7 +178,7 @@ In this case query must be changed for supporting pagination. Trucated flag in r
   }
   ```
 - in `v3`:
-  By default, truncated result wraps as non-retryable error
+  By default, truncated result wraps as non-retryable error for `session.Execute` and retryable error for `session.StreamExecuteScanQuery`
   ```go
   import (
     "github.com/ydb-platform/ydb-go-sdk/v3/table/result"  
diff --git a/internal/table/scanner/result.go b/internal/table/scanner/result.go
index e2abbc576..a9d8244aa 100644
--- a/internal/table/scanner/result.go
+++ b/internal/table/scanner/result.go
@@ -78,6 +78,12 @@ func WithIgnoreTruncated(ignoreTruncated bool) option {
 	}
 }
 
+func WithMarkTruncatedAsRetryable() option {
+	return func(r *baseResult) {
+		r.scanner.markTruncatedAsRetryable = true
+	}
+}
+
 func NewStream(
 	recv func(ctx context.Context) (*Ydb.ResultSet, *Ydb_TableStats.QueryStats, error),
 	onClose func(error) error,
diff --git a/internal/table/scanner/scanner.go b/internal/table/scanner/scanner.go
index b236fa49c..df0a994d2 100644
--- a/internal/table/scanner/scanner.go
+++ b/internal/table/scanner/scanner.go
@@ -23,16 +23,17 @@ import (
 	"github.com/ydb-platform/ydb-go-sdk/v3/table/types"
 )
 
-var errTruncated = xerrors.Retryable(errors.New("truncated result"))
+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
-	ignoreTruncated bool
+	set                      *Ydb.ResultSet
+	row                      *Ydb.Value
+	converter                *rawConverter
+	stack                    scanStack
+	nextRow                  int
+	nextItem                 int
+	ignoreTruncated          bool
+	markTruncatedAsRetryable bool
 
 	columnIndexes []int
 
@@ -226,6 +227,9 @@ func (s *scanner) Err() error {
 		return s.err
 	}
 	if !s.ignoreTruncated && s.truncated() {
+		if s.markTruncatedAsRetryable {
+			return xerrors.WithStackTrace(xerrors.Retryable(errTruncated))
+		}
 		return xerrors.WithStackTrace(errTruncated)
 	}
 	return nil
diff --git a/internal/table/session.go b/internal/table/session.go
index f45394196..a41929c1e 100644
--- a/internal/table/session.go
+++ b/internal/table/session.go
@@ -1084,6 +1084,7 @@ func (s *session) StreamExecuteScanQuery(
 			return err
 		},
 		scanner.WithIgnoreTruncated(s.config.IgnoreTruncated()),
+		scanner.WithMarkTruncatedAsRetryable(),
 	), nil
 }