From 5810f96eb6537bd88d5a59de2f202f7ac75c0ab0 Mon Sep 17 00:00:00 2001
From: dengliming <liming.d.pro@gmail.com>
Date: Sat, 10 Oct 2020 14:44:20 +0800
Subject: [PATCH] Add support for ZADD GT/LT options #1451

Original pull request: #1455.
---
 src/main/java/io/lettuce/core/ZAddArgs.java   | 59 +++++++++++++++++++
 .../SortedSetCommandIntegrationTests.java     | 32 ++++++++++
 2 files changed, 91 insertions(+)

diff --git a/src/main/java/io/lettuce/core/ZAddArgs.java b/src/main/java/io/lettuce/core/ZAddArgs.java
index 972a48f032..7d37e38df6 100644
--- 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
index eac688946a..e0b078a15e 100644
--- 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")));
+    }
 }