github.com/ydb-platform/ydb-go-sdk/v3@v3.57.0/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(
    30  						info trace.RetryLoopStartInfo,
    31  					) func(trace.RetryLoopIntermediateInfo) func(trace.RetryLoopDoneInfo) {
    32  						retryCalled[info.Label] = true
    33  						return nil
    34  					},
    35  				}),
    36  			)
    37  		)
    38  
    39  		require.NoError(t, db.Table().Do(ctx,
    40  			func(ctx context.Context, s table.Session) error {
    41  				return nil
    42  			},
    43  			table.WithLabel("db.Table().Do"),
    44  		))
    45  
    46  		require.NoError(t, db.Table().DoTx(ctx,
    47  			func(ctx context.Context, tx table.TransactionActor) error {
    48  				return nil
    49  			},
    50  			table.WithLabel("db.Table().DoTx"),
    51  		))
    52  
    53  		for _, key := range []string{
    54  			"db.Table().Do",
    55  			"db.Table().DoTx",
    56  		} {
    57  			require.True(t, retryCalled[key], key)
    58  		}
    59  	})
    60  
    61  	t.Run("database/sql", func(t *testing.T) {
    62  		var (
    63  			retryCalled = make(map[string]bool, 2)
    64  			scope       = newScope(t)
    65  			nativeDb    = scope.Driver(
    66  				ydb.WithTraceRetry(trace.Retry{
    67  					OnRetry: func(
    68  						info trace.RetryLoopStartInfo,
    69  					) func(trace.RetryLoopIntermediateInfo) func(trace.RetryLoopDoneInfo) {
    70  						retryCalled[info.Label] = true
    71  						return nil
    72  					},
    73  				}),
    74  			)
    75  			db = sql.OpenDB(ydb.MustConnector(nativeDb))
    76  		)
    77  		require.NoError(t, retry.Do(ctx, db,
    78  			func(ctx context.Context, cc *sql.Conn) error {
    79  				return nil
    80  			},
    81  			retry.WithLabel("retry.Do"),
    82  		))
    83  
    84  		require.NoError(t, retry.DoTx(ctx, db,
    85  			func(ctx context.Context, tx *sql.Tx) error {
    86  				return nil
    87  			},
    88  			retry.WithLabel("retry.DoTx"),
    89  		))
    90  
    91  		for _, key := range []string{
    92  			"retry.Do",
    93  			"retry.DoTx",
    94  		} {
    95  			require.True(t, retryCalled[key], key)
    96  		}
    97  	})
    98  }