mirror of
https://github.com/gofiber/storage.git
synced 2025-10-05 16:48:25 +08:00
mongodb: add support for context management
This commit is contained in:
@@ -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
|
||||
|
@@ -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())
|
||||
|
Reference in New Issue
Block a user