github.com/ydb-platform/ydb-go-sdk/v3@v3.57.0/internal/query/result_set_test.go (about)

     1  package query
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"io"
     7  	"testing"
     8  
     9  	"github.com/stretchr/testify/require"
    10  	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb"
    11  	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb_Query"
    12  	"go.uber.org/mock/gomock"
    13  	grpcCodes "google.golang.org/grpc/codes"
    14  	grpcStatus "google.golang.org/grpc/status"
    15  
    16  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
    17  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xtest"
    18  )
    19  
    20  func TestResultSetNext(t *testing.T) {
    21  	ctx := xtest.Context(t)
    22  	ctrl := gomock.NewController(t)
    23  	t.Run("OverTwoParts", func(t *testing.T) {
    24  		stream := NewMockQueryService_ExecuteQueryClient(ctrl)
    25  		stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{
    26  			Status:         Ydb.StatusIds_SUCCESS,
    27  			ResultSetIndex: 0,
    28  			ResultSet: &Ydb.ResultSet{
    29  				Columns: []*Ydb.Column{
    30  					{
    31  						Name: "a",
    32  						Type: &Ydb.Type{
    33  							Type: &Ydb.Type_TypeId{
    34  								TypeId: Ydb.Type_UINT64,
    35  							},
    36  						},
    37  					},
    38  					{
    39  						Name: "b",
    40  						Type: &Ydb.Type{
    41  							Type: &Ydb.Type_TypeId{
    42  								TypeId: Ydb.Type_UTF8,
    43  							},
    44  						},
    45  					},
    46  				},
    47  				Rows: []*Ydb.Value{
    48  					{
    49  						Items: []*Ydb.Value{{
    50  							Value: &Ydb.Value_Uint64Value{
    51  								Uint64Value: 1,
    52  							},
    53  						}, {
    54  							Value: &Ydb.Value_TextValue{
    55  								TextValue: "1",
    56  							},
    57  						}},
    58  					},
    59  					{
    60  						Items: []*Ydb.Value{{
    61  							Value: &Ydb.Value_Uint64Value{
    62  								Uint64Value: 2,
    63  							},
    64  						}, {
    65  							Value: &Ydb.Value_TextValue{
    66  								TextValue: "2",
    67  							},
    68  						}},
    69  					},
    70  					{
    71  						Items: []*Ydb.Value{{
    72  							Value: &Ydb.Value_Uint64Value{
    73  								Uint64Value: 3,
    74  							},
    75  						}, {
    76  							Value: &Ydb.Value_TextValue{
    77  								TextValue: "3",
    78  							},
    79  						}},
    80  					},
    81  				},
    82  			},
    83  		}, nil)
    84  		stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{
    85  			Status:         Ydb.StatusIds_SUCCESS,
    86  			ResultSetIndex: 0,
    87  			ResultSet: &Ydb.ResultSet{
    88  				Rows: []*Ydb.Value{
    89  					{
    90  						Items: []*Ydb.Value{{
    91  							Value: &Ydb.Value_Uint64Value{
    92  								Uint64Value: 4,
    93  							},
    94  						}, {
    95  							Value: &Ydb.Value_TextValue{
    96  								TextValue: "4",
    97  							},
    98  						}},
    99  					},
   100  					{
   101  						Items: []*Ydb.Value{{
   102  							Value: &Ydb.Value_Uint64Value{
   103  								Uint64Value: 5,
   104  							},
   105  						}, {
   106  							Value: &Ydb.Value_TextValue{
   107  								TextValue: "5",
   108  							},
   109  						}},
   110  					},
   111  				},
   112  			},
   113  		}, nil)
   114  		stream.EXPECT().Recv().Return(nil, io.EOF)
   115  		recv, err := stream.Recv()
   116  		require.NoError(t, err)
   117  		rs := newResultSet(func() (*Ydb_Query.ExecuteQueryResponsePart, error) {
   118  			part, err := stream.Recv()
   119  			if err != nil {
   120  				return nil, xerrors.WithStackTrace(err)
   121  			}
   122  
   123  			return part, nil
   124  		}, recv)
   125  		require.EqualValues(t, 0, rs.index)
   126  		{
   127  			_, err := rs.next(ctx)
   128  			require.NoError(t, err)
   129  			require.EqualValues(t, 0, rs.rowIndex)
   130  		}
   131  		{
   132  			_, err := rs.next(ctx)
   133  			require.NoError(t, err)
   134  			require.EqualValues(t, 1, rs.rowIndex)
   135  		}
   136  		{
   137  			_, err := rs.next(ctx)
   138  			require.NoError(t, err)
   139  			require.EqualValues(t, 2, rs.rowIndex)
   140  		}
   141  		{
   142  			_, err := rs.next(ctx)
   143  			require.NoError(t, err)
   144  			require.EqualValues(t, 0, rs.rowIndex)
   145  		}
   146  		{
   147  			_, err := rs.next(ctx)
   148  			require.NoError(t, err)
   149  			require.EqualValues(t, 1, rs.rowIndex)
   150  		}
   151  		{
   152  			_, err := rs.next(ctx)
   153  			require.ErrorIs(t, err, io.EOF)
   154  		}
   155  	})
   156  	t.Run("CanceledContext", func(t *testing.T) {
   157  		ctx, cancel := context.WithCancel(xtest.Context(t))
   158  		ctrl := gomock.NewController(t)
   159  		stream := NewMockQueryService_ExecuteQueryClient(ctrl)
   160  		stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{
   161  			Status:         Ydb.StatusIds_SUCCESS,
   162  			ResultSetIndex: 0,
   163  			ResultSet: &Ydb.ResultSet{
   164  				Columns: []*Ydb.Column{
   165  					{
   166  						Name: "a",
   167  						Type: &Ydb.Type{
   168  							Type: &Ydb.Type_TypeId{
   169  								TypeId: Ydb.Type_UINT64,
   170  							},
   171  						},
   172  					},
   173  					{
   174  						Name: "b",
   175  						Type: &Ydb.Type{
   176  							Type: &Ydb.Type_TypeId{
   177  								TypeId: Ydb.Type_UTF8,
   178  							},
   179  						},
   180  					},
   181  				},
   182  				Rows: []*Ydb.Value{
   183  					{
   184  						Items: []*Ydb.Value{{
   185  							Value: &Ydb.Value_Uint64Value{
   186  								Uint64Value: 1,
   187  							},
   188  						}, {
   189  							Value: &Ydb.Value_TextValue{
   190  								TextValue: "1",
   191  							},
   192  						}},
   193  					},
   194  					{
   195  						Items: []*Ydb.Value{{
   196  							Value: &Ydb.Value_Uint64Value{
   197  								Uint64Value: 2,
   198  							},
   199  						}, {
   200  							Value: &Ydb.Value_TextValue{
   201  								TextValue: "2",
   202  							},
   203  						}},
   204  					},
   205  					{
   206  						Items: []*Ydb.Value{{
   207  							Value: &Ydb.Value_Uint64Value{
   208  								Uint64Value: 3,
   209  							},
   210  						}, {
   211  							Value: &Ydb.Value_TextValue{
   212  								TextValue: "3",
   213  							},
   214  						}},
   215  					},
   216  				},
   217  			},
   218  		}, nil)
   219  		recv, err := stream.Recv()
   220  		require.NoError(t, err)
   221  		rs := newResultSet(func() (*Ydb_Query.ExecuteQueryResponsePart, error) {
   222  			part, err := stream.Recv()
   223  			if err != nil {
   224  				return nil, xerrors.WithStackTrace(err)
   225  			}
   226  
   227  			return part, nil
   228  		}, recv)
   229  		require.EqualValues(t, 0, rs.index)
   230  		{
   231  			_, err := rs.next(ctx)
   232  			require.NoError(t, err)
   233  			require.EqualValues(t, 0, rs.rowIndex)
   234  		}
   235  		cancel()
   236  		{
   237  			_, err := rs.next(ctx)
   238  			require.ErrorIs(t, err, context.Canceled)
   239  		}
   240  	})
   241  	t.Run("OperationError", func(t *testing.T) {
   242  		ctx := xtest.Context(t)
   243  		ctrl := gomock.NewController(t)
   244  		stream := NewMockQueryService_ExecuteQueryClient(ctrl)
   245  		stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{
   246  			Status:         Ydb.StatusIds_SUCCESS,
   247  			ResultSetIndex: 0,
   248  			ResultSet: &Ydb.ResultSet{
   249  				Columns: []*Ydb.Column{
   250  					{
   251  						Name: "a",
   252  						Type: &Ydb.Type{
   253  							Type: &Ydb.Type_TypeId{
   254  								TypeId: Ydb.Type_UINT64,
   255  							},
   256  						},
   257  					},
   258  					{
   259  						Name: "b",
   260  						Type: &Ydb.Type{
   261  							Type: &Ydb.Type_TypeId{
   262  								TypeId: Ydb.Type_UTF8,
   263  							},
   264  						},
   265  					},
   266  				},
   267  				Rows: []*Ydb.Value{
   268  					{
   269  						Items: []*Ydb.Value{{
   270  							Value: &Ydb.Value_Uint64Value{
   271  								Uint64Value: 1,
   272  							},
   273  						}, {
   274  							Value: &Ydb.Value_TextValue{
   275  								TextValue: "1",
   276  							},
   277  						}},
   278  					},
   279  					{
   280  						Items: []*Ydb.Value{{
   281  							Value: &Ydb.Value_Uint64Value{
   282  								Uint64Value: 2,
   283  							},
   284  						}, {
   285  							Value: &Ydb.Value_TextValue{
   286  								TextValue: "2",
   287  							},
   288  						}},
   289  					},
   290  					{
   291  						Items: []*Ydb.Value{{
   292  							Value: &Ydb.Value_Uint64Value{
   293  								Uint64Value: 3,
   294  							},
   295  						}, {
   296  							Value: &Ydb.Value_TextValue{
   297  								TextValue: "3",
   298  							},
   299  						}},
   300  					},
   301  				},
   302  			},
   303  		}, nil)
   304  		stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{
   305  			Status:         Ydb.StatusIds_OVERLOADED,
   306  			ResultSetIndex: 0,
   307  		}, nil)
   308  		recv, err := stream.Recv()
   309  		require.NoError(t, err)
   310  		rs := newResultSet(func() (*Ydb_Query.ExecuteQueryResponsePart, error) {
   311  			part, err := nextPart(stream)
   312  			if err != nil {
   313  				return nil, xerrors.WithStackTrace(err)
   314  			}
   315  			if resultSetIndex := part.GetResultSetIndex(); resultSetIndex != 0 {
   316  				return nil, xerrors.WithStackTrace(fmt.Errorf(
   317  					"critical violation of the logic: wrong result set index: %d != %d",
   318  					resultSetIndex, 0,
   319  				))
   320  			}
   321  
   322  			return part, nil
   323  		}, recv)
   324  		require.EqualValues(t, 0, rs.index)
   325  		{
   326  			_, err := rs.next(ctx)
   327  			require.NoError(t, err)
   328  			require.EqualValues(t, 0, rs.rowIndex)
   329  		}
   330  		{
   331  			_, err := rs.next(ctx)
   332  			require.NoError(t, err)
   333  			require.EqualValues(t, 1, rs.rowIndex)
   334  		}
   335  		{
   336  			_, err := rs.next(ctx)
   337  			require.NoError(t, err)
   338  			require.EqualValues(t, 2, rs.rowIndex)
   339  		}
   340  		{
   341  			_, err := rs.next(ctx)
   342  			require.True(t, xerrors.IsOperationError(err, Ydb.StatusIds_OVERLOADED))
   343  		}
   344  	})
   345  	t.Run("TransportError", func(t *testing.T) {
   346  		ctx := xtest.Context(t)
   347  		ctrl := gomock.NewController(t)
   348  		stream := NewMockQueryService_ExecuteQueryClient(ctrl)
   349  		stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{
   350  			Status:         Ydb.StatusIds_SUCCESS,
   351  			ResultSetIndex: 0,
   352  			ResultSet: &Ydb.ResultSet{
   353  				Columns: []*Ydb.Column{
   354  					{
   355  						Name: "a",
   356  						Type: &Ydb.Type{
   357  							Type: &Ydb.Type_TypeId{
   358  								TypeId: Ydb.Type_UINT64,
   359  							},
   360  						},
   361  					},
   362  					{
   363  						Name: "b",
   364  						Type: &Ydb.Type{
   365  							Type: &Ydb.Type_TypeId{
   366  								TypeId: Ydb.Type_UTF8,
   367  							},
   368  						},
   369  					},
   370  				},
   371  				Rows: []*Ydb.Value{
   372  					{
   373  						Items: []*Ydb.Value{{
   374  							Value: &Ydb.Value_Uint64Value{
   375  								Uint64Value: 1,
   376  							},
   377  						}, {
   378  							Value: &Ydb.Value_TextValue{
   379  								TextValue: "1",
   380  							},
   381  						}},
   382  					},
   383  					{
   384  						Items: []*Ydb.Value{{
   385  							Value: &Ydb.Value_Uint64Value{
   386  								Uint64Value: 2,
   387  							},
   388  						}, {
   389  							Value: &Ydb.Value_TextValue{
   390  								TextValue: "2",
   391  							},
   392  						}},
   393  					},
   394  					{
   395  						Items: []*Ydb.Value{{
   396  							Value: &Ydb.Value_Uint64Value{
   397  								Uint64Value: 3,
   398  							},
   399  						}, {
   400  							Value: &Ydb.Value_TextValue{
   401  								TextValue: "3",
   402  							},
   403  						}},
   404  					},
   405  				},
   406  			},
   407  		}, nil)
   408  		stream.EXPECT().Recv().Return(nil, grpcStatus.Error(grpcCodes.Unavailable, ""))
   409  		recv, err := stream.Recv()
   410  		require.NoError(t, err)
   411  		rs := newResultSet(func() (*Ydb_Query.ExecuteQueryResponsePart, error) {
   412  			part, err := nextPart(stream)
   413  			if err != nil {
   414  				return nil, xerrors.WithStackTrace(err)
   415  			}
   416  			if resultSetIndex := part.GetResultSetIndex(); resultSetIndex != 0 {
   417  				return nil, xerrors.WithStackTrace(fmt.Errorf(
   418  					"critical violation of the logic: wrong result set index: %d != %d",
   419  					resultSetIndex, 0,
   420  				))
   421  			}
   422  
   423  			return part, nil
   424  		}, recv)
   425  		require.EqualValues(t, 0, rs.index)
   426  		{
   427  			_, err := rs.next(ctx)
   428  			require.NoError(t, err)
   429  			require.EqualValues(t, 0, rs.rowIndex)
   430  		}
   431  		{
   432  			_, err := rs.next(ctx)
   433  			require.NoError(t, err)
   434  			require.EqualValues(t, 1, rs.rowIndex)
   435  		}
   436  		{
   437  			_, err := rs.next(ctx)
   438  			require.NoError(t, err)
   439  			require.EqualValues(t, 2, rs.rowIndex)
   440  		}
   441  		{
   442  			_, err := rs.next(ctx)
   443  			require.True(t, xerrors.IsTransportError(err, grpcCodes.Unavailable))
   444  		}
   445  	})
   446  	t.Run("WrongResultSetIndex", func(t *testing.T) {
   447  		ctx := xtest.Context(t)
   448  		ctrl := gomock.NewController(t)
   449  		stream := NewMockQueryService_ExecuteQueryClient(ctrl)
   450  		stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{
   451  			Status:         Ydb.StatusIds_SUCCESS,
   452  			ResultSetIndex: 0,
   453  			ResultSet: &Ydb.ResultSet{
   454  				Columns: []*Ydb.Column{
   455  					{
   456  						Name: "a",
   457  						Type: &Ydb.Type{
   458  							Type: &Ydb.Type_TypeId{
   459  								TypeId: Ydb.Type_UINT64,
   460  							},
   461  						},
   462  					},
   463  					{
   464  						Name: "b",
   465  						Type: &Ydb.Type{
   466  							Type: &Ydb.Type_TypeId{
   467  								TypeId: Ydb.Type_UTF8,
   468  							},
   469  						},
   470  					},
   471  				},
   472  				Rows: []*Ydb.Value{
   473  					{
   474  						Items: []*Ydb.Value{{
   475  							Value: &Ydb.Value_Uint64Value{
   476  								Uint64Value: 1,
   477  							},
   478  						}, {
   479  							Value: &Ydb.Value_TextValue{
   480  								TextValue: "1",
   481  							},
   482  						}},
   483  					},
   484  					{
   485  						Items: []*Ydb.Value{{
   486  							Value: &Ydb.Value_Uint64Value{
   487  								Uint64Value: 2,
   488  							},
   489  						}, {
   490  							Value: &Ydb.Value_TextValue{
   491  								TextValue: "2",
   492  							},
   493  						}},
   494  					},
   495  					{
   496  						Items: []*Ydb.Value{{
   497  							Value: &Ydb.Value_Uint64Value{
   498  								Uint64Value: 3,
   499  							},
   500  						}, {
   501  							Value: &Ydb.Value_TextValue{
   502  								TextValue: "3",
   503  							},
   504  						}},
   505  					},
   506  				},
   507  			},
   508  		}, nil)
   509  		stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{
   510  			Status:         Ydb.StatusIds_SUCCESS,
   511  			ResultSetIndex: 1,
   512  			ResultSet: &Ydb.ResultSet{
   513  				Columns: []*Ydb.Column{
   514  					{
   515  						Name: "a",
   516  						Type: &Ydb.Type{
   517  							Type: &Ydb.Type_TypeId{
   518  								TypeId: Ydb.Type_UINT64,
   519  							},
   520  						},
   521  					},
   522  					{
   523  						Name: "b",
   524  						Type: &Ydb.Type{
   525  							Type: &Ydb.Type_TypeId{
   526  								TypeId: Ydb.Type_UTF8,
   527  							},
   528  						},
   529  					},
   530  				},
   531  				Rows: []*Ydb.Value{
   532  					{
   533  						Items: []*Ydb.Value{{
   534  							Value: &Ydb.Value_Uint64Value{
   535  								Uint64Value: 1,
   536  							},
   537  						}, {
   538  							Value: &Ydb.Value_TextValue{
   539  								TextValue: "1",
   540  							},
   541  						}},
   542  					},
   543  					{
   544  						Items: []*Ydb.Value{{
   545  							Value: &Ydb.Value_Uint64Value{
   546  								Uint64Value: 2,
   547  							},
   548  						}, {
   549  							Value: &Ydb.Value_TextValue{
   550  								TextValue: "2",
   551  							},
   552  						}},
   553  					},
   554  					{
   555  						Items: []*Ydb.Value{{
   556  							Value: &Ydb.Value_Uint64Value{
   557  								Uint64Value: 3,
   558  							},
   559  						}, {
   560  							Value: &Ydb.Value_TextValue{
   561  								TextValue: "3",
   562  							},
   563  						}},
   564  					},
   565  				},
   566  			},
   567  		}, nil)
   568  		recv, err := stream.Recv()
   569  		require.NoError(t, err)
   570  		rs := newResultSet(func() (*Ydb_Query.ExecuteQueryResponsePart, error) {
   571  			part, err := nextPart(stream)
   572  			if err != nil {
   573  				return nil, xerrors.WithStackTrace(err)
   574  			}
   575  
   576  			return part, nil
   577  		}, recv)
   578  		require.EqualValues(t, 0, rs.index)
   579  		{
   580  			_, err := rs.next(ctx)
   581  			require.NoError(t, err)
   582  			require.EqualValues(t, 0, rs.rowIndex)
   583  		}
   584  		{
   585  			_, err := rs.next(ctx)
   586  			require.NoError(t, err)
   587  			require.EqualValues(t, 1, rs.rowIndex)
   588  		}
   589  		{
   590  			_, err := rs.next(ctx)
   591  			require.NoError(t, err)
   592  			require.EqualValues(t, 2, rs.rowIndex)
   593  		}
   594  		{
   595  			_, err := rs.next(ctx)
   596  			require.ErrorIs(t, err, errWrongResultSetIndex)
   597  		}
   598  		{
   599  			_, err := rs.next(ctx)
   600  			require.ErrorIs(t, err, io.EOF)
   601  		}
   602  	})
   603  }