Add support for ZADD options (NX, XX, CH, INCR)

Issue #545

Author: Joao Fernandes

lib/redis.rb

@@ -1410,20 +1410,53 @@ def zcard(key)
   # @param [[Float, String], Array<[Float, String]>] args
   #   - a single `[score, member]` pair
   #   - an array of `[score, member]` pairs
-  #
-  # @return [Boolean, Fixnum]
+  # @param [Hash] options
+  #   - `:xx => true`: Only update elements that already exist (never
+  #   add elements)
+  #   - `:nx => true`: Don't update already existing elements (always
+  #   add new elements)
+  #   - `:ch => true`: Modify the return value from the number of new
+  #   elements added, to the total number of elements changed (CH is an
+  #   abbreviation of changed); changed elements are new elements added
+  #   and elements already existing for which the score was updated
+  #   - `:incr => true`: When this option is specified ZADD acts like
+  #   ZINCRBY; only one score-element pair can be specified in this mode
+  #
+  # @return [Boolean, Fixnum, Float]
   #   - `Boolean` when a single pair is specified, holding whether or not it was
-  #   **added** to the sorted set
+  #   **added** to the sorted set.
   #   - `Fixnum` when an array of pairs is specified, holding the number of
-  #   pairs that were **added** to the sorted set
-  def zadd(key, *args)
+  #   pairs that were **added** to the sorted set.
+  #   - `Float` when option :incr is specified, holding the score of the member
+  #   after incrementing it.
+  def zadd(key, *args, options)
+    zadd_options = []
+    if options.is_a?(Hash)
+      nx = options[:nx]
+      zadd_options << "NX" if nx
+
+      xx = options[:xx]
+      zadd_options << "XX" if xx
+
+      raise ArgumentError, "XX and NX options at the same time are not compatible" if nx && xx
+
+      ch = options[:ch]
+      zadd_options << "CH" if ch
+
+      incr = options[:incr]
+      zadd_options << "INCR" if incr
+    else
+      args << options
+    end
+
     synchronize do |client|
       if args.size == 1 && args[0].is_a?(Array)
         # Variadic: return integer
-        client.call([:zadd, key] + args[0])
+        raise ArgumentError, "only one score-element pair can be specified with the :incr option" if incr
+        client.call([:zadd, key] + zadd_options + args[0])
       elsif args.size == 2
-        # Single pair: return boolean
-        client.call([:zadd, key, args[0], args[1]], &_boolify)
+        # Single pair: return float if INCR, boolean if !INCR
+        client.call([:zadd, key, *zadd_options, args[0], args[1]], &(incr ? _floatify : _boolify))
       else
         raise ArgumentError, "wrong number of arguments"
       end

test/lint/sorted_sets.rb

@@ -9,6 +9,50 @@ def test_zadd
       assert_equal true, r.zadd("foo", 1, "s1")
       assert_equal false, r.zadd("foo", 1, "s1")
       assert_equal 1, r.zcard("foo")
+      r.del "foo"
+
+      target_version "3.0.2" do
+        # XX option
+        assert_equal 0, r.zcard("foo")
+        assert_equal false, r.zadd("foo", 1, "s1", xx: true)
+        r.zadd("foo", 1, "s1")
+        assert_equal false, r.zadd("foo", 2, "s1", xx: true)
+        assert_equal 2, r.zscore("foo", "s1")
+        r.del "foo"
+
+        # NX option
+        assert_equal 0, r.zcard("foo")
+        assert_equal true, r.zadd("foo", 1, "s1", nx: true)
+        assert_equal false, r.zadd("foo", 2, "s1", nx: true)
+        assert_equal 1, r.zscore("foo", "s1")
+        assert_equal 1, r.zcard("foo")
+        r.del "foo"
+
+        # CH option
+        assert_equal 0, r.zcard("foo")
+        assert_equal true, r.zadd("foo", 1, "s1", ch: true)
+        assert_equal false, r.zadd("foo", 1, "s1", ch: true)
+        assert_equal true, r.zadd("foo", 2, "s1", ch: true)
+        assert_equal 1, r.zcard("foo")
+        r.del "foo"
+
+        # INCR option
+        rv = r.zadd("foo", 1, "s1", incr: true)
+        assert_equal 1.0, rv
+
+        rv = r.zadd("foo", 10, "s1", incr: true)
+        assert_equal 11.0, rv
+
+        rv = r.zadd("bar", "-inf", "s1", incr: true)
+        assert_equal(-Infinity, rv)
+
+        rv = r.zadd("bar", "+inf", "s2", incr: true)
+        assert_equal(+Infinity, rv)
+        r.del "foo", "bar"
+
+        # Incompatible options combinations
+        assert_raise(ArgumentError) { r.zadd("foo", 1, "s1", xx: true, nx: true) }
+      end
     end
 
     def test_variadic_zadd
@@ -31,6 +75,42 @@ def test_variadic_zadd
         assert_raise(Redis::CommandError) { r.zadd("foo", ["bar"]) }
         assert_raise(Redis::CommandError) { r.zadd("foo", ["bar", "qux", "zap"]) }
       end
+
+      target_version "3.0.2" do
+        # XX option
+        assert_equal 0, r.zcard("foo")
+        assert_equal 0, r.zadd("foo", [1, "s1", 2, "s2"], xx: true)
+        r.zadd("foo", [1, "s1", 2, "s2"])
+        assert_equal 0, r.zadd("foo", [2, "s1", 3, "s2", 4, "s3"], xx: true)
+        assert_equal 2, r.zscore("foo", "s1")
+        assert_equal 3, r.zscore("foo", "s2")
+        assert_equal nil, r.zscore("foo", "s3")
+        assert_equal 2, r.zcard("foo")
+        r.del "foo"
+
+        # NX option
+        assert_equal 0, r.zcard("foo")
+        assert_equal 2, r.zadd("foo", [1, "s1", 2, "s2"], nx: true)
+        assert_equal 1, r.zadd("foo", [2, "s1", 3, "s2", 4, "s3"], nx: true)
+        assert_equal 1, r.zscore("foo", "s1")
+        assert_equal 2, r.zscore("foo", "s2")
+        assert_equal 4, r.zscore("foo", "s3")
+        assert_equal 3, r.zcard("foo")
+        r.del "foo"
+
+        # CH option
+        assert_equal 0, r.zcard("foo")
+        assert_equal 2, r.zadd("foo", [1, "s1", 2, "s2"], ch: true)
+        assert_equal 2, r.zadd("foo", [1, "s1", 3, "s2", 4, "s3"], ch: true)
+        assert_equal 3, r.zcard("foo")
+        r.del "foo"
+
+        # INCR option
+        assert_raise(ArgumentError) { r.zadd("foo", [1, "s1"], incr: true) }
+
+        # Incompatible options combinations
+        assert_raise(ArgumentError) { r.zadd("foo", [1, "s1"], xx: true, nx: true) }
+      end
     end
 
     def test_zrem