dynamodb: add support for context management

dynamodb/dynamodb.go

@@ -19,7 +19,6 @@ import (
 type Storage struct {
 	db    *awsdynamodb.Client
 	table string
-	requestTimeout time.Duration
 }
 
 // "k" is used as table column name for the key.
@@ -77,11 +76,8 @@ func New(config Config) *Storage {
 	return store
 }
 
-// Get value by key
-func (s *Storage) Get(key string) ([]byte, error) {
-	ctx, cancel := s.requestContext()
-	defer cancel()
-
+// GetWithContext value by key with context
+func (s *Storage) GetWithContext(ctx context.Context, key string) ([]byte, error) {
 	k := make(map[string]types.AttributeValue)
 	k[keyAttrName] = &types.AttributeValueMemberS{
 		Value: key,
@@ -108,11 +104,12 @@ func (s *Storage) Get(key string) ([]byte, error) {
 	return item.V, err
 }
 
-// Set key with value
-func (s *Storage) Set(key string, val []byte, exp time.Duration) error {
-	ctx, cancel := s.requestContext()
-	defer cancel()
+func (s *Storage) Get(key string) ([]byte, error) {
+	return s.GetWithContext(context.Background(), key)
+}
 
+// Set key with value
+func (s *Storage) SetWithContext(ctx context.Context, key string, val []byte, exp time.Duration) error {
 	// Ain't Nobody Got Time For That
 	if len(key) <= 0 || len(val) <= 0 {
 		return nil
@@ -134,11 +131,12 @@ func (s *Storage) Set(key string, val []byte, exp time.Duration) error {
 	return err
 }
 
-// Delete entry by key
-func (s *Storage) Delete(key string) error {
-	ctx, cancel := s.requestContext()
-	defer cancel()
+func (s *Storage) Set(key string, val []byte, exp time.Duration) error {
+	return s.SetWithContext(context.Background(), key, val, exp)
+}
 
+// Delete entry by key
+func (s *Storage) DeleteWithContext(ctx context.Context, key string) error {
 	// Ain't Nobody Got Time For That
 	if len(key) <= 0 {
 		return nil
@@ -157,11 +155,12 @@ func (s *Storage) Delete(key string) error {
 	return err
 }
 
-// Reset all entries, including unexpired
-func (s *Storage) Reset() error {
-	ctx, cancel := s.requestContext()
-	defer cancel()
+func (s *Storage) Delete(key string) error {
+	return s.DeleteWithContext(context.Background(), key)
+}
 
+// Reset all entries, including unexpired
+func (s *Storage) ResetWithContext(ctx context.Context) error {
 	deleteTableInput := awsdynamodb.DeleteTableInput{
 		TableName: &s.table,
 	}
@@ -169,14 +168,17 @@ func (s *Storage) Reset() error {
 	return err
 }
 
+func (s *Storage) Reset() error {
+	return s.ResetWithContext(context.Background())
+}
+
 // Close the database
 func (s *Storage) Close() error {
 	return nil
 }
 
 func (s *Storage) createTable(cfg Config, describeTableInput awsdynamodb.DescribeTableInput) error {
-	ctx, cancel := s.requestContext()
-	defer cancel()
+	ctx := context.Background()
 
 	keyAttrType := "S" // For "string"
 	keyType := "HASH"  // As opposed to "RANGE"
@@ -225,14 +227,6 @@ func (s *Storage) createTable(cfg Config, describeTableInput awsdynamodb.Describ
 	return nil
 }
 
-// Context for making requests will timeout if a non-zero timeout is configured
-func (s *Storage) requestContext() (context.Context, context.CancelFunc) {
-	if s.requestTimeout > 0 {
-		return context.WithTimeout(context.Background(), s.requestTimeout)
-	}
-	return context.Background(), func() {}
-}
-
 func returnAWSConfig(cfg Config) (aws.Config, error) {
 	if cfg.Credentials != (Credentials{}) {
 		credentials := credentials.NewStaticCredentialsProvider(cfg.Credentials.AccessKey, cfg.Credentials.SecretAccessKey, "")

dynamodb/dynamodb_test.go

@@ -60,6 +60,19 @@ func Test_DynamoDB_Set(t *testing.T) {
 	require.NoError(t, err)
 }
 
+func Test_DynamoDB_SetWithContext(t *testing.T) {
+	var (
+		key = "john"
+		val = []byte("doe")
+	)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	err := testStore.SetWithContext(ctx, key, val, 0)
+	require.ErrorIs(t, err, context.Canceled)
+}
+
 func Test_DynamoDB_Set_Override(t *testing.T) {
 	var (
 		key = "john"
@@ -93,6 +106,19 @@ func Test_DynamoDB_Get(t *testing.T) {
 	require.Equal(t, val, result)
 }
 
+func Test_DynamoDB_GetWithContext(t *testing.T) {
+	var (
+		key = "john"
+	)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	result, err := testStore.GetWithContext(ctx, key)
+	require.ErrorIs(t, err, context.Canceled)
+	require.Zero(t, len(result))
+}
+
 func Test_DynamoDB_Get_NotExist(t *testing.T) {
 	testStore := newTestStore(t)
 	defer testStore.Close()
@@ -122,6 +148,26 @@ func Test_DynamoDB_Delete(t *testing.T) {
 	require.Zero(t, len(result))
 }
 
+func Test_DynamoDB_DeleteWithContext(t *testing.T) {
+	var (
+		key = "john"
+		val = []byte("doe")
+	)
+
+	err := testStore.Set(key, val, 0)
+	require.NoError(t, err)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	err = testStore.DeleteWithContext(ctx, key)
+	require.ErrorIs(t, err, context.Canceled)
+
+	result, err := testStore.Get(key)
+	require.NoError(t, err)
+	require.Equal(t, val, result)
+}
+
 func Test_DynamoDB_Reset(t *testing.T) {
 	val := []byte("doe")
 
@@ -146,6 +192,30 @@ func Test_DynamoDB_Reset(t *testing.T) {
 	require.Zero(t, len(result))
 }
 
+func Test_DynamoDB_ResetWithContext(t *testing.T) {
+	val := []byte("doe")
+
+	err := testStore.Set("john1", val, 0)
+	require.NoError(t, err)
+
+	err = testStore.Set("john2", val, 0)
+	require.NoError(t, err)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	err = testStore.ResetWithContext(ctx)
+	require.ErrorIs(t, err, context.Canceled)
+
+	result, err := testStore.Get("john1")
+	require.NoError(t, err)
+	require.Equal(t, val, result)
+
+	result, err = testStore.Get("john2")
+	require.NoError(t, err)
+	require.Equal(t, val, result)
+}
+
 func Test_DynamoDB_Close(t *testing.T) {
 	testStore := newTestStore(t)
 	require.Nil(t, testStore.Close())