Add SetNX SetXX (#236)

* Add SetNX * Add SetXX
trustwallet · Sep 14, 2023 · 1d4cb3d · 1d4cb3d
1 parent 6022cf2
commit 1d4cb3d
Show file tree

Hide file tree

Showing 3 changed files with 80 additions and 0 deletions.
diff --git a/cache/redis/client_interface.go b/cache/redis/client_interface.go
@@ -16,6 +16,8 @@ type redisClient interface {
 	Del(ctx context.Context, keys ...string) *redis.IntCmd
 	Pipeline() redis.Pipeliner
 	Watch(ctx context.Context, fn func(tx *redis.Tx) error, keys ...string) error
+	SetNX(ctx context.Context, key string, value interface{}, expiration time.Duration) *redis.BoolCmd
+	SetXX(ctx context.Context, key string, value interface{}, expiration time.Duration) *redis.BoolCmd
 
 	Ping(ctx context.Context) *redis.StatusCmd
 	Scan(ctx context.Context, cursor uint64, match string, count int64) *redis.ScanCmd

diff --git a/cache/redis/redis.go b/cache/redis/redis.go
@@ -265,6 +265,34 @@ func (r *Redis) Reconnect(ctx context.Context, host string) error {
 	return nil
 }
 
+func (r *Redis) SetNX(ctx context.Context, key string, value interface{}, expiration time.Duration) (bool, error) {
+	data, err := json.Marshal(value)
+	if err != nil {
+		return false, err
+	}
+
+	cmd := r.client.SetNX(ctx, key, data, expiration)
+	if cmd.Err() != nil {
+		return false, cmd.Err()
+	}
+
+	return cmd.Val(), nil
+}
+
+func (r *Redis) SetXX(ctx context.Context, key string, value interface{}, expiration time.Duration) (bool, error) {
+	data, err := json.Marshal(value)
+	if err != nil {
+		return false, err
+	}
+
+	cmd := r.client.SetXX(ctx, key, data, expiration)
+	if cmd.Err() != nil {
+		return false, cmd.Err()
+	}
+
+	return cmd.Val(), nil
+}
+
 func (r *Redis) reconnectCluster(ctx context.Context, redisURL string) error {
 	options, err := redis.ParseClusterURL(redisURL)
 	if err != nil {

diff --git a/cache/redis/redis_test.go b/cache/redis/redis_test.go
@@ -280,6 +280,56 @@ func TestRedis_Reconnect(t *testing.T) {
 
 }
 
+func TestRedis_SetNX(t *testing.T) {
+	redisInitFns := []redisInitFn{redisInit}
+	for _, redisInit := range redisInitFns {
+		t.Run("", func(t *testing.T) {
+			r, err := redisInit(t)
+			assert.Nil(t, err)
+
+			b1, err := r.SetNX(context.TODO(), "test", 1, time.Second)
+			assert.Nil(t, err)
+			assert.Equal(t, true, b1)
+
+			b2, err := r.SetNX(context.TODO(), "test", 1, time.Second)
+			assert.Nil(t, err)
+			assert.Equal(t, false, b2)
+
+			var v []byte
+			err = r.Get(context.TODO(), "test", &v)
+			assert.NotNil(t, err)
+			assert.Equal(t, 1, 1)
+		})
+	}
+}
+
+func TestRedis_SetXX(t *testing.T) {
+	redisInitFns := []redisInitFn{redisInit}
+	for _, redisInit := range redisInitFns {
+		t.Run("", func(t *testing.T) {
+			r, err := redisInit(t)
+			assert.Nil(t, err)
+
+			notExist, err := r.SetXX(context.TODO(), "test", 1, time.Minute)
+			assert.Nil(t, err)
+			assert.Equal(t, false, notExist)
+
+			b, err := r.SetNX(context.TODO(), "test", 1, time.Minute)
+			assert.Nil(t, err)
+			assert.Equal(t, true, b)
+
+			exist, err := r.SetXX(context.TODO(), "test", 1, time.Minute)
+			assert.Nil(t, err)
+			assert.Equal(t, true, exist)
+
+			var v []byte
+			err = r.Get(context.TODO(), "test", &v)
+			assert.NotNil(t, err)
+			assert.Equal(t, 1, 1)
+		})
+	}
+}
+
 func redisInit(t *testing.T) (*Redis, error) {
 	mr, err := miniredis.Run()
 	assert.NotNil(t, mr)