* Added receiving first result set on construct internal/table/scanner.NewStream()

Author: asmyasnikov

CHANGELOG.md

@@ -1,3 +1,4 @@
+* Added receiving first result set on construct `internal/table/scanner.NewStream()`
 * Added experimental package `metrics` with SDK metrics
 * Fixed redundant trace call for finished `database/sql` transactions
 * Added repeater event type to wake-up func context

internal/scripting/client.go

@@ -222,7 +222,7 @@ func (c *Client) streamExecute(
 		return nil, xerrors.WithStackTrace(err)
 	}
 
-	return scanner.NewStream(
+	return scanner.NewStream(ctx,
 		func(ctx context.Context) (
 			set *Ydb.ResultSet,
 			stats *Ydb_TableStats.QueryStats,
@@ -249,7 +249,7 @@ func (c *Client) streamExecute(
 			onIntermediate(xerrors.HideEOF(err))(xerrors.HideEOF(err))
 			return err
 		},
-	), nil
+	)
 }
 
 func (c *Client) Close(ctx context.Context) (err error) {

internal/table/scanner/result.go

@@ -20,8 +20,9 @@ var errAlreadyClosed = xerrors.Wrap(errors.New("result closed early"))
 type baseResult struct {
 	scanner
 
-	statsMtx xsync.RWMutex
-	stats    *Ydb_TableStats.QueryStats
+	nextResultSetCounter xatomic.Uint64
+	statsMtx             xsync.RWMutex
+	stats                *Ydb_TableStats.QueryStats
 
 	closed xatomic.Bool
 }
@@ -103,10 +104,11 @@ func WithMarkTruncatedAsRetryable() option {
 }
 
 func NewStream(
+	ctx context.Context,
 	recv func(ctx context.Context) (*Ydb.ResultSet, *Ydb_TableStats.QueryStats, error),
 	onClose func(error) error,
 	opts ...option,
-) StreamResult {
+) (StreamResult, error) {
 	r := &streamResult{
 		recv:  recv,
 		close: onClose,
@@ -116,7 +118,10 @@ func NewStream(
 			o(&r.baseResult)
 		}
 	}
-	return r
+	if err := r.nextResultSetErr(ctx); err != nil {
+		return nil, xerrors.WithStackTrace(err)
+	}
+	return r, nil
 }
 
 func NewUnary(sets []*Ydb.ResultSet, stats *Ydb_TableStats.QueryStats, opts ...option) UnaryResult {
@@ -157,20 +162,18 @@ func (r *unaryResult) NextResultSet(ctx context.Context, columns ...string) bool
 	return r.NextResultSetErr(ctx, columns...) == nil
 }
 
-func (r *streamResult) NextResultSetErr(ctx context.Context, columns ...string) (err error) {
-	if r.isClosed() {
-		return xerrors.WithStackTrace(errAlreadyClosed)
-	}
-	if err = r.Err(); err != nil {
-		return xerrors.WithStackTrace(err)
+func (r *streamResult) nextResultSetErr(ctx context.Context, columns ...string) (err error) {
+	if r.nextResultSetCounter.Add(1) == 2 {
+		r.setColumnIndexes(columns)
+		return ctx.Err()
 	}
 	s, stats, err := r.recv(ctx)
 	if err != nil {
 		r.Reset(nil)
 		if xerrors.Is(err, io.EOF) {
 			return err
 		}
-		return r.errorf(0, "streamResult.NextResultSetErr(): %w", err)
+		return r.errorf(1, "streamResult.NextResultSetErr(): %w", err)
 	}
 	r.Reset(s, columns...)
 	if stats != nil {
@@ -181,6 +184,22 @@ func (r *streamResult) NextResultSetErr(ctx context.Context, columns ...string)
 	return ctx.Err()
 }
 
+func (r *streamResult) NextResultSetErr(ctx context.Context, columns ...string) (err error) {
+	if r.isClosed() {
+		return xerrors.WithStackTrace(errAlreadyClosed)
+	}
+	if err = r.Err(); err != nil {
+		return xerrors.WithStackTrace(err)
+	}
+	if err := r.nextResultSetErr(ctx, columns...); err != nil {
+		if xerrors.Is(err, io.EOF) {
+			return io.EOF
+		}
+		return xerrors.WithStackTrace(err)
+	}
+	return nil
+}
+
 func (r *streamResult) NextResultSet(ctx context.Context, columns ...string) bool {
 	return r.NextResultSetErr(ctx, columns...) == nil
 }

internal/table/scanner/result_test.go

@@ -3,10 +3,14 @@ package scanner
 import (
 	"context"
 	"fmt"
+	"io"
 	"reflect"
 	"testing"
+	"time"
 
+	"github.com/stretchr/testify/require"
 	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb"
+	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb_TableStats"
 
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/allocator"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/value"
@@ -199,3 +203,75 @@ func NewResultSet(a *allocator.Allocator, opts ...ResultSetOption) *Ydb.ResultSe
 	}
 	return (*Ydb.ResultSet)(&d)
 }
+
+func TestNewStreamWithRecvFirstResultSet(t *testing.T) {
+	for _, tt := range []struct {
+		ctx         context.Context
+		recvCounter int
+		err         error
+	}{
+		{
+			ctx: context.Background(),
+			err: nil,
+		},
+		{
+			ctx: func() context.Context {
+				ctx, cancel := context.WithCancel(context.Background())
+				cancel()
+				return ctx
+			}(),
+			err: context.Canceled,
+		},
+		{
+			ctx: func() context.Context {
+				ctx, cancel := context.WithTimeout(context.Background(), 0)
+				cancel()
+				return ctx
+			}(),
+			err: context.DeadlineExceeded,
+		},
+		{
+			ctx: func() context.Context {
+				ctx, cancel := context.WithTimeout(context.Background(), time.Hour)
+				cancel()
+				return ctx
+			}(),
+			err: context.Canceled,
+		},
+	} {
+		t.Run("", func(t *testing.T) {
+			result, err := NewStream(tt.ctx,
+				func(ctx context.Context) (*Ydb.ResultSet, *Ydb_TableStats.QueryStats, error) {
+					tt.recvCounter++
+					if tt.recvCounter > 1000 {
+						return nil, nil, io.EOF
+					}
+					return &Ydb.ResultSet{}, nil, ctx.Err()
+				},
+				func(err error) error {
+					return err
+				},
+			)
+			if tt.err != nil {
+				require.ErrorIs(t, err, tt.err)
+				require.Nil(t, result)
+			} else {
+				require.NoError(t, err)
+				require.NotNil(t, result)
+				require.EqualValues(t, 1, tt.recvCounter)
+				require.EqualValues(t, 1, result.(*streamResult).nextResultSetCounter.Load())
+				for i := range make([]struct{}, 1000) {
+					err = result.NextResultSetErr(tt.ctx)
+					require.NoError(t, err)
+					require.Equal(t, i+1, tt.recvCounter)
+					require.Equal(t, i+2, int(result.(*streamResult).nextResultSetCounter.Load()))
+				}
+				err = result.NextResultSetErr(tt.ctx)
+				require.ErrorIs(t, err, io.EOF)
+				require.True(t, err == io.EOF) //nolint:errorlint
+				require.Equal(t, 1001, tt.recvCounter)
+				require.Equal(t, 1002, int(result.(*streamResult).nextResultSetCounter.Load()))
+			}
+		})
+	}
+}

internal/table/session.go

@@ -996,7 +996,7 @@ func (s *session) StreamReadTable(
 		return nil, xerrors.WithStackTrace(err)
 	}
 
-	return scanner.NewStream(
+	return scanner.NewStream(ctx,
 		func(ctx context.Context) (
 			set *Ydb.ResultSet,
 			stats *Ydb_TableStats.QueryStats,
@@ -1024,7 +1024,7 @@ func (s *session) StreamReadTable(
 			return err
 		},
 		scanner.WithIgnoreTruncated(true), // stream read table always returns truncated flag on last result set
-	), nil
+	)
 }
 
 func (s *session) ReadRows(
@@ -1123,7 +1123,7 @@ func (s *session) StreamExecuteScanQuery(
 		return nil, xerrors.WithStackTrace(err)
 	}
 
-	return scanner.NewStream(
+	return scanner.NewStream(ctx,
 		func(ctx context.Context) (
 			set *Ydb.ResultSet,
 			stats *Ydb_TableStats.QueryStats,
@@ -1152,7 +1152,7 @@ func (s *session) StreamExecuteScanQuery(
 		},
 		scanner.WithIgnoreTruncated(s.config.IgnoreTruncated()),
 		scanner.WithMarkTruncatedAsRetryable(),
-	), nil
+	)
 }
 
 // BulkUpsert uploads given list of ydb struct values to the table.

tests/integration/table_scan_error_test.go

@@ -71,13 +71,8 @@ func TestIssue847ScanError(t *testing.T) {
 		if err != nil {
 			return err
 		}
-		err = res.NextResultSetErr(ctx)
-		if err != nil {
-			return err
-		}
 		return res.Err()
 	}, table.WithTxSettings(table.TxSettings(table.WithSnapshotReadOnly())))
 	require.Error(t, err)
-	t.Log(err)
 	require.ErrorContains(t, err, "Unexpected token 'SELICT'")
 }