github.com/ydb-platform/ydb-go-sdk/v3@v3.89.2/internal/query/result_set_go1.23_test.go (about)

     1  //go:build go1.23
     2  
     3  package query
     4  
     5  import (
     6  	"context"
     7  	"fmt"
     8  	"io"
     9  	"testing"
    10  
    11  	"github.com/stretchr/testify/require"
    12  	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb"
    13  	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb_Query"
    14  	"go.uber.org/mock/gomock"
    15  	grpcCodes "google.golang.org/grpc/codes"
    16  	grpcStatus "google.golang.org/grpc/status"
    17  
    18  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
    19  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xtest"
    20  )
    21  
    22  func TestResultSetRangeRows(t *testing.T) {
    23  	ctx := xtest.Context(t)
    24  	ctrl := gomock.NewController(t)
    25  	t.Run("EmptyResultSet", func(t *testing.T) {
    26  		stream := NewMockQueryService_ExecuteQueryClient(ctrl)
    27  		stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{
    28  			Status:         Ydb.StatusIds_SUCCESS,
    29  			ResultSetIndex: 0,
    30  			ResultSet: &Ydb.ResultSet{
    31  				Columns: []*Ydb.Column{
    32  					{
    33  						Name: "a",
    34  						Type: &Ydb.Type{
    35  							Type: &Ydb.Type_TypeId{
    36  								TypeId: Ydb.Type_UINT64,
    37  							},
    38  						},
    39  					},
    40  					{
    41  						Name: "b",
    42  						Type: &Ydb.Type{
    43  							Type: &Ydb.Type_TypeId{
    44  								TypeId: Ydb.Type_UTF8,
    45  							},
    46  						},
    47  					},
    48  				},
    49  				Rows: []*Ydb.Value{},
    50  			},
    51  		}, nil)
    52  		stream.EXPECT().Recv().Return(nil, io.EOF)
    53  		recv, err := stream.Recv()
    54  		require.NoError(t, err)
    55  		rs := newResultSet(func() (*Ydb_Query.ExecuteQueryResponsePart, error) {
    56  			part, err := stream.Recv()
    57  			if err != nil {
    58  				return nil, xerrors.WithStackTrace(err)
    59  			}
    60  
    61  			return part, nil
    62  		}, recv)
    63  		require.EqualValues(t, 0, rs.index)
    64  		count := 0
    65  		for _, err := range rs.Rows(ctx) {
    66  			require.NoError(t, err)
    67  			count++
    68  		}
    69  		require.EqualValues(t, 0, count)
    70  	})
    71  	t.Run("SecondResultSetEmpty", func(t *testing.T) {
    72  		stream := NewMockQueryService_ExecuteQueryClient(ctrl)
    73  		stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{
    74  			Status:         Ydb.StatusIds_SUCCESS,
    75  			ResultSetIndex: 0,
    76  			ResultSet: &Ydb.ResultSet{
    77  				Columns: []*Ydb.Column{
    78  					{
    79  						Name: "a",
    80  						Type: &Ydb.Type{
    81  							Type: &Ydb.Type_TypeId{
    82  								TypeId: Ydb.Type_UINT64,
    83  							},
    84  						},
    85  					},
    86  					{
    87  						Name: "b",
    88  						Type: &Ydb.Type{
    89  							Type: &Ydb.Type_TypeId{
    90  								TypeId: Ydb.Type_UTF8,
    91  							},
    92  						},
    93  					},
    94  				},
    95  				Rows: []*Ydb.Value{
    96  					{
    97  						Items: []*Ydb.Value{{
    98  							Value: &Ydb.Value_Uint64Value{
    99  								Uint64Value: 1,
   100  							},
   101  						}, {
   102  							Value: &Ydb.Value_TextValue{
   103  								TextValue: "1",
   104  							},
   105  						}},
   106  					},
   107  					{
   108  						Items: []*Ydb.Value{{
   109  							Value: &Ydb.Value_Uint64Value{
   110  								Uint64Value: 2,
   111  							},
   112  						}, {
   113  							Value: &Ydb.Value_TextValue{
   114  								TextValue: "2",
   115  							},
   116  						}},
   117  					},
   118  					{
   119  						Items: []*Ydb.Value{{
   120  							Value: &Ydb.Value_Uint64Value{
   121  								Uint64Value: 3,
   122  							},
   123  						}, {
   124  							Value: &Ydb.Value_TextValue{
   125  								TextValue: "3",
   126  							},
   127  						}},
   128  					},
   129  				},
   130  			},
   131  		}, nil)
   132  		stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{
   133  			Status:         Ydb.StatusIds_SUCCESS,
   134  			ResultSetIndex: 0,
   135  			ResultSet: &Ydb.ResultSet{
   136  				Rows: []*Ydb.Value{},
   137  			},
   138  		}, nil)
   139  		stream.EXPECT().Recv().Return(nil, io.EOF)
   140  		recv, err := stream.Recv()
   141  		require.NoError(t, err)
   142  		rs := newResultSet(func() (*Ydb_Query.ExecuteQueryResponsePart, error) {
   143  			part, err := stream.Recv()
   144  			if err != nil {
   145  				return nil, xerrors.WithStackTrace(err)
   146  			}
   147  
   148  			return part, nil
   149  		}, recv)
   150  		require.EqualValues(t, 0, rs.index)
   151  		count := 0
   152  		for row, err := range rs.Rows(ctx) {
   153  			require.NoError(t, err)
   154  			require.EqualValues(t, count, rs.rowIndex)
   155  			var (
   156  				a uint64
   157  				b string
   158  			)
   159  			err := row.Scan(&a, &b)
   160  			require.NoError(t, err)
   161  			count++
   162  			require.EqualValues(t, count, a)
   163  			require.EqualValues(t, fmt.Sprintf("%v", count), b)
   164  		}
   165  		require.EqualValues(t, count, 3)
   166  	})
   167  	t.Run("BreakIterate", func(t *testing.T) {
   168  		stream := NewMockQueryService_ExecuteQueryClient(ctrl)
   169  		stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{
   170  			Status:         Ydb.StatusIds_SUCCESS,
   171  			ResultSetIndex: 0,
   172  			ResultSet: &Ydb.ResultSet{
   173  				Columns: []*Ydb.Column{
   174  					{
   175  						Name: "a",
   176  						Type: &Ydb.Type{
   177  							Type: &Ydb.Type_TypeId{
   178  								TypeId: Ydb.Type_UINT64,
   179  							},
   180  						},
   181  					},
   182  					{
   183  						Name: "b",
   184  						Type: &Ydb.Type{
   185  							Type: &Ydb.Type_TypeId{
   186  								TypeId: Ydb.Type_UTF8,
   187  							},
   188  						},
   189  					},
   190  				},
   191  				Rows: []*Ydb.Value{
   192  					{
   193  						Items: []*Ydb.Value{{
   194  							Value: &Ydb.Value_Uint64Value{
   195  								Uint64Value: 1,
   196  							},
   197  						}, {
   198  							Value: &Ydb.Value_TextValue{
   199  								TextValue: "1",
   200  							},
   201  						}},
   202  					},
   203  					{
   204  						Items: []*Ydb.Value{{
   205  							Value: &Ydb.Value_Uint64Value{
   206  								Uint64Value: 2,
   207  							},
   208  						}, {
   209  							Value: &Ydb.Value_TextValue{
   210  								TextValue: "2",
   211  							},
   212  						}},
   213  					},
   214  					{
   215  						Items: []*Ydb.Value{{
   216  							Value: &Ydb.Value_Uint64Value{
   217  								Uint64Value: 3,
   218  							},
   219  						}, {
   220  							Value: &Ydb.Value_TextValue{
   221  								TextValue: "3",
   222  							},
   223  						}},
   224  					},
   225  				},
   226  			},
   227  		}, nil)
   228  		recv, err := stream.Recv()
   229  		require.NoError(t, err)
   230  		rs := newResultSet(func() (*Ydb_Query.ExecuteQueryResponsePart, error) {
   231  			part, err := stream.Recv()
   232  			if err != nil {
   233  				return nil, xerrors.WithStackTrace(err)
   234  			}
   235  
   236  			return part, nil
   237  		}, recv)
   238  		require.EqualValues(t, 0, rs.index)
   239  		count := 0
   240  		for _, err := range rs.Rows(ctx) {
   241  			require.NoError(t, err)
   242  			require.EqualValues(t, count, rs.rowIndex)
   243  			if count > 0 {
   244  				break
   245  			}
   246  			count++
   247  		}
   248  		require.EqualValues(t, count, 1)
   249  	})
   250  	t.Run("IntermediateResultSetEmpty", func(t *testing.T) {
   251  		stream := NewMockQueryService_ExecuteQueryClient(ctrl)
   252  		stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{
   253  			Status:         Ydb.StatusIds_SUCCESS,
   254  			ResultSetIndex: 0,
   255  			ResultSet: &Ydb.ResultSet{
   256  				Columns: []*Ydb.Column{
   257  					{
   258  						Name: "a",
   259  						Type: &Ydb.Type{
   260  							Type: &Ydb.Type_TypeId{
   261  								TypeId: Ydb.Type_UINT64,
   262  							},
   263  						},
   264  					},
   265  					{
   266  						Name: "b",
   267  						Type: &Ydb.Type{
   268  							Type: &Ydb.Type_TypeId{
   269  								TypeId: Ydb.Type_UTF8,
   270  							},
   271  						},
   272  					},
   273  				},
   274  				Rows: []*Ydb.Value{
   275  					{
   276  						Items: []*Ydb.Value{{
   277  							Value: &Ydb.Value_Uint64Value{
   278  								Uint64Value: 1,
   279  							},
   280  						}, {
   281  							Value: &Ydb.Value_TextValue{
   282  								TextValue: "1",
   283  							},
   284  						}},
   285  					},
   286  					{
   287  						Items: []*Ydb.Value{{
   288  							Value: &Ydb.Value_Uint64Value{
   289  								Uint64Value: 2,
   290  							},
   291  						}, {
   292  							Value: &Ydb.Value_TextValue{
   293  								TextValue: "2",
   294  							},
   295  						}},
   296  					},
   297  					{
   298  						Items: []*Ydb.Value{{
   299  							Value: &Ydb.Value_Uint64Value{
   300  								Uint64Value: 3,
   301  							},
   302  						}, {
   303  							Value: &Ydb.Value_TextValue{
   304  								TextValue: "3",
   305  							},
   306  						}},
   307  					},
   308  				},
   309  			},
   310  		}, nil)
   311  		stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{
   312  			Status:         Ydb.StatusIds_SUCCESS,
   313  			ResultSetIndex: 0,
   314  			ResultSet: &Ydb.ResultSet{
   315  				Rows: []*Ydb.Value{},
   316  			},
   317  		}, nil)
   318  		stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{
   319  			Status:         Ydb.StatusIds_SUCCESS,
   320  			ResultSetIndex: 0,
   321  			ResultSet: &Ydb.ResultSet{
   322  				Columns: []*Ydb.Column{
   323  					{
   324  						Name: "a",
   325  						Type: &Ydb.Type{
   326  							Type: &Ydb.Type_TypeId{
   327  								TypeId: Ydb.Type_UINT64,
   328  							},
   329  						},
   330  					},
   331  					{
   332  						Name: "b",
   333  						Type: &Ydb.Type{
   334  							Type: &Ydb.Type_TypeId{
   335  								TypeId: Ydb.Type_UTF8,
   336  							},
   337  						},
   338  					},
   339  				},
   340  				Rows: []*Ydb.Value{
   341  					{
   342  						Items: []*Ydb.Value{{
   343  							Value: &Ydb.Value_Uint64Value{
   344  								Uint64Value: 4,
   345  							},
   346  						}, {
   347  							Value: &Ydb.Value_TextValue{
   348  								TextValue: "4",
   349  							},
   350  						}},
   351  					},
   352  					{
   353  						Items: []*Ydb.Value{{
   354  							Value: &Ydb.Value_Uint64Value{
   355  								Uint64Value: 5,
   356  							},
   357  						}, {
   358  							Value: &Ydb.Value_TextValue{
   359  								TextValue: "5",
   360  							},
   361  						}},
   362  					},
   363  					{
   364  						Items: []*Ydb.Value{{
   365  							Value: &Ydb.Value_Uint64Value{
   366  								Uint64Value: 6,
   367  							},
   368  						}, {
   369  							Value: &Ydb.Value_TextValue{
   370  								TextValue: "6",
   371  							},
   372  						}},
   373  					},
   374  				},
   375  			},
   376  		}, nil)
   377  		stream.EXPECT().Recv().Return(nil, io.EOF)
   378  		recv, err := stream.Recv()
   379  		require.NoError(t, err)
   380  		rs := newResultSet(func() (*Ydb_Query.ExecuteQueryResponsePart, error) {
   381  			part, err := stream.Recv()
   382  			if err != nil {
   383  				return nil, xerrors.WithStackTrace(err)
   384  			}
   385  
   386  			return part, nil
   387  		}, recv)
   388  		require.EqualValues(t, 0, rs.index)
   389  		count := 0
   390  		for row, err := range rs.Rows(ctx) {
   391  			require.NoError(t, err)
   392  			var (
   393  				a uint64
   394  				b string
   395  			)
   396  			err := row.Scan(&a, &b)
   397  			require.NoError(t, err)
   398  			count++
   399  			require.EqualValues(t, count, a)
   400  			require.EqualValues(t, fmt.Sprintf("%v", count), b)
   401  		}
   402  		require.EqualValues(t, count, 6)
   403  	})
   404  	t.Run("OverTwoParts", func(t *testing.T) {
   405  		stream := NewMockQueryService_ExecuteQueryClient(ctrl)
   406  		stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{
   407  			Status:         Ydb.StatusIds_SUCCESS,
   408  			ResultSetIndex: 0,
   409  			ResultSet: &Ydb.ResultSet{
   410  				Columns: []*Ydb.Column{
   411  					{
   412  						Name: "a",
   413  						Type: &Ydb.Type{
   414  							Type: &Ydb.Type_TypeId{
   415  								TypeId: Ydb.Type_UINT64,
   416  							},
   417  						},
   418  					},
   419  					{
   420  						Name: "b",
   421  						Type: &Ydb.Type{
   422  							Type: &Ydb.Type_TypeId{
   423  								TypeId: Ydb.Type_UTF8,
   424  							},
   425  						},
   426  					},
   427  				},
   428  				Rows: []*Ydb.Value{
   429  					{
   430  						Items: []*Ydb.Value{{
   431  							Value: &Ydb.Value_Uint64Value{
   432  								Uint64Value: 1,
   433  							},
   434  						}, {
   435  							Value: &Ydb.Value_TextValue{
   436  								TextValue: "1",
   437  							},
   438  						}},
   439  					},
   440  					{
   441  						Items: []*Ydb.Value{{
   442  							Value: &Ydb.Value_Uint64Value{
   443  								Uint64Value: 2,
   444  							},
   445  						}, {
   446  							Value: &Ydb.Value_TextValue{
   447  								TextValue: "2",
   448  							},
   449  						}},
   450  					},
   451  					{
   452  						Items: []*Ydb.Value{{
   453  							Value: &Ydb.Value_Uint64Value{
   454  								Uint64Value: 3,
   455  							},
   456  						}, {
   457  							Value: &Ydb.Value_TextValue{
   458  								TextValue: "3",
   459  							},
   460  						}},
   461  					},
   462  				},
   463  			},
   464  		}, nil)
   465  		stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{
   466  			Status:         Ydb.StatusIds_SUCCESS,
   467  			ResultSetIndex: 0,
   468  			ResultSet: &Ydb.ResultSet{
   469  				Rows: []*Ydb.Value{
   470  					{
   471  						Items: []*Ydb.Value{{
   472  							Value: &Ydb.Value_Uint64Value{
   473  								Uint64Value: 4,
   474  							},
   475  						}, {
   476  							Value: &Ydb.Value_TextValue{
   477  								TextValue: "4",
   478  							},
   479  						}},
   480  					},
   481  					{
   482  						Items: []*Ydb.Value{{
   483  							Value: &Ydb.Value_Uint64Value{
   484  								Uint64Value: 5,
   485  							},
   486  						}, {
   487  							Value: &Ydb.Value_TextValue{
   488  								TextValue: "5",
   489  							},
   490  						}},
   491  					},
   492  				},
   493  			},
   494  		}, nil)
   495  		stream.EXPECT().Recv().Return(nil, io.EOF)
   496  		recv, err := stream.Recv()
   497  		require.NoError(t, err)
   498  		rs := newResultSet(func() (*Ydb_Query.ExecuteQueryResponsePart, error) {
   499  			part, err := stream.Recv()
   500  			if err != nil {
   501  				return nil, xerrors.WithStackTrace(err)
   502  			}
   503  
   504  			return part, nil
   505  		}, recv)
   506  		require.EqualValues(t, 0, rs.index)
   507  		count := 0
   508  		for row, err := range rs.Rows(ctx) {
   509  			require.NoError(t, err)
   510  			var (
   511  				a uint64
   512  				b string
   513  			)
   514  			err := row.Scan(&a, &b)
   515  			require.NoError(t, err)
   516  			count++
   517  			require.EqualValues(t, count, a)
   518  			require.EqualValues(t, fmt.Sprintf("%v", count), b)
   519  		}
   520  		require.EqualValues(t, count, 5)
   521  	})
   522  	t.Run("CanceledContext", func(t *testing.T) {
   523  		childCtx, cancel := context.WithCancel(xtest.Context(t))
   524  		defer cancel()
   525  		stream := NewMockQueryService_ExecuteQueryClient(ctrl)
   526  		stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{
   527  			Status:         Ydb.StatusIds_SUCCESS,
   528  			ResultSetIndex: 0,
   529  			ResultSet: &Ydb.ResultSet{
   530  				Columns: []*Ydb.Column{
   531  					{
   532  						Name: "a",
   533  						Type: &Ydb.Type{
   534  							Type: &Ydb.Type_TypeId{
   535  								TypeId: Ydb.Type_UINT64,
   536  							},
   537  						},
   538  					},
   539  					{
   540  						Name: "b",
   541  						Type: &Ydb.Type{
   542  							Type: &Ydb.Type_TypeId{
   543  								TypeId: Ydb.Type_UTF8,
   544  							},
   545  						},
   546  					},
   547  				},
   548  				Rows: []*Ydb.Value{
   549  					{
   550  						Items: []*Ydb.Value{{
   551  							Value: &Ydb.Value_Uint64Value{
   552  								Uint64Value: 1,
   553  							},
   554  						}, {
   555  							Value: &Ydb.Value_TextValue{
   556  								TextValue: "1",
   557  							},
   558  						}},
   559  					},
   560  					{
   561  						Items: []*Ydb.Value{{
   562  							Value: &Ydb.Value_Uint64Value{
   563  								Uint64Value: 2,
   564  							},
   565  						}, {
   566  							Value: &Ydb.Value_TextValue{
   567  								TextValue: "2",
   568  							},
   569  						}},
   570  					},
   571  					{
   572  						Items: []*Ydb.Value{{
   573  							Value: &Ydb.Value_Uint64Value{
   574  								Uint64Value: 3,
   575  							},
   576  						}, {
   577  							Value: &Ydb.Value_TextValue{
   578  								TextValue: "3",
   579  							},
   580  						}},
   581  					},
   582  				},
   583  			},
   584  		}, nil)
   585  		recv, err := stream.Recv()
   586  		require.NoError(t, err)
   587  		rs := newResultSet(func() (*Ydb_Query.ExecuteQueryResponsePart, error) {
   588  			part, err := stream.Recv()
   589  			if err != nil {
   590  				return nil, xerrors.WithStackTrace(err)
   591  			}
   592  
   593  			return part, nil
   594  		}, recv)
   595  		require.EqualValues(t, 0, rs.index)
   596  		var (
   597  			count     = 0
   598  			cancelled = false
   599  		)
   600  		for _, err := range rs.Rows(childCtx) {
   601  			count++
   602  			if !cancelled {
   603  				require.NoError(t, err)
   604  				cancel()
   605  				cancelled = true
   606  			} else {
   607  				require.ErrorIs(t, err, context.Canceled)
   608  			}
   609  		}
   610  		require.EqualValues(t, count, 2)
   611  	})
   612  	t.Run("OperationError", func(t *testing.T) {
   613  		stream := NewMockQueryService_ExecuteQueryClient(ctrl)
   614  		stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{
   615  			Status:         Ydb.StatusIds_SUCCESS,
   616  			ResultSetIndex: 0,
   617  			ResultSet: &Ydb.ResultSet{
   618  				Columns: []*Ydb.Column{
   619  					{
   620  						Name: "a",
   621  						Type: &Ydb.Type{
   622  							Type: &Ydb.Type_TypeId{
   623  								TypeId: Ydb.Type_UINT64,
   624  							},
   625  						},
   626  					},
   627  					{
   628  						Name: "b",
   629  						Type: &Ydb.Type{
   630  							Type: &Ydb.Type_TypeId{
   631  								TypeId: Ydb.Type_UTF8,
   632  							},
   633  						},
   634  					},
   635  				},
   636  				Rows: []*Ydb.Value{
   637  					{
   638  						Items: []*Ydb.Value{{
   639  							Value: &Ydb.Value_Uint64Value{
   640  								Uint64Value: 1,
   641  							},
   642  						}, {
   643  							Value: &Ydb.Value_TextValue{
   644  								TextValue: "1",
   645  							},
   646  						}},
   647  					},
   648  					{
   649  						Items: []*Ydb.Value{{
   650  							Value: &Ydb.Value_Uint64Value{
   651  								Uint64Value: 2,
   652  							},
   653  						}, {
   654  							Value: &Ydb.Value_TextValue{
   655  								TextValue: "2",
   656  							},
   657  						}},
   658  					},
   659  					{
   660  						Items: []*Ydb.Value{{
   661  							Value: &Ydb.Value_Uint64Value{
   662  								Uint64Value: 3,
   663  							},
   664  						}, {
   665  							Value: &Ydb.Value_TextValue{
   666  								TextValue: "3",
   667  							},
   668  						}},
   669  					},
   670  				},
   671  			},
   672  		}, nil)
   673  		stream.EXPECT().Recv().Return(nil, xerrors.Operation(xerrors.WithStatusCode(
   674  			Ydb.StatusIds_OVERLOADED,
   675  		)))
   676  		recv, err := stream.Recv()
   677  		require.NoError(t, err)
   678  		rs := newResultSet(func() (*Ydb_Query.ExecuteQueryResponsePart, error) {
   679  			part, err := nextPart(stream)
   680  			if err != nil {
   681  				return nil, xerrors.WithStackTrace(err)
   682  			}
   683  			if resultSetIndex := part.GetResultSetIndex(); resultSetIndex != 0 {
   684  				return nil, xerrors.WithStackTrace(fmt.Errorf(
   685  					"%w: %d != %d",
   686  					errWrongNextResultSetIndex,
   687  					resultSetIndex, 0,
   688  				))
   689  			}
   690  
   691  			return part, nil
   692  		}, recv)
   693  		require.EqualValues(t, 0, rs.index)
   694  		count := 0
   695  		for _, err := range rs.Rows(ctx) {
   696  			if count < 3 {
   697  				require.NoError(t, err)
   698  			} else {
   699  				require.True(t, xerrors.IsOperationError(err, Ydb.StatusIds_OVERLOADED))
   700  			}
   701  			count++
   702  		}
   703  		require.EqualValues(t, count, 4)
   704  	})
   705  	t.Run("TransportError", func(t *testing.T) {
   706  		stream := NewMockQueryService_ExecuteQueryClient(ctrl)
   707  		stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{
   708  			Status:         Ydb.StatusIds_SUCCESS,
   709  			ResultSetIndex: 0,
   710  			ResultSet: &Ydb.ResultSet{
   711  				Columns: []*Ydb.Column{
   712  					{
   713  						Name: "a",
   714  						Type: &Ydb.Type{
   715  							Type: &Ydb.Type_TypeId{
   716  								TypeId: Ydb.Type_UINT64,
   717  							},
   718  						},
   719  					},
   720  					{
   721  						Name: "b",
   722  						Type: &Ydb.Type{
   723  							Type: &Ydb.Type_TypeId{
   724  								TypeId: Ydb.Type_UTF8,
   725  							},
   726  						},
   727  					},
   728  				},
   729  				Rows: []*Ydb.Value{
   730  					{
   731  						Items: []*Ydb.Value{{
   732  							Value: &Ydb.Value_Uint64Value{
   733  								Uint64Value: 1,
   734  							},
   735  						}, {
   736  							Value: &Ydb.Value_TextValue{
   737  								TextValue: "1",
   738  							},
   739  						}},
   740  					},
   741  					{
   742  						Items: []*Ydb.Value{{
   743  							Value: &Ydb.Value_Uint64Value{
   744  								Uint64Value: 2,
   745  							},
   746  						}, {
   747  							Value: &Ydb.Value_TextValue{
   748  								TextValue: "2",
   749  							},
   750  						}},
   751  					},
   752  					{
   753  						Items: []*Ydb.Value{{
   754  							Value: &Ydb.Value_Uint64Value{
   755  								Uint64Value: 3,
   756  							},
   757  						}, {
   758  							Value: &Ydb.Value_TextValue{
   759  								TextValue: "3",
   760  							},
   761  						}},
   762  					},
   763  				},
   764  			},
   765  		}, nil)
   766  		stream.EXPECT().Recv().Return(nil, grpcStatus.Error(grpcCodes.Unavailable, ""))
   767  		recv, err := stream.Recv()
   768  		require.NoError(t, err)
   769  		rs := newResultSet(func() (*Ydb_Query.ExecuteQueryResponsePart, error) {
   770  			part, err := nextPart(stream)
   771  			if err != nil {
   772  				return nil, xerrors.WithStackTrace(err)
   773  			}
   774  			if resultSetIndex := part.GetResultSetIndex(); resultSetIndex != 0 {
   775  				return nil, xerrors.WithStackTrace(fmt.Errorf(
   776  					"%w: %d != %d",
   777  					errWrongNextResultSetIndex,
   778  					resultSetIndex, 0,
   779  				))
   780  			}
   781  
   782  			return part, nil
   783  		}, recv)
   784  		require.EqualValues(t, 0, rs.index)
   785  		count := 0
   786  		for _, err := range rs.Rows(ctx) {
   787  			if count < 3 {
   788  				require.NoError(t, err)
   789  			} else {
   790  				require.True(t, xerrors.IsTransportError(err, grpcCodes.Unavailable))
   791  			}
   792  			count++
   793  		}
   794  		require.EqualValues(t, count, 4)
   795  	})
   796  	t.Run("WrongResultSetIndex", func(t *testing.T) {
   797  		stream := NewMockQueryService_ExecuteQueryClient(ctrl)
   798  		stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{
   799  			Status:         Ydb.StatusIds_SUCCESS,
   800  			ResultSetIndex: 0,
   801  			ResultSet: &Ydb.ResultSet{
   802  				Columns: []*Ydb.Column{
   803  					{
   804  						Name: "a",
   805  						Type: &Ydb.Type{
   806  							Type: &Ydb.Type_TypeId{
   807  								TypeId: Ydb.Type_UINT64,
   808  							},
   809  						},
   810  					},
   811  					{
   812  						Name: "b",
   813  						Type: &Ydb.Type{
   814  							Type: &Ydb.Type_TypeId{
   815  								TypeId: Ydb.Type_UTF8,
   816  							},
   817  						},
   818  					},
   819  				},
   820  				Rows: []*Ydb.Value{
   821  					{
   822  						Items: []*Ydb.Value{{
   823  							Value: &Ydb.Value_Uint64Value{
   824  								Uint64Value: 1,
   825  							},
   826  						}, {
   827  							Value: &Ydb.Value_TextValue{
   828  								TextValue: "1",
   829  							},
   830  						}},
   831  					},
   832  					{
   833  						Items: []*Ydb.Value{{
   834  							Value: &Ydb.Value_Uint64Value{
   835  								Uint64Value: 2,
   836  							},
   837  						}, {
   838  							Value: &Ydb.Value_TextValue{
   839  								TextValue: "2",
   840  							},
   841  						}},
   842  					},
   843  					{
   844  						Items: []*Ydb.Value{{
   845  							Value: &Ydb.Value_Uint64Value{
   846  								Uint64Value: 3,
   847  							},
   848  						}, {
   849  							Value: &Ydb.Value_TextValue{
   850  								TextValue: "3",
   851  							},
   852  						}},
   853  					},
   854  				},
   855  			},
   856  		}, nil)
   857  		stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{
   858  			Status:         Ydb.StatusIds_SUCCESS,
   859  			ResultSetIndex: 1,
   860  			ResultSet: &Ydb.ResultSet{
   861  				Columns: []*Ydb.Column{
   862  					{
   863  						Name: "a",
   864  						Type: &Ydb.Type{
   865  							Type: &Ydb.Type_TypeId{
   866  								TypeId: Ydb.Type_UINT64,
   867  							},
   868  						},
   869  					},
   870  					{
   871  						Name: "b",
   872  						Type: &Ydb.Type{
   873  							Type: &Ydb.Type_TypeId{
   874  								TypeId: Ydb.Type_UTF8,
   875  							},
   876  						},
   877  					},
   878  				},
   879  				Rows: []*Ydb.Value{
   880  					{
   881  						Items: []*Ydb.Value{{
   882  							Value: &Ydb.Value_Uint64Value{
   883  								Uint64Value: 1,
   884  							},
   885  						}, {
   886  							Value: &Ydb.Value_TextValue{
   887  								TextValue: "1",
   888  							},
   889  						}},
   890  					},
   891  					{
   892  						Items: []*Ydb.Value{{
   893  							Value: &Ydb.Value_Uint64Value{
   894  								Uint64Value: 2,
   895  							},
   896  						}, {
   897  							Value: &Ydb.Value_TextValue{
   898  								TextValue: "2",
   899  							},
   900  						}},
   901  					},
   902  					{
   903  						Items: []*Ydb.Value{{
   904  							Value: &Ydb.Value_Uint64Value{
   905  								Uint64Value: 3,
   906  							},
   907  						}, {
   908  							Value: &Ydb.Value_TextValue{
   909  								TextValue: "3",
   910  							},
   911  						}},
   912  					},
   913  				},
   914  			},
   915  		}, nil)
   916  		recv, err := stream.Recv()
   917  		require.NoError(t, err)
   918  		rs := newResultSet(func() (*Ydb_Query.ExecuteQueryResponsePart, error) {
   919  			part, err := nextPart(stream)
   920  			if err != nil {
   921  				return nil, xerrors.WithStackTrace(err)
   922  			}
   923  
   924  			return part, nil
   925  		}, recv)
   926  		require.EqualValues(t, 0, rs.index)
   927  		count := 0
   928  		for _, err := range rs.Rows(ctx) {
   929  			if count < 3 {
   930  				require.NoError(t, err)
   931  			} else {
   932  				require.ErrorIs(t, err, errWrongResultSetIndex)
   933  			}
   934  			count++
   935  		}
   936  		require.EqualValues(t, count, 4)
   937  	})
   938  }