valkey: add support for context management

valkey/valkey.go

@@ -85,41 +85,61 @@ func New(config ...Config) *Storage {
 	}
 }
 
-// Get value by key
+// GetWithContext gets value by key with context
-func (s *Storage) Get(key string) ([]byte, error) {
+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
+// Get gets value by key
-func (s *Storage) Set(key string, val []byte, exp time.Duration) error {
+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
+// Set sets key with value and expiration
-func (s *Storage) Delete(key string) error {
+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

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