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))
batch.put(KEY_CODEC.encode_zset_value(key, value),
to_float_string(score))
batch.put(KEY_CODEC.encode_zset_score(key, value, score), bytes(''))
batch.write()
return result
def zrem(self, key, *members):
"""
see zadd() for information about score and value structures
"""
result = 0
zset_length = int(self._db.get(KEY_CODEC.encode_zset(key), '0'))
# safe guard
if zset_length == 0:
return result
batch = self._db.write_batch()
for member in members:
score = self._db.get(KEY_CODEC.encode_zset_value(key, member))
if score is None:
continue
result += 1
zset_length -= 1
batch.delete(KEY_CODEC.encode_zset_value(key, member))
batch.delete(KEY_CODEC.encode_zset_score(key, member, score))
# empty zset should be removed from keyspace
if zset_length == 0:
self.delete(key)
else:
batch.put(KEY_CODEC.encode_zset(key), bytes(zset_length))
batch.write()
return result
def zrange(self, key, start, stop, with_scores):
result = []
zset_length = int(self._db.get(KEY_CODEC.encode_zset(key), '0'))
if stop < 0:
end = zset_length + stop
else:
end = stop
if start < 0:
begin = max(0, zset_length + start)
else:
begin = start
for i, (db_key, _) in enumerate(
self._get_db_iterator(KEY_CODEC.get_min_zset_score(key))):
if i < begin:
continue
if i > end:
break
db_score = KEY_CODEC.decode_zset_score(db_key)
db_value = KEY_CODEC.decode_zset_value(db_key)
result.append(db_value)
if with_scores:
result.append(db_score)
return result
def type(self, key):
if self._db.get(KEY_CODEC.encode_string(key)) is not None:
return 'string'
if self._db.get(KEY_CODEC.encode_set(key)) is not None:
return 'set'
if self._db.get(KEY_CODEC.encode_hash(key)) is not None:
return 'hash'
if self._db.get(KEY_CODEC.encode_zset(key)) is not None:
return 'zset'
return 'none'
def _delete_db_zset(self, key):
# there are three sets of db keys for zsets:
# * zset
# * zset scores
# * zset values
with self._db.write_batch() as batch:
batch.delete(KEY_CODEC.encode_zset(key))
for db_key, _ in self._get_db_iterator(
KEY_CODEC.get_min_zset_score(key)):
batch.delete(db_key)
for db_key, _ in self._get_db_iterator(
KEY_CODEC.get_min_zset_value(key)):
batch.delete(db_key)
def _delete_db_zset(self, key):
# there are three sets of db keys for zsets:
# * zset
# * zset scores
# * zset values
#
# currently the `zset` key is immediately deleted and the other keys
# will be collected by gc.KeyGarbageCollector()
key_id, _ = self._get_zset_key_id_and_length(key)
with self._db.write_batch() as batch:
batch.delete(KEY_CODEC.encode_zset(key))
batch.put(KEY_CODEC.encode_deleted_zset_score(key_id), bytes(''))
batch.put(KEY_CODEC.encode_deleted_zset_value(key_id), bytes(''))
def rename(self, old_name, new_name):
if self.exists(old_name):
if old_name == new_name:
return
# replace the key that holds the key ID and don't touch the rest
key_type = self.type(old_name)
if key_type == 'zset':
old_db_key = KEY_CODEC.encode_zset(old_name)
new_db_key = KEY_CODEC.encode_zset(new_name)
elif key_type == 'hash':
old_db_key = KEY_CODEC.encode_hash(old_name)
new_db_key = KEY_CODEC.encode_hash(new_name)
elif key_type == 'set':
old_db_key = KEY_CODEC.encode_set(old_name)
new_db_key = KEY_CODEC.encode_set(new_name)
elif key_type == 'string':
old_db_key = KEY_CODEC.encode_string(old_name)
new_db_key = KEY_CODEC.encode_string(new_name)
else:
raise DredisError("invalid key type")
self._replace_db_key(new_db_key, old_db_key)
else:
raise NoKeyError()
def delete(self, *keys):
result = 0
for key in keys:
if self._db.get(KEY_CODEC.encode_string(key)) is not None:
self._delete_db_string(key)
result += 1
elif self._db.get(KEY_CODEC.encode_set(key)) is not None:
self._delete_db_set(key)
result += 1
elif self._db.get(KEY_CODEC.encode_hash(key)) is not None:
self._delete_db_hash(key)
result += 1
elif self._db.get(KEY_CODEC.encode_zset(key)) is not None:
self._delete_db_zset(key)
result += 1
return result
def zcard(self, key):
return int(self._db.get(KEY_CODEC.encode_zset(key), '0'))
def _get_zset_key_id_and_length(self, key):
db_value = self._db.get(KEY_CODEC.encode_zset(key))
return KEY_CODEC.decode_key_id_and_length(key, db_value)
