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

     1  package query
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"fmt"
     7  	"io"
     8  	"sync"
     9  	"testing"
    10  	"time"
    11  
    12  	"github.com/rekby/fixenv"
    13  	"github.com/rekby/fixenv/sf"
    14  	"github.com/stretchr/testify/require"
    15  	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb"
    16  	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb_Query"
    17  	"go.uber.org/mock/gomock"
    18  	"google.golang.org/grpc"
    19  	grpcCodes "google.golang.org/grpc/codes"
    20  	grpcStatus "google.golang.org/grpc/status"
    21  
    22  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/query/options"
    23  	baseTx "github.com/ydb-platform/ydb-go-sdk/v3/internal/tx"
    24  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
    25  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xtest"
    26  	"github.com/ydb-platform/ydb-go-sdk/v3/query"
    27  )
    28  
    29  var _ baseTx.Transaction = &Transaction{}
    30  
    31  func TestBegin(t *testing.T) {
    32  	t.Run("HappyWay", func(t *testing.T) {
    33  		ctx := xtest.Context(t)
    34  		ctrl := gomock.NewController(t)
    35  		client := NewMockQueryServiceClient(ctrl)
    36  		client.EXPECT().BeginTransaction(gomock.Any(), gomock.Any()).Return(&Ydb_Query.BeginTransactionResponse{
    37  			Status: Ydb.StatusIds_SUCCESS,
    38  			TxMeta: &Ydb_Query.TransactionMeta{
    39  				Id: "123",
    40  			},
    41  		}, nil)
    42  		t.Log("begin")
    43  		txID, err := begin(ctx, client, "123", query.TxSettings())
    44  		require.NoError(t, err)
    45  		require.Equal(t, "123", txID)
    46  	})
    47  	t.Run("TransportError", func(t *testing.T) {
    48  		ctx := xtest.Context(t)
    49  		ctrl := gomock.NewController(t)
    50  		client := NewMockQueryServiceClient(ctrl)
    51  		client.EXPECT().BeginTransaction(gomock.Any(), gomock.Any()).Return(nil, grpcStatus.Error(grpcCodes.Unavailable, ""))
    52  		t.Log("begin")
    53  		_, err := begin(ctx, client, "123", query.TxSettings())
    54  		require.Error(t, err)
    55  		require.True(t, xerrors.IsTransportError(err, grpcCodes.Unavailable))
    56  	})
    57  	t.Run("OperationError", func(t *testing.T) {
    58  		ctx := xtest.Context(t)
    59  		ctrl := gomock.NewController(t)
    60  		client := NewMockQueryServiceClient(ctrl)
    61  		client.EXPECT().BeginTransaction(gomock.Any(), gomock.Any()).Return(nil,
    62  			xerrors.Operation(xerrors.WithStatusCode(Ydb.StatusIds_UNAVAILABLE)),
    63  		)
    64  		t.Log("begin")
    65  		_, err := begin(ctx, client, "123", query.TxSettings())
    66  		require.Error(t, err)
    67  		require.True(t, xerrors.IsOperationError(err, Ydb.StatusIds_UNAVAILABLE))
    68  	})
    69  }
    70  
    71  func TestCommitTx(t *testing.T) {
    72  	t.Run("HappyWay", func(t *testing.T) {
    73  		ctx := xtest.Context(t)
    74  		ctrl := gomock.NewController(t)
    75  		service := NewMockQueryServiceClient(ctrl)
    76  		service.EXPECT().CommitTransaction(gomock.Any(), gomock.Any()).Return(
    77  			&Ydb_Query.CommitTransactionResponse{
    78  				Status: Ydb.StatusIds_SUCCESS,
    79  			}, nil,
    80  		)
    81  		t.Log("commit")
    82  		err := commitTx(ctx, service, "123", "456")
    83  		require.NoError(t, err)
    84  	})
    85  	t.Run("TransportError", func(t *testing.T) {
    86  		ctx := xtest.Context(t)
    87  		ctrl := gomock.NewController(t)
    88  		service := NewMockQueryServiceClient(ctrl)
    89  		service.EXPECT().CommitTransaction(gomock.Any(), gomock.Any()).Return(
    90  			nil, grpcStatus.Error(grpcCodes.Unavailable, ""),
    91  		)
    92  		t.Log("commit")
    93  		err := commitTx(ctx, service, "123", "456")
    94  		require.Error(t, err)
    95  		require.True(t, xerrors.IsTransportError(err, grpcCodes.Unavailable))
    96  	})
    97  	t.Run("OperationError", func(t *testing.T) {
    98  		ctx := xtest.Context(t)
    99  		ctrl := gomock.NewController(t)
   100  		service := NewMockQueryServiceClient(ctrl)
   101  		service.EXPECT().CommitTransaction(gomock.Any(), gomock.Any()).Return(nil,
   102  			xerrors.Operation(xerrors.WithStatusCode(Ydb.StatusIds_UNAVAILABLE)),
   103  		)
   104  		t.Log("commit")
   105  		err := commitTx(ctx, service, "123", "456")
   106  		require.Error(t, err)
   107  		require.True(t, xerrors.IsOperationError(err, Ydb.StatusIds_UNAVAILABLE))
   108  	})
   109  }
   110  
   111  func TestTxOnCompleted(t *testing.T) {
   112  	t.Run("OnCommitTxSuccess", func(t *testing.T) {
   113  		e := fixenv.New(t)
   114  
   115  		QueryGrpcMock(e).EXPECT().CommitTransaction(gomock.Any(), gomock.Any()).Return(
   116  			&Ydb_Query.CommitTransactionResponse{
   117  				Status: Ydb.StatusIds_SUCCESS,
   118  			}, nil,
   119  		)
   120  
   121  		tx := TransactionOverGrpcMock(e)
   122  
   123  		var completed []error
   124  		tx.OnCompleted(func(transactionResult error) {
   125  			completed = append(completed, transactionResult)
   126  		})
   127  		err := tx.CommitTx(sf.Context(e))
   128  		require.NoError(t, err)
   129  		require.Equal(t, []error{nil}, completed)
   130  	})
   131  	t.Run("OnCommitTxFailed", func(t *testing.T) {
   132  		e := fixenv.New(t)
   133  
   134  		testError := errors.New("test-error")
   135  
   136  		QueryGrpcMock(e).EXPECT().CommitTransaction(gomock.Any(), gomock.Any()).Return(nil,
   137  			testError,
   138  		)
   139  
   140  		tx := TransactionOverGrpcMock(e)
   141  
   142  		var completed []error
   143  		tx.OnCompleted(func(transactionResult error) {
   144  			completed = append(completed, transactionResult)
   145  		})
   146  		err := tx.CommitTx(sf.Context(e))
   147  		require.ErrorIs(t, err, testError)
   148  		require.Len(t, completed, 1)
   149  		require.ErrorIs(t, completed[0], err)
   150  	})
   151  	t.Run("OnRollback", func(t *testing.T) {
   152  		e := fixenv.New(t)
   153  
   154  		rollbackCalled := false
   155  		QueryGrpcMock(e).EXPECT().RollbackTransaction(gomock.Any(), gomock.Any()).DoAndReturn(
   156  			func(
   157  				ctx context.Context,
   158  				request *Ydb_Query.RollbackTransactionRequest,
   159  				option ...grpc.CallOption,
   160  			) (
   161  				*Ydb_Query.RollbackTransactionResponse,
   162  				error,
   163  			) {
   164  				rollbackCalled = true
   165  
   166  				return &Ydb_Query.RollbackTransactionResponse{
   167  					Status: Ydb.StatusIds_SUCCESS,
   168  				}, nil
   169  			})
   170  
   171  		tx := TransactionOverGrpcMock(e)
   172  		var completed error
   173  
   174  		tx.OnCompleted(func(transactionResult error) {
   175  			// notification before call to the server
   176  			require.False(t, rollbackCalled)
   177  			completed = transactionResult
   178  		})
   179  
   180  		_ = tx.Rollback(sf.Context(e))
   181  		require.ErrorIs(t, completed, ErrTransactionRollingBack)
   182  	})
   183  	t.Run("OnExecWithoutCommitTxSuccess", func(t *testing.T) {
   184  		e := fixenv.New(t)
   185  
   186  		responseStream := NewMockQueryService_ExecuteQueryClient(MockController(e))
   187  		responseStream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{
   188  			Status: Ydb.StatusIds_SUCCESS,
   189  		}, nil)
   190  		responseStream.EXPECT().Recv().Return(nil, io.EOF)
   191  
   192  		QueryGrpcMock(e).EXPECT().ExecuteQuery(gomock.Any(), gomock.Any()).Return(responseStream, nil)
   193  
   194  		tx := TransactionOverGrpcMock(e)
   195  		var completed []error
   196  
   197  		tx.OnCompleted(func(transactionResult error) {
   198  			completed = append(completed, transactionResult)
   199  		})
   200  
   201  		err := tx.Exec(sf.Context(e), "")
   202  		require.NoError(t, err)
   203  		require.Empty(t, completed)
   204  	})
   205  	t.Run("OnQueryWithoutCommitTxSuccess", func(t *testing.T) {
   206  		e := fixenv.New(t)
   207  
   208  		responseStream := NewMockQueryService_ExecuteQueryClient(MockController(e))
   209  		responseStream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{
   210  			Status: Ydb.StatusIds_SUCCESS,
   211  		}, nil)
   212  
   213  		QueryGrpcMock(e).EXPECT().ExecuteQuery(gomock.Any(), gomock.Any()).Return(responseStream, nil)
   214  
   215  		tx := TransactionOverGrpcMock(e)
   216  		var completed []error
   217  
   218  		tx.OnCompleted(func(transactionResult error) {
   219  			completed = append(completed, transactionResult)
   220  		})
   221  
   222  		_, err := tx.Query(sf.Context(e), "")
   223  		require.NoError(t, err)
   224  		require.Empty(t, completed)
   225  	})
   226  	t.Run("OnExecWithoutTxSuccess", func(t *testing.T) {
   227  		xtest.TestManyTimes(t, func(t testing.TB) {
   228  			e := fixenv.New(t)
   229  
   230  			responseStream := NewMockQueryService_ExecuteQueryClient(MockController(e))
   231  			responseStream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{
   232  				Status: Ydb.StatusIds_SUCCESS,
   233  			}, nil)
   234  			responseStream.EXPECT().Recv().Return(nil, io.EOF)
   235  
   236  			QueryGrpcMock(e).EXPECT().ExecuteQuery(gomock.Any(), gomock.Any()).Return(responseStream, nil)
   237  
   238  			tx := TransactionOverGrpcMock(e)
   239  			var completedMutex sync.Mutex
   240  			var completed []error
   241  
   242  			tx.OnCompleted(func(transactionResult error) {
   243  				completedMutex.Lock()
   244  				completed = append(completed, transactionResult)
   245  				completedMutex.Unlock()
   246  			})
   247  
   248  			err := tx.Exec(sf.Context(e), "")
   249  			require.NoError(t, err)
   250  			require.Empty(t, completed)
   251  		})
   252  	})
   253  	t.Run("OnQueryWithoutTxSuccess", func(t *testing.T) {
   254  		xtest.TestManyTimes(t, func(t testing.TB) {
   255  			e := fixenv.New(t)
   256  
   257  			responseStream := NewMockQueryService_ExecuteQueryClient(MockController(e))
   258  			responseStream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{
   259  				Status: Ydb.StatusIds_SUCCESS,
   260  			}, nil)
   261  			responseStream.EXPECT().Recv().Return(nil, io.EOF)
   262  
   263  			QueryGrpcMock(e).EXPECT().ExecuteQuery(gomock.Any(), gomock.Any()).Return(responseStream, nil)
   264  
   265  			tx := TransactionOverGrpcMock(e)
   266  			var completedMutex sync.Mutex
   267  			var completed []error
   268  
   269  			tx.OnCompleted(func(transactionResult error) {
   270  				completedMutex.Lock()
   271  				completed = append(completed, transactionResult)
   272  				completedMutex.Unlock()
   273  			})
   274  
   275  			res, err := tx.Query(sf.Context(e), "")
   276  			require.NoError(t, err)
   277  			_ = res.Close(sf.Context(e))
   278  			time.Sleep(time.Millisecond) // time for reaction for closing channel
   279  			require.Empty(t, completed)
   280  		})
   281  	})
   282  	t.Run("OnExecWithTxSuccess", func(t *testing.T) {
   283  		xtest.TestManyTimes(t, func(t testing.TB) {
   284  			e := fixenv.New(t)
   285  
   286  			responseStream := NewMockQueryService_ExecuteQueryClient(MockController(e))
   287  			responseStream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{
   288  				Status: Ydb.StatusIds_SUCCESS,
   289  			}, nil)
   290  			responseStream.EXPECT().Recv().Return(nil, io.EOF)
   291  
   292  			QueryGrpcMock(e).EXPECT().ExecuteQuery(gomock.Any(), gomock.Any()).Return(responseStream, nil)
   293  
   294  			tx := TransactionOverGrpcMock(e)
   295  			var completedMutex sync.Mutex
   296  			var completed []error
   297  
   298  			tx.OnCompleted(func(transactionResult error) {
   299  				completedMutex.Lock()
   300  				completed = append(completed, transactionResult)
   301  				completedMutex.Unlock()
   302  			})
   303  
   304  			err := tx.Exec(sf.Context(e), "", options.WithCommit())
   305  			require.NoError(t, err)
   306  			xtest.SpinWaitCondition(t, &completedMutex, func() bool {
   307  				return len(completed) != 0
   308  			})
   309  			require.Equal(t, []error{nil}, completed)
   310  		})
   311  	})
   312  	t.Run("OnQueryWithTxSuccess", func(t *testing.T) {
   313  		xtest.TestManyTimes(t, func(t testing.TB) {
   314  			e := fixenv.New(t)
   315  
   316  			responseStream := NewMockQueryService_ExecuteQueryClient(MockController(e))
   317  			responseStream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{
   318  				Status: Ydb.StatusIds_SUCCESS,
   319  			}, nil)
   320  			responseStream.EXPECT().Recv().Return(nil, io.EOF)
   321  
   322  			QueryGrpcMock(e).EXPECT().ExecuteQuery(gomock.Any(), gomock.Any()).Return(responseStream, nil)
   323  
   324  			tx := TransactionOverGrpcMock(e)
   325  			var completedMutex sync.Mutex
   326  			var completed []error
   327  
   328  			tx.OnCompleted(func(transactionResult error) {
   329  				completedMutex.Lock()
   330  				completed = append(completed, transactionResult)
   331  				completedMutex.Unlock()
   332  			})
   333  
   334  			res, err := tx.Query(sf.Context(e), "", options.WithCommit())
   335  			_ = res.Close(sf.Context(e))
   336  			require.NoError(t, err)
   337  			xtest.SpinWaitCondition(t, &completedMutex, func() bool {
   338  				return len(completed) != 0
   339  			})
   340  			require.Equal(t, []error{nil}, completed)
   341  		})
   342  	})
   343  	t.Run("OnQueryWithTxSuccessWithTwoResultSet", func(t *testing.T) {
   344  		xtest.TestManyTimes(t, func(t testing.TB) {
   345  			e := fixenv.New(t)
   346  
   347  			responseStream := NewMockQueryService_ExecuteQueryClient(MockController(e))
   348  			responseStream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{
   349  				ResultSetIndex: 0,
   350  				ResultSet:      &Ydb.ResultSet{},
   351  			}, nil)
   352  			responseStream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{
   353  				Status:         Ydb.StatusIds_SUCCESS,
   354  				ResultSetIndex: 1,
   355  				ResultSet:      &Ydb.ResultSet{},
   356  			}, nil)
   357  			responseStream.EXPECT().Recv().Return(nil, io.EOF)
   358  
   359  			QueryGrpcMock(e).EXPECT().ExecuteQuery(gomock.Any(), gomock.Any()).Return(responseStream, nil)
   360  
   361  			tx := TransactionOverGrpcMock(e)
   362  			var completedMutex sync.Mutex
   363  			var completed []error
   364  
   365  			tx.OnCompleted(func(transactionResult error) {
   366  				completedMutex.Lock()
   367  				completed = append(completed, transactionResult)
   368  				completedMutex.Unlock()
   369  			})
   370  
   371  			res, err := tx.Query(sf.Context(e), "", options.WithCommit())
   372  			require.NoError(t, err)
   373  
   374  			// time for event happened if is
   375  			time.Sleep(time.Millisecond)
   376  			require.Empty(t, completed)
   377  
   378  			_, err = res.NextResultSet(sf.Context(e))
   379  			require.NoError(t, err)
   380  			// time for event happened if is
   381  			time.Sleep(time.Millisecond)
   382  			require.Empty(t, completed)
   383  
   384  			_, err = res.NextResultSet(sf.Context(e))
   385  			require.NoError(t, err)
   386  
   387  			_ = res.Close(sf.Context(e))
   388  			require.NoError(t, err)
   389  			xtest.SpinWaitCondition(t, &completedMutex, func() bool {
   390  				return len(completed) != 0
   391  			})
   392  			require.Equal(t, []error{nil}, completed)
   393  		})
   394  	})
   395  	t.Run("OnExecuteFailedOnInitResponse", func(t *testing.T) {
   396  		for _, commit := range []bool{true, false} {
   397  			t.Run(fmt.Sprint("commit:", commit), func(t *testing.T) {
   398  				e := fixenv.New(t)
   399  
   400  				testErr := errors.New("test")
   401  				responseStream := NewMockQueryService_ExecuteQueryClient(MockController(e))
   402  				responseStream.EXPECT().Recv().Return(nil, testErr)
   403  
   404  				QueryGrpcMock(e).EXPECT().ExecuteQuery(gomock.Any(), gomock.Any()).Return(responseStream, nil)
   405  
   406  				tx := TransactionOverGrpcMock(e)
   407  				var transactionResult error
   408  
   409  				tx.OnCompleted(func(err error) {
   410  					transactionResult = err
   411  				})
   412  
   413  				err := tx.Exec(sf.Context(e), "", query.WithCommit())
   414  				require.ErrorIs(t, err, testErr)
   415  				require.Error(t, transactionResult)
   416  				require.ErrorIs(t, transactionResult, testErr)
   417  			})
   418  		}
   419  	})
   420  	t.Run("OnExecuteFailedInResponsePart", func(t *testing.T) {
   421  		for _, commit := range []bool{true, false} {
   422  			t.Run(fmt.Sprint("commit:", commit), func(t *testing.T) {
   423  				xtest.TestManyTimes(t, func(t testing.TB) {
   424  					e := fixenv.New(t)
   425  
   426  					errorReturned := false
   427  					responseStream := NewMockQueryService_ExecuteQueryClient(MockController(e))
   428  					responseStream.EXPECT().Recv().DoAndReturn(func() (*Ydb_Query.ExecuteQueryResponsePart, error) {
   429  						errorReturned = true
   430  
   431  						return nil, xerrors.Operation(xerrors.WithStatusCode(Ydb.StatusIds_BAD_SESSION))
   432  					})
   433  
   434  					QueryGrpcMock(e).EXPECT().ExecuteQuery(gomock.Any(), gomock.Any()).Return(responseStream, nil)
   435  
   436  					tx := TransactionOverGrpcMock(e)
   437  					var transactionResult error
   438  
   439  					tx.OnCompleted(func(err error) {
   440  						transactionResult = err
   441  					})
   442  
   443  					err := tx.Exec(sf.Context(e), "", query.WithCommit())
   444  					require.True(t, errorReturned)
   445  					require.True(t, xerrors.IsOperationError(err, Ydb.StatusIds_BAD_SESSION))
   446  					require.Error(t, transactionResult)
   447  					require.True(t, xerrors.IsOperationError(transactionResult, Ydb.StatusIds_BAD_SESSION))
   448  				})
   449  			})
   450  		}
   451  	})
   452  }
   453  
   454  func TestRollback(t *testing.T) {
   455  	t.Run("HappyWay", func(t *testing.T) {
   456  		ctx := xtest.Context(t)
   457  		ctrl := gomock.NewController(t)
   458  		service := NewMockQueryServiceClient(ctrl)
   459  		service.EXPECT().RollbackTransaction(gomock.Any(), gomock.Any()).Return(
   460  			&Ydb_Query.RollbackTransactionResponse{
   461  				Status: Ydb.StatusIds_SUCCESS,
   462  			}, nil,
   463  		)
   464  		t.Log("rollback")
   465  		err := rollback(ctx, service, "123", "456")
   466  		require.NoError(t, err)
   467  	})
   468  	t.Run("TransportError", func(t *testing.T) {
   469  		ctx := xtest.Context(t)
   470  		ctrl := gomock.NewController(t)
   471  		service := NewMockQueryServiceClient(ctrl)
   472  		service.EXPECT().RollbackTransaction(gomock.Any(), gomock.Any()).Return(
   473  			nil, grpcStatus.Error(grpcCodes.Unavailable, ""),
   474  		)
   475  		t.Log("rollback")
   476  		err := rollback(ctx, service, "123", "456")
   477  		require.Error(t, err)
   478  		require.True(t, xerrors.IsTransportError(err, grpcCodes.Unavailable))
   479  	})
   480  	t.Run("OperationError", func(t *testing.T) {
   481  		ctx := xtest.Context(t)
   482  		ctrl := gomock.NewController(t)
   483  		service := NewMockQueryServiceClient(ctrl)
   484  		service.EXPECT().RollbackTransaction(gomock.Any(), gomock.Any()).Return(nil,
   485  			xerrors.Operation(xerrors.WithStatusCode(Ydb.StatusIds_UNAVAILABLE)),
   486  		)
   487  		t.Log("rollback")
   488  		err := rollback(ctx, service, "123", "456")
   489  		require.Error(t, err)
   490  		require.True(t, xerrors.IsOperationError(err, Ydb.StatusIds_UNAVAILABLE))
   491  	})
   492  }