enhancement: add context management support to storages

aerospike/aerospike.go

@@ -1,6 +1,7 @@
 package aerospike
 
 import (
+	"context"
 	"log"
 	"time"
 
@@ -216,6 +217,11 @@ func (s *Storage) Get(key string) ([]byte, error) {
 	return data, nil
 }
 
+// GetWithContext gets value by key (dummy context support)
+func (s *Storage) GetWithContext(ctx context.Context, key string) ([]byte, error) {
+	return s.Get(key)
+}
+
 // Set key with value
 func (s *Storage) Set(key string, val []byte, exp time.Duration) error {
 	k, err := aerospike.NewKey(s.namespace, s.setName, key)
@@ -242,6 +248,11 @@ func (s *Storage) Set(key string, val []byte, exp time.Duration) error {
 	return s.client.Put(writePolicy, k, bins)
 }
 
+// SetWithContext sets value by key (dummy context support)
+func (s *Storage) SetWithContext(ctx context.Context, key string, val []byte, exp time.Duration) error {
+	return s.Set(key, val, exp)
+}
+
 // Delete key
 func (s *Storage) Delete(key string) error {
 	k, err := aerospike.NewKey(s.namespace, s.setName, key)
@@ -253,6 +264,11 @@ func (s *Storage) Delete(key string) error {
 	return err
 }
 
+// DeleteWithContext deletes key (dummy context support)
+func (s *Storage) DeleteWithContext(ctx context.Context, key string) error {
+	return s.Delete(key)
+}
+
 // Reset all keys
 func (s *Storage) Reset() error {
 	// Use ScanAll which returns a Recordset
@@ -293,6 +309,11 @@ func (s *Storage) Reset() error {
 	return nil
 }
 
+// ResetWithContext resets all keys (dummy context support)
+func (s *Storage) ResetWithContext(ctx context.Context) error {
+	return s.Reset()
+}
+
 // Close the storage
 func (s *Storage) Close() error {
 	s.client.Close()

arangodb/README.md

@@ -21,9 +21,13 @@ A ArangoDB storage driver using `arangodb/go-driver` and [arangodb/go-driver](ht
 ### Signatures
 ```go
 func New(config ...Config) Storage
+func (s *Storage) GetWithContext(ctx context.Context, key string) ([]byte, error)
 func (s *Storage) Get(key string) ([]byte, error)
+func (s *Storage) SetWithContext(ctx context.Context, key string, val []byte, exp time.Duration) error
 func (s *Storage) Set(key string, val []byte, exp time.Duration) error
+func (s *Storage) DeleteWithContext(ctx context.Context, key string) error
 func (s *Storage) Delete(key string) error
+func (s *Storage) ResetWithContext(ctx context.Context) error
 func (s *Storage) Reset() error
 func (s *Storage) Close() error
 func (s *Storage) Conn() driver.Client

arangodb/arangodb.go

@@ -116,14 +116,12 @@ func New(config ...Config) *Storage {
 	return store
 }
 
-// Get value by key
-func (s *Storage) Get(key string) ([]byte, error) {
+// GetWithContext value by key with given context
+func (s *Storage) GetWithContext(ctx context.Context, key string) ([]byte, error) {
 	if len(key) <= 0 {
 		return nil, nil
 	}
 
-	ctx := context.Background()
-
 	// Check if the document exists
 	// to avoid errors later
 	exists, err := s.collection.DocumentExists(ctx, key)
@@ -151,8 +149,13 @@ func (s *Storage) Get(key string) ([]byte, error) {
 	return utils.UnsafeBytes(model.Val), nil
 }
 
-// Set key with value
-func (s *Storage) Set(key string, val []byte, exp time.Duration) error {
+// Get value by key
+func (s *Storage) Get(key string) ([]byte, error) {
+	return s.GetWithContext(context.Background(), key)
+}
+
+// SetWithContext key with value with given context
+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
@@ -169,7 +172,6 @@ func (s *Storage) Set(key string, val []byte, exp time.Duration) error {
 		Val: valStr,
 		Exp: expireAt,
 	}
-	ctx := context.Background()
 
 	// Arango does not support documents with the same key
 	// So we need to check if the document exists
@@ -188,20 +190,35 @@ func (s *Storage) Set(key string, val []byte, exp time.Duration) error {
 	return err
 }
 
-// Delete value by key
-func (s *Storage) Delete(key string) error {
+// Set key with value
+func (s *Storage) Set(key string, val []byte, exp time.Duration) error {
+	return s.SetWithContext(context.Background(), key, val, exp)
+}
+
+// DeleteWithContext value by key with given context
+func (s *Storage) DeleteWithContext(ctx context.Context, key string) error {
 	// Ain't Nobody Got Time For That
 	if len(key) <= 0 {
 		return nil
 	}
-	_, err := s.collection.RemoveDocument(context.Background(), key)
+	_, err := s.collection.RemoveDocument(ctx, key)
 	return err
 }
 
+// Delete value by key
+func (s *Storage) Delete(key string) error {
+	return s.DeleteWithContext(context.Background(), key)
+}
+
+// ResetWithContext all keys with given context
+func (s *Storage) ResetWithContext(ctx context.Context) error {
+	return s.collection.Truncate(ctx)
+}
+
 // Reset all keys
 // truncate the collection
 func (s *Storage) Reset() error {
-	return s.collection.Truncate(context.Background())
+	return s.ResetWithContext(context.Background())
 }
 
 // Close the database

arangodb/arangodb_test.go

@@ -67,6 +67,22 @@ func Test_ArangoDB_Set(t *testing.T) {
 	require.NoError(t, err)
 }
 
+func Test_ArangoDB_SetWithContext(t *testing.T) {
+	var (
+		key = "john"
+		val = []byte("doe")
+	)
+
+	testStore := newTestStore(t)
+	defer testStore.Close()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	err := testStore.SetWithContext(ctx, key, val, 0)
+	require.ErrorIs(t, err, context.Canceled)
+}
+
 func Test_ArangoDB_Upsert(t *testing.T) {
 	var (
 		key = "john"
@@ -100,6 +116,26 @@ func Test_ArangoDB_Get(t *testing.T) {
 	require.Equal(t, val, result)
 }
 
+func Test_ArangoDB_GetWithContext(t *testing.T) {
+	var (
+		key = "john"
+		val = []byte("doe")
+	)
+
+	testStore := newTestStore(t)
+	defer testStore.Close()
+
+	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_ArangoDB_Set_Expiration(t *testing.T) {
 	var (
 		key = "john"
@@ -156,6 +192,29 @@ func Test_ArangoDB_Delete(t *testing.T) {
 	require.Zero(t, len(result))
 }
 
+func Test_ArangoDB_DeleteWithContext(t *testing.T) {
+	var (
+		key = "john"
+		val = []byte("doe")
+	)
+
+	testStore := newTestStore(t)
+	defer testStore.Close()
+
+	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_ArangoDB_Reset(t *testing.T) {
 	val := []byte("doe")
 
@@ -180,6 +239,33 @@ func Test_ArangoDB_Reset(t *testing.T) {
 	require.Zero(t, len(result))
 }
 
+func Test_ArangoDB_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.Equal(t, err, nil)
+
+	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_ArangoDB_Non_UTF8(t *testing.T) {
 	val := []byte("0xF5")
 

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