Merge branch 'master' into fix-pool-item-close

Author: 4el0ve4ek

CHANGELOG.md

@@ -1,5 +1,10 @@
 * Fixed connections pool leak on closing
 
+## v3.92.3
+* Fixed error with incompleted data returen from transaction.ReadQueryResult method
+* Added option `query/WithResponsePartLimitSizeBytes(...)` for queries with query service
+
+
 ## v3.92.2
 * Added `table/options.WithShardNodesInfo()` experimental option to get shard nodeId for describe table call
 

internal/query/execute_query.go

@@ -31,6 +31,7 @@ type executeSettings interface {
 	CallOptions() []grpc.CallOption
 	RetryOpts() []retry.Option
 	ResourcePool() string
+	ResponsePartLimitSizeBytes() int64
 }
 
 type executeScriptConfig interface {
@@ -77,6 +78,7 @@ func executeQueryRequest(a *allocator.Allocator, sessionID, q string, cfg execut
 	request.StatsMode = Ydb_Query.StatsMode(cfg.StatsMode())
 	request.ConcurrentResultSets = false
 	request.PoolId = cfg.ResourcePool()
+	request.ResponsePartLimitBytes = cfg.ResponsePartLimitSizeBytes()
 
 	return request, cfg.CallOptions()
 }
@@ -146,14 +148,7 @@ func readResultSet(ctx context.Context, r *streamResult) (_ *resultSetWithClose,
 	if err != nil {
 		return nil, xerrors.WithStackTrace(err)
 	}
-
-	_, err = r.nextResultSet(ctx)
-	if err == nil {
-		return nil, xerrors.WithStackTrace(errMoreThanOneResultSet)
-	}
-	if !xerrors.Is(err, io.EOF) {
-		return nil, xerrors.WithStackTrace(err)
-	}
+	rs.mustBeLastResultSet = true
 
 	return &resultSetWithClose{
 		resultSet: rs,

internal/query/options/execute.go

@@ -27,15 +27,16 @@ type (
 
 	// executeSettings is a holder for execute settings
 	executeSettings struct {
-		syntax        Syntax
-		params        params.Parameters
-		execMode      ExecMode
-		statsMode     StatsMode
-		resourcePool  string
-		statsCallback func(queryStats stats.QueryStats)
-		callOptions   []grpc.CallOption
-		txControl     *tx.Control
-		retryOptions  []retry.Option
+		syntax                 Syntax
+		params                 params.Parameters
+		execMode               ExecMode
+		statsMode              StatsMode
+		resourcePool           string
+		statsCallback          func(queryStats stats.QueryStats)
+		callOptions            []grpc.CallOption
+		txControl              *tx.Control
+		retryOptions           []retry.Option
+		responsePartLimitBytes int64
 	}
 
 	// Execute is an interface for execute method options
@@ -58,7 +59,8 @@ type (
 		mode     StatsMode
 		callback func(stats.QueryStats)
 	}
-	execModeOption = ExecMode
+	execModeOption         = ExecMode
+	responsePartLimitBytes int64
 )
 
 func (poolID resourcePool) applyExecuteOption(s *executeSettings) {
@@ -175,6 +177,10 @@ func (s *executeSettings) Params() *params.Parameters {
 	return &s.params
 }
 
+func (s *executeSettings) ResponsePartLimitSizeBytes() int64 {
+	return s.responsePartLimitBytes
+}
+
 func WithParameters(parameters *params.Parameters) parametersOption {
 	return parametersOption(*parameters)
 }
@@ -201,6 +207,14 @@ func WithExecMode(mode ExecMode) execModeOption {
 	return mode
 }
 
+func WithResponsePartLimitSizeBytes(size int64) responsePartLimitBytes {
+	return responsePartLimitBytes(size)
+}
+
+func (size responsePartLimitBytes) applyExecuteOption(s *executeSettings) {
+	s.responsePartLimitBytes = int64(size)
+}
+
 func WithSyntax(syntax Syntax) syntaxOption {
 	return syntax
 }

internal/query/result.go

@@ -19,6 +19,8 @@ import (
 	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
 )
 
+var errReadNextResultSet = xerrors.Wrap(errors.New("ydb: stop read the result set because see part of next result set"))
+
 var (
 	_ result.Result = (*streamResult)(nil)
 	_ result.Result = (*materializedResult)(nil)
@@ -294,8 +296,8 @@ func (r *streamResult) nextPartFunc(
 			}
 			if part.GetResultSetIndex() > nextResultSetIndex {
 				return nil, xerrors.WithStackTrace(fmt.Errorf(
-					"result set (index=%d) receive part (index=%d) for next result set: %w",
-					nextResultSetIndex, part.GetResultSetIndex(), io.EOF,
+					"result set (index=%d) receive part (index=%d) for next result set: %w (%w)",
+					nextResultSetIndex, part.GetResultSetIndex(), io.EOF, errReadNextResultSet,
 				))
 			}
 

internal/query/result_set.go

@@ -2,6 +2,7 @@ package query
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"io"
 
@@ -29,12 +30,13 @@ type (
 		rowIndex    int
 	}
 	resultSet struct {
-		index       int64
-		recv        func() (*Ydb_Query.ExecuteQueryResponsePart, error)
-		columns     []*Ydb.Column
-		currentPart *Ydb_Query.ExecuteQueryResponsePart
-		rowIndex    int
-		done        chan struct{}
+		index               int64
+		recv                func() (*Ydb_Query.ExecuteQueryResponsePart, error)
+		columns             []*Ydb.Column
+		currentPart         *Ydb_Query.ExecuteQueryResponsePart
+		rowIndex            int
+		done                chan struct{}
+		mustBeLastResultSet bool
 	}
 	resultSetWithClose struct {
 		*resultSet
@@ -158,13 +160,19 @@ func (rs *resultSet) nextRow(ctx context.Context) (*Row, error) {
 		case <-ctx.Done():
 			return nil, xerrors.WithStackTrace(ctx.Err())
 		default:
+			//nolint:nestif
 			if rs.rowIndex == len(rs.currentPart.GetResultSet().GetRows()) {
 				part, err := rs.recv()
 				if err != nil {
 					if xerrors.Is(err, io.EOF) {
 						close(rs.done)
 					}
 
+					if rs.mustBeLastResultSet && errors.Is(err, errReadNextResultSet) {
+						// prevent detect io.EOF in the error
+						return nil, xerrors.WithStackTrace(xerrors.Wrap(errors.New(err.Error())))
+					}
+
 					return nil, xerrors.WithStackTrace(err)
 				}
 				rs.rowIndex = 0

internal/query/result_test.go

@@ -1520,22 +1520,22 @@ func TestCloseResultOnCloseClosableResultSet(t *testing.T) {
 				{
 					Items: []*Ydb.Value{{
 						Value: &Ydb.Value_Uint64Value{
-							Uint64Value: 1,
+							Uint64Value: 3,
 						},
 					}, {
 						Value: &Ydb.Value_TextValue{
-							TextValue: "1",
+							TextValue: "3",
 						},
 					}},
 				},
 				{
 					Items: []*Ydb.Value{{
 						Value: &Ydb.Value_Uint64Value{
-							Uint64Value: 2,
+							Uint64Value: 4,
 						},
 					}, {
 						Value: &Ydb.Value_TextValue{
-							TextValue: "2",
+							TextValue: "4",
 						},
 					}},
 				},
@@ -1576,8 +1576,20 @@ func TestCloseResultOnCloseClosableResultSet(t *testing.T) {
 	require.EqualValues(t, 2, a)
 	require.EqualValues(t, "2", b)
 	r3, err3 := rs.NextRow(ctx)
-	require.ErrorIs(t, err3, io.EOF)
-	require.Nil(t, r3)
+	require.NoError(t, err3)
+	scanErr3 := r3.Scan(&a, &b)
+	require.EqualValues(t, 3, a)
+	require.EqualValues(t, "3", b)
+	require.NoError(t, scanErr3)
+	r4, err4 := rs.NextRow(ctx)
+	require.NoError(t, err4)
+	scanErr4 := r4.Scan(&a, &b)
+	require.EqualValues(t, 4, a)
+	require.EqualValues(t, "4", b)
+	require.NoError(t, scanErr4)
+	r5, err5 := rs.NextRow(ctx)
+	require.ErrorIs(t, err5, io.EOF)
+	require.Nil(t, r5)
 	err = rs.Close(ctx)
 	require.NoError(t, err)
 	require.True(t, closed)

internal/version/version.go

@@ -3,7 +3,7 @@ package version
 const (
 	Major = "3"
 	Minor = "92"
-	Patch = "2"
+	Patch = "3"
 
 	Package = "ydb-go-sdk"
 )

query/execute_options.go

@@ -55,6 +55,12 @@ func WithStatsMode(mode options.StatsMode, callback func(Stats)) options.Execute
 	return options.WithStatsMode(mode, callback)
 }
 
+// WithResponsePartLimitSizeBytes limit size of each part (data portion) in stream for query service resoponse
+// it isn't limit total size of answer
+func WithResponsePartLimitSizeBytes(size int64) options.Execute {
+	return options.WithResponsePartLimitSizeBytes(size)
+}
+
 func WithCallOptions(opts ...grpc.CallOption) options.Execute {
 	return options.WithCallOptions(opts...)
 }

tests/integration/query_execute_test.go

@@ -409,3 +409,144 @@ func TestIssue1456TooManyUnknownTransactions(t *testing.T) {
 		wg.Wait()
 	})
 }
+
+func TestQueryResultSet(t *testing.T) {
+	if version.Lt(os.Getenv("YDB_VERSION"), "24.1") {
+		t.Skip("query service not allowed in YDB version '" + os.Getenv("YDB_VERSION") + "'")
+	}
+
+	t.Run("OK", func(t *testing.T) {
+		scope := newScope(t)
+
+		partSizeBytes := 1000
+		targetCount := partSizeBytes * 10 // for guarantee size of response will contain many parts
+		items := make([]types.Value, 0, targetCount)
+		for i := 0; i < targetCount; i++ {
+			item := types.StructValue(
+				types.StructFieldValue("val", types.Int64Value(int64(i))),
+			)
+			items = append(items, item)
+		}
+
+		err := scope.Driver().Query().DoTx(scope.Ctx, func(ctx context.Context, tx query.TxActor) error {
+			rs, err := tx.QueryResultSet(ctx, `
+DECLARE $arg AS List<Struct<val: Int64>>;
+
+SELECT * FROM AS_TABLE($arg);
+`,
+				query.WithParameters(ydb.ParamsBuilder().Param("$arg").Any(types.ListValue(items...)).Build()),
+				query.WithResponsePartLimitSizeBytes(int64(partSizeBytes)),
+			)
+			if err != nil {
+				return err
+			}
+
+			for i := 0; i < targetCount; i++ {
+				row, err := rs.NextRow(ctx)
+				if err != nil {
+					return err
+				}
+
+				var val int64
+				err = row.Scan(&val)
+				require.NoError(t, err)
+				require.Equal(t, int64(i), val)
+			}
+
+			return nil
+		})
+		require.NoError(t, err)
+	})
+	t.Run("FailOnSecondResultSet", func(t *testing.T) {
+		scope := newScope(t)
+
+		var secondRowError error
+		err := scope.Driver().Query().DoTx(scope.Ctx, func(ctx context.Context, tx query.TxActor) error {
+			rs, err := tx.QueryResultSet(ctx, "SELECT 1; SELECT 2")
+			if err != nil {
+				return err
+			}
+
+			_, err = rs.NextRow(ctx)
+			if err != nil {
+				return err
+			}
+
+			_, secondRowError = rs.NextRow(ctx)
+
+			return nil
+		})
+		require.NoError(t, err)
+		require.Error(t, secondRowError)
+		require.NotErrorIs(t, secondRowError, io.EOF)
+	})
+}
+
+func TestQueryPartLimiter(t *testing.T) {
+	if os.Getenv("YDB_VERSION") != "nightly" && version.Lt(os.Getenv("YDB_VERSION"), "25.0") {
+		t.Skip("require enables transactions for topics")
+	}
+
+	scope := newScope(t)
+
+	var readPartCount int
+	scope.Driver(ydb.WithTraceQuery(trace.Query{
+		OnResultNextPart: func(info trace.QueryResultNextPartStartInfo) func(info trace.QueryResultNextPartDoneInfo) {
+			return func(info trace.QueryResultNextPartDoneInfo) {
+				if info.Error == nil {
+					readPartCount++
+				}
+			}
+		},
+	}))
+
+	targetCount := 1000
+	items := make([]types.Value, 0, targetCount)
+	for i := 0; i < targetCount; i++ {
+		item := types.StructValue(
+			types.StructFieldValue("val", types.Int64Value(int64(i))),
+		)
+		items = append(items, item)
+	}
+
+	getPartCount := func(partSize int64) int {
+		partCount := 0
+		err := scope.Driver().Query().DoTx(scope.Ctx, func(ctx context.Context, tx query.TxActor) error {
+			oldParts := readPartCount
+			rs, err := tx.QueryResultSet(ctx, `
+DECLARE $arg AS List<Struct<val: Int64>>;
+
+SELECT * FROM AS_TABLE($arg);
+`,
+				query.WithParameters(ydb.ParamsBuilder().Param("$arg").Any(types.ListValue(items...)).Build()),
+				query.WithResponsePartLimitSizeBytes(partSize),
+			)
+			if err != nil {
+				return err
+			}
+
+			rowCount := 0
+			for {
+				_, err = rs.NextRow(scope.Ctx)
+				if errors.Is(err, io.EOF) {
+					break
+				}
+				require.NoError(t, err)
+				rowCount++
+			}
+			require.Equal(t, targetCount, rowCount)
+
+			partCount = readPartCount - oldParts
+			return nil
+		})
+
+		require.NoError(t, err)
+		return partCount
+	}
+
+	partsWithBigSize := getPartCount(1000000)
+	partsWithLittleSize := getPartCount(100)
+
+	require.Equal(t, 1, partsWithBigSize)
+	require.Greater(t, partsWithLittleSize, 1)
+}