mirror of
https://github.com/gofiber/storage.git
synced 2025-10-26 18:00:26 +08:00
add context support to more storages
This commit is contained in:
@@ -1,6 +1,7 @@
|
|||||||
package badger
|
package badger
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"context"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
"github.com/dgraph-io/badger/v3"
|
"github.com/dgraph-io/badger/v3"
|
||||||
@@ -72,6 +73,12 @@ func (s *Storage) Get(key string) ([]byte, error) {
|
|||||||
return data, err
|
return data, err
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// GetWithContext gets value by key.
|
||||||
|
// Note: This method is not used in the current implementation, but is included to satisfy the Storage interface.
|
||||||
|
func (s *Storage) GetWithContext(ctx context.Context, key string) ([]byte, error) {
|
||||||
|
return s.Get(key)
|
||||||
|
}
|
||||||
|
|
||||||
// Set key with value
|
// Set key with value
|
||||||
func (s *Storage) Set(key string, val []byte, exp time.Duration) error {
|
func (s *Storage) Set(key string, val []byte, exp time.Duration) error {
|
||||||
// Ain't Nobody Got Time For That
|
// Ain't Nobody Got Time For That
|
||||||
@@ -88,6 +95,12 @@ func (s *Storage) Set(key string, val []byte, exp time.Duration) error {
|
|||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// SetWithContext sets key with value.
|
||||||
|
// Note: This method is not used in the current implementation, but is included to satisfy the Storage interface.
|
||||||
|
func (s *Storage) SetWithContext(ctx context.Context, key string, val []byte, exp time.Duration) error {
|
||||||
|
return s.Set(key, val, exp)
|
||||||
|
}
|
||||||
|
|
||||||
// Delete key by key
|
// Delete key by key
|
||||||
func (s *Storage) Delete(key string) error {
|
func (s *Storage) Delete(key string) error {
|
||||||
// Ain't Nobody Got Time For That
|
// Ain't Nobody Got Time For That
|
||||||
@@ -99,11 +112,23 @@ func (s *Storage) Delete(key string) error {
|
|||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// DeleteWithContext deletes key by key.
|
||||||
|
// Note: This method is not used in the current implementation, but is included to satisfy the Storage interface.
|
||||||
|
func (s *Storage) DeleteWithContext(ctx context.Context, key string) error {
|
||||||
|
return s.Delete(key)
|
||||||
|
}
|
||||||
|
|
||||||
// Reset all keys
|
// Reset all keys
|
||||||
func (s *Storage) Reset() error {
|
func (s *Storage) Reset() error {
|
||||||
return s.db.DropAll()
|
return s.db.DropAll()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// ResetWithContext resets all keys.
|
||||||
|
// Note: This method is not used in the current implementation, but is included to satisfy the Storage interface.
|
||||||
|
func (s *Storage) ResetWithContext(ctx context.Context) error {
|
||||||
|
return s.db.DropAll()
|
||||||
|
}
|
||||||
|
|
||||||
// Close the memory storage
|
// Close the memory storage
|
||||||
func (s *Storage) Close() error {
|
func (s *Storage) Close() error {
|
||||||
s.done <- struct{}{}
|
s.done <- struct{}{}
|
||||||
|
|||||||
@@ -1,6 +1,7 @@
|
|||||||
package bbolt
|
package bbolt
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"context"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
"github.com/gofiber/utils/v2"
|
"github.com/gofiber/utils/v2"
|
||||||
@@ -62,6 +63,11 @@ func (s *Storage) Get(key string) ([]byte, error) {
|
|||||||
return value, err
|
return value, err
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// 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
|
// Set key with value
|
||||||
func (s *Storage) Set(key string, value []byte, exp time.Duration) error {
|
func (s *Storage) Set(key string, value []byte, exp time.Duration) error {
|
||||||
if len(key) <= 0 || len(value) <= 0 {
|
if len(key) <= 0 || len(value) <= 0 {
|
||||||
@@ -70,11 +76,15 @@ func (s *Storage) Set(key string, value []byte, exp time.Duration) error {
|
|||||||
|
|
||||||
return s.conn.Update(func(tx *bbolt.Tx) error {
|
return s.conn.Update(func(tx *bbolt.Tx) error {
|
||||||
b := tx.Bucket(utils.UnsafeBytes(s.bucket))
|
b := tx.Bucket(utils.UnsafeBytes(s.bucket))
|
||||||
|
|
||||||
return b.Put(utils.UnsafeBytes(key), value)
|
return b.Put(utils.UnsafeBytes(key), value)
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// SetWithContext sets key with value (dummy context support)
|
||||||
|
func (s *Storage) SetWithContext(ctx context.Context, key string, value []byte, exp time.Duration) error {
|
||||||
|
return s.Set(key, value, exp)
|
||||||
|
}
|
||||||
|
|
||||||
// Delete entry by key
|
// Delete entry by key
|
||||||
func (s *Storage) Delete(key string) error {
|
func (s *Storage) Delete(key string) error {
|
||||||
if len(key) <= 0 {
|
if len(key) <= 0 {
|
||||||
@@ -83,28 +93,36 @@ func (s *Storage) Delete(key string) error {
|
|||||||
|
|
||||||
return s.conn.Update(func(tx *bbolt.Tx) error {
|
return s.conn.Update(func(tx *bbolt.Tx) error {
|
||||||
b := tx.Bucket(utils.UnsafeBytes(s.bucket))
|
b := tx.Bucket(utils.UnsafeBytes(s.bucket))
|
||||||
|
|
||||||
return b.Delete(utils.UnsafeBytes(key))
|
return b.Delete(utils.UnsafeBytes(key))
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// DeleteWithContext deletes key by key (dummy context support)
|
||||||
|
func (s *Storage) DeleteWithContext(ctx context.Context, key string) error {
|
||||||
|
return s.Delete(key)
|
||||||
|
}
|
||||||
|
|
||||||
// Reset all entries
|
// Reset all entries
|
||||||
func (s *Storage) Reset() error {
|
func (s *Storage) Reset() error {
|
||||||
return s.conn.Update(func(tx *bbolt.Tx) error {
|
return s.conn.Update(func(tx *bbolt.Tx) error {
|
||||||
b := tx.Bucket(utils.UnsafeBytes(s.bucket))
|
b := tx.Bucket(utils.UnsafeBytes(s.bucket))
|
||||||
|
|
||||||
return b.ForEach(func(k, _ []byte) error {
|
return b.ForEach(func(k, _ []byte) error {
|
||||||
return b.Delete(k)
|
return b.Delete(k)
|
||||||
})
|
})
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// ResetWithContext resets all entries (dummy context support)
|
||||||
|
func (s *Storage) ResetWithContext(ctx context.Context) error {
|
||||||
|
return s.Reset()
|
||||||
|
}
|
||||||
|
|
||||||
// Close the database
|
// Close the database
|
||||||
func (s *Storage) Close() error {
|
func (s *Storage) Close() error {
|
||||||
return s.conn.Close()
|
return s.conn.Close()
|
||||||
}
|
}
|
||||||
|
|
||||||
// Return database client
|
// Conn returns the database client
|
||||||
func (s *Storage) Conn() *bbolt.DB {
|
func (s *Storage) Conn() *bbolt.DB {
|
||||||
return s.conn
|
return s.conn
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1,6 +1,7 @@
|
|||||||
package memcache
|
package memcache
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"context"
|
||||||
"strings"
|
"strings"
|
||||||
"sync"
|
"sync"
|
||||||
"time"
|
"time"
|
||||||
@@ -68,10 +69,13 @@ func (s *Storage) Get(key string) ([]byte, error) {
|
|||||||
return item.Value, nil
|
return item.Value, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
// Set key with value
|
// 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
|
// Set key with value
|
||||||
func (s *Storage) Set(key string, val []byte, exp time.Duration) error {
|
func (s *Storage) Set(key string, val []byte, exp time.Duration) error {
|
||||||
// Ain't Nobody Got Time For That
|
|
||||||
if len(key) <= 0 || len(val) <= 0 {
|
if len(key) <= 0 || len(val) <= 0 {
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
@@ -87,20 +91,34 @@ func (s *Storage) Set(key string, val []byte, exp time.Duration) error {
|
|||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// SetWithContext sets key with value (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 by key
|
// Delete key by key
|
||||||
func (s *Storage) Delete(key string) error {
|
func (s *Storage) Delete(key string) error {
|
||||||
// Ain't Nobody Got Time For That
|
|
||||||
if len(key) <= 0 {
|
if len(key) <= 0 {
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
return s.db.Delete(key)
|
return s.db.Delete(key)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// DeleteWithContext deletes key by key (dummy context support)
|
||||||
|
func (s *Storage) DeleteWithContext(ctx context.Context, key string) error {
|
||||||
|
return s.Delete(key)
|
||||||
|
}
|
||||||
|
|
||||||
// Reset all keys
|
// Reset all keys
|
||||||
func (s *Storage) Reset() error {
|
func (s *Storage) Reset() error {
|
||||||
return s.db.DeleteAll()
|
return s.db.DeleteAll()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// ResetWithContext resets all keys (dummy context support)
|
||||||
|
func (s *Storage) ResetWithContext(ctx context.Context) error {
|
||||||
|
return s.Reset()
|
||||||
|
}
|
||||||
|
|
||||||
// Close the database
|
// Close the database
|
||||||
func (s *Storage) Close() error {
|
func (s *Storage) Close() error {
|
||||||
return nil
|
return nil
|
||||||
@@ -111,7 +129,7 @@ func (s *Storage) acquireItem() *mc.Item {
|
|||||||
return s.items.Get().(*mc.Item)
|
return s.items.Get().(*mc.Item)
|
||||||
}
|
}
|
||||||
|
|
||||||
// Release item from pool
|
// Release item back to pool
|
||||||
func (s *Storage) releaseItem(item *mc.Item) {
|
func (s *Storage) releaseItem(item *mc.Item) {
|
||||||
if item != nil {
|
if item != nil {
|
||||||
item.Key = ""
|
item.Key = ""
|
||||||
|
|||||||
@@ -1,6 +1,7 @@
|
|||||||
package memory
|
package memory
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"context"
|
||||||
"sync"
|
"sync"
|
||||||
"sync/atomic"
|
"sync/atomic"
|
||||||
"time"
|
"time"
|
||||||
@@ -17,9 +18,8 @@ type Storage struct {
|
|||||||
}
|
}
|
||||||
|
|
||||||
type entry struct {
|
type entry struct {
|
||||||
data []byte
|
data []byte
|
||||||
// max value is 4294967295 -> Sun Feb 07 2106 06:28:15 GMT+0000
|
expiry uint32 // max value is 4294967295 -> Sun Feb 07 2106 06:28:15 GMT+0000
|
||||||
expiry uint32
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// New creates a new memory storage
|
// New creates a new memory storage
|
||||||
@@ -49,16 +49,20 @@ func (s *Storage) Get(key string) ([]byte, error) {
|
|||||||
s.mux.RLock()
|
s.mux.RLock()
|
||||||
v, ok := s.db[key]
|
v, ok := s.db[key]
|
||||||
s.mux.RUnlock()
|
s.mux.RUnlock()
|
||||||
if !ok || v.expiry != 0 && v.expiry <= atomic.LoadUint32(&internal.Timestamp) {
|
if !ok || (v.expiry != 0 && v.expiry <= atomic.LoadUint32(&internal.Timestamp)) {
|
||||||
return nil, nil
|
return nil, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
return v.data, nil
|
return v.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
|
// Set key with value
|
||||||
func (s *Storage) Set(key string, val []byte, exp time.Duration) error {
|
func (s *Storage) Set(key string, val []byte, exp time.Duration) error {
|
||||||
// Ain't Nobody Got Time For That
|
|
||||||
if len(key) <= 0 || len(val) <= 0 {
|
if len(key) <= 0 || len(val) <= 0 {
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
@@ -75,9 +79,13 @@ func (s *Storage) Set(key string, val []byte, exp time.Duration) error {
|
|||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// 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 by key
|
// Delete key by key
|
||||||
func (s *Storage) Delete(key string) error {
|
func (s *Storage) Delete(key string) error {
|
||||||
// Ain't Nobody Got Time For That
|
|
||||||
if len(key) <= 0 {
|
if len(key) <= 0 {
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
@@ -87,6 +95,11 @@ func (s *Storage) Delete(key string) error {
|
|||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// DeleteWithContext deletes key (dummy context support)
|
||||||
|
func (s *Storage) DeleteWithContext(ctx context.Context, key string) error {
|
||||||
|
return s.Delete(key)
|
||||||
|
}
|
||||||
|
|
||||||
// Reset all keys
|
// Reset all keys
|
||||||
func (s *Storage) Reset() error {
|
func (s *Storage) Reset() error {
|
||||||
ndb := make(map[string]entry)
|
ndb := make(map[string]entry)
|
||||||
@@ -96,6 +109,11 @@ func (s *Storage) Reset() error {
|
|||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// ResetWithContext resets all keys (dummy context support)
|
||||||
|
func (s *Storage) ResetWithContext(ctx context.Context) error {
|
||||||
|
return s.Reset()
|
||||||
|
}
|
||||||
|
|
||||||
// Close the memory storage
|
// Close the memory storage
|
||||||
func (s *Storage) Close() error {
|
func (s *Storage) Close() error {
|
||||||
s.done <- struct{}{}
|
s.done <- struct{}{}
|
||||||
@@ -122,8 +140,6 @@ func (s *Storage) gc() {
|
|||||||
}
|
}
|
||||||
s.mux.RUnlock()
|
s.mux.RUnlock()
|
||||||
s.mux.Lock()
|
s.mux.Lock()
|
||||||
// Double-checked locking.
|
|
||||||
// We might have replaced the item in the meantime.
|
|
||||||
for i := range expired {
|
for i := range expired {
|
||||||
v := s.db[expired[i]]
|
v := s.db[expired[i]]
|
||||||
if v.expiry != 0 && v.expiry <= ts {
|
if v.expiry != 0 && v.expiry <= ts {
|
||||||
@@ -135,14 +151,14 @@ func (s *Storage) gc() {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// Return database client
|
// Conn returns database client
|
||||||
func (s *Storage) Conn() map[string]entry {
|
func (s *Storage) Conn() map[string]entry {
|
||||||
s.mux.RLock()
|
s.mux.RLock()
|
||||||
defer s.mux.RUnlock()
|
defer s.mux.RUnlock()
|
||||||
return s.db
|
return s.db
|
||||||
}
|
}
|
||||||
|
|
||||||
// Return all the keys
|
// Keys returns all the keys
|
||||||
func (s *Storage) Keys() ([][]byte, error) {
|
func (s *Storage) Keys() ([][]byte, error) {
|
||||||
s.mux.RLock()
|
s.mux.RLock()
|
||||||
defer s.mux.RUnlock()
|
defer s.mux.RUnlock()
|
||||||
@@ -154,13 +170,11 @@ func (s *Storage) Keys() ([][]byte, error) {
|
|||||||
ts := atomic.LoadUint32(&internal.Timestamp)
|
ts := atomic.LoadUint32(&internal.Timestamp)
|
||||||
keys := make([][]byte, 0, len(s.db))
|
keys := make([][]byte, 0, len(s.db))
|
||||||
for key, v := range s.db {
|
for key, v := range s.db {
|
||||||
// Filter out the expired keys
|
|
||||||
if v.expiry == 0 || v.expiry > ts {
|
if v.expiry == 0 || v.expiry > ts {
|
||||||
keys = append(keys, []byte(key))
|
keys = append(keys, []byte(key))
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// Double check if no valid keys were found
|
|
||||||
if len(keys) == 0 {
|
if len(keys) == 0 {
|
||||||
return nil, nil
|
return nil, nil
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1,6 +1,7 @@
|
|||||||
package mockstorage
|
package mockstorage
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"context"
|
||||||
"errors"
|
"errors"
|
||||||
"sync"
|
"sync"
|
||||||
"time"
|
"time"
|
||||||
@@ -26,31 +27,26 @@ type Entry struct {
|
|||||||
|
|
||||||
// CustomFuncs allows injecting custom behaviors for testing.
|
// CustomFuncs allows injecting custom behaviors for testing.
|
||||||
type CustomFuncs struct {
|
type CustomFuncs struct {
|
||||||
GetFunc func(key string) ([]byte, error)
|
GetFunc func(key string) ([]byte, error)
|
||||||
SetFunc func(key string, val []byte, exp time.Duration) error
|
GetWithContext func(ctx context.Context, key string) ([]byte, error)
|
||||||
DeleteFunc func(key string) error
|
SetFunc func(key string, val []byte, exp time.Duration) error
|
||||||
ResetFunc func() error
|
SetWithContext func(ctx context.Context, key string, val []byte, exp time.Duration) error
|
||||||
CloseFunc func() error
|
DeleteFunc func(key string) error
|
||||||
ConnFunc func() map[string]Entry
|
DeleteWithContext func(ctx context.Context, key string) error
|
||||||
KeysFunc func() ([][]byte, error)
|
ResetFunc func() error
|
||||||
|
ResetWithContext func(ctx context.Context) error
|
||||||
|
CloseFunc func() error
|
||||||
|
ConnFunc func() map[string]Entry
|
||||||
|
KeysFunc func() ([][]byte, error)
|
||||||
}
|
}
|
||||||
|
|
||||||
// New creates a new mock storage with optional configuration.
|
// New creates a new mock storage with optional configuration.
|
||||||
func New(config ...Config) *Storage {
|
func New(config ...Config) *Storage {
|
||||||
s := &Storage{
|
s := &Storage{
|
||||||
data: make(map[string]Entry),
|
data: make(map[string]Entry),
|
||||||
custom: &CustomFuncs{
|
custom: &CustomFuncs{}, // default no-op
|
||||||
GetFunc: nil,
|
|
||||||
SetFunc: nil,
|
|
||||||
DeleteFunc: nil,
|
|
||||||
ResetFunc: nil,
|
|
||||||
CloseFunc: nil,
|
|
||||||
ConnFunc: nil,
|
|
||||||
KeysFunc: nil,
|
|
||||||
},
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// If a config is provided and it has CustomFuncs, use them
|
|
||||||
if len(config) > 0 && config[0].CustomFuncs != nil {
|
if len(config) > 0 && config[0].CustomFuncs != nil {
|
||||||
s.custom = config[0].CustomFuncs
|
s.custom = config[0].CustomFuncs
|
||||||
}
|
}
|
||||||
@@ -78,7 +74,15 @@ func (s *Storage) Get(key string) ([]byte, error) {
|
|||||||
return e.Value, nil
|
return e.Value, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
// Set sets the value for a given key with an expiration time.
|
// GetWithContext retrieves value by key using a context (functional or fallback)
|
||||||
|
func (s *Storage) GetWithContext(ctx context.Context, key string) ([]byte, error) {
|
||||||
|
if s.custom.GetWithContext != nil {
|
||||||
|
return s.custom.GetWithContext(ctx, key)
|
||||||
|
}
|
||||||
|
return s.Get(key)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Set sets the value for a given key with expiration.
|
||||||
func (s *Storage) Set(key string, val []byte, exp time.Duration) error {
|
func (s *Storage) Set(key string, val []byte, exp time.Duration) error {
|
||||||
if s.custom.SetFunc != nil {
|
if s.custom.SetFunc != nil {
|
||||||
return s.custom.SetFunc(key, val, exp)
|
return s.custom.SetFunc(key, val, exp)
|
||||||
@@ -96,6 +100,14 @@ func (s *Storage) Set(key string, val []byte, exp time.Duration) error {
|
|||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// SetWithContext sets value using context.
|
||||||
|
func (s *Storage) SetWithContext(ctx context.Context, key string, val []byte, exp time.Duration) error {
|
||||||
|
if s.custom.SetWithContext != nil {
|
||||||
|
return s.custom.SetWithContext(ctx, key, val, exp)
|
||||||
|
}
|
||||||
|
return s.Set(key, val, exp)
|
||||||
|
}
|
||||||
|
|
||||||
// Delete removes a key from the storage.
|
// Delete removes a key from the storage.
|
||||||
func (s *Storage) Delete(key string) error {
|
func (s *Storage) Delete(key string) error {
|
||||||
if s.custom.DeleteFunc != nil {
|
if s.custom.DeleteFunc != nil {
|
||||||
@@ -104,12 +116,19 @@ func (s *Storage) Delete(key string) error {
|
|||||||
|
|
||||||
s.mu.Lock()
|
s.mu.Lock()
|
||||||
defer s.mu.Unlock()
|
defer s.mu.Unlock()
|
||||||
|
|
||||||
delete(s.data, key)
|
delete(s.data, key)
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
// Reset clears all keys from the storage.
|
// DeleteWithContext deletes key using context.
|
||||||
|
func (s *Storage) DeleteWithContext(ctx context.Context, key string) error {
|
||||||
|
if s.custom.DeleteWithContext != nil {
|
||||||
|
return s.custom.DeleteWithContext(ctx, key)
|
||||||
|
}
|
||||||
|
return s.Delete(key)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Reset clears all keys.
|
||||||
func (s *Storage) Reset() error {
|
func (s *Storage) Reset() error {
|
||||||
if s.custom.ResetFunc != nil {
|
if s.custom.ResetFunc != nil {
|
||||||
return s.custom.ResetFunc()
|
return s.custom.ResetFunc()
|
||||||
@@ -117,22 +136,27 @@ func (s *Storage) Reset() error {
|
|||||||
|
|
||||||
s.mu.Lock()
|
s.mu.Lock()
|
||||||
defer s.mu.Unlock()
|
defer s.mu.Unlock()
|
||||||
|
|
||||||
s.data = make(map[string]Entry)
|
s.data = make(map[string]Entry)
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
// Close closes the storage (no-op for mock).
|
// ResetWithContext resets storage using context.
|
||||||
|
func (s *Storage) ResetWithContext(ctx context.Context) error {
|
||||||
|
if s.custom.ResetWithContext != nil {
|
||||||
|
return s.custom.ResetWithContext(ctx)
|
||||||
|
}
|
||||||
|
return s.Reset()
|
||||||
|
}
|
||||||
|
|
||||||
|
// Close closes the mock storage (no-op).
|
||||||
func (s *Storage) Close() error {
|
func (s *Storage) Close() error {
|
||||||
if s.custom.CloseFunc != nil {
|
if s.custom.CloseFunc != nil {
|
||||||
return s.custom.CloseFunc()
|
return s.custom.CloseFunc()
|
||||||
}
|
}
|
||||||
|
|
||||||
// No resources to clean up in mock
|
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
// Conn returns the internal data map (for testing purposes).
|
// Conn returns internal map.
|
||||||
func (s *Storage) Conn() map[string]Entry {
|
func (s *Storage) Conn() map[string]Entry {
|
||||||
if s.custom.ConnFunc != nil {
|
if s.custom.ConnFunc != nil {
|
||||||
return s.custom.ConnFunc()
|
return s.custom.ConnFunc()
|
||||||
@@ -148,7 +172,7 @@ func (s *Storage) Conn() map[string]Entry {
|
|||||||
return copyData
|
return copyData
|
||||||
}
|
}
|
||||||
|
|
||||||
// Keys returns all keys in the storage.
|
// Keys returns all keys.
|
||||||
func (s *Storage) Keys() ([][]byte, error) {
|
func (s *Storage) Keys() ([][]byte, error) {
|
||||||
if s.custom.KeysFunc != nil {
|
if s.custom.KeysFunc != nil {
|
||||||
return s.custom.KeysFunc()
|
return s.custom.KeysFunc()
|
||||||
@@ -164,7 +188,7 @@ func (s *Storage) Keys() ([][]byte, error) {
|
|||||||
return keys, nil
|
return keys, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
// SetCustomFuncs allows setting custom function implementations.
|
// SetCustomFuncs allows runtime injection of function implementations.
|
||||||
func (s *Storage) SetCustomFuncs(custom *CustomFuncs) {
|
func (s *Storage) SetCustomFuncs(custom *CustomFuncs) {
|
||||||
s.custom = custom
|
s.custom = custom
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -2,6 +2,7 @@ package mockstorage
|
|||||||
|
|
||||||
import (
|
import (
|
||||||
"bytes"
|
"bytes"
|
||||||
|
"context"
|
||||||
"errors"
|
"errors"
|
||||||
"testing"
|
"testing"
|
||||||
"time"
|
"time"
|
||||||
@@ -274,3 +275,111 @@ func TestStorageConnAndKeys(t *testing.T) {
|
|||||||
t.Errorf("Keys() = %v, want %v", keys, [][]byte{[]byte("key1")})
|
t.Errorf("Keys() = %v, want %v", keys, [][]byte{[]byte("key1")})
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestGetWithContext(t *testing.T) {
|
||||||
|
store := New()
|
||||||
|
|
||||||
|
// fallback to Get
|
||||||
|
_ = store.Set("key1", []byte("val1"), 0)
|
||||||
|
val, err := store.GetWithContext(context.Background(), "key1")
|
||||||
|
if err != nil || !bytes.Equal(val, []byte("val1")) {
|
||||||
|
t.Errorf("GetWithContext fallback failed: got %v, err %v", val, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// custom override
|
||||||
|
store.SetCustomFuncs(&CustomFuncs{
|
||||||
|
GetWithContext: func(ctx context.Context, key string) ([]byte, error) {
|
||||||
|
if key == "override" {
|
||||||
|
return []byte("ctx-value"), nil
|
||||||
|
}
|
||||||
|
return nil, errors.New("not found")
|
||||||
|
},
|
||||||
|
})
|
||||||
|
val, err = store.GetWithContext(context.TODO(), "override")
|
||||||
|
if err != nil || !bytes.Equal(val, []byte("ctx-value")) {
|
||||||
|
t.Errorf("GetWithContext custom failed: got %v, err %v", val, err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSetWithContext(t *testing.T) {
|
||||||
|
store := New()
|
||||||
|
|
||||||
|
// fallback to Set
|
||||||
|
err := store.SetWithContext(context.TODO(), "key2", []byte("val2"), 0)
|
||||||
|
if err != nil {
|
||||||
|
t.Errorf("SetWithContext fallback failed: %v", err)
|
||||||
|
}
|
||||||
|
val, _ := store.Get("key2")
|
||||||
|
if !bytes.Equal(val, []byte("val2")) {
|
||||||
|
t.Errorf("SetWithContext fallback mismatch: got %v", val)
|
||||||
|
}
|
||||||
|
|
||||||
|
// custom override
|
||||||
|
store.SetCustomFuncs(&CustomFuncs{
|
||||||
|
SetWithContext: func(ctx context.Context, key string, val []byte, exp time.Duration) error {
|
||||||
|
if key == "readonly" {
|
||||||
|
return errors.New("forbidden")
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
},
|
||||||
|
})
|
||||||
|
err = store.SetWithContext(context.TODO(), "readonly", []byte("fail"), 0)
|
||||||
|
if err == nil || err.Error() != "forbidden" {
|
||||||
|
t.Errorf("SetWithContext custom override failed: err=%v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestDeleteWithContext(t *testing.T) {
|
||||||
|
store := New()
|
||||||
|
|
||||||
|
// fallback to Delete
|
||||||
|
_ = store.Set("key3", []byte("val3"), 0)
|
||||||
|
err := store.DeleteWithContext(context.TODO(), "key3")
|
||||||
|
if err != nil {
|
||||||
|
t.Errorf("DeleteWithContext fallback failed: %v", err)
|
||||||
|
}
|
||||||
|
val, err := store.Get("key3")
|
||||||
|
if err == nil {
|
||||||
|
t.Errorf("expected deletion, but got value: %v", val)
|
||||||
|
}
|
||||||
|
|
||||||
|
// custom override
|
||||||
|
store.SetCustomFuncs(&CustomFuncs{
|
||||||
|
DeleteWithContext: func(ctx context.Context, key string) error {
|
||||||
|
if key == "undeletable" {
|
||||||
|
return errors.New("blocked")
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
},
|
||||||
|
})
|
||||||
|
err = store.DeleteWithContext(context.TODO(), "undeletable")
|
||||||
|
if err == nil || err.Error() != "blocked" {
|
||||||
|
t.Errorf("DeleteWithContext custom override failed: err=%v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestResetWithContext(t *testing.T) {
|
||||||
|
store := New()
|
||||||
|
|
||||||
|
// fallback to Reset
|
||||||
|
_ = store.Set("key4", []byte("val4"), 0)
|
||||||
|
err := store.ResetWithContext(context.TODO())
|
||||||
|
if err != nil {
|
||||||
|
t.Errorf("ResetWithContext fallback failed: %v", err)
|
||||||
|
}
|
||||||
|
val, err := store.Get("key4")
|
||||||
|
if err == nil {
|
||||||
|
t.Errorf("expected reset to remove key, but got value: %v", val)
|
||||||
|
}
|
||||||
|
|
||||||
|
// custom override
|
||||||
|
store.SetCustomFuncs(&CustomFuncs{
|
||||||
|
ResetWithContext: func(ctx context.Context) error {
|
||||||
|
return errors.New("custom reset error")
|
||||||
|
},
|
||||||
|
})
|
||||||
|
err = store.ResetWithContext(context.Background())
|
||||||
|
if err == nil || err.Error() != "custom reset error" {
|
||||||
|
t.Errorf("ResetWithContext custom override failed: err=%v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -1,6 +1,7 @@
|
|||||||
package pebble
|
package pebble
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"context"
|
||||||
"encoding/json"
|
"encoding/json"
|
||||||
"errors"
|
"errors"
|
||||||
"log"
|
"log"
|
||||||
@@ -39,8 +40,7 @@ func New(config ...Config) *Storage {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// // Implement the logic to retrieve the value for the given key from the storage provider
|
// Get retrieves the value by key.
|
||||||
// // Return nil, nil if the key does not exist
|
|
||||||
func (s *Storage) Get(key string) ([]byte, error) {
|
func (s *Storage) Get(key string) ([]byte, error) {
|
||||||
if len(key) <= 0 {
|
if len(key) <= 0 {
|
||||||
return nil, nil
|
return nil, nil
|
||||||
@@ -60,7 +60,6 @@ func (s *Storage) Get(key string) ([]byte, error) {
|
|||||||
|
|
||||||
var cache CacheType
|
var cache CacheType
|
||||||
err = json.Unmarshal(data, &cache)
|
err = json.Unmarshal(data, &cache)
|
||||||
|
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, nil
|
return nil, nil
|
||||||
}
|
}
|
||||||
@@ -71,18 +70,26 @@ func (s *Storage) Get(key string) ([]byte, error) {
|
|||||||
err = s.db.Delete([]byte(key), nil)
|
err = s.db.Delete([]byte(key), nil)
|
||||||
return nil, err
|
return nil, err
|
||||||
}
|
}
|
||||||
|
|
||||||
return cache.Data, nil
|
return cache.Data, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
// // Implement the logic to store the given value for the given key in the storage provider
|
// GetWithContext retrieves value by key (dummy context support)
|
||||||
// // Use the provided expiration value (0 means no expiration)
|
func (s *Storage) GetWithContext(ctx context.Context, key string) ([]byte, error) {
|
||||||
// // Ignore empty key or value without returning an error
|
return s.Get(key)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Set stores the given value with optional expiration
|
||||||
func (s *Storage) Set(key string, val []byte, exp time.Duration) error {
|
func (s *Storage) Set(key string, val []byte, exp time.Duration) error {
|
||||||
if len(key) <= 0 || len(val) <= 0 {
|
if len(key) <= 0 || len(val) <= 0 {
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
cache := CacheType{Data: []byte(val), Created: time.Now().Unix(), Expires: 0}
|
cache := CacheType{
|
||||||
|
Data: val,
|
||||||
|
Created: time.Now().Unix(),
|
||||||
|
Expires: 0,
|
||||||
|
}
|
||||||
|
|
||||||
if exp > 0 {
|
if exp > 0 {
|
||||||
cache.Expires = cache.Created + int64(exp.Seconds())
|
cache.Expires = cache.Created + int64(exp.Seconds())
|
||||||
@@ -95,8 +102,12 @@ func (s *Storage) Set(key string, val []byte, exp time.Duration) error {
|
|||||||
return s.db.Set([]byte(key), jsonString, s.writeOptions)
|
return s.db.Set([]byte(key), jsonString, s.writeOptions)
|
||||||
}
|
}
|
||||||
|
|
||||||
// // Implement the logic to delete the value for the given key from the storage provider
|
// SetWithContext sets value by key (dummy context support)
|
||||||
// // Return no error if the key does not exist in the storage
|
func (s *Storage) SetWithContext(ctx context.Context, key string, val []byte, exp time.Duration) error {
|
||||||
|
return s.Set(key, val, exp)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Delete removes a value by key
|
||||||
func (s *Storage) Delete(key string) error {
|
func (s *Storage) Delete(key string) error {
|
||||||
if len(key) <= 0 {
|
if len(key) <= 0 {
|
||||||
return nil
|
return nil
|
||||||
@@ -104,26 +115,36 @@ func (s *Storage) Delete(key string) error {
|
|||||||
return s.db.Delete([]byte(key), s.writeOptions)
|
return s.db.Delete([]byte(key), s.writeOptions)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// DeleteWithContext deletes key (dummy context support)
|
||||||
|
func (s *Storage) DeleteWithContext(ctx context.Context, key string) error {
|
||||||
|
return s.Delete(key)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Reset flushes the DB
|
||||||
func (s *Storage) Reset() error {
|
func (s *Storage) Reset() error {
|
||||||
return s.db.Flush()
|
return s.db.Flush()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// ResetWithContext resets storage (dummy context support)
|
||||||
|
func (s *Storage) ResetWithContext(ctx context.Context) error {
|
||||||
|
return s.Reset()
|
||||||
|
}
|
||||||
|
|
||||||
|
// Close closes the database
|
||||||
func (s *Storage) Close() error {
|
func (s *Storage) Close() error {
|
||||||
return s.db.Close()
|
return s.db.Close()
|
||||||
}
|
}
|
||||||
|
|
||||||
// // Return database client
|
// Conn returns the database client
|
||||||
func (s *Storage) Conn() *pebble.DB {
|
func (s *Storage) Conn() *pebble.DB {
|
||||||
return s.db
|
return s.db
|
||||||
}
|
}
|
||||||
|
|
||||||
func isValid(fp string) bool {
|
func isValid(fp string) bool {
|
||||||
// Check if file already exists
|
|
||||||
if _, err := os.Stat(fp); err == nil {
|
if _, err := os.Stat(fp); err == nil {
|
||||||
return true
|
return true
|
||||||
}
|
}
|
||||||
|
|
||||||
// Attempt to create it
|
|
||||||
var d []byte
|
var d []byte
|
||||||
err := os.WriteFile(fp, d, 0o600)
|
err := os.WriteFile(fp, d, 0o600)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
|
|||||||
@@ -1,6 +1,7 @@
|
|||||||
package ristretto
|
package ristretto
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"context"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
"github.com/dgraph-io/ristretto"
|
"github.com/dgraph-io/ristretto"
|
||||||
@@ -52,6 +53,11 @@ func (s *Storage) Get(key string) ([]byte, error) {
|
|||||||
return buf, nil
|
return buf, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// GetWithContext gets the value by key (dummy context support)
|
||||||
|
func (s *Storage) GetWithContext(ctx context.Context, key string) ([]byte, error) {
|
||||||
|
return s.Get(key)
|
||||||
|
}
|
||||||
|
|
||||||
// Set stores the given value for the given key along
|
// Set stores the given value for the given key along
|
||||||
// with an expiration value, time.Time{} means no expiration.
|
// with an expiration value, time.Time{} means no expiration.
|
||||||
// Empty key or value will be ignored without an error.
|
// Empty key or value will be ignored without an error.
|
||||||
@@ -66,8 +72,12 @@ func (s *Storage) Set(key string, val []byte, exp time.Duration) error {
|
|||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// 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 deletes the value for the given key.
|
// Delete deletes the value for the given key.
|
||||||
// It returns no error if the storage does not contain the key,
|
|
||||||
func (s *Storage) Delete(key string) error {
|
func (s *Storage) Delete(key string) error {
|
||||||
if len(key) <= 0 {
|
if len(key) <= 0 {
|
||||||
return nil
|
return nil
|
||||||
@@ -76,12 +86,22 @@ func (s *Storage) Delete(key string) error {
|
|||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
// Reset resets the storage and delete all keys.
|
// DeleteWithContext deletes key (dummy context support)
|
||||||
|
func (s *Storage) DeleteWithContext(ctx context.Context, key string) error {
|
||||||
|
return s.Delete(key)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Reset resets the storage and deletes all keys.
|
||||||
func (s *Storage) Reset() error {
|
func (s *Storage) Reset() error {
|
||||||
s.cache.Clear()
|
s.cache.Clear()
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// ResetWithContext resets storage (dummy context support)
|
||||||
|
func (s *Storage) ResetWithContext(ctx context.Context) error {
|
||||||
|
return s.Reset()
|
||||||
|
}
|
||||||
|
|
||||||
// Close closes the storage and will stop any running garbage
|
// Close closes the storage and will stop any running garbage
|
||||||
// collectors and open connections.
|
// collectors and open connections.
|
||||||
func (s *Storage) Close() error {
|
func (s *Storage) Close() error {
|
||||||
@@ -89,7 +109,7 @@ func (s *Storage) Close() error {
|
|||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
// Return database client
|
// Conn returns the database client
|
||||||
func (s *Storage) Conn() *ristretto.Cache {
|
func (s *Storage) Conn() *ristretto.Cache {
|
||||||
return s.cache
|
return s.cache
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user