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

     1  package query
     2  
     3  import (
     4  	"context"
     5  	"io"
     6  	"testing"
     7  
     8  	"github.com/stretchr/testify/require"
     9  	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb"
    10  	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb_Query"
    11  	"go.uber.org/mock/gomock"
    12  	"google.golang.org/grpc"
    13  	grpcCodes "google.golang.org/grpc/codes"
    14  	"google.golang.org/grpc/metadata"
    15  	grpcStatus "google.golang.org/grpc/status"
    16  
    17  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/allocator"
    18  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/params"
    19  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
    20  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xtest"
    21  	"github.com/ydb-platform/ydb-go-sdk/v3/query"
    22  )
    23  
    24  func TestExecute(t *testing.T) {
    25  	t.Run("HappyWay", func(t *testing.T) {
    26  		ctx := xtest.Context(t)
    27  		ctrl := gomock.NewController(t)
    28  		stream := NewMockQueryService_ExecuteQueryClient(ctrl)
    29  		stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{
    30  			Status: Ydb.StatusIds_SUCCESS,
    31  			TxMeta: &Ydb_Query.TransactionMeta{
    32  				Id: "456",
    33  			},
    34  			ResultSetIndex: 0,
    35  			ResultSet: &Ydb.ResultSet{
    36  				Columns: []*Ydb.Column{
    37  					{
    38  						Name: "a",
    39  						Type: &Ydb.Type{
    40  							Type: &Ydb.Type_TypeId{
    41  								TypeId: Ydb.Type_UINT64,
    42  							},
    43  						},
    44  					},
    45  					{
    46  						Name: "b",
    47  						Type: &Ydb.Type{
    48  							Type: &Ydb.Type_TypeId{
    49  								TypeId: Ydb.Type_UTF8,
    50  							},
    51  						},
    52  					},
    53  				},
    54  				Rows: []*Ydb.Value{
    55  					{
    56  						Items: []*Ydb.Value{{
    57  							Value: &Ydb.Value_Uint64Value{
    58  								Uint64Value: 1,
    59  							},
    60  						}, {
    61  							Value: &Ydb.Value_TextValue{
    62  								TextValue: "1",
    63  							},
    64  						}},
    65  					},
    66  					{
    67  						Items: []*Ydb.Value{{
    68  							Value: &Ydb.Value_Uint64Value{
    69  								Uint64Value: 2,
    70  							},
    71  						}, {
    72  							Value: &Ydb.Value_TextValue{
    73  								TextValue: "2",
    74  							},
    75  						}},
    76  					},
    77  					{
    78  						Items: []*Ydb.Value{{
    79  							Value: &Ydb.Value_Uint64Value{
    80  								Uint64Value: 3,
    81  							},
    82  						}, {
    83  							Value: &Ydb.Value_TextValue{
    84  								TextValue: "3",
    85  							},
    86  						}},
    87  					},
    88  				},
    89  			},
    90  		}, nil)
    91  		stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{
    92  			Status:         Ydb.StatusIds_SUCCESS,
    93  			ResultSetIndex: 0,
    94  			ResultSet: &Ydb.ResultSet{
    95  				Rows: []*Ydb.Value{
    96  					{
    97  						Items: []*Ydb.Value{{
    98  							Value: &Ydb.Value_Uint64Value{
    99  								Uint64Value: 4,
   100  							},
   101  						}, {
   102  							Value: &Ydb.Value_TextValue{
   103  								TextValue: "4",
   104  							},
   105  						}},
   106  					},
   107  					{
   108  						Items: []*Ydb.Value{{
   109  							Value: &Ydb.Value_Uint64Value{
   110  								Uint64Value: 5,
   111  							},
   112  						}, {
   113  							Value: &Ydb.Value_TextValue{
   114  								TextValue: "5",
   115  							},
   116  						}},
   117  					},
   118  				},
   119  			},
   120  		}, nil)
   121  		stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{
   122  			Status:         Ydb.StatusIds_SUCCESS,
   123  			ResultSetIndex: 1,
   124  			ResultSet: &Ydb.ResultSet{
   125  				Columns: []*Ydb.Column{
   126  					{
   127  						Name: "c",
   128  						Type: &Ydb.Type{
   129  							Type: &Ydb.Type_TypeId{
   130  								TypeId: Ydb.Type_UINT64,
   131  							},
   132  						},
   133  					},
   134  					{
   135  						Name: "d",
   136  						Type: &Ydb.Type{
   137  							Type: &Ydb.Type_TypeId{
   138  								TypeId: Ydb.Type_UTF8,
   139  							},
   140  						},
   141  					},
   142  					{
   143  						Name: "e",
   144  						Type: &Ydb.Type{
   145  							Type: &Ydb.Type_TypeId{
   146  								TypeId: Ydb.Type_BOOL,
   147  							},
   148  						},
   149  					},
   150  				},
   151  				Rows: []*Ydb.Value{
   152  					{
   153  						Items: []*Ydb.Value{{
   154  							Value: &Ydb.Value_Uint64Value{
   155  								Uint64Value: 1,
   156  							},
   157  						}, {
   158  							Value: &Ydb.Value_TextValue{
   159  								TextValue: "1",
   160  							},
   161  						}, {
   162  							Value: &Ydb.Value_BoolValue{
   163  								BoolValue: true,
   164  							},
   165  						}},
   166  					},
   167  					{
   168  						Items: []*Ydb.Value{{
   169  							Value: &Ydb.Value_Uint64Value{
   170  								Uint64Value: 2,
   171  							},
   172  						}, {
   173  							Value: &Ydb.Value_TextValue{
   174  								TextValue: "2",
   175  							},
   176  						}, {
   177  							Value: &Ydb.Value_BoolValue{
   178  								BoolValue: false,
   179  							},
   180  						}},
   181  					},
   182  				},
   183  			},
   184  		}, nil)
   185  		stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{
   186  			Status:         Ydb.StatusIds_SUCCESS,
   187  			ResultSetIndex: 1,
   188  			ResultSet: &Ydb.ResultSet{
   189  				Rows: []*Ydb.Value{
   190  					{
   191  						Items: []*Ydb.Value{{
   192  							Value: &Ydb.Value_Uint64Value{
   193  								Uint64Value: 3,
   194  							},
   195  						}, {
   196  							Value: &Ydb.Value_TextValue{
   197  								TextValue: "3",
   198  							},
   199  						}, {
   200  							Value: &Ydb.Value_BoolValue{
   201  								BoolValue: true,
   202  							},
   203  						}},
   204  					},
   205  					{
   206  						Items: []*Ydb.Value{{
   207  							Value: &Ydb.Value_Uint64Value{
   208  								Uint64Value: 4,
   209  							},
   210  						}, {
   211  							Value: &Ydb.Value_TextValue{
   212  								TextValue: "4",
   213  							},
   214  						}, {
   215  							Value: &Ydb.Value_BoolValue{
   216  								BoolValue: false,
   217  							},
   218  						}},
   219  					},
   220  					{
   221  						Items: []*Ydb.Value{{
   222  							Value: &Ydb.Value_Uint64Value{
   223  								Uint64Value: 5,
   224  							},
   225  						}, {
   226  							Value: &Ydb.Value_TextValue{
   227  								TextValue: "5",
   228  							},
   229  						}, {
   230  							Value: &Ydb.Value_BoolValue{
   231  								BoolValue: false,
   232  							},
   233  						}},
   234  					},
   235  				},
   236  			},
   237  		}, nil)
   238  		stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{
   239  			Status:         Ydb.StatusIds_SUCCESS,
   240  			ResultSetIndex: 2,
   241  			ResultSet: &Ydb.ResultSet{
   242  				Columns: []*Ydb.Column{
   243  					{
   244  						Name: "c",
   245  						Type: &Ydb.Type{
   246  							Type: &Ydb.Type_TypeId{
   247  								TypeId: Ydb.Type_UINT64,
   248  							},
   249  						},
   250  					},
   251  					{
   252  						Name: "d",
   253  						Type: &Ydb.Type{
   254  							Type: &Ydb.Type_TypeId{
   255  								TypeId: Ydb.Type_UTF8,
   256  							},
   257  						},
   258  					},
   259  					{
   260  						Name: "e",
   261  						Type: &Ydb.Type{
   262  							Type: &Ydb.Type_TypeId{
   263  								TypeId: Ydb.Type_BOOL,
   264  							},
   265  						},
   266  					},
   267  				},
   268  				Rows: []*Ydb.Value{
   269  					{
   270  						Items: []*Ydb.Value{{
   271  							Value: &Ydb.Value_Uint64Value{
   272  								Uint64Value: 1,
   273  							},
   274  						}, {
   275  							Value: &Ydb.Value_TextValue{
   276  								TextValue: "1",
   277  							},
   278  						}, {
   279  							Value: &Ydb.Value_BoolValue{
   280  								BoolValue: true,
   281  							},
   282  						}},
   283  					},
   284  					{
   285  						Items: []*Ydb.Value{{
   286  							Value: &Ydb.Value_Uint64Value{
   287  								Uint64Value: 2,
   288  							},
   289  						}, {
   290  							Value: &Ydb.Value_TextValue{
   291  								TextValue: "2",
   292  							},
   293  						}, {
   294  							Value: &Ydb.Value_BoolValue{
   295  								BoolValue: false,
   296  							},
   297  						}},
   298  					},
   299  				},
   300  			},
   301  		}, nil)
   302  		stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{
   303  			Status:         Ydb.StatusIds_SUCCESS,
   304  			ResultSetIndex: 2,
   305  			ResultSet: &Ydb.ResultSet{
   306  				Rows: []*Ydb.Value{
   307  					{
   308  						Items: []*Ydb.Value{{
   309  							Value: &Ydb.Value_Uint64Value{
   310  								Uint64Value: 3,
   311  							},
   312  						}, {
   313  							Value: &Ydb.Value_TextValue{
   314  								TextValue: "3",
   315  							},
   316  						}, {
   317  							Value: &Ydb.Value_BoolValue{
   318  								BoolValue: true,
   319  							},
   320  						}},
   321  					},
   322  					{
   323  						Items: []*Ydb.Value{{
   324  							Value: &Ydb.Value_Uint64Value{
   325  								Uint64Value: 4,
   326  							},
   327  						}, {
   328  							Value: &Ydb.Value_TextValue{
   329  								TextValue: "4",
   330  							},
   331  						}, {
   332  							Value: &Ydb.Value_BoolValue{
   333  								BoolValue: false,
   334  							},
   335  						}},
   336  					},
   337  					{
   338  						Items: []*Ydb.Value{{
   339  							Value: &Ydb.Value_Uint64Value{
   340  								Uint64Value: 5,
   341  							},
   342  						}, {
   343  							Value: &Ydb.Value_TextValue{
   344  								TextValue: "5",
   345  							},
   346  						}, {
   347  							Value: &Ydb.Value_BoolValue{
   348  								BoolValue: false,
   349  							},
   350  						}},
   351  					},
   352  				},
   353  			},
   354  		}, nil)
   355  		stream.EXPECT().Recv().Return(nil, io.EOF)
   356  		service := NewMockQueryServiceClient(ctrl)
   357  		service.EXPECT().ExecuteQuery(gomock.Any(), gomock.Any()).Return(stream, nil)
   358  		t.Log("execute")
   359  		tx, r, err := execute(ctx, &Session{id: "123"}, service, "", query.ExecuteSettings())
   360  		require.NoError(t, err)
   361  		defer r.Close(ctx)
   362  		require.EqualValues(t, "456", tx.id)
   363  		require.EqualValues(t, "123", tx.s.id)
   364  		require.EqualValues(t, -1, r.resultSetIndex)
   365  		{
   366  			t.Log("nextResultSet")
   367  			rs, err := r.nextResultSet(ctx)
   368  			require.NoError(t, err)
   369  			require.EqualValues(t, 0, rs.index)
   370  			{
   371  				t.Log("next (row=1)")
   372  				_, err := rs.next(ctx)
   373  				require.NoError(t, err)
   374  				require.EqualValues(t, 0, rs.rowIndex)
   375  			}
   376  			{
   377  				t.Log("next (row=2)")
   378  				_, err := rs.next(ctx)
   379  				require.NoError(t, err)
   380  				require.EqualValues(t, 1, rs.rowIndex)
   381  			}
   382  			{
   383  				t.Log("next (row=3)")
   384  				_, err := rs.next(ctx)
   385  				require.NoError(t, err)
   386  				require.EqualValues(t, 2, rs.rowIndex)
   387  			}
   388  			{
   389  				t.Log("next (row=4)")
   390  				_, err := rs.next(ctx)
   391  				require.NoError(t, err)
   392  				require.EqualValues(t, 0, rs.rowIndex)
   393  			}
   394  			{
   395  				t.Log("next (row=5)")
   396  				_, err := rs.next(ctx)
   397  				require.NoError(t, err)
   398  				require.EqualValues(t, 1, rs.rowIndex)
   399  			}
   400  			{
   401  				t.Log("next (row=6)")
   402  				_, err := rs.next(ctx)
   403  				require.ErrorIs(t, err, io.EOF)
   404  			}
   405  		}
   406  		{
   407  			t.Log("nextResultSet")
   408  			rs, err := r.nextResultSet(ctx)
   409  			require.NoError(t, err)
   410  			require.EqualValues(t, 1, rs.index)
   411  		}
   412  		{
   413  			t.Log("nextResultSet")
   414  			rs, err := r.nextResultSet(ctx)
   415  			require.NoError(t, err)
   416  			require.EqualValues(t, 2, rs.index)
   417  			{
   418  				t.Log("next (row=1)")
   419  				_, err := rs.next(ctx)
   420  				require.NoError(t, err)
   421  				require.EqualValues(t, 0, rs.rowIndex)
   422  			}
   423  			{
   424  				t.Log("next (row=2)")
   425  				_, err := rs.next(ctx)
   426  				require.NoError(t, err)
   427  				require.EqualValues(t, 1, rs.rowIndex)
   428  			}
   429  			{
   430  				t.Log("next (row=3)")
   431  				_, err := rs.next(ctx)
   432  				require.NoError(t, err)
   433  				require.EqualValues(t, 0, rs.rowIndex)
   434  			}
   435  			{
   436  				t.Log("next (row=4)")
   437  				_, err := rs.next(ctx)
   438  				require.NoError(t, err)
   439  				require.EqualValues(t, 1, rs.rowIndex)
   440  			}
   441  			{
   442  				t.Log("next (row=5)")
   443  				_, err := rs.next(ctx)
   444  				require.NoError(t, err)
   445  				require.EqualValues(t, 2, rs.rowIndex)
   446  			}
   447  			{
   448  				t.Log("next (row=6)")
   449  				_, err := rs.next(ctx)
   450  				require.ErrorIs(t, err, io.EOF)
   451  			}
   452  		}
   453  		{
   454  			t.Log("close result")
   455  			r.Close(context.Background())
   456  		}
   457  		{
   458  			t.Log("nextResultSet")
   459  			_, err := r.nextResultSet(context.Background())
   460  			require.ErrorIs(t, err, errClosedResult)
   461  		}
   462  		t.Log("check final error")
   463  		require.NoError(t, r.Err())
   464  	})
   465  	t.Run("TransportError", func(t *testing.T) {
   466  		t.Run("OnCall", func(t *testing.T) {
   467  			ctx := xtest.Context(t)
   468  			ctrl := gomock.NewController(t)
   469  			service := NewMockQueryServiceClient(ctrl)
   470  			service.EXPECT().ExecuteQuery(gomock.Any(), gomock.Any()).Return(nil, grpcStatus.Error(grpcCodes.Unavailable, ""))
   471  			t.Log("execute")
   472  			_, _, err := execute(ctx, &Session{id: "123"}, service, "", query.ExecuteSettings())
   473  			require.Error(t, err)
   474  			require.True(t, xerrors.IsTransportError(err, grpcCodes.Unavailable))
   475  		})
   476  		t.Run("OnStream", func(t *testing.T) {
   477  			ctx := xtest.Context(t)
   478  			ctrl := gomock.NewController(t)
   479  			stream := NewMockQueryService_ExecuteQueryClient(ctrl)
   480  			stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{
   481  				Status: Ydb.StatusIds_SUCCESS,
   482  				TxMeta: &Ydb_Query.TransactionMeta{
   483  					Id: "456",
   484  				},
   485  				ResultSetIndex: 0,
   486  				ResultSet: &Ydb.ResultSet{
   487  					Columns: []*Ydb.Column{
   488  						{
   489  							Name: "a",
   490  							Type: &Ydb.Type{
   491  								Type: &Ydb.Type_TypeId{
   492  									TypeId: Ydb.Type_UINT64,
   493  								},
   494  							},
   495  						},
   496  						{
   497  							Name: "b",
   498  							Type: &Ydb.Type{
   499  								Type: &Ydb.Type_TypeId{
   500  									TypeId: Ydb.Type_UTF8,
   501  								},
   502  							},
   503  						},
   504  					},
   505  					Rows: []*Ydb.Value{
   506  						{
   507  							Items: []*Ydb.Value{{
   508  								Value: &Ydb.Value_Uint64Value{
   509  									Uint64Value: 1,
   510  								},
   511  							}, {
   512  								Value: &Ydb.Value_TextValue{
   513  									TextValue: "1",
   514  								},
   515  							}},
   516  						},
   517  						{
   518  							Items: []*Ydb.Value{{
   519  								Value: &Ydb.Value_Uint64Value{
   520  									Uint64Value: 2,
   521  								},
   522  							}, {
   523  								Value: &Ydb.Value_TextValue{
   524  									TextValue: "2",
   525  								},
   526  							}},
   527  						},
   528  						{
   529  							Items: []*Ydb.Value{{
   530  								Value: &Ydb.Value_Uint64Value{
   531  									Uint64Value: 3,
   532  								},
   533  							}, {
   534  								Value: &Ydb.Value_TextValue{
   535  									TextValue: "3",
   536  								},
   537  							}},
   538  						},
   539  					},
   540  				},
   541  			}, nil)
   542  			stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{
   543  				Status:         Ydb.StatusIds_SUCCESS,
   544  				ResultSetIndex: 0,
   545  				ResultSet: &Ydb.ResultSet{
   546  					Rows: []*Ydb.Value{
   547  						{
   548  							Items: []*Ydb.Value{{
   549  								Value: &Ydb.Value_Uint64Value{
   550  									Uint64Value: 4,
   551  								},
   552  							}, {
   553  								Value: &Ydb.Value_TextValue{
   554  									TextValue: "4",
   555  								},
   556  							}},
   557  						},
   558  						{
   559  							Items: []*Ydb.Value{{
   560  								Value: &Ydb.Value_Uint64Value{
   561  									Uint64Value: 5,
   562  								},
   563  							}, {
   564  								Value: &Ydb.Value_TextValue{
   565  									TextValue: "5",
   566  								},
   567  							}},
   568  						},
   569  					},
   570  				},
   571  			}, nil)
   572  			stream.EXPECT().Recv().Return(nil, grpcStatus.Error(grpcCodes.Unavailable, ""))
   573  			service := NewMockQueryServiceClient(ctrl)
   574  			service.EXPECT().ExecuteQuery(gomock.Any(), gomock.Any()).Return(stream, nil)
   575  			t.Log("execute")
   576  			tx, r, err := execute(ctx, &Session{id: "123"}, service, "", query.ExecuteSettings())
   577  			require.NoError(t, err)
   578  			defer r.Close(ctx)
   579  			require.EqualValues(t, "456", tx.id)
   580  			require.EqualValues(t, "123", tx.s.id)
   581  			require.EqualValues(t, -1, r.resultSetIndex)
   582  			{
   583  				t.Log("nextResultSet")
   584  				rs, err := r.nextResultSet(ctx)
   585  				require.NoError(t, err)
   586  				require.EqualValues(t, 0, rs.index)
   587  				{
   588  					t.Log("next (row=1)")
   589  					_, err := rs.next(ctx)
   590  					require.NoError(t, err)
   591  					require.EqualValues(t, 0, rs.rowIndex)
   592  				}
   593  				{
   594  					t.Log("next (row=2)")
   595  					_, err := rs.next(ctx)
   596  					require.NoError(t, err)
   597  					require.EqualValues(t, 1, rs.rowIndex)
   598  				}
   599  				{
   600  					t.Log("next (row=3)")
   601  					_, err := rs.next(ctx)
   602  					require.NoError(t, err)
   603  					require.EqualValues(t, 2, rs.rowIndex)
   604  				}
   605  				{
   606  					t.Log("next (row=4)")
   607  					_, err := rs.next(ctx)
   608  					require.NoError(t, err)
   609  					require.EqualValues(t, 0, rs.rowIndex)
   610  				}
   611  				{
   612  					t.Log("next (row=5)")
   613  					_, err := rs.next(ctx)
   614  					require.NoError(t, err)
   615  					require.EqualValues(t, 1, rs.rowIndex)
   616  				}
   617  				{
   618  					t.Log("next (row=6)")
   619  					_, err := rs.next(ctx)
   620  					require.Error(t, err)
   621  					require.True(t, xerrors.IsTransportError(err, grpcCodes.Unavailable))
   622  				}
   623  			}
   624  			t.Log("check final error")
   625  			require.Error(t, r.Err())
   626  			require.True(t, xerrors.IsTransportError(r.Err(), grpcCodes.Unavailable))
   627  		})
   628  	})
   629  	t.Run("OperationError", func(t *testing.T) {
   630  		t.Run("OnCall", func(t *testing.T) {
   631  			ctx := xtest.Context(t)
   632  			ctrl := gomock.NewController(t)
   633  			stream := NewMockQueryService_ExecuteQueryClient(ctrl)
   634  			stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{
   635  				Status: Ydb.StatusIds_UNAVAILABLE,
   636  			}, nil)
   637  			service := NewMockQueryServiceClient(ctrl)
   638  			service.EXPECT().ExecuteQuery(gomock.Any(), gomock.Any()).Return(stream, nil)
   639  			t.Log("execute")
   640  			_, _, err := execute(ctx, &Session{id: "123"}, service, "", query.ExecuteSettings())
   641  			require.Error(t, err)
   642  			require.True(t, xerrors.IsOperationError(err, Ydb.StatusIds_UNAVAILABLE))
   643  		})
   644  		t.Run("OnStream", func(t *testing.T) {
   645  			ctx := xtest.Context(t)
   646  			ctrl := gomock.NewController(t)
   647  			stream := NewMockQueryService_ExecuteQueryClient(ctrl)
   648  			stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{
   649  				Status: Ydb.StatusIds_SUCCESS,
   650  				TxMeta: &Ydb_Query.TransactionMeta{
   651  					Id: "456",
   652  				},
   653  				ResultSetIndex: 0,
   654  				ResultSet: &Ydb.ResultSet{
   655  					Columns: []*Ydb.Column{
   656  						{
   657  							Name: "a",
   658  							Type: &Ydb.Type{
   659  								Type: &Ydb.Type_TypeId{
   660  									TypeId: Ydb.Type_UINT64,
   661  								},
   662  							},
   663  						},
   664  						{
   665  							Name: "b",
   666  							Type: &Ydb.Type{
   667  								Type: &Ydb.Type_TypeId{
   668  									TypeId: Ydb.Type_UTF8,
   669  								},
   670  							},
   671  						},
   672  					},
   673  					Rows: []*Ydb.Value{
   674  						{
   675  							Items: []*Ydb.Value{{
   676  								Value: &Ydb.Value_Uint64Value{
   677  									Uint64Value: 1,
   678  								},
   679  							}, {
   680  								Value: &Ydb.Value_TextValue{
   681  									TextValue: "1",
   682  								},
   683  							}},
   684  						},
   685  						{
   686  							Items: []*Ydb.Value{{
   687  								Value: &Ydb.Value_Uint64Value{
   688  									Uint64Value: 2,
   689  								},
   690  							}, {
   691  								Value: &Ydb.Value_TextValue{
   692  									TextValue: "2",
   693  								},
   694  							}},
   695  						},
   696  						{
   697  							Items: []*Ydb.Value{{
   698  								Value: &Ydb.Value_Uint64Value{
   699  									Uint64Value: 3,
   700  								},
   701  							}, {
   702  								Value: &Ydb.Value_TextValue{
   703  									TextValue: "3",
   704  								},
   705  							}},
   706  						},
   707  					},
   708  				},
   709  			}, nil)
   710  			stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{
   711  				Status: Ydb.StatusIds_UNAVAILABLE,
   712  			}, nil)
   713  			service := NewMockQueryServiceClient(ctrl)
   714  			service.EXPECT().ExecuteQuery(gomock.Any(), gomock.Any()).Return(stream, nil)
   715  			t.Log("execute")
   716  			tx, r, err := execute(ctx, &Session{id: "123"}, service, "", query.ExecuteSettings())
   717  			require.NoError(t, err)
   718  			defer r.Close(ctx)
   719  			require.EqualValues(t, "456", tx.id)
   720  			require.EqualValues(t, "123", tx.s.id)
   721  			require.EqualValues(t, -1, r.resultSetIndex)
   722  			{
   723  				t.Log("nextResultSet")
   724  				rs, err := r.nextResultSet(ctx)
   725  				require.NoError(t, err)
   726  				require.EqualValues(t, 0, rs.index)
   727  				{
   728  					t.Log("next (row=1)")
   729  					_, err := rs.next(ctx)
   730  					require.NoError(t, err)
   731  					require.EqualValues(t, 0, rs.rowIndex)
   732  				}
   733  				{
   734  					t.Log("next (row=2)")
   735  					_, err := rs.next(ctx)
   736  					require.NoError(t, err)
   737  					require.EqualValues(t, 1, rs.rowIndex)
   738  				}
   739  				{
   740  					t.Log("next (row=3)")
   741  					_, err := rs.next(ctx)
   742  					require.NoError(t, err)
   743  					require.EqualValues(t, 2, rs.rowIndex)
   744  				}
   745  				{
   746  					t.Log("next (row=4)")
   747  					_, err := rs.next(ctx)
   748  					require.Error(t, err)
   749  					require.True(t, xerrors.IsOperationError(err, Ydb.StatusIds_UNAVAILABLE))
   750  				}
   751  			}
   752  			t.Log("check final error")
   753  			require.Error(t, r.Err())
   754  			require.True(t, xerrors.IsOperationError(r.Err(), Ydb.StatusIds_UNAVAILABLE))
   755  		})
   756  	})
   757  }
   758  
   759  func TestExecuteQueryRequest(t *testing.T) {
   760  	a := allocator.New()
   761  	for _, tt := range []struct {
   762  		name        string
   763  		opts        []query.ExecuteOption
   764  		request     *Ydb_Query.ExecuteQueryRequest
   765  		callOptions []grpc.CallOption
   766  	}{
   767  		{
   768  			name: "WithoutOptions",
   769  			request: &Ydb_Query.ExecuteQueryRequest{
   770  				SessionId: "WithoutOptions",
   771  				ExecMode:  Ydb_Query.ExecMode_EXEC_MODE_EXECUTE,
   772  				TxControl: &Ydb_Query.TransactionControl{
   773  					TxSelector: &Ydb_Query.TransactionControl_BeginTx{
   774  						BeginTx: &Ydb_Query.TransactionSettings{
   775  							TxMode: &Ydb_Query.TransactionSettings_SerializableReadWrite{
   776  								SerializableReadWrite: &Ydb_Query.SerializableModeSettings{},
   777  							},
   778  						},
   779  					},
   780  					CommitTx: true,
   781  				},
   782  				Query: &Ydb_Query.ExecuteQueryRequest_QueryContent{
   783  					QueryContent: &Ydb_Query.QueryContent{
   784  						Syntax: Ydb_Query.Syntax_SYNTAX_YQL_V1,
   785  						Text:   "WithoutOptions",
   786  					},
   787  				},
   788  				StatsMode:            Ydb_Query.StatsMode_STATS_MODE_NONE,
   789  				ConcurrentResultSets: false,
   790  			},
   791  		},
   792  		{
   793  			name: "WithParams",
   794  			opts: []query.ExecuteOption{
   795  				query.WithParameters(
   796  					params.Builder{}.
   797  						Param("$a").Text("A").
   798  						Param("$b").Text("B").
   799  						Param("$c").Text("C").
   800  						Build(),
   801  				),
   802  			},
   803  			request: &Ydb_Query.ExecuteQueryRequest{
   804  				SessionId: "WithParams",
   805  				ExecMode:  Ydb_Query.ExecMode_EXEC_MODE_EXECUTE,
   806  				TxControl: &Ydb_Query.TransactionControl{
   807  					TxSelector: &Ydb_Query.TransactionControl_BeginTx{
   808  						BeginTx: &Ydb_Query.TransactionSettings{
   809  							TxMode: &Ydb_Query.TransactionSettings_SerializableReadWrite{
   810  								SerializableReadWrite: &Ydb_Query.SerializableModeSettings{},
   811  							},
   812  						},
   813  					},
   814  					CommitTx: true,
   815  				},
   816  				Query: &Ydb_Query.ExecuteQueryRequest_QueryContent{
   817  					QueryContent: &Ydb_Query.QueryContent{
   818  						Syntax: Ydb_Query.Syntax_SYNTAX_YQL_V1,
   819  						Text:   "WithParams",
   820  					},
   821  				},
   822  				Parameters: map[string]*Ydb.TypedValue{
   823  					"$a": {
   824  						Type: &Ydb.Type{
   825  							Type: &Ydb.Type_TypeId{
   826  								TypeId: Ydb.Type_UTF8,
   827  							},
   828  						},
   829  						Value: &Ydb.Value{
   830  							Value: &Ydb.Value_TextValue{
   831  								TextValue: "A",
   832  							},
   833  						},
   834  					},
   835  					"$b": {
   836  						Type: &Ydb.Type{
   837  							Type: &Ydb.Type_TypeId{
   838  								TypeId: Ydb.Type_UTF8,
   839  							},
   840  						},
   841  						Value: &Ydb.Value{
   842  							Value: &Ydb.Value_TextValue{
   843  								TextValue: "B",
   844  							},
   845  						},
   846  					},
   847  					"$c": {
   848  						Type: &Ydb.Type{
   849  							Type: &Ydb.Type_TypeId{
   850  								TypeId: Ydb.Type_UTF8,
   851  							},
   852  						},
   853  						Value: &Ydb.Value{
   854  							Value: &Ydb.Value_TextValue{
   855  								TextValue: "C",
   856  							},
   857  						},
   858  					},
   859  				},
   860  				StatsMode:            Ydb_Query.StatsMode_STATS_MODE_NONE,
   861  				ConcurrentResultSets: false,
   862  			},
   863  		},
   864  		{
   865  			name: "WithExplain",
   866  			opts: []query.ExecuteOption{
   867  				query.WithExecMode(query.ExecModeExplain),
   868  			},
   869  			request: &Ydb_Query.ExecuteQueryRequest{
   870  				SessionId: "WithExplain",
   871  				ExecMode:  Ydb_Query.ExecMode_EXEC_MODE_EXPLAIN,
   872  				TxControl: &Ydb_Query.TransactionControl{
   873  					TxSelector: &Ydb_Query.TransactionControl_BeginTx{
   874  						BeginTx: &Ydb_Query.TransactionSettings{
   875  							TxMode: &Ydb_Query.TransactionSettings_SerializableReadWrite{
   876  								SerializableReadWrite: &Ydb_Query.SerializableModeSettings{},
   877  							},
   878  						},
   879  					},
   880  					CommitTx: true,
   881  				},
   882  				Query: &Ydb_Query.ExecuteQueryRequest_QueryContent{
   883  					QueryContent: &Ydb_Query.QueryContent{
   884  						Syntax: Ydb_Query.Syntax_SYNTAX_YQL_V1,
   885  						Text:   "WithExplain",
   886  					},
   887  				},
   888  				StatsMode:            Ydb_Query.StatsMode_STATS_MODE_NONE,
   889  				ConcurrentResultSets: false,
   890  			},
   891  		},
   892  		{
   893  			name: "WithValidate",
   894  			opts: []query.ExecuteOption{
   895  				query.WithExecMode(query.ExecModeValidate),
   896  			},
   897  			request: &Ydb_Query.ExecuteQueryRequest{
   898  				SessionId: "WithValidate",
   899  				ExecMode:  Ydb_Query.ExecMode_EXEC_MODE_VALIDATE,
   900  				TxControl: &Ydb_Query.TransactionControl{
   901  					TxSelector: &Ydb_Query.TransactionControl_BeginTx{
   902  						BeginTx: &Ydb_Query.TransactionSettings{
   903  							TxMode: &Ydb_Query.TransactionSettings_SerializableReadWrite{
   904  								SerializableReadWrite: &Ydb_Query.SerializableModeSettings{},
   905  							},
   906  						},
   907  					},
   908  					CommitTx: true,
   909  				},
   910  				Query: &Ydb_Query.ExecuteQueryRequest_QueryContent{
   911  					QueryContent: &Ydb_Query.QueryContent{
   912  						Syntax: Ydb_Query.Syntax_SYNTAX_YQL_V1,
   913  						Text:   "WithValidate",
   914  					},
   915  				},
   916  				StatsMode:            Ydb_Query.StatsMode_STATS_MODE_NONE,
   917  				ConcurrentResultSets: false,
   918  			},
   919  		},
   920  		{
   921  			name: "WithValidate",
   922  			opts: []query.ExecuteOption{
   923  				query.WithExecMode(query.ExecModeParse),
   924  			},
   925  			request: &Ydb_Query.ExecuteQueryRequest{
   926  				SessionId: "WithValidate",
   927  				ExecMode:  Ydb_Query.ExecMode_EXEC_MODE_PARSE,
   928  				TxControl: &Ydb_Query.TransactionControl{
   929  					TxSelector: &Ydb_Query.TransactionControl_BeginTx{
   930  						BeginTx: &Ydb_Query.TransactionSettings{
   931  							TxMode: &Ydb_Query.TransactionSettings_SerializableReadWrite{
   932  								SerializableReadWrite: &Ydb_Query.SerializableModeSettings{},
   933  							},
   934  						},
   935  					},
   936  					CommitTx: true,
   937  				},
   938  				Query: &Ydb_Query.ExecuteQueryRequest_QueryContent{
   939  					QueryContent: &Ydb_Query.QueryContent{
   940  						Syntax: Ydb_Query.Syntax_SYNTAX_YQL_V1,
   941  						Text:   "WithValidate",
   942  					},
   943  				},
   944  				StatsMode:            Ydb_Query.StatsMode_STATS_MODE_NONE,
   945  				ConcurrentResultSets: false,
   946  			},
   947  		},
   948  		{
   949  			name: "WithStatsFull",
   950  			opts: []query.ExecuteOption{
   951  				query.WithStatsMode(query.StatsModeFull),
   952  			},
   953  			request: &Ydb_Query.ExecuteQueryRequest{
   954  				SessionId: "WithStatsFull",
   955  				ExecMode:  Ydb_Query.ExecMode_EXEC_MODE_EXECUTE,
   956  				TxControl: &Ydb_Query.TransactionControl{
   957  					TxSelector: &Ydb_Query.TransactionControl_BeginTx{
   958  						BeginTx: &Ydb_Query.TransactionSettings{
   959  							TxMode: &Ydb_Query.TransactionSettings_SerializableReadWrite{
   960  								SerializableReadWrite: &Ydb_Query.SerializableModeSettings{},
   961  							},
   962  						},
   963  					},
   964  					CommitTx: true,
   965  				},
   966  				Query: &Ydb_Query.ExecuteQueryRequest_QueryContent{
   967  					QueryContent: &Ydb_Query.QueryContent{
   968  						Syntax: Ydb_Query.Syntax_SYNTAX_YQL_V1,
   969  						Text:   "WithStatsFull",
   970  					},
   971  				},
   972  				StatsMode:            Ydb_Query.StatsMode_STATS_MODE_FULL,
   973  				ConcurrentResultSets: false,
   974  			},
   975  		},
   976  		{
   977  			name: "WithStatsBasic",
   978  			opts: []query.ExecuteOption{
   979  				query.WithStatsMode(query.StatsModeBasic),
   980  			},
   981  			request: &Ydb_Query.ExecuteQueryRequest{
   982  				SessionId: "WithStatsBasic",
   983  				ExecMode:  Ydb_Query.ExecMode_EXEC_MODE_EXECUTE,
   984  				TxControl: &Ydb_Query.TransactionControl{
   985  					TxSelector: &Ydb_Query.TransactionControl_BeginTx{
   986  						BeginTx: &Ydb_Query.TransactionSettings{
   987  							TxMode: &Ydb_Query.TransactionSettings_SerializableReadWrite{
   988  								SerializableReadWrite: &Ydb_Query.SerializableModeSettings{},
   989  							},
   990  						},
   991  					},
   992  					CommitTx: true,
   993  				},
   994  				Query: &Ydb_Query.ExecuteQueryRequest_QueryContent{
   995  					QueryContent: &Ydb_Query.QueryContent{
   996  						Syntax: Ydb_Query.Syntax_SYNTAX_YQL_V1,
   997  						Text:   "WithStatsBasic",
   998  					},
   999  				},
  1000  				StatsMode:            Ydb_Query.StatsMode_STATS_MODE_BASIC,
  1001  				ConcurrentResultSets: false,
  1002  			},
  1003  		},
  1004  		{
  1005  			name: "WithStatsProfile",
  1006  			opts: []query.ExecuteOption{
  1007  				query.WithStatsMode(query.StatsModeProfile),
  1008  			},
  1009  			request: &Ydb_Query.ExecuteQueryRequest{
  1010  				SessionId: "WithStatsProfile",
  1011  				ExecMode:  Ydb_Query.ExecMode_EXEC_MODE_EXECUTE,
  1012  				TxControl: &Ydb_Query.TransactionControl{
  1013  					TxSelector: &Ydb_Query.TransactionControl_BeginTx{
  1014  						BeginTx: &Ydb_Query.TransactionSettings{
  1015  							TxMode: &Ydb_Query.TransactionSettings_SerializableReadWrite{
  1016  								SerializableReadWrite: &Ydb_Query.SerializableModeSettings{},
  1017  							},
  1018  						},
  1019  					},
  1020  					CommitTx: true,
  1021  				},
  1022  				Query: &Ydb_Query.ExecuteQueryRequest_QueryContent{
  1023  					QueryContent: &Ydb_Query.QueryContent{
  1024  						Syntax: Ydb_Query.Syntax_SYNTAX_YQL_V1,
  1025  						Text:   "WithStatsProfile",
  1026  					},
  1027  				},
  1028  				StatsMode:            Ydb_Query.StatsMode_STATS_MODE_PROFILE,
  1029  				ConcurrentResultSets: false,
  1030  			},
  1031  		},
  1032  		{
  1033  			name: "WithGrpcCallOptions",
  1034  			opts: []query.ExecuteOption{
  1035  				query.WithCallOptions(grpc.Header(&metadata.MD{
  1036  					"ext-header": []string{"test"},
  1037  				})),
  1038  			},
  1039  			request: &Ydb_Query.ExecuteQueryRequest{
  1040  				SessionId: "WithGrpcCallOptions",
  1041  				ExecMode:  Ydb_Query.ExecMode_EXEC_MODE_EXECUTE,
  1042  				TxControl: &Ydb_Query.TransactionControl{
  1043  					TxSelector: &Ydb_Query.TransactionControl_BeginTx{
  1044  						BeginTx: &Ydb_Query.TransactionSettings{
  1045  							TxMode: &Ydb_Query.TransactionSettings_SerializableReadWrite{
  1046  								SerializableReadWrite: &Ydb_Query.SerializableModeSettings{},
  1047  							},
  1048  						},
  1049  					},
  1050  					CommitTx: true,
  1051  				},
  1052  				Query: &Ydb_Query.ExecuteQueryRequest_QueryContent{
  1053  					QueryContent: &Ydb_Query.QueryContent{
  1054  						Syntax: Ydb_Query.Syntax_SYNTAX_YQL_V1,
  1055  						Text:   "WithGrpcCallOptions",
  1056  					},
  1057  				},
  1058  				StatsMode:            Ydb_Query.StatsMode_STATS_MODE_NONE,
  1059  				ConcurrentResultSets: false,
  1060  			},
  1061  			callOptions: []grpc.CallOption{
  1062  				grpc.Header(&metadata.MD{
  1063  					"ext-header": []string{"test"},
  1064  				}),
  1065  			},
  1066  		},
  1067  	} {
  1068  		t.Run(tt.name, func(t *testing.T) {
  1069  			request, callOptions := executeQueryRequest(a, tt.name, tt.name, query.ExecuteSettings(tt.opts...))
  1070  			require.Equal(t, request.String(), tt.request.String())
  1071  			require.Equal(t, tt.callOptions, callOptions)
  1072  		})
  1073  	}
  1074  }