diff --git a/mongodb/mongodb.go b/mongodb/mongodb.go index 4164551f..f424ea3a 100644 --- a/mongodb/mongodb.go +++ b/mongodb/mongodb.go @@ -121,12 +121,12 @@ func New(config ...Config) *Storage { return store } -// 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) { if len(key) <= 0 { return nil, nil } - res := s.col.FindOne(context.Background(), bson.M{"key": key}) + res := s.col.FindOne(ctx, bson.M{"key": key}) item := s.acquireItem() if err := res.Err(); err != nil { @@ -149,11 +149,16 @@ func (s *Storage) Get(key string) ([]byte, error) { return item.Value, nil } -// Set key with value, replace if document exits +// Get gets value by key +func (s *Storage) Get(key string) ([]byte, error) { + return s.GetWithContext(context.Background(), key) +} + +// SetWithContext sets key with value, replace if document exits with context // // document will be remove automatically if exp is set, based on MongoDB TTL Indexes // Set key with value -func (s *Storage) Set(key string, val []byte, exp time.Duration) error { +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 @@ -167,25 +172,43 @@ func (s *Storage) Set(key string, val []byte, exp time.Duration) error { if exp != 0 { item.Expiration = time.Now().Add(exp).UTC() } - _, err := s.col.ReplaceOne(context.Background(), filter, item, options.Replace().SetUpsert(true)) + _, err := s.col.ReplaceOne(ctx, filter, item, options.Replace().SetUpsert(true)) s.releaseItem(item) return err } -// Delete document by key -func (s *Storage) Delete(key string) error { +// Set sets key with value, replace if document exits +// +// document will be remove automatically if exp is set, based on MongoDB TTL Indexes +// 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 deletes document by key with 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.col.DeleteOne(context.Background(), bson.M{"key": key}) + _, err := s.col.DeleteOne(ctx, bson.M{"key": key}) return err } +// Delete deletes document by key +func (s *Storage) Delete(key string) error { + return s.DeleteWithContext(context.Background(), key) +} + +// Reset all keys by drop collection with context +func (s *Storage) ResetWithContext(ctx context.Context) error { + return s.col.Drop(ctx) +} + // Reset all keys by drop collection func (s *Storage) Reset() error { - return s.col.Drop(context.Background()) + return s.ResetWithContext(context.Background()) } // Close the database diff --git a/mongodb/mongodb_test.go b/mongodb/mongodb_test.go index 287519d7..edd90091 100644 --- a/mongodb/mongodb_test.go +++ b/mongodb/mongodb_test.go @@ -66,6 +66,19 @@ func Test_MongoDB_Set(t *testing.T) { require.NoError(t, err) } +func Test_MongoDB_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_MongoDB_Set_Override(t *testing.T) { var ( key = "john" @@ -99,6 +112,23 @@ func Test_MongoDB_Get(t *testing.T) { require.Equal(t, val, result) } +func Test_MongoDB_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_MongoDB_Set_Expiration(t *testing.T) { var ( key = "john" @@ -159,6 +189,26 @@ func Test_MongoDB_Delete(t *testing.T) { require.Zero(t, len(result)) } +func Test_MongoDB_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_MongoDB_Reset(t *testing.T) { val := []byte("doe") @@ -183,6 +233,30 @@ func Test_MongoDB_Reset(t *testing.T) { require.Zero(t, len(result)) } +func Test_MongoDB_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_MongoDB_Close(t *testing.T) { testStore := newTestStore(t) require.NoError(t, testStore.Close())