mongodb: add support for context management

This commit is contained in:
Muhammed Efe Cetin
2025-02-09 21:45:08 +03:00
parent 134e774e80
commit c1e526d31d
2 changed files with 107 additions and 10 deletions

View File

@@ -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

View File

@@ -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())