From 2192ac6cfe41990142b61b3839727d0681e6ff44 Mon Sep 17 00:00:00 2001 From: Muhammed Efe Cetin Date: Tue, 11 Feb 2025 15:27:15 +0300 Subject: [PATCH] azureblob: add support for context management --- azureblob/azureblob.go | 65 +++++++++++++++++++++---------- azureblob/azureblob_test.go | 77 +++++++++++++++++++++++++++++++++++++ 2 files changed, 122 insertions(+), 20 deletions(-) diff --git a/azureblob/azureblob.go b/azureblob/azureblob.go index cae22a61..b6ba7bcb 100644 --- a/azureblob/azureblob.go +++ b/azureblob/azureblob.go @@ -48,13 +48,12 @@ func New(config ...Config) *Storage { return storage } -// Get value by key -func (s *Storage) Get(key string) ([]byte, error) { +// GetWithContext gets value by key +func (s *Storage) GetWithContext(ctx context.Context, key string) ([]byte, error) { if len(key) <= 0 { return nil, nil } - ctx, cancel := s.requestContext() - defer cancel() + resp, err := s.client.DownloadStream(ctx, s.container, key, nil) if err != nil { return []byte{}, err @@ -63,55 +62,81 @@ func (s *Storage) Get(key string) ([]byte, error) { return data, 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 sets key with value +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.client.UploadBuffer(ctx, s.container, key, val, nil) return err } -// Delete entry by key -func (s *Storage) Delete(key string) error { +// Set sets key with value +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 +func (s *Storage) DeleteWithContext(ctx context.Context, key string) error { if len(key) <= 0 { return nil } - ctx, cancel := s.requestContext() - defer cancel() _, err := s.client.DeleteBlob(ctx, s.container, key, nil) return err } -// Reset all entries -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 +func (s *Storage) ResetWithContext(ctx context.Context) error { //_, err := s.client.DeleteContainer(ctx, s.container, nil) //return err pager := s.client.NewListBlobsFlatPager(s.container, nil) - errCounter := 0 + for pager.More() { resp, err := pager.NextPage(ctx) if err != nil { - errCounter = errCounter + 1 + return err } + for _, v := range resp.Segment.BlobItems { _, err = s.client.DeleteBlob(ctx, s.container, *v.Name, nil) if err != nil { - errCounter = errCounter + 1 + return err } } } - if errCounter > 0 { - return fmt.Errorf("%d errors occured while resetting", errCounter) - } + return nil } +// Reset resets all entries +func (s *Storage) Reset() error { + ctx, cancel := s.requestContext() + defer cancel() + + return s.ResetWithContext(ctx) +} + // Conn returns storage client func (s *Storage) Conn() *azblob.Client { return s.client diff --git a/azureblob/azureblob_test.go b/azureblob/azureblob_test.go index ddbdfb9e..0abf2c2c 100644 --- a/azureblob/azureblob_test.go +++ b/azureblob/azureblob_test.go @@ -63,6 +63,23 @@ func Test_AzureBlob_Get(t *testing.T) { require.Equal(t, val, result) } +func Test_AzureBlob_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_AzureBlob_Set(t *testing.T) { var ( key = "john" @@ -76,6 +93,19 @@ func Test_AzureBlob_Set(t *testing.T) { require.NoError(t, err) } +func Test_AzureBlob_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_AzureBlob_Delete(t *testing.T) { var ( key = "john" @@ -101,6 +131,26 @@ func Test_AzureBlob_Delete(t *testing.T) { require.Zero(t, len(result)) } +func Test_AzureBlob_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_AzureBlob_Override(t *testing.T) { var ( key = "john" @@ -172,6 +222,33 @@ func Test_AzureBlob_Conn(t *testing.T) { require.True(t, testStore.Conn() != nil) } +func Test_AzureBlob_ResetWithContext(t *testing.T) { + val := []byte("doe") + + testStore := newTestStore(t) + defer testStore.Close() + + 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_AzureBlob_Close(t *testing.T) { testStore := newTestStore(t) require.NoError(t, testStore.Close())