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

     1  //go:build integration
     2  // +build integration
     3  
     4  package integration
     5  
     6  import (
     7  	"context"
     8  	"database/sql"
     9  	"errors"
    10  	"os"
    11  	"testing"
    12  	"time"
    13  
    14  	"github.com/stretchr/testify/require"
    15  
    16  	"github.com/ydb-platform/ydb-go-sdk/v3"
    17  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/version"
    18  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xtest"
    19  	"github.com/ydb-platform/ydb-go-sdk/v3/query"
    20  	"github.com/ydb-platform/ydb-go-sdk/v3/retry"
    21  	"github.com/ydb-platform/ydb-go-sdk/v3/retry/budget"
    22  	"github.com/ydb-platform/ydb-go-sdk/v3/table"
    23  )
    24  
    25  type noQuota struct{}
    26  
    27  func (n noQuota) Acquire(ctx context.Context) error {
    28  	return errors.New("no quota")
    29  }
    30  
    31  func TestRetryBudget(t *testing.T) {
    32  	ctx := xtest.Context(t)
    33  
    34  	defaultLimiter := budget.Limited(1)
    35  	defer defaultLimiter.Stop()
    36  
    37  	nativeDriver, err := ydb.Open(ctx, os.Getenv("YDB_CONNECTION_STRING"),
    38  		ydb.WithDiscoveryInterval(time.Second),
    39  		ydb.WithRetryBudget(defaultLimiter),
    40  	)
    41  	require.NoError(t, err)
    42  
    43  	defer func() {
    44  		// cleanup
    45  		_ = nativeDriver.Close(ctx)
    46  	}()
    47  
    48  	c, err := ydb.Connector(nativeDriver)
    49  	require.NoError(t, err)
    50  
    51  	defer func() {
    52  		// cleanup
    53  		_ = c.Close()
    54  	}()
    55  
    56  	db := sql.OpenDB(c)
    57  	defer func() {
    58  		// cleanup
    59  		_ = db.Close()
    60  	}()
    61  
    62  	retryBudget := noQuota{}
    63  
    64  	t.Run("retry.Retry", func(t *testing.T) {
    65  		err := retry.Retry(ctx, func(ctx context.Context) (err error) {
    66  			return retry.RetryableError(errors.New("custom error"))
    67  		}, retry.WithBudget(retryBudget))
    68  		require.ErrorIs(t, err, budget.ErrNoQuota)
    69  	})
    70  	t.Run("retry.Do", func(t *testing.T) {
    71  		err := retry.Do(ctx, db, func(ctx context.Context, cc *sql.Conn) (err error) {
    72  			return retry.RetryableError(errors.New("custom error"))
    73  		}, retry.WithBudget(retryBudget))
    74  		require.ErrorIs(t, err, budget.ErrNoQuota)
    75  	})
    76  	t.Run("retry.DoTx", func(t *testing.T) {
    77  		err := retry.DoTx(ctx, db, func(ctx context.Context, tx *sql.Tx) (err error) {
    78  			return retry.RetryableError(errors.New("custom error"))
    79  		}, retry.WithBudget(retryBudget))
    80  		require.ErrorIs(t, err, budget.ErrNoQuota)
    81  	})
    82  	t.Run("db.Table().Do", func(t *testing.T) {
    83  		err := nativeDriver.Table().Do(ctx, func(ctx context.Context, s table.Session) error {
    84  			return retry.RetryableError(errors.New("custom error"))
    85  		}, table.WithRetryBudget(retryBudget))
    86  		require.ErrorIs(t, err, budget.ErrNoQuota)
    87  	})
    88  	t.Run("db.Table().DoTx", func(t *testing.T) {
    89  		err := nativeDriver.Table().DoTx(ctx, func(ctx context.Context, tx table.TransactionActor) error {
    90  			return retry.RetryableError(errors.New("custom error"))
    91  		}, table.WithRetryBudget(retryBudget))
    92  		require.ErrorIs(t, err, budget.ErrNoQuota)
    93  	})
    94  	if version.Gte(os.Getenv("YDB_VERSION"), "24.1") {
    95  		t.Run("db.Query().Do", func(t *testing.T) {
    96  			err := nativeDriver.Query().Do(ctx, func(ctx context.Context, s query.Session) error {
    97  				return retry.RetryableError(errors.New("custom error"))
    98  			}, query.WithRetryBudget(retryBudget))
    99  			require.ErrorIs(t, err, budget.ErrNoQuota)
   100  		})
   101  		t.Run("db.Query().DoTx", func(t *testing.T) {
   102  			err := nativeDriver.Query().DoTx(ctx, func(ctx context.Context, tx query.TxActor) error {
   103  				return retry.RetryableError(errors.New("custom error"))
   104  			}, query.WithRetryBudget(retryBudget))
   105  			require.ErrorIs(t, err, budget.ErrNoQuota)
   106  		})
   107  	}
   108  }