Add ZADD [NX|XX] support

dredis/commands.py

@@ -294,15 +294,32 @@ def cmd_eval(keyspace, script, numkeys, *args):
 
 
 @command('ZADD', arity=-4, flags=CMD_WRITE)
-def cmd_zadd(keyspace, key, *flat_pairs):
- if len(flat_pairs) % 2 != 0:
+def cmd_zadd(keyspace, key, *args):
+ nx = False
+ xx = False
+ start_index = 0
+
+ for i, arg in enumerate(args):
+ if arg.lower() == 'nx':
+ nx = True
+ elif arg.lower() == 'xx':
+ xx = True
+ else:
+ break
+ start_index = i + 1
+
+ if nx and xx:
+ raise DredisError("XX and NX options at the same time are not compatible")
+
+ args = args[start_index:]
+ if len(args) % 2 != 0:
  raise DredisSyntaxError()
 
  count = 0
- pairs = zip(flat_pairs[0::2], flat_pairs[1::2]) # [1, 2, 3, 4] -> [(1,2), (3,4)]
+ pairs = zip(args[0::2], args[1::2]) # [1, 2, 3, 4] -> [(1,2), (3,4)]
  for score, value in pairs:
  _validate_zset_score(score)
- count += keyspace.zadd(key, score, value)
+ count += keyspace.zadd(key, score, value, nx=nx, xx=xx)
  return count
 
 

dredis/keyspace.py

@@ -187,19 +187,23 @@ def _get_db_iterator(self, key_prefix=None, start=None):
  for db_key, db_value in self._db.iterator(prefix=key_prefix, start=start):
  yield db_key, db_value
 
- def zadd(self, key, score, value):
+ def zadd(self, key, score, value, nx=False, xx=False):
  zset_length = int(self._db.get(KEY_CODEC.encode_zset(key), '0'))
 
  batch = self._db.write_batch()
  db_score = self._db.get(KEY_CODEC.encode_zset_value(key, value))
  if db_score is not None:
+ if nx:
+ return 0
  result = 0
  previous_score = db_score
  if float(previous_score) == float(score):
  return result
  else:
  batch.delete(KEY_CODEC.encode_zset_score(key, value, previous_score))
  else:
+ if xx:
+ return 0
  result = 1
  zset_length += 1
  batch.put(KEY_CODEC.encode_zset(key), bytes(zset_length))

tests/integration/test_zset.py

@@ -477,3 +477,40 @@ def test_zscan_with_a_cursor_that_doesnt_exist():
  r = fresh_redis()
 
  assert r.zscan('myzset', 123) == (0, [])
+
+
+def test_zadd_nx_should_only_add_if_element_isnt_present():
+ r = fresh_redis()
+
+ r.zadd('myzset', 0, 'test1')
+
+ # the current version of redis-py doesn't support r.zadd(..., nx=True)
+ assert r.execute_command('ZADD', 'myzset', 'NX', 10, 'test1') == 0
+ assert r.execute_command('ZADD', 'myzset', 'NX', 10, 'test2') == 1
+ assert r.zrange('myzset', 0, -1, withscores=True) == [
+ ('test1', 0),
+ ('test2', 10),
+ ]
+
+
+def test_zadd_xx_should_never_add_new_elements():
+ r = fresh_redis()
+
+ r.zadd('myzset', 0, 'test1')
+
+ # the current version of redis-py doesn't support r.zadd(..., xx=True)
+ assert r.execute_command('ZADD', 'myzset', 'XX', 10, 'test1') == 0
+ assert r.execute_command('ZADD', 'myzset', 'XX', 10, 'test2') == 0
+ assert r.zrange('myzset', 0, -1, withscores=True) == [
+ ('test1', 10),
+ ]
+
+
+def test_zadd_should_not_accept_nx_and_xx_at_the_same_time():
+ r = fresh_redis()
+
+ with pytest.raises(redis.ResponseError) as exc:
+ # the current version of redis-py doesn't support r.zadd(..., nx=True, xx=True)
+ r.execute_command('ZADD', 'myzset', 'XX', 'NX', 0, 'test1')
+
+ assert str(exc.value) == "XX and NX options at the same time are not compatible"