diff --git a/valkey/valkey.go b/valkey/valkey.go index 2bd8831b..fd345d9a 100644 --- a/valkey/valkey.go +++ b/valkey/valkey.go @@ -85,41 +85,61 @@ func New(config ...Config) *Storage { } } -// 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 } - val, err := s.db.DoCache(context.Background(), s.db.B().Get().Key(key).Cache(), cacheTTL).AsBytes() + val, err := s.db.DoCache(ctx, s.db.B().Get().Key(key).Cache(), cacheTTL).AsBytes() if err != nil && valkey.IsValkeyNil(err) { return nil, nil } return val, err } -// Set key with value -func (s *Storage) Set(key string, val []byte, exp time.Duration) error { +// Get gets value by key +func (s *Storage) Get(key string) ([]byte, error) { + return s.GetWithContext(context.Background(), key) +} + +// SetWithContext sets key with value and expiration with context +func (s *Storage) SetWithContext(ctx context.Context, key string, val []byte, exp time.Duration) error { if len(key) <= 0 || len(val) <= 0 { return nil } if exp > 0 { - return s.db.Do(context.Background(), s.db.B().Set().Key(key).Value(string(val)).Ex(exp).Build()).Error() + return s.db.Do(ctx, s.db.B().Set().Key(key).Value(string(val)).Ex(exp).Build()).Error() } else { - return s.db.Do(context.Background(), s.db.B().Set().Key(key).Value(string(val)).Build()).Error() + return s.db.Do(ctx, s.db.B().Set().Key(key).Value(string(val)).Build()).Error() } } -// Delete key by key -func (s *Storage) Delete(key string) error { +// Set sets key with value and expiration +func (s *Storage) Set(key string, val []byte, exp time.Duration) error { + return s.SetWithContext(context.Background(), key, val, exp) +} + +// DeleteWithContext deletes key by key with context +func (s *Storage) DeleteWithContext(ctx context.Context, key string) error { if len(key) <= 0 { return nil } - return s.db.Do(context.Background(), s.db.B().Del().Key(key).Build()).Error() + return s.db.Do(ctx, s.db.B().Del().Key(key).Build()).Error() } -// Reset all keys +// Delete deletes key by key +func (s *Storage) Delete(key string) error { + return s.DeleteWithContext(context.Background(), key) +} + +// ResetWithContext resets all keys with context +func (s *Storage) ResetWithContext(ctx context.Context) error { + return s.db.Do(ctx, s.db.B().Flushdb().Build()).Error() +} + +// Reset resets all keys func (s *Storage) Reset() error { - return s.db.Do(context.Background(), s.db.B().Flushdb().Build()).Error() + return s.ResetWithContext(context.Background()) } // Close the database diff --git a/valkey/valkey_test.go b/valkey/valkey_test.go index cf59cf75..2d9c9aeb 100644 --- a/valkey/valkey_test.go +++ b/valkey/valkey_test.go @@ -1,6 +1,7 @@ package valkey import ( + "context" "crypto/tls" "log" "testing" @@ -23,6 +24,19 @@ func Test_Valkey_Set(t *testing.T) { require.NoError(t, err) } +func Test_Valkey_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_Valkey_Set_Override(t *testing.T) { var ( key = "john" @@ -50,6 +64,23 @@ func Test_Valkey_Get(t *testing.T) { require.Equal(t, val, result) } +func Test_Valkey_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_Valkey_Set_Expiration(t *testing.T) { var ( key = "john" @@ -94,6 +125,26 @@ func Test_Valkey_Delete(t *testing.T) { require.Zero(t, len(result)) } +func Test_Valkey_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_Valkey_Reset(t *testing.T) { val := []byte("doe") @@ -115,6 +166,30 @@ func Test_Valkey_Reset(t *testing.T) { require.Zero(t, len(result)) } +func Test_Valkey_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_Valkey_Close(t *testing.T) { require.Nil(t, testStore.Close()) }