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