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 hdel(self, key, *fields):
result = 0
key_id, hash_length = self._get_hash_key_id_and_length(key)
# safe guard
if hash_length == 0:
return result
batch = self._db.write_batch()
for field in fields:
if self._db.get(KEY_CODEC.encode_hash_field(key_id,
field)) is not None:
result += 1
hash_length -= 1
batch.delete(KEY_CODEC.encode_hash_field(key_id, field))
if hash_length == 0:
# remove empty hashes from keyspace
self.delete(key)
else:
batch.put(
KEY_CODEC.encode_hash(key),
KEY_CODEC.encode_key_id_and_length(key, key_id, hash_length))
batch.write()
return result
def zrangebyscore(self,
key,
min_score,
max_score,
withscores=False,
offset=0,
count=float('+inf')):
result = []
num_elems_read = 0
if withscores:
num_elems_per_entry = 2
else:
num_elems_per_entry = 1
score_range = ScoreRange(min_score, max_score)
for db_key in self._db.iterator(
prefix=KEY_CODEC.get_min_zset_score(key), include_value=False):
if len(result) / num_elems_per_entry >= count:
return result
db_score = KEY_CODEC.decode_zset_score(db_key)
if score_range.above_max(db_score):
break
if score_range.check(db_score):
num_elems_read += 1
if len(result) / num_elems_per_entry >= count:
return result
if num_elems_read > offset:
db_value = KEY_CODEC.decode_zset_value(db_key)
result.append(db_value)
if withscores:
result.append(db_score)
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 hvals(self, key):
result = []
if self._db.get(KEY_CODEC.encode_hash(key)) is not None:
for _, db_value in self._get_db_iterator(
KEY_CODEC.get_min_hash_field(key)):
result.append(db_value)
return result
def smembers(self, key):
result = set()
if self._db.get(KEY_CODEC.encode_set(key)):
for db_key, _ in self._get_db_iterator(
KEY_CODEC.get_min_set_member(key)):
_, length, member_key = KEY_CODEC.decode_key(db_key)
member_value = member_key[length:]
result.add(member_value)
return result
def _delete_db_hash(self, key):
# there are two sets of db keys for hashes:
# * hash
# * hash fields
with self._db.write_batch() as batch:
batch.delete(KEY_CODEC.encode_hash(key))
for db_key, _ in self._get_db_iterator(
KEY_CODEC.get_min_hash_field(key)):
batch.delete(db_key)
def hset(self, key, field, value):
result = 0
if self._db.get(KEY_CODEC.encode_hash_field(key, field)) is None:
result = 1
hash_length = int(self._db.get(KEY_CODEC.encode_hash(key), '0'))
with self._db.write_batch() as batch:
batch.put(KEY_CODEC.encode_hash(key), bytes(hash_length + 1))
batch.put(KEY_CODEC.encode_hash_field(key, field), value)
return result
def sadd(self, key, value):
if self._db.get(KEY_CODEC.encode_set_member(key, value)) is None:
length = int(self._db.get(KEY_CODEC.encode_set(key)) or b'0')
with self._db.write_batch() as batch:
batch.put(KEY_CODEC.encode_set(key), bytes(length + 1))
batch.put(KEY_CODEC.encode_set_member(key, value), bytes(''))
return 1
else:
return 0
def _delete_db_set(self, key):
# there are two sets of db keys for sets:
# * set
# * set members
with self._db.write_batch() as batch:
batch.delete(KEY_CODEC.encode_set(key))
for db_key, _ in self._get_db_iterator(
KEY_CODEC.get_min_set_member(key)):
batch.delete(db_key)
def smembers(self, key):
result = set()
key_id, length = self._get_set_key_id_and_length(key)
if length > 0:
for db_key, _ in self._get_db_iterator(
KEY_CODEC.get_min_set_member(key_id)):
_, length, member_key = KEY_CODEC.decode_key(db_key)
member_value = member_key[length:]
result.add(member_value)
return result
def hgetall(self, key):
result = []
if self._db.get(KEY_CODEC.encode_hash(key)) is not None:
for db_key, db_value in self._get_db_iterator(
KEY_CODEC.get_min_hash_field(key)):
_, length, field_key = KEY_CODEC.decode_key(db_key)
field = field_key[length:]
result.append(field)
result.append(db_value)
return result
def hsetnx(self, key, field, value):
# only set if not set before
if self._db.get(KEY_CODEC.encode_hash_field(key, field)) is None:
hash_length = int(self._db.get(KEY_CODEC.encode_hash(key), '0'))
with self._db.write_batch() as batch:
batch.put(KEY_CODEC.encode_hash(key), bytes(hash_length + 1))
batch.put(KEY_CODEC.encode_hash_field(key, field), value)
return 1
else:
return 0
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 hgetall(self, key):
result = []
key_id, hash_length = self._get_hash_key_id_and_length(key)
if hash_length > 0:
for db_key, db_value in self._get_db_iterator(
KEY_CODEC.get_min_hash_field(key_id)):
_, length, field_key = KEY_CODEC.decode_key(db_key)
field = field_key[length:]
result.append(field)
result.append(db_value)
return result
def _delete_db_hash(self, key):
# there are two sets of db keys for hashes:
# * hash
# * hash fields
#
# currently the `hash` key is immediately deleted and the other keys
# will be collected by gc.KeyGarbageCollector()
key_id, _ = self._get_hash_key_id_and_length(key)
with self._db.write_batch() as batch:
batch.delete(KEY_CODEC.encode_hash(key))
batch.put(KEY_CODEC.encode_deleted_hash(key_id), bytes(''))
def _delete_db_set(self, key):
# there are two sets of db keys for sets:
# * set
# * set members
#
# currently the `set` key is immediately deleted and the other keys
# will be collected by gc.KeyGarbageCollector()
key_id, _ = self._get_set_key_id_and_length(key)
with self._db.write_batch() as batch:
batch.delete(KEY_CODEC.encode_set(key))
batch.put(KEY_CODEC.encode_deleted_set(key_id), bytes(''))
def hset(self, key, field, value):
result = 0
key_id, hash_length = self._get_hash_key_id_and_length(key)
if self._db.get(KEY_CODEC.encode_hash_field(key_id, field)) is None:
result = 1
with self._db.write_batch() as batch:
batch.put(
KEY_CODEC.encode_hash(key),
KEY_CODEC.encode_key_id_and_length(key, key_id,
hash_length + result))
batch.put(KEY_CODEC.encode_hash_field(key_id, field), value)
return result
def hsetnx(self, key, field, value):
key_id, hash_length = self._get_hash_key_id_and_length(key)
# only set if not set before
if self._db.get(KEY_CODEC.encode_hash_field(key_id, field)) is None:
with self._db.write_batch() as batch:
batch.put(
KEY_CODEC.encode_hash(key),
KEY_CODEC.encode_key_id_and_length(key, key_id,
hash_length + 1))
batch.put(KEY_CODEC.encode_hash_field(key_id, field), value)
return 1
else:
return 0
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 sadd(self, key, value):
key_id, length = self._get_set_key_id_and_length(key)
if self._db.get(KEY_CODEC.encode_set_member(key_id, value)) is None:
with self._db.write_batch() as batch:
batch.put(
KEY_CODEC.encode_set(key),
KEY_CODEC.encode_key_id_and_length(key, key_id,
length + 1))
batch.put(KEY_CODEC.encode_set_member(key_id, value),
bytes(''))
return 1
else:
return 0
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 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 zrank(self, key, member):
score = self._db.get(KEY_CODEC.encode_zset_value(key, member))
if score is None:
return None
rank = 0
for db_key, _ in self._get_db_iterator(
KEY_CODEC.get_min_zset_score(key)):
db_score = KEY_CODEC.decode_zset_score(db_key)
db_value = KEY_CODEC.decode_zset_value(db_key)
if db_score < float(score):
rank += 1
elif db_score == float(score) and db_value < member:
rank += 1
else:
break
return rank
def hvals(self, key):
result = []
key_id, hash_length = self._get_hash_key_id_and_length(key)
if hash_length > 0:
for _, db_value in self._get_db_iterator(
KEY_CODEC.get_min_hash_field(key_id)):
result.append(db_value)
return result
def keys(self, pattern):
db_keys = set()
for key_type in KEY_CODEC.KEY_TYPES:
for key in self._db.iterator(prefix=chr(key_type),
include_value=False):
_, _, key_value = KEY_CODEC.decode_key(key)
if pattern is None or fnmatch.fnmatch(key_value, pattern):
db_keys.add(key_value)
return db_keys
def _convert_key(batch, db_key, db_value):
type_id, _, key = KEY_CODEC.decode_key(db_key)
if len(db_value) < UUID_LENGTH_IN_BYTES:
# older schema before uuid
key_id = key
length = db_value
else:
# new schema with uuid
key_id = db_value[:UUID_LENGTH_IN_BYTES]
length = db_value[UUID_LENGTH_IN_BYTES:]
print('batch.put({!r}, {!r})'.format(KEY_CODEC.get_key(key_id, type_id),
length))
batch.put(KEY_CODEC.get_key(key_id, type_id), length)
if key != key_id:
batch.delete(db_key)
print('batch.delete({!r})'.format(db_key))
print('[TO RUN] RENAME {} {}'.format(_rediscli_str(key_id),
_rediscli_str(key)))
def _collect(self, db):
deleted = 0
with db.write_batch() as batch:
for deleted_db_key, _ in db.iterator(
prefix=KEY_CODEC.MIN_DELETED_VALUE):
_, _, deleted_key_value = KEY_CODEC.decode_key(deleted_db_key)
for db_key, _ in db.iterator(prefix=deleted_key_value):
deleted += 1
batch.delete(db_key)
if deleted == self._batch_size:
return
batch.delete(deleted_db_key)
def zcount(self, key, min_score, max_score):
# TODO: optimize for performance. it's probably possible to create a new entry only for scores
# like:
# <prefix>myzset<score> = number of elements with that score
#
# ZADD myzset 10 a
# <prefix>_myzset_10 = 1 ; one element with score 10
#
# ZADD myzset 10 b
# <prefix>_myzset_10 = 2 ; two elements with score 10
score_range = ScoreRange(min_score, max_score)
count = 0
for db_key, _ in self._get_db_iterator(
KEY_CODEC.get_min_zset_score(key)):
db_score = KEY_CODEC.decode_zset_score(db_key)
if score_range.check(db_score):
count += 1
if score_range.above_max(db_score):
break
return count
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
