add retry client for DeleteItem

Thumbscrew · Jun 18, 2024 · ce8dac6 · ce8dac6
1 parent 25af5c7
commit ce8dac6
Show file tree

Hide file tree

Showing 5 changed files with 344 additions and 0 deletions.
diff --git a/ddbretry/ddbretry.go b/ddbretry/ddbretry.go
@@ -0,0 +1,60 @@
+package ddbretry
+
+import (
+	"context"
+	"time"
+
+	"github.com/Thumbscrew/aws-go-tools/ddbretry/errors"
+	"github.com/aws/aws-sdk-go-v2/service/dynamodb"
+	"github.com/aws/aws-sdk-go-v2/service/dynamodb/types"
+)
+
+type DynamoDBClient interface {
+	DeleteItem(context.Context, *dynamodb.DeleteItemInput, ...func(*dynamodb.Options)) (*dynamodb.DeleteItemOutput, error)
+}
+
+type RetryDynamoDBClient struct {
+	DynamoDBClient
+	Retries     int
+	BackOffTime time.Duration
+}
+
+func NewRetryDynamoDBClient(client DynamoDBClient, retries int, backOff time.Duration) *RetryDynamoDBClient {
+	return &RetryDynamoDBClient{
+		DynamoDBClient: client,
+		Retries:        retries,
+		BackOffTime:    backOff,
+	}
+}
+
+func (c *RetryDynamoDBClient) DeleteItem(ctx context.Context, input *dynamodb.DeleteItemInput, o ...func(*dynamodb.Options)) (output *dynamodb.DeleteItemOutput, err error) {
+	retries := c.Retries
+	infinite := retries == -1
+	for retries >= 0 || infinite {
+		output, err = c.DynamoDBClient.DeleteItem(ctx, input, o...)
+		if err != nil {
+			if IsProvisionedThroughputExceededException(err) {
+				if retries > 0 {
+					retries--
+					time.Sleep(c.BackOffTime)
+				} else if infinite {
+					time.Sleep(c.BackOffTime)
+				} else {
+					return
+				}
+			} else {
+				return
+			}
+		} else {
+			return
+		}
+	}
+
+	return nil, errors.NewInvalidRetryError(retries)
+}
+
+func IsProvisionedThroughputExceededException(err error) bool {
+	_, ok := err.(*types.ProvisionedThroughputExceededException)
+
+	return ok
+}
diff --git a/ddbretry/ddbretry_test.go b/ddbretry/ddbretry_test.go
@@ -0,0 +1,208 @@
+package ddbretry
+
+import (
+	"context"
+	"errors"
+	"testing"
+	"time"
+
+	ierrors "github.com/Thumbscrew/aws-go-tools/ddbretry/errors"
+	"github.com/aws/aws-sdk-go-v2/service/dynamodb"
+	"github.com/aws/aws-sdk-go-v2/service/dynamodb/types"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestIsProvisionedThroughputExceededException(t *testing.T) {
+	type args struct {
+		err error
+	}
+	tests := []struct {
+		name string
+		args args
+		want bool
+	}{
+		{
+			name: "should return true when error is ProvisionedThroughputExceededException",
+			args: args{
+				err: &types.ProvisionedThroughputExceededException{},
+			},
+			want: true,
+		},
+		{
+			name: "should return false when error is not ProvisionedThroughputExceededException",
+			args: args{
+				err: errors.New("foo"),
+			},
+			want: false,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			assert.Equal(t, tt.want, IsProvisionedThroughputExceededException(tt.args.err))
+		})
+	}
+}
+
+type SuccessfulDynamoDBClient struct {
+	ThroughputExceededCount int
+}
+
+func (c *SuccessfulDynamoDBClient) DeleteItem(ctx context.Context, input *dynamodb.DeleteItemInput, o ...func(*dynamodb.Options)) (*dynamodb.DeleteItemOutput, error) {
+	for c.ThroughputExceededCount > 0 {
+		c.ThroughputExceededCount--
+		return nil, &types.ProvisionedThroughputExceededException{}
+	}
+
+	return &dynamodb.DeleteItemOutput{}, nil
+}
+
+type FailingDynamoDBClient struct {
+	ThroughputExceededCount int
+	Err                     error
+}
+
+func (c *FailingDynamoDBClient) DeleteItem(ctx context.Context, input *dynamodb.DeleteItemInput, o ...func(*dynamodb.Options)) (*dynamodb.DeleteItemOutput, error) {
+	for c.ThroughputExceededCount > 0 {
+		c.ThroughputExceededCount--
+		return nil, &types.ProvisionedThroughputExceededException{}
+	}
+
+	return nil, c.Err
+}
+
+func TestRetryDynamoDBClient_DeleteItem(t *testing.T) {
+	ctx := context.Background()
+
+	type fields struct {
+		DynamoDBClient DynamoDBClient
+		Retries        int
+		BackOffTime    time.Duration
+	}
+	type args struct {
+		ctx   context.Context
+		input *dynamodb.DeleteItemInput
+		o     []func(*dynamodb.Options)
+	}
+	tests := []struct {
+		name       string
+		fields     fields
+		args       args
+		wantOutput *dynamodb.DeleteItemOutput
+		wantErr    error
+	}{
+		{
+			name: "should receive output from successful call in DynamoDBClient",
+			fields: fields{
+				DynamoDBClient: &SuccessfulDynamoDBClient{},
+			},
+			args: args{
+				ctx:   ctx,
+				input: &dynamodb.DeleteItemInput{},
+			},
+			wantOutput: &dynamodb.DeleteItemOutput{},
+			wantErr:    nil,
+		},
+		{
+			name: "should receive error from failed call in DynamoDBClient",
+			fields: fields{
+				DynamoDBClient: &FailingDynamoDBClient{
+					Err: errors.New("foo"),
+				},
+			},
+			args: args{
+				ctx:   ctx,
+				input: &dynamodb.DeleteItemInput{},
+			},
+			wantOutput: nil,
+			wantErr:    errors.New("foo"),
+		},
+		{
+			name: "should receive output when retries is higher than number of throughput exceptions",
+			fields: fields{
+				DynamoDBClient: &SuccessfulDynamoDBClient{
+					ThroughputExceededCount: 2,
+				},
+				Retries: 3,
+			},
+			args: args{
+				ctx:   ctx,
+				input: &dynamodb.DeleteItemInput{},
+			},
+			wantOutput: &dynamodb.DeleteItemOutput{},
+			wantErr:    nil,
+		},
+		{
+			name: "should receive throughput exception when number of throughput exceptions is higher than retries",
+			fields: fields{
+				DynamoDBClient: &SuccessfulDynamoDBClient{
+					ThroughputExceededCount: 3,
+				},
+				Retries: 2,
+			},
+			args: args{
+				ctx:   ctx,
+				input: &dynamodb.DeleteItemInput{},
+			},
+			wantOutput: nil,
+			wantErr:    &types.ProvisionedThroughputExceededException{},
+		},
+		{
+			name: "should receive error after throughput exceptions when retries is higher",
+			fields: fields{
+				DynamoDBClient: &FailingDynamoDBClient{
+					ThroughputExceededCount: 2,
+					Err:                     errors.New("foo"),
+				},
+				Retries: 3,
+			},
+			args: args{
+				ctx:   ctx,
+				input: &dynamodb.DeleteItemInput{},
+			},
+			wantOutput: nil,
+			wantErr:    errors.New("foo"),
+		},
+		{
+			name: "should receive output after throughput exceptions when retries is infinite",
+			fields: fields{
+				DynamoDBClient: &SuccessfulDynamoDBClient{
+					ThroughputExceededCount: 10,
+				},
+				Retries: -1,
+			},
+			args: args{
+				ctx:   ctx,
+				input: &dynamodb.DeleteItemInput{},
+			},
+			wantOutput: &dynamodb.DeleteItemOutput{},
+			wantErr:    nil,
+		},
+		{
+			name: "should receive InvalidRetryError when retries value is invalid",
+			fields: fields{
+				DynamoDBClient: &SuccessfulDynamoDBClient{
+					ThroughputExceededCount: 10,
+				},
+				Retries: -2,
+			},
+			args: args{
+				ctx:   ctx,
+				input: &dynamodb.DeleteItemInput{},
+			},
+			wantOutput: nil,
+			wantErr:    ierrors.NewInvalidRetryError(-2),
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			c := &RetryDynamoDBClient{
+				DynamoDBClient: tt.fields.DynamoDBClient,
+				Retries:        tt.fields.Retries,
+				BackOffTime:    tt.fields.BackOffTime,
+			}
+			gotOutput, err := c.DeleteItem(tt.args.ctx, tt.args.input, tt.args.o...)
+			assert.Equal(t, tt.wantOutput, gotOutput)
+			assert.Equal(t, tt.wantErr, err)
+		})
+	}
+}
diff --git a/ddbretry/errors/errors.go b/ddbretry/errors/errors.go
@@ -0,0 +1,23 @@
+package errors
+
+import "fmt"
+
+type InvalidRetryError struct {
+	Retries int
+}
+
+func (e *InvalidRetryError) Error() string {
+	return fmt.Sprintf("invalid value for retries: %d", e.Retries)
+}
+
+func NewInvalidRetryError(retries int) *InvalidRetryError {
+	return &InvalidRetryError{
+		Retries: retries,
+	}
+}
+
+func IsInvalidRetryError(err error) bool {
+	_, ok := err.(*InvalidRetryError)
+
+	return ok
+}
diff --git a/ddbretry/go.mod b/ddbretry/go.mod
@@ -0,0 +1,21 @@
+module github.com/Thumbscrew/aws-go-tools/ddbretry
+
+go 1.21
+
+require (
+	github.com/aws/aws-sdk-go-v2/service/dynamodb v1.32.9
+	github.com/stretchr/testify v1.9.0
+)
+
+require (
+	github.com/aws/aws-sdk-go-v2 v1.28.0 // indirect
+	github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.10 // indirect
+	github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.10 // indirect
+	github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.11.2 // indirect
+	github.com/aws/aws-sdk-go-v2/service/internal/endpoint-discovery v1.9.11 // indirect
+	github.com/aws/smithy-go v1.20.2 // indirect
+	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/jmespath/go-jmespath v0.4.0 // indirect
+	github.com/pmezard/go-difflib v1.0.0 // indirect
+	gopkg.in/yaml.v3 v3.0.1 // indirect
+)
diff --git a/ddbretry/go.sum b/ddbretry/go.sum
@@ -0,0 +1,32 @@
+github.com/aws/aws-sdk-go-v2 v1.28.0 h1:ne6ftNhY0lUvlazMUQF15FF6NH80wKmPRFG7g2q6TCw=
+github.com/aws/aws-sdk-go-v2 v1.28.0/go.mod h1:ffIFB97e2yNsv4aTSGkqtHnppsIJzw7G7BReUZ3jCXM=
+github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.10 h1:LZIUb8sQG2cb89QaVFtMSnER10gyKkqU1k3hP3g9das=
+github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.10/go.mod h1:BRIqay//vnIOCZjoXWSLffL2uzbtxEmnSlfbvVh7Z/4=
+github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.10 h1:HY7CXLA0GiQUo3WYxOP7WYkLcwvRX4cLPf5joUcrQGk=
+github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.10/go.mod h1:kfRBSxRa+I+VyON7el3wLZdrO91oxUxEwdAaWgFqN90=
+github.com/aws/aws-sdk-go-v2/service/dynamodb v1.32.9 h1:EQ6Th8HvCAaVDGVTSpGHP+aGhOI77ANNW/RByMWY2eU=
+github.com/aws/aws-sdk-go-v2/service/dynamodb v1.32.9/go.mod h1:9WEu5LY+YUn9hvsnw89QdlCc5tpwo9mrJ5RQooMV7t4=
+github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.11.2 h1:Ji0DY1xUsUr3I8cHps0G+XM3WWU16lP6yG8qu1GAZAs=
+github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.11.2/go.mod h1:5CsjAbs3NlGQyZNFACh+zztPDI7fU6eW9QsxjfnuBKg=
+github.com/aws/aws-sdk-go-v2/service/internal/endpoint-discovery v1.9.11 h1:F5o2FRQkUByNwIhkU3xPl8jmsnA2i6+cX7aJt1qJpBM=
+github.com/aws/aws-sdk-go-v2/service/internal/endpoint-discovery v1.9.11/go.mod h1:oKamKUpKwRMfg3o6yMyUXbKcDcvdnsvkJW+euxc3jPk=
+github.com/aws/smithy-go v1.20.2 h1:tbp628ireGtzcHDDmLT/6ADHidqnwgF57XOXZe6tp4Q=
+github.com/aws/smithy-go v1.20.2/go.mod h1:krry+ya/rV9RDcV/Q16kpu6ypI4K2czasz0NC3qS14E=
+github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
+github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/jmespath/go-jmespath v0.4.0 h1:BEgLn5cpjn8UN1mAw4NjwDrS35OdebyEtFe+9YPoQUg=
+github.com/jmespath/go-jmespath v0.4.0/go.mod h1:T8mJZnbsbmF+m6zOOFylbeCJqk5+pHWvzYPziyZiYoo=
+github.com/jmespath/go-jmespath/internal/testify v1.5.1 h1:shLQSRRSCCPj3f2gpwzGwWFoC7ycTf1rcQZHOlsJ6N8=
+github.com/jmespath/go-jmespath/internal/testify v1.5.1/go.mod h1:L3OGu8Wl2/fWfCI6z80xFu9LTZmf1ZRjMHUOPmWr69U=
+github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
+github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
+github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/yaml.v2 v2.2.8 h1:obN1ZagJSUGI0Ek/LBmuj4SNLPfIny3KsKFopxRdj10=
+gopkg.in/yaml.v2 v2.2.8/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
+gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
+gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=