Add support for ZADD GT/LT options #1451

Original pull request: #1455.
redis · Oct 26, 2020 · 5810f96 · 5810f96
1 parent d032662
commit 5810f96
Show file tree

Hide file tree

Showing 2 changed files with 91 additions and 0 deletions.
diff --git a/src/main/java/io/lettuce/core/ZAddArgs.java b/src/main/java/io/lettuce/core/ZAddArgs.java
@@ -24,6 +24,7 @@
  * {@link ZAddArgs} is a mutable object and instances should be used only once to avoid shared mutable state.
  *
  * @author Mark Paluch
+ * @author dengliming
  */
 public class ZAddArgs implements CompositeArgument {
 
@@ -33,6 +34,10 @@ public class ZAddArgs implements CompositeArgument {
 
     private boolean ch = false;
 
+    private boolean lt = false;
+
+    private boolean gt = false;
+
     /**
      * Builder entry points for {@link ScanArgs}.
      */
@@ -74,6 +79,27 @@ public static ZAddArgs ch() {
             return new ZAddArgs().ch();
         }
 
+        /**
+         * Creates new {@link ZAddArgs} and enabling {@literal GT}.
+         *
+         * @return new {@link ZAddArgs} with {@literal GT} enabled.
+         * @see ZAddArgs#gt()
+         * @since 6.1
+         */
+        public static ZAddArgs gt() {
+            return new ZAddArgs().gt();
+        }
+
+        /**
+         * Creates new {@link ZAddArgs} and enabling {@literal LT}.
+         *
+         * @return new {@link ZAddArgs} with {@literal LT} enabled.
+         * @see ZAddArgs#lt()
+         * @since 6.1
+         */
+        public static ZAddArgs lt() {
+            return new ZAddArgs().lt();
+        }
     }
 
     /**
@@ -109,6 +135,32 @@ public ZAddArgs ch() {
         return this;
     }
 
+    /**
+     * Only update existing elements if the new score is greater than the current score. This flag doesn't prevent adding new elements.
+     *
+     * @return {@code this} {@link ZAddArgs}.
+     * @since 6.1
+     */
+    public ZAddArgs gt() {
+
+        this.gt = true;
+        this.lt = false;
+        return this;
+    }
+
+    /**
+     * Only update existing elements if the new score is less than the current score. This flag doesn't prevent adding new elements.
+     *
+     * @return {@code this} {@link ZAddArgs}.
+     * @since 6.1
+     */
+    public ZAddArgs lt() {
+
+        this.lt = true;
+        this.gt = false;
+        return this;
+    }
+
     public <K, V> void build(CommandArgs<K, V> args) {
 
         if (nx) {
@@ -119,6 +171,13 @@ public <K, V> void build(CommandArgs<K, V> args) {
             args.add("XX");
         }
 
+        if (gt) {
+            args.add("GT");
+        }
+        if (lt) {
+            args.add("LT");
+        }
+
         if (ch) {
             args.add("CH");
         }

diff --git a/src/test/java/io/lettuce/core/commands/SortedSetCommandIntegrationTests.java b/src/test/java/io/lettuce/core/commands/SortedSetCommandIntegrationTests.java
@@ -795,4 +795,36 @@ void setup100KeyValues(Set<String> expect) {
             expect.add(value + i);
         }
     }
+
+    @Test
+    void zaddgt() {
+        assertThat(redis.zadd(key, 1.0, "a")).isEqualTo(1);
+        // new score less than the current score
+        assertThat(redis.zadd(key, ZAddArgs.Builder.gt(), 0.0, "a")).isEqualTo(0);
+        assertThat(redis.zrangeWithScores(key, 0, -1)).isEqualTo(svlist(sv(1.0, "a")));
+
+        // new score greater than the current score
+        assertThat(redis.zadd(key, ZAddArgs.Builder.gt(), 2.0, "a")).isEqualTo(0);
+        assertThat(redis.zrangeWithScores(key, 0, -1)).isEqualTo(svlist(sv(2.0, "a")));
+
+        // add new element
+        assertThat(redis.zadd(key, ZAddArgs.Builder.gt(), 0.0, "b")).isEqualTo(1);
+        assertThat(redis.zrangeWithScores(key, 0, -1)).isEqualTo(svlist(sv(0.0, "b"), sv(2.0, "a")));
+    }
+
+    @Test
+    void zaddlt() {
+        assertThat(redis.zadd(key, 2.0, "a")).isEqualTo(1);
+        // new score greater than the current score
+        assertThat(redis.zadd(key, ZAddArgs.Builder.lt(), 3.0, "a")).isEqualTo(0);
+        assertThat(redis.zrangeWithScores(key, 0, -1)).isEqualTo(svlist(sv(2.0, "a")));
+
+        // new score less than the current score
+        assertThat(redis.zadd(key, ZAddArgs.Builder.lt(), 1.0, "a")).isEqualTo(0);
+        assertThat(redis.zrangeWithScores(key, 0, -1)).isEqualTo(svlist(sv(1.0, "a")));
+
+        // add new element
+        assertThat(redis.zadd(key, ZAddArgs.Builder.lt(), 0.0, "b")).isEqualTo(1);
+        assertThat(redis.zrangeWithScores(key, 0, -1)).isEqualTo(svlist(sv(0.0, "b"), sv(1.0, "a")));
+    }
 }