Merge pull request #547 from Talkdesk/545-zadd-options

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

Author: djanowski

lib/redis.rb

@@ -1410,20 +1410,50 @@ 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 args.last.is_a?(Hash)
+      options = args.pop
+
+      nx = options[:nx]
+      zadd_options << "NX" if nx
+
+      xx = options[:xx]
+      zadd_options << "XX" if xx
+
+      ch = options[:ch]
+      zadd_options << "CH" if ch
+
+      incr = options[:incr]
+      zadd_options << "INCR" if incr
+    end
+
     synchronize do |client|
       if args.size == 1 && args[0].is_a?(Array)
-        # Variadic: return integer
-        client.call([:zadd, key] + args[0])
+        # Variadic: return float if INCR, integer if !INCR
+        client.call([:zadd, key] + zadd_options + args[0], &(incr ? _floatify : _identity))
       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, &(incr ? _floatify : _boolify))
       else
         raise ArgumentError, "wrong number of arguments"
       end
@@ -2614,6 +2644,12 @@ def _pairify(array)
     array.each_slice(2).to_a
   end
 
+  def _identity
+    lambda { |value|
+      value
+    }
+  end
+
   def _subscription(method, channels, block)
     return @client.call([method] + channels) if subscribed?
 

test/lint/sorted_sets.rb

@@ -9,6 +9,43 @@ 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
+        assert_equal 1.0, r.zadd("foo", 1, "s1", :incr => true)
+        assert_equal 11.0, r.zadd("foo", 10, "s1", :incr => true)
+        assert_equal -Infinity, r.zadd("bar", "-inf", "s1", :incr => true)
+        assert_equal +Infinity, r.zadd("bar", "+inf", "s2", :incr => true)
+        r.del "foo", "bar"
+
+        # Incompatible options combination
+        assert_raise(Redis::CommandError) { r.zadd("foo", 1, "s1", :xx => true, :nx => true) }
+      end
     end
 
     def test_variadic_zadd
@@ -31,6 +68,47 @@ 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_equal 1.0, r.zadd("foo", [1, "s1"], :incr => true)
+        assert_equal 11.0, r.zadd("foo", [10, "s1"], :incr => true)
+        assert_equal -Infinity, r.zadd("bar", ["-inf", "s1"], :incr => true)
+        assert_equal +Infinity, r.zadd("bar", ["+inf", "s2"], :incr => true)
+        assert_raise(Redis::CommandError) { r.zadd("foo", [1, "s1", 2, "s2"], :incr => true) }
+        r.del "foo", "bar"
+
+        # Incompatible options combination
+        assert_raise(Redis::CommandError) { r.zadd("foo", [1, "s1"], :xx => true, :nx => true) }
+      end
     end
 
     def test_zrem