async def test_script_object_in_pipeline(self, r):
multiply = await r.register_script(multiply_script)
assert not multiply.sha
pipe = r.pipeline()
await pipe.set('a', 2)
await pipe.get('a')
multiply(keys=['a'], args=[3], client=pipe)
# even though the pipeline wasn't executed yet, we made sure the
# script was loaded and got a valid sha
assert multiply.sha
assert await r.script_exists(multiply.sha) == [True]
# [SET worked, GET 'a', result of multiple script]
assert await pipe.execute() == [True, b('2'), 6]
# purge the script from redis's cache and re-run the pipeline
# the multiply script object knows it's sha, so it shouldn't get
# reloaded until pipe.execute()
await r.script_flush()
pipe = await r.pipeline()
await pipe.set('a', 2)
await pipe.get('a')
assert multiply.sha
multiply(keys=['a'], args=[3], client=pipe)
assert await r.script_exists(multiply.sha) == [False]
# [SET worked, GET 'a', result of multiple script]
assert await pipe.execute() == [True, b('2'), 6]
async def test_access_correct_slave_with_readonly_mode_client(sr):
"""
Test that the client can get value normally with readonly mode
when we connect to correct slave.
"""
# we assume this key is set on 127.0.0.1:7000(7003)
await sr.set('foo16706', 'foo')
await asyncio.sleep(1)
with patch.object(ClusterConnectionPool,
'get_node_by_slot') as return_slave_mock:
return_slave_mock.return_value = {
'name': '127.0.0.1:7004',
'host': '127.0.0.1',
'port': 7004,
'server_type': 'slave',
}
master_value = {
'host': '127.0.0.1',
'name': '127.0.0.1:7000',
'port': 7000,
'server_type': 'master'
}
with patch.object(ClusterConnectionPool,
'get_master_node_by_slot',
return_value=master_value) as return_master_mock:
readonly_client = StrictRedisCluster(host="127.0.0.1",
port=7000,
readonly=True)
assert b('foo') == await readonly_client.get('foo16706')
readonly_client = StrictRedisCluster.from_url(
url="redis://127.0.0.1:7000/0", readonly=True)
assert b('foo') == await readonly_client.get('foo16706')
async def test_exec_error_in_response(self, r):
"""
an invalid pipeline command at exec time adds the exception instance
to the list of returned values
"""
await r.flushdb()
await r.set('c', 'a')
async with await r.pipeline() as pipe:
await pipe.set('a', 1)
await pipe.set('b', 2)
await pipe.lpush('c', 3)
await pipe.set('d', 4)
result = await pipe.execute(raise_on_error=False)
assert result[0]
assert await r.get('a') == b('1')
assert result[1]
assert await r.get('b') == b('2')
# we can't lpush to a key that's a string value, so this should
# be a ResponseError exception
assert isinstance(result[2], ResponseError)
assert await r.get('c') == b('a')
# since this isn't a transaction, the other commands after the
# error are still executed
assert result[3]
assert await r.get('d') == b('4')
# make sure the pipe was restored to a working state
await pipe.set('z', 'zzz')
assert await pipe.execute() == [True]
assert await r.get('z') == b('zzz')
async def test_script_object_in_pipeline(self, r):
await r.script_flush()
multiply = r.register_script(multiply_script)
precalculated_sha = multiply.sha
assert precalculated_sha
pipe = await r.pipeline()
await pipe.set('a', 2)
await pipe.get('a')
await multiply.execute(keys=['a'], args=[3], client=pipe)
assert await r.script_exists(multiply.sha) == [False]
# [SET worked, GET 'a', result of multiple script]
assert await pipe.execute() == [True, b('2'), 6]
# The script should have been loaded by pipe.execute()
assert await r.script_exists(multiply.sha) == [True]
# The precalculated sha should have been the correct one
assert multiply.sha == precalculated_sha
# purge the script from redis's cache and re-run the pipeline
# the multiply script should be reloaded by pipe.execute()
await r.script_flush()
pipe = await r.pipeline()
await pipe.set('a', 2)
await pipe.get('a')
await multiply.execute(keys=['a'], args=[3], client=pipe)
assert await r.script_exists(multiply.sha) == [False]
# [SET worked, GET 'a', result of multiple script]
assert await pipe.execute() == [True, b('2'), 6]
assert await r.script_exists(multiply.sha) == [True]
async def _georadiusgeneric(self, command, *args, **kwargs):
pieces = list(args)
if kwargs['unit'] and kwargs['unit'] not in ('m', 'km', 'mi', 'ft'):
raise RedisError("GEORADIUS invalid unit")
elif kwargs['unit']:
pieces.append(kwargs['unit'])
else:
pieces.append('m', )
for token in ('withdist', 'withcoord', 'withhash'):
if kwargs[token]:
pieces.append(b(token.upper()))
if kwargs['count']:
pieces.extend([b('COUNT'), kwargs['count']])
if kwargs['sort'] and kwargs['sort'] not in ('ASC', 'DESC'):
raise RedisError("GEORADIUS invalid sort")
elif kwargs['sort']:
pieces.append(b(kwargs['sort']))
if kwargs['store'] and kwargs['store_dist']:
raise RedisError("GEORADIUS store and store_dist cant be set"
" together")
if kwargs['store']:
pieces.extend([b('STORE'), kwargs['store']])
if kwargs['store_dist']:
pieces.extend([b('STOREDIST'), kwargs['store_dist']])
return await self.execute_command(command, *pieces, **kwargs)
async def zrevrangebyscore(self,
name,
max,
min,
start=None,
num=None,
withscores=False,
score_cast_func=float):
"""
Returns a range of values from the sorted set ``name`` with scores
between ``min`` and ``max`` in descending order.
If ``start`` and ``num`` are specified, then return a slice
of the range.
``withscores`` indicates to return the scores along with the values.
The return type is a list of (value, score) pairs
``score_cast_func`` a callable used to cast the score return value
"""
if (start is not None and num is None) or \
(num is not None and start is None):
raise RedisError("``start`` and ``num`` must both be specified")
pieces = ['ZREVRANGEBYSCORE', name, max, min]
if start is not None and num is not None:
pieces.extend([b('LIMIT'), start, num])
if withscores:
pieces.append(b('WITHSCORES'))
options = {
'withscores': withscores,
'score_cast_func': score_cast_func
}
return await self.execute_command(*pieces, **options)
def _trans_type(self, content):
if isinstance(content, str):
content = content.encode(self.encoding)
elif isinstance(content, int):
content = b(str(content))
elif isinstance(content, float):
content = b(repr(content))
return content
async def slaveof(self, host=None, port=None):
"""
Set the server to be a replicated slave of the instance identified
by the ``host`` and ``port``. If called without arguments, the
instance is promoted to a master instead.
"""
if host is None and port is None:
return await self.execute_command('SLAVEOF', b('NO'), b('ONE'))
return await self.execute_command('SLAVEOF', host, port)
async def test_pipeline_no_transaction(self, r):
await r.flushdb()
async with await r.pipeline(transaction=False) as pipe:
await (await (await pipe.set('a', 'a1')).set('b',
'b1')).set('c', 'c1')
assert await pipe.execute() == [True, True, True]
assert await r.get('a') == b('a1')
assert await r.get('b') == b('b1')
assert await r.get('c') == b('c1')
async def test_pipeline_eval(self, r):
await r.flushdb()
async with await r.pipeline(transaction=False) as pipe:
await pipe.eval("return {KEYS[1],KEYS[2],ARGV[1],ARGV[2]}", 2,
"A{foo}", "B{foo}", "first", "second")
res = (await pipe.execute())[0]
assert res[0] == b('A{foo}')
assert res[1] == b('B{foo}')
assert res[2] == b('first')
assert res[3] == b('second')
def encode(self, value):
"Return a bytestring representation of the value"
if isinstance(value, bytes):
return value
elif isinstance(value, int):
value = b(str(value))
elif isinstance(value, float):
value = b(repr(value))
elif not isinstance(value, str):
value = str(value).encode()
return value
def _trans_type(self, content):
if isinstance(content, str):
content = content.encode(self.encoding)
elif isinstance(content, int):
content = b(str(content))
elif isinstance(content, float):
content = b(repr(content))
if not isinstance(content, bytes):
raise TypeError('Wrong data type({}) to compress'.format(
type(content)))
return content
async def test_pubsub_numsub(self, r):
p1 = r.pubsub(ignore_subscribe_messages=True)
await p1.subscribe('foo', 'bar', 'baz')
p2 = r.pubsub(ignore_subscribe_messages=True)
await p2.subscribe('bar', 'baz')
p3 = r.pubsub(ignore_subscribe_messages=True)
await p3.subscribe('baz')
channels = [(b('foo'), 1), (b('bar'), 2), (b('baz'), 3)]
assert channels == await r.pubsub_numsub('foo', 'bar', 'baz')
await p1.unsubscribe()
await p2.unsubscribe()
await p3.unsubscribe()
async def _zaggregate(self, command, dest, keys, aggregate=None):
pieces = [command, dest, len(keys)]
if isinstance(keys, dict):
keys, weights = iterkeys(keys), itervalues(keys)
else:
weights = None
pieces.extend(keys)
if weights:
pieces.append(b('WEIGHTS'))
pieces.extend(weights)
if aggregate:
pieces.append(b('AGGREGATE'))
pieces.append(aggregate)
return await self.execute_command(*pieces)
async def my_transaction(pipe):
a_value = await pipe.get('a')
assert a_value in (b('1'), b('2'))
b_value = await pipe.get('b')
assert b_value == b('2')
# silly run-once code... incr's "a" so WatchError should be raised
# forcing this all to run again. this should incr "a" once to "2"
if not has_run:
await r.incr('a')
has_run.append('it has')
pipe.multi()
await pipe.set('c', int(a_value) + int(b_value))
async def hscan(self, name, cursor=0, match=None, count=None):
"""
Incrementallys return key/value slices in a hash. Also returns a
cursor pointing to the scan position.
``match`` allows for filtering the keys by pattern
``count`` allows for hint the minimum number of returns
"""
pieces = [name, cursor]
if match is not None:
pieces.extend([b('MATCH'), match])
if count is not None:
pieces.extend([b('COUNT'), count])
return await self.execute_command('HSCAN', *pieces)
async def sscan(self, name, cursor=0, match=None, count=None):
"""
Incrementally return lists of elements in a set. Also return a cursor
indicating the scan position.
``match`` allows for filtering the keys by pattern
``count`` allows for hint the minimum number of returns
"""
pieces = [name, cursor]
if match is not None:
pieces.extend([b('MATCH'), match])
if count is not None:
pieces.extend([b('COUNT'), count])
return await self.execute_command('SSCAN', *pieces)
async def acquire(self, blocking=None, blocking_timeout=None):
"""
Use Redis to hold a shared, distributed lock named ``name``.
Returns True once the lock is acquired.
If ``blocking`` is False, always return immediately. If the lock
was acquired, return True, otherwise return False.
``blocking_timeout`` specifies the maximum number of seconds to
wait trying to acquire the lock.
"""
sleep = self.sleep
token = b(uuid.uuid1().hex)
if blocking is None:
blocking = self.blocking
if blocking_timeout is None:
blocking_timeout = self.blocking_timeout
stop_trying_at = None
if blocking_timeout is not None:
stop_trying_at = mod_time.time() + blocking_timeout
while True:
if await self.do_acquire(token):
self.local.token = token
return True
if not blocking:
return False
if stop_trying_at is not None and mod_time.time() > stop_trying_at:
return False
await asyncio.sleep(sleep, loop=self.redis.connection_pool.loop)
def parse_slowlog_get(response, **options):
return [{
'id': item[0],
'start_time': int(item[1]),
'duration': int(item[2]),
'command': b(' ').join(item[3])
} for item in response]
async def zrevrange(self,
name,
start,
end,
withscores=False,
score_cast_func=float):
"""
Returns a range of values from sorted set ``name`` between
``start`` and ``end`` sorted in descending order.
``start`` and ``end`` can be negative, indicating the end of the range.
``withscores`` indicates to return the scores along with the values
The return type is a list of (value, score) pairs
``score_cast_func`` a callable used to cast the score return value
"""
pieces = ['ZREVRANGE', name, start, end]
if withscores:
pieces.append(b('WITHSCORES'))
options = {
'withscores': withscores,
'score_cast_func': score_cast_func
}
return await self.execute_command(*pieces, **options)
async def assert_moved_redirection_on_slave(sr, connection_pool_cls,
cluster_obj):
"""
"""
# we assume this key is set on 127.0.0.1:7000(7003)
await sr.set('foo16706', 'foo')
asyncio.sleep(1)
with patch.object(connection_pool_cls,
'get_node_by_slot') as return_slave_mock:
return_slave_mock.return_value = {
'name': '127.0.0.1:7004',
'host': '127.0.0.1',
'port': 7004,
'server_type': 'slave',
}
master_value = {
'host': '127.0.0.1',
'name': '127.0.0.1:7000',
'port': 7000,
'server_type': 'master'
}
with patch.object(ClusterConnectionPool,
'get_master_node_by_slot') as return_master_mock:
return_master_mock.return_value = master_value
assert await cluster_obj.get('foo16706') == b('foo')
assert return_master_mock.call_count == 1
async def test_transaction_callable(self, r):
await r.flushdb()
await r.set('a', 1)
await r.set('b', 2)
has_run = []
async def my_transaction(pipe):
a_value = await pipe.get('a')
assert a_value in (b('1'), b('2'))
b_value = await pipe.get('b')
assert b_value == b('2')
# silly run-once code... incr's "a" so WatchError should be raised
# forcing this all to run again. this should incr "a" once to "2"
if not has_run:
await r.incr('a')
has_run.append('it has')
pipe.multi()
await pipe.set('c', int(a_value) + int(b_value))
result = await r.transaction(my_transaction,
'a',
'b',
watch_delay=0.01)
assert result == [True]
assert await r.get('c') == b('4')
async def test_watch_succeed(self, r):
await r.flushdb()
await r.set('a', 1)
await r.set('b', 2)
async with await r.pipeline() as pipe:
await pipe.watch('a', 'b')
assert pipe.watching
a_value = await pipe.get('a')
b_value = await pipe.get('b')
assert a_value == b('1')
assert b_value == b('2')
pipe.multi()
await pipe.set('c', 3)
assert await pipe.execute() == [True]
assert not pipe.watching
def pack_command(self, *args):
"Pack a series of arguments into the Redis protocol"
output = []
# the client might have included 1 or more literal arguments in
# the command name, e.g., 'CONFIG GET'. The Redis server expects these
# arguments to be sent separately, so split the first argument
# manually. All of these arguements get wrapped in the Token class
# to prevent them from being encoded.
command = args[0]
if ' ' in command:
args = tuple([b(s) for s in command.split()]) + args[1:]
else:
args = (b(command), ) + args[1:]
buff = SYM_EMPTY.join((SYM_STAR, b(str(len(args))), SYM_CRLF))
for arg in map(self.encode, args):
# to avoid large string mallocs, chunk the command into the
# output list if we're sending large values
if len(buff) > 6000 or len(arg) > 6000:
buff = SYM_EMPTY.join(
(buff, SYM_DOLLAR, b(str(len(arg))), SYM_CRLF))
output.append(buff)
output.append(b(arg))
buff = SYM_CRLF
else:
buff = SYM_EMPTY.join((buff, SYM_DOLLAR, b(str(len(arg))),
SYM_CRLF, b(arg), SYM_CRLF))
output.append(buff)
return output
async def test_pipeline(self, r):
await r.flushdb()
async with await r.pipeline() as pipe:
await pipe.set('a', 'a1')
await pipe.get('a')
await pipe.zadd('z', z1=1)
await pipe.zadd('z', z2=4)
await pipe.zincrby('z', 'z1')
await pipe.zrange('z', 0, 5, withscores=True)
assert await pipe.execute() == \
[
True,
b('a1'),
1,
1,
2.0,
[(b('z1'), 2.0), (b('z2'), 4)],
]
async def zscan(self, name, cursor=0, match=None, count=None,
score_cast_func=float):
"""
Incrementally return lists of elements in a sorted set. Also return a
cursor indicating the scan position.
``match`` allows for filtering the keys by pattern
``count`` allows for hint the minimum number of returns
``score_cast_func`` a callable used to cast the score return value
"""
pieces = [name, cursor]
if match is not None:
pieces.extend([b('MATCH'), match])
if count is not None:
pieces.extend([b('COUNT'), count])
options = {'score_cast_func': score_cast_func}
return await self.execute_command('ZSCAN', *pieces, **options)
async def test_exec_error_in_no_transaction_pipeline(self, r):
await r.flushdb()
await r.set('a', 1)
async with await r.pipeline(transaction=False) as pipe:
await pipe.llen('a')
await pipe.expire('a', 100)
with pytest.raises(ResponseError) as ex:
await pipe.execute()
assert await r.get('a') == b('1')
async def test_exec_error_in_no_transaction_pipeline_unicode_command(self, r):
key = chr(11) + 'abcd' + chr(23)
await r.set(key, 1)
async with await r.pipeline(transaction=False) as pipe:
await pipe.llen(key)
await pipe.expire(key, 100)
with pytest.raises(ResponseError) as ex:
await pipe.execute()
assert await r.get(key) == b('1')
async def test_unwatch(self, r):
await r.flushdb()
await r.set('a', 1)
await r.set('b', 2)
async with await r.pipeline() as pipe:
await pipe.watch('a', 'b')
await r.set('b', 3)
await pipe.unwatch()
assert not pipe.watching
await pipe.get('a')
assert await pipe.execute() == [b('1')]
async def test_parse_error_raised(self, r):
async with await r.pipeline() as pipe:
# the zrem is invalid because we don't pass any keys to it
await (await (await pipe.set('a', 1)).zrem('b')).set('b', 2)
with pytest.raises(ResponseError) as ex:
await pipe.execute()
assert str(ex.value).startswith(
'Command # 2 (ZREM b) of pipeline caused error: ')
# make sure the pipe was restored to a working state
assert await (await pipe.set('z', 'zzz')).execute() == [True]
assert await r.get('z') == b('zzz')
