github.com/ydb-platform/ydb-go-sdk/v3@v3.89.2/tests/integration/with_trace_retry_test.go (about)

     1  //go:build integration
     2  // +build integration
     3  
     4  package integration
     5  
     6  import (
     7  	"context"
     8  	"database/sql"
     9  	"testing"
    10  
    11  	"github.com/stretchr/testify/require"
    12  
    13  	"github.com/ydb-platform/ydb-go-sdk/v3"
    14  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xtest"
    15  	"github.com/ydb-platform/ydb-go-sdk/v3/retry"
    16  	"github.com/ydb-platform/ydb-go-sdk/v3/table"
    17  	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
    18  )
    19  
    20  func TestWithTraceRetry(t *testing.T) {
    21  	ctx := xtest.Context(t)
    22  
    23  	t.Run("table", func(t *testing.T) {
    24  		var (
    25  			retryCalled = make(map[string]bool, 2)
    26  			scope       = newScope(t)
    27  			db          = scope.Driver(
    28  				ydb.WithTraceRetry(trace.Retry{
    29  					OnRetry: func(info trace.RetryLoopStartInfo) func(trace.RetryLoopDoneInfo) {
    30  						retryCalled[info.Label] = true
    31  						return nil
    32  					},
    33  				}),
    34  			)
    35  		)
    36  
    37  		require.NoError(t, db.Table().Do(ctx,
    38  			func(ctx context.Context, s table.Session) error {
    39  				return nil
    40  			},
    41  			table.WithLabel("db.Table().Do"),
    42  		))
    43  
    44  		require.NoError(t, db.Table().DoTx(ctx,
    45  			func(ctx context.Context, tx table.TransactionActor) error {
    46  				return nil
    47  			},
    48  			table.WithLabel("db.Table().DoTx"),
    49  		))
    50  
    51  		for _, key := range []string{
    52  			"db.Table().Do",
    53  			"db.Table().DoTx",
    54  		} {
    55  			require.True(t, retryCalled[key], key)
    56  		}
    57  	})
    58  
    59  	t.Run("database/sql", func(t *testing.T) {
    60  		var (
    61  			retryCalled = make(map[string]bool, 2)
    62  			scope       = newScope(t)
    63  			nativeDb    = scope.Driver(
    64  				ydb.WithTraceRetry(trace.Retry{
    65  					OnRetry: func(info trace.RetryLoopStartInfo) func(trace.RetryLoopDoneInfo) {
    66  						retryCalled[info.Label] = true
    67  						return nil
    68  					},
    69  				}),
    70  			)
    71  			db = sql.OpenDB(ydb.MustConnector(nativeDb))
    72  		)
    73  		require.NoError(t, retry.Do(ctx, db,
    74  			func(ctx context.Context, cc *sql.Conn) error {
    75  				return nil
    76  			},
    77  			retry.WithLabel("retry.Do"),
    78  		))
    79  
    80  		require.NoError(t, retry.DoTx(ctx, db,
    81  			func(ctx context.Context, tx *sql.Tx) error {
    82  				return nil
    83  			},
    84  			retry.WithLabel("retry.DoTx"),
    85  		))
    86  
    87  		for _, key := range []string{
    88  			"retry.Do",
    89  			"retry.DoTx",
    90  		} {
    91  			require.True(t, retryCalled[key], key)
    92  		}
    93  	})
    94  }