diff --git a/arangodb/README.md b/arangodb/README.md index ae3a6c39..74f5a680 100644 --- a/arangodb/README.md +++ b/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 diff --git a/arangodb/arangodb.go b/arangodb/arangodb.go index 19a72300..ce3b59fa 100644 --- a/arangodb/arangodb.go +++ b/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 diff --git a/arangodb/arangodb_test.go b/arangodb/arangodb_test.go index ef0ef4f1..d5ad6d31 100644 --- a/arangodb/arangodb_test.go +++ b/arangodb/arangodb_test.go @@ -1,6 +1,7 @@ package arangodb import ( + "context" "testing" "time" @@ -21,6 +22,19 @@ func Test_ArangoDB_Set(t *testing.T) { require.NoError(t, err) } +func Test_ArangoDB_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_ArangoDB_Upsert(t *testing.T) { var ( key = "john" @@ -48,6 +62,23 @@ 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") + ) + + 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" @@ -92,6 +123,26 @@ 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") + ) + + 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") @@ -113,6 +164,30 @@ func Test_ArangoDB_Reset(t *testing.T) { require.Zero(t, len(result)) } +func Test_ArangoDB_ResetWithContext(t *testing.T) { + val := []byte("doe") + + 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")