287 lines
9.2 KiB
Python
287 lines
9.2 KiB
Python
|
from collections import OrderedDict
|
||
|
|
|
||
|
|
from aioredis.util import wait_convert, wait_make_dict, wait_ok
|
||
|
|
|
||
|
|
|
||
|
|
def fields_to_dict(fields, type_=OrderedDict):
|
||
|
|
"""Convert a flat list of key/values into an OrderedDict"""
|
||
|
|
fields_iterator = iter(fields)
|
||
|
|
return type_(zip(fields_iterator, fields_iterator))
|
||
|
|
|
||
|
|
|
||
|
|
def parse_messages(messages):
|
||
|
|
""" Parse messages as returned by Redis into something useful
|
||
|
|
|
||
|
|
Messages returned by XRANGE arrive in the form:
|
||
|
|
|
||
|
|
[
|
||
|
|
[message_id, [key1, value1, key2, value2, ...]],
|
||
|
|
...
|
||
|
|
]
|
||
|
|
|
||
|
|
Here we parse this into:
|
||
|
|
|
||
|
|
[
|
||
|
|
[message_id, OrderedDict(
|
||
|
|
(key1, value1),
|
||
|
|
(key2, value2),
|
||
|
|
...
|
||
|
|
)],
|
||
|
|
...
|
||
|
|
]
|
||
|
|
|
||
|
|
"""
|
||
|
|
if messages is None:
|
||
|
|
return []
|
||
|
|
|
||
|
|
messages = (message for message in messages if message is not None)
|
||
|
|
return [
|
||
|
|
(mid, fields_to_dict(values))
|
||
|
|
for mid, values
|
||
|
|
in messages if values is not None
|
||
|
|
]
|
||
|
|
|
||
|
|
|
||
|
|
def parse_messages_by_stream(messages_by_stream):
|
||
|
|
""" Parse messages returned by stream
|
||
|
|
|
||
|
|
Messages returned by XREAD arrive in the form:
|
||
|
|
[stream_name,
|
||
|
|
[
|
||
|
|
[message_id, [key1, value1, key2, value2, ...]],
|
||
|
|
...
|
||
|
|
],
|
||
|
|
...
|
||
|
|
]
|
||
|
|
|
||
|
|
Here we parse this into (with the help of the above parse_messages()
|
||
|
|
function):
|
||
|
|
|
||
|
|
[
|
||
|
|
[stream_name, message_id, OrderedDict(
|
||
|
|
(key1, value1),
|
||
|
|
(key2, value2),.
|
||
|
|
...
|
||
|
|
)],
|
||
|
|
...
|
||
|
|
]
|
||
|
|
|
||
|
|
"""
|
||
|
|
if messages_by_stream is None:
|
||
|
|
return []
|
||
|
|
|
||
|
|
parsed = []
|
||
|
|
for stream, messages in messages_by_stream:
|
||
|
|
for message_id, fields in parse_messages(messages):
|
||
|
|
parsed.append((stream, message_id, fields))
|
||
|
|
return parsed
|
||
|
|
|
||
|
|
|
||
|
|
def parse_lists_to_dicts(lists):
|
||
|
|
""" Convert [[a, 1, b, 2], ...] into [{a:1, b: 2}, ...]"""
|
||
|
|
return [fields_to_dict(l, type_=dict) for l in lists]
|
||
|
|
|
||
|
|
|
||
|
|
class StreamCommandsMixin:
|
||
|
|
"""Stream commands mixin
|
||
|
|
|
||
|
|
Streams are available in Redis since v5.0
|
||
|
|
"""
|
||
|
|
|
||
|
|
def xadd(self, stream, fields, message_id=b'*', max_len=None,
|
||
|
|
exact_len=False):
|
||
|
|
"""Add a message to a stream."""
|
||
|
|
args = []
|
||
|
|
if max_len is not None:
|
||
|
|
if exact_len:
|
||
|
|
args.extend((b'MAXLEN', max_len))
|
||
|
|
else:
|
||
|
|
args.extend((b'MAXLEN', b'~', max_len))
|
||
|
|
|
||
|
|
args.append(message_id)
|
||
|
|
|
||
|
|
for k, v in fields.items():
|
||
|
|
args.extend([k, v])
|
||
|
|
return self.execute(b'XADD', stream, *args)
|
||
|
|
|
||
|
|
def xrange(self, stream, start='-', stop='+', count=None):
|
||
|
|
"""Retrieve messages from a stream."""
|
||
|
|
if count is not None:
|
||
|
|
extra = ['COUNT', count]
|
||
|
|
else:
|
||
|
|
extra = []
|
||
|
|
fut = self.execute(b'XRANGE', stream, start, stop, *extra)
|
||
|
|
return wait_convert(fut, parse_messages)
|
||
|
|
|
||
|
|
def xrevrange(self, stream, start='+', stop='-', count=None):
|
||
|
|
"""Retrieve messages from a stream in reverse order."""
|
||
|
|
if count is not None:
|
||
|
|
extra = ['COUNT', count]
|
||
|
|
else:
|
||
|
|
extra = []
|
||
|
|
fut = self.execute(b'XREVRANGE', stream, start, stop, *extra)
|
||
|
|
return wait_convert(fut, parse_messages)
|
||
|
|
|
||
|
|
def xread(self, streams, timeout=0, count=None, latest_ids=None):
|
||
|
|
"""Perform a blocking read on the given stream
|
||
|
|
|
||
|
|
:raises ValueError: if the length of streams and latest_ids do
|
||
|
|
not match
|
||
|
|
"""
|
||
|
|
args = self._xread(streams, timeout, count, latest_ids)
|
||
|
|
fut = self.execute(b'XREAD', *args)
|
||
|
|
return wait_convert(fut, parse_messages_by_stream)
|
||
|
|
|
||
|
|
def xread_group(self, group_name, consumer_name, streams, timeout=0,
|
||
|
|
count=None, latest_ids=None, no_ack=False):
|
||
|
|
"""Perform a blocking read on the given stream as part of a consumer group
|
||
|
|
|
||
|
|
:raises ValueError: if the length of streams and latest_ids do
|
||
|
|
not match
|
||
|
|
"""
|
||
|
|
args = self._xread(
|
||
|
|
streams, timeout, count, latest_ids, no_ack
|
||
|
|
)
|
||
|
|
fut = self.execute(
|
||
|
|
b'XREADGROUP', b'GROUP', group_name, consumer_name, *args
|
||
|
|
)
|
||
|
|
return wait_convert(fut, parse_messages_by_stream)
|
||
|
|
|
||
|
|
def xgroup_create(self, stream, group_name, latest_id='$', mkstream=False):
|
||
|
|
"""Create a consumer group"""
|
||
|
|
args = [b'CREATE', stream, group_name, latest_id]
|
||
|
|
if mkstream:
|
||
|
|
args.append(b'MKSTREAM')
|
||
|
|
fut = self.execute(b'XGROUP', *args)
|
||
|
|
return wait_ok(fut)
|
||
|
|
|
||
|
|
def xgroup_setid(self, stream, group_name, latest_id='$'):
|
||
|
|
"""Set the latest ID for a consumer group"""
|
||
|
|
fut = self.execute(b'XGROUP', b'SETID', stream, group_name, latest_id)
|
||
|
|
return wait_ok(fut)
|
||
|
|
|
||
|
|
def xgroup_destroy(self, stream, group_name):
|
||
|
|
"""Delete a consumer group"""
|
||
|
|
fut = self.execute(b'XGROUP', b'DESTROY', stream, group_name)
|
||
|
|
return wait_ok(fut)
|
||
|
|
|
||
|
|
def xgroup_delconsumer(self, stream, group_name, consumer_name):
|
||
|
|
"""Delete a specific consumer from a group"""
|
||
|
|
fut = self.execute(
|
||
|
|
b'XGROUP', b'DELCONSUMER', stream, group_name, consumer_name
|
||
|
|
)
|
||
|
|
return wait_convert(fut, int)
|
||
|
|
|
||
|
|
def xpending(self, stream, group_name, start=None, stop=None, count=None,
|
||
|
|
consumer=None):
|
||
|
|
"""Get information on pending messages for a stream
|
||
|
|
|
||
|
|
Returned data will vary depending on the presence (or not)
|
||
|
|
of the start/stop/count parameters. For more details see:
|
||
|
|
https://redis.io/commands/xpending
|
||
|
|
|
||
|
|
:raises ValueError: if the start/stop/count parameters are only
|
||
|
|
partially specified
|
||
|
|
"""
|
||
|
|
# Returns: total pel messages, min id, max id, count
|
||
|
|
ssc = [start, stop, count]
|
||
|
|
ssc_count = len([v for v in ssc if v is not None])
|
||
|
|
if ssc_count != 3 and ssc_count != 0:
|
||
|
|
raise ValueError(
|
||
|
|
'Either specify non or all of the start/stop/count arguments'
|
||
|
|
)
|
||
|
|
if not any(ssc):
|
||
|
|
ssc = []
|
||
|
|
|
||
|
|
args = [stream, group_name] + ssc
|
||
|
|
if consumer:
|
||
|
|
args.append(consumer)
|
||
|
|
return self.execute(b'XPENDING', *args)
|
||
|
|
|
||
|
|
def xclaim(self, stream, group_name, consumer_name, min_idle_time,
|
||
|
|
id, *ids):
|
||
|
|
"""Claim a message for a given consumer"""
|
||
|
|
fut = self.execute(
|
||
|
|
b'XCLAIM', stream, group_name, consumer_name, min_idle_time,
|
||
|
|
id, *ids
|
||
|
|
)
|
||
|
|
return wait_convert(fut, parse_messages)
|
||
|
|
|
||
|
|
def xack(self, stream, group_name, id, *ids):
|
||
|
|
"""Acknowledge a message for a given consumer group"""
|
||
|
|
return self.execute(b'XACK', stream, group_name, id, *ids)
|
||
|
|
|
||
|
|
def xdel(self, stream, id):
|
||
|
|
"""Removes the specified entries(IDs) from a stream"""
|
||
|
|
return self.execute(b'XDEL', stream, id)
|
||
|
|
|
||
|
|
def xtrim(self, stream, max_len, exact_len=False):
|
||
|
|
"""trims the stream to a given number of items, evicting older items"""
|
||
|
|
args = []
|
||
|
|
if exact_len:
|
||
|
|
args.extend((b'MAXLEN', max_len))
|
||
|
|
else:
|
||
|
|
args.extend((b'MAXLEN', b'~', max_len))
|
||
|
|
return self.execute(b'XTRIM', stream, *args)
|
||
|
|
|
||
|
|
def xlen(self, stream):
|
||
|
|
"""Returns the number of entries inside a stream"""
|
||
|
|
return self.execute(b'XLEN', stream)
|
||
|
|
|
||
|
|
def xinfo(self, stream):
|
||
|
|
"""Retrieve information about the given stream.
|
||
|
|
|
||
|
|
An alias for ``xinfo_stream()``
|
||
|
|
"""
|
||
|
|
return self.xinfo_stream(stream)
|
||
|
|
|
||
|
|
def xinfo_consumers(self, stream, group_name):
|
||
|
|
"""Retrieve consumers of a consumer group"""
|
||
|
|
fut = self.execute(b'XINFO', b'CONSUMERS', stream, group_name)
|
||
|
|
|
||
|
|
return wait_convert(fut, parse_lists_to_dicts)
|
||
|
|
|
||
|
|
def xinfo_groups(self, stream):
|
||
|
|
"""Retrieve the consumer groups for a stream"""
|
||
|
|
fut = self.execute(b'XINFO', b'GROUPS', stream)
|
||
|
|
return wait_convert(fut, parse_lists_to_dicts)
|
||
|
|
|
||
|
|
def xinfo_stream(self, stream):
|
||
|
|
"""Retrieve information about the given stream."""
|
||
|
|
fut = self.execute(b'XINFO', b'STREAM', stream)
|
||
|
|
return wait_make_dict(fut)
|
||
|
|
|
||
|
|
def xinfo_help(self):
|
||
|
|
"""Retrieve help regarding the ``XINFO`` sub-commands"""
|
||
|
|
fut = self.execute(b'XINFO', b'HELP')
|
||
|
|
return wait_convert(fut, lambda l: b'\n'.join(l))
|
||
|
|
|
||
|
|
def _xread(self, streams, timeout=0, count=None, latest_ids=None,
|
||
|
|
no_ack=False):
|
||
|
|
"""Wraps up common functionality between ``xread()``
|
||
|
|
and ``xread_group()``
|
||
|
|
|
||
|
|
You should probably be using ``xread()`` or ``xread_group()`` directly.
|
||
|
|
"""
|
||
|
|
if latest_ids is None:
|
||
|
|
latest_ids = ['$'] * len(streams)
|
||
|
|
if len(streams) != len(latest_ids):
|
||
|
|
raise ValueError(
|
||
|
|
'The streams and latest_ids parameters must be of the '
|
||
|
|
'same length'
|
||
|
|
)
|
||
|
|
|
||
|
|
count_args = [b'COUNT', count] if count else []
|
||
|
|
if timeout is None:
|
||
|
|
block_args = []
|
||
|
|
elif not isinstance(timeout, int):
|
||
|
|
raise TypeError(
|
||
|
|
"timeout argument must be int, not {!r}".format(timeout))
|
||
|
|
else:
|
||
|
|
block_args = [b'BLOCK', timeout]
|
||
|
|
|
||
|
|
noack_args = [b'NOACK'] if no_ack else []
|
||
|
|
|
||
|
|
return count_args + block_args + noack_args + [b'STREAMS'] + streams \
|
||
|
|
+ latest_ids
|