s3: add support for context management

s3/s3.go

@@ -78,17 +78,14 @@ func New(config ...Config) *Storage {
 	return storage
 }
 
-// Get value by key
-func (s *Storage) Get(key string) ([]byte, error) {
+// GetWithContext gets value by key with context
+func (s *Storage) GetWithContext(ctx context.Context, key string) ([]byte, error) {
 	var nsk *types.NoSuchKey
 
 	if len(key) <= 0 {
 		return nil, nil
 	}
 
-	ctx, cancel := s.requestContext()
-	defer cancel()
-
 	buf := manager.NewWriteAtBuffer([]byte{})
 
 	_, err := s.downloader.Download(ctx, buf, &s3.GetObjectInput{
@@ -102,15 +99,20 @@ func (s *Storage) Get(key string) ([]byte, error) {
 	return buf.Bytes(), err
 }
 
-// Set key with value
-func (s *Storage) Set(key string, val []byte, exp time.Duration) error {
+// Get gets value by key
+func (s *Storage) Get(key string) ([]byte, error) {
+	ctx, cancel := s.requestContext()
+	defer cancel()
+
+	return s.GetWithContext(ctx, key)
+}
+
+// SetWithContext key with value and expiration time with context
+func (s *Storage) SetWithContext(ctx context.Context, key string, val []byte, exp time.Duration) error {
 	if len(key) <= 0 {
 		return nil
 	}
 
-	ctx, cancel := s.requestContext()
-	defer cancel()
-
 	_, err := s.uploader.Upload(ctx, &s3.PutObjectInput{
 		Bucket: &s.bucket,
 		Key:    aws.String(key),
@@ -120,15 +122,20 @@ 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 {
+// Set key with value and expiration time
+func (s *Storage) Set(key string, val []byte, exp time.Duration) error {
+	ctx, cancel := s.requestContext()
+	defer cancel()
+
+	return s.SetWithContext(ctx, key, val, exp)
+}
+
+// DeleteWithContext deletes entry by key with context
+func (s *Storage) DeleteWithContext(ctx context.Context, key string) error {
 	if len(key) <= 0 {
 		return nil
 	}
 
-	ctx, cancel := s.requestContext()
-	defer cancel()
-
 	_, err := s.svc.DeleteObject(ctx, &s3.DeleteObjectInput{
 		Bucket: &s.bucket,
 		Key:    aws.String(key),
@@ -137,11 +144,16 @@ func (s *Storage) Delete(key string) error {
 	return err
 }
 
-// Reset all entries, including unexpired
-func (s *Storage) Reset() error {
+// Delete deletes entry by key
+func (s *Storage) Delete(key string) error {
 	ctx, cancel := s.requestContext()
 	defer cancel()
 
+	return s.DeleteWithContext(ctx, key)
+}
+
+// ResetWithContext resets all entries, including unexpired ones with context
+func (s *Storage) ResetWithContext(ctx context.Context) error {
 	paginator := s3.NewListObjectsV2Paginator(s.svc, &s3.ListObjectsV2Input{
 		Bucket: &s.bucket,
 	})
@@ -173,6 +185,14 @@ func (s *Storage) Reset() error {
 	return nil
 }
 
+// Reset resets all entries, including unexpired ones
+func (s *Storage) Reset() error {
+	ctx, cancel := s.requestContext()
+	defer cancel()
+
+	return s.ResetWithContext(ctx)
+}
+
 // Close the database
 func (s *Storage) Close() error {
 	return nil

s3/s3_test.go

@@ -1,6 +1,7 @@
 package s3
 
 import (
+	"context"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -19,6 +20,19 @@ func Test_S3_Set(t *testing.T) {
 	require.NoError(t, err)
 }
 
+func Test_S3_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_S3_Set_Override(t *testing.T) {
 	var (
 		key = "john"
@@ -52,6 +66,23 @@ func Test_S3_Get(t *testing.T) {
 	require.Equal(t, val, result)
 }
 
+func Test_S3_GetWithContext(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()
+
+	result, err := testStore.GetWithContext(ctx, key)
+	require.ErrorIs(t, err, context.Canceled)
+	require.Zero(t, len(result))
+}
+
 func Test_S3_Get_NotExist(t *testing.T) {
 	testStore := newTestStore(t)
 	defer testStore.Close()
@@ -81,6 +112,26 @@ func Test_S3_Delete(t *testing.T) {
 	require.Zero(t, len(result))
 }
 
+func Test_S3_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_S3_Reset(t *testing.T) {
 	val := []byte("doe")
 
@@ -105,6 +156,30 @@ func Test_S3_Reset(t *testing.T) {
 	require.Zero(t, len(result))
 }
 
+func Test_S3_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_S3_Close(t *testing.T) {
 	testStore := newTestStore(t)
 	require.NoError(t, testStore.Close())