diff --git a/s3/s3.go b/s3/s3.go index 2e189f4c..6d77ba11 100644 --- a/s3/s3.go +++ b/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 diff --git a/s3/s3_test.go b/s3/s3_test.go index fe2ad8f2..1ae3b70c 100644 --- a/s3/s3_test.go +++ b/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())