Beispiel #1
0
class ClusterSentinelCommands(SentinelCommandMixin):
    NODES_FLAGS = dict_merge(
        list_keys_to_dict([
            'SENTINEL GET-MASTER-ADDR-BY-NAME', 'SENTINEL MASTER',
            'SENTINEL MASTERS', 'SENTINEL MONITOR', 'SENTINEL REMOVE',
            'SENTINEL SENTINELS', 'SENTINEL SET', 'SENTINEL SLAVES'
        ], NodeFlag.BLOCKED))
Beispiel #2
0
class ClusterServerCommandMixin(ServerCommandMixin):
    NODES_FLAGS = dict_merge(
        list_keys_to_dict(['SHUTDOWN', 'SLAVEOF', 'CLIENT SETNAME'],
                          NodeFlag.BLOCKED),
        list_keys_to_dict(['FLUSHALL', 'FLUSHDB'], NodeFlag.ALL_MASTERS),
        list_keys_to_dict([
            'SLOWLOG LEN', 'SLOWLOG RESET', 'SLOWLOG GET', 'TIME', 'SAVE',
            'LASTSAVE', 'DBSIZE', 'CONFIG RESETSTAT', 'CONFIG REWRITE',
            'CONFIG GET', 'CONFIG SET', 'CLIENT KILL', 'CLIENT LIST',
            'CLIENT GETNAME', 'INFO', 'BGSAVE', 'BGREWRITEAOF'
        ], NodeFlag.ALL_NODES))

    RESULT_CALLBACKS = dict_merge(
        list_keys_to_dict([
            'CONFIG GET', 'CONFIG SET', 'SLOWLOG GET', 'CLIENT KILL', 'INFO',
            'BGREWRITEAOF', 'BGSAVE', 'CLIENT LIST', 'CLIENT GETNAME',
            'CONFIG RESETSTAT', 'CONFIG REWRITE', 'DBSIZE', 'LASTSAVE', 'SAVE',
            'SLOWLOG LEN', 'SLOWLOG RESET', 'TIME', 'FLUSHALL', 'FLUSHDB'
        ], lambda res: res))
Beispiel #3
0
class ClusterScriptingCommandMixin(ScriptingCommandMixin):

    NODES_FLAGS = dict_merge({'SCRIPT KILL': NodeFlag.BLOCKED},
                             list_keys_to_dict([
                                 "SCRIPT LOAD",
                                 "SCRIPT FLUSH",
                                 "SCRIPT EXISTS",
                             ], NodeFlag.ALL_MASTERS))

    RESULT_CALLBACKS = {
        "SCRIPT LOAD": lambda res: list(res.values()).pop(),
        "SCRIPT EXISTS": lambda res: [all(k) for k in zip(*res.values())],
        "SCRIPT FLUSH": lambda res: all(res.values())
    }
Beispiel #4
0
class CLusterPubSubCommandMixin(PubSubCommandMixin):
    
    NODES_FLAGS = dict_merge(
        list_keys_to_dict(
            ['PUBSUB CHANNELS', 'PUBSUB NUMSUB', 'PUBSUB NUMPAT'],
            NodeFlag.ALL_NODES
        )
    )

    RESULT_CALLBACKS = dict_merge(
        list_keys_to_dict([
            "PUBSUB CHANNELS",
        ], parse_cluster_pubsub_channels),
        list_keys_to_dict([
            "PUBSUB NUMSUB",
        ], parse_cluster_pubsub_numsub),
        list_keys_to_dict([
            "PUBSUB NUMPAT",
        ], parse_cluster_pubsub_numpat),
    )

    def pubsub(self, **kwargs):
        return ClusterPubSub(self.connection_pool, **kwargs)
Beispiel #5
0
 def __init__(self,
              connection_pool,
              result_callbacks=None,
              response_callbacks=None,
              startup_nodes=None):
     """
     """
     self.command_stack = []
     self.refresh_table_asap = False
     self.connection_pool = connection_pool
     self.result_callbacks = result_callbacks or self.__class__.RESULT_CALLBACKS.copy(
     )
     self.startup_nodes = startup_nodes if startup_nodes else []
     self.nodes_flags = self.__class__.NODES_FLAGS.copy()
     self.response_callbacks = dict_merge(
         response_callbacks or self.__class__.RESPONSE_CALLBACKS.copy())
Beispiel #6
0
class HyperLogCommandMixin:

    RESPONSE_CALLBACKS = dict_merge(string_keys_to_dict('PFADD PFCOUNT', int),
                                    {
                                        'PFMERGE': bool_ok,
                                    })

    async def pfadd(self, name, *values):
        "Adds the specified elements to the specified HyperLogLog."
        return await self.execute_command('PFADD', name, *values)

    async def pfcount(self, *sources):
        """
        Return the approximated cardinality of
        the set observed by the HyperLogLog at key(s).
        """
        return await self.execute_command('PFCOUNT', *sources)

    async def pfmerge(self, dest, *sources):
        "Merge N different HyperLogLogs into a single one."
        return await self.execute_command('PFMERGE', dest, *sources)
Beispiel #7
0
class StrictRedis(*mixins):
    """
    Implementation of the Redis protocol.

    This abstract class provides a Python interface to all Redis commands
    and an implementation of the Redis protocol.

    Connection and Pipeline derive from this, implementing how
    the commands are sent and received to the Redis server
    """

    RESPONSE_CALLBACKS = dict_merge(
        *[mixin.RESPONSE_CALLBACKS for mixin in mixins]
    )

    @classmethod
    def from_url(cls, url, db=None, **kwargs):
        """
        Return a Redis client object configured from the given URL, which must
        use either `the ``redis://`` scheme
        <http://www.iana.org/assignments/uri-schemes/prov/redis>`_ for RESP
        connections or the ``unix://`` scheme for Unix domain sockets.

        For example::

            redis://[:password]@localhost:6379/0
            unix://[:password]@/path/to/socket.sock?db=0

        There are several ways to specify a database number. The parse function
        will return the first specified option:
            1. A ``db`` querystring option, e.g. redis://localhost?db=0
            2. If using the redis:// scheme, the path argument of the url, e.g.
               redis://localhost/0
            3. The ``db`` argument to this function.

        If none of these options are specified, db=0 is used.

        Any additional querystring arguments and keyword arguments will be
        passed along to the ConnectionPool class's initializer. In the case
        of conflicting arguments, querystring arguments always win.
        """
        connection_pool = ConnectionPool.from_url(url, db=db, **kwargs)
        return cls(connection_pool=connection_pool)

    def __init__(self, host='localhost', port=6379,
                 db=0, password=None, stream_timeout=None,
                 connect_timeout=None, connection_pool=None,
                 unix_socket_path=None,
                 ssl=False, ssl_keyfile=None, ssl_certfile=None,
                 ssl_cert_reqs=None, ssl_ca_certs=None,
                 max_connections=None, retry_on_timeout=False):
        if not connection_pool:
            kwargs = {
                'db': db,
                'password': password,
                'stream_timeout': stream_timeout,
                'connect_timeout': connect_timeout,
                'max_connections': max_connections,
                'retry_on_timeout': retry_on_timeout
            }
            # based on input, setup appropriate connection args
            if unix_socket_path is not None:
                kwargs.update({
                    'path': unix_socket_path,
                    'connection_class': UnixDomainSocketConnection
                })
            else:
                # TCP specific options
                kwargs.update({
                    'host': host,
                    'port': port
                })

                if ssl:
                    kwargs.update({
                        'ssl_keyfile': ssl_keyfile,
                        'ssl_certfile': ssl_certfile,
                        'ssl_cert_reqs': ssl_cert_reqs,
                        'ssl_ca_certs': ssl_ca_certs,
                    })
            connection_pool = ConnectionPool(**kwargs)
        self.connection_pool = connection_pool
        self._use_lua_lock = None

        self.response_callbacks = self.__class__.RESPONSE_CALLBACKS.copy()

    def __repr__(self):
        return "{}<{}>".format(type(self).__name__, repr(self.connection_pool))

    def set_response_callback(self, command, callback):
        "Set a custom Response Callback"
        self.response_callbacks[command] = callback

    # COMMAND EXECUTION AND PROTOCOL PARSING
    async def execute_command(self, *args, **options):
        "Execute a command and return a parsed response"
        pool = self.connection_pool
        command_name = args[0]
        connection = pool.get_connection()
        try:
            await connection.send_command(*args)
            return await self.parse_response(connection, command_name, **options)
        except (ConnectionError, TimeoutError) as e:
            connection.disconnect()
            if not connection.retry_on_timeout and isinstance(e, TimeoutError):
                raise
            await connection.send_command(*args)
            return await self.parse_response(connection, command_name, **options)
        finally:
            pool.release(connection)

    async def parse_response(self, connection, command_name, **options):
        "Parses a response from the Redis server"
        response = await connection.read_response()
        if command_name in self.response_callbacks:
            callback = self.response_callbacks[command_name]
            return callback(response, **options)
        return response
Beispiel #8
0
class StrictRedisCluster(StrictRedis, *cluster_mixins):
    """
    If a command is implemented over the one in StrictRedis then it requires some changes compared to
    the regular implementation of the method.
    """
    RedisClusterRequestTTL = 16
    NODES_FLAGS = dict_merge(*[mixin.NODES_FLAGS
                               for mixin in cluster_mixins
                               if hasattr(mixin, 'NODES_FLAGS')])
    RESULT_CALLBACKS = dict_merge(*[mixin.RESULT_CALLBACKS
                                    for mixin in cluster_mixins
                                    if hasattr(mixin, 'RESULT_CALLBACKS')])

    def __init__(self, host=None, port=None, startup_nodes=None, max_connections=32,
                 max_connections_per_node=False, readonly=False,
                 reinitialize_steps=None, skip_full_coverage_check=False,
                 nodemanager_follow_cluster=False, **kwargs):
        """
        :startup_nodes:
        List of nodes that initial bootstrapping can be done from
        :host:
        Can be used to point to a startup node
        :port:
        Can be used to point to a startup node
        :max_connections:
        Maximum number of connections that should be kept open at one time
        :readonly:
        enable READONLY mode. You can read possibly stale data from slave.
        :skip_full_coverage_check:
        Skips the check of cluster-require-full-coverage config, useful for clusters
        without the CONFIG command (like aws)
        :nodemanager_follow_cluster:
        The node manager will during initialization try the last set of nodes that
        it was operating on. This will allow the client to drift along side the cluster
        if the cluster nodes move around alot.
        :**kwargs:
        Extra arguments that will be sent into StrictRedis instance when created
        (See Official redis-py doc for supported kwargs
        [https://github.com/andymccurdy/redis-py/blob/master/redis/client.py])
        Some kwargs is not supported and will raise RedisClusterException
        - db (Redis do not support database SELECT in cluster mode)
        """
        # Tweaks to StrictRedis client arguments when running in cluster mode
        if "db" in kwargs:
            raise RedisClusterException("Argument 'db' is not possible to use in cluster mode")
        if "connection_pool" in kwargs:
            pool = kwargs.pop('connection_pool')
        else:
            startup_nodes = [] if startup_nodes is None else startup_nodes

            # Support host/port as argument
            if host:
                startup_nodes.append({"host": host, "port": port if port else 7000})
            pool = ClusterConnectionPool(
                startup_nodes=startup_nodes,
                max_connections=max_connections,
                reinitialize_steps=reinitialize_steps,
                max_connections_per_node=max_connections_per_node,
                skip_full_coverage_check=skip_full_coverage_check,
                nodemanager_follow_cluster=nodemanager_follow_cluster,
                readonly=readonly,
                **kwargs
            )

        super(StrictRedisCluster, self).__init__(connection_pool=pool, **kwargs)

        self.refresh_table_asap = False
        self.nodes_flags = self.__class__.NODES_FLAGS.copy()
        self.result_callbacks = self.__class__.RESULT_CALLBACKS.copy()
        self.response_callbacks = self.__class__.RESPONSE_CALLBACKS.copy()

    @classmethod
    def from_url(cls, url, db=None, skip_full_coverage_check=False, **kwargs):
        """
        Return a Redis client object configured from the given URL, which must
        use either `the ``redis://`` scheme
        <http://www.iana.org/assignments/uri-schemes/prov/redis>`_ for RESP
        connections or the ``unix://`` scheme for Unix domain sockets.
        For example::
        redis://[:password]@localhost:6379/0
        unix://[:password]@/path/to/socket.sock?db=0
        There are several ways to specify a database number. The parse function
        will return the first specified option:
        1. A ``db`` querystring option, e.g. redis://localhost?db=0
        2. If using the redis:// scheme, the path argument of the url, e.g.
        redis://localhost/0
        3. The ``db`` argument to this function.
        If none of these options are specified, db=0 is used.
        Any additional querystring arguments and keyword arguments will be
        passed along to the ConnectionPool class's initializer. In the case
        of conflicting arguments, querystring arguments always win.
        """
        connection_pool = ClusterConnectionPool.from_url(url, db=db, **kwargs)
        return cls(connection_pool=connection_pool, skip_full_coverage_check=skip_full_coverage_check)

    def __repr__(self):
        """
        """
        servers = list({'{0}:{1}'.format(info['host'], info['port'])
                        for info in self.connection_pool.nodes.startup_nodes})
        servers.sort()
        return "{0}<{1}>".format(type(self).__name__, ', '.join(servers))

    def set_result_callback(self, command, callback):
        "Set a custom Result Callback"
        self.result_callbacks[command] = callback

    def _determine_slot(self, *args):
        """
        figure out what slot based on command and args
        """
        if len(args) <= 1:
            raise RedisClusterException("No way to dispatch this command to Redis Cluster. Missing key.")
        command = args[0]

        if command in ['EVAL', 'EVALSHA']:
            numkeys = args[2]
            keys = args[3: 3 + numkeys]
            slots = {self.connection_pool.nodes.keyslot(key) for key in keys}
            if len(slots) != 1:
                raise RedisClusterException("{0} - all keys must map to the same key slot".format(command))
            return slots.pop()

        key = args[1]

        return self.connection_pool.nodes.keyslot(key)

    def _merge_result(self, command, res, **kwargs):
        """
        `res` is a dict with the following structure Dict(NodeName, CommandResult)
        """
        if command in self.result_callbacks:
            return self.result_callbacks[command](res, **kwargs)

        # Default way to handle result
        return first_key(res)

    def determine_node(self, *args, **kwargs):
        """
        """
        command = args[0]
        node_flag = self.nodes_flags.get(command)

        if node_flag == NodeFlag.BLOCKED:
            return blocked_command(self, command)
        elif node_flag == NodeFlag.RANDOM:
            return [self.connection_pool.nodes.random_node()]
        elif node_flag == NodeFlag.ALL_MASTERS:
            return self.connection_pool.nodes.all_masters()
        elif node_flag == NodeFlag.ALL_NODES:
            return self.connection_pool.nodes.all_nodes()
        elif node_flag == NodeFlag.SLOT_ID:
            # if node flag of command is SLOT_ID
            # `slot_id` should is assumed in kwargs
            slot = kwargs.get('slot_id')
            if not slot:
                raise RedisClusterException('slot_id is needed to execute command {}'
                                            .format(command))
            return [self.connection_pool.nodes.node_from_slot(slot)]
        else:
            return None

    @clusterdown_wrapper
    async def execute_command(self, *args, **kwargs):
        """
        Send a command to a node in the cluster
        """
        if not self.connection_pool.initialized:
            await self.connection_pool.initialize()
        if not args:
            raise RedisClusterException("Unable to determine command to use")

        command = args[0]

        node = self.determine_node(*args, **kwargs)
        if node:
            return await self._execute_command_on_nodes(node, *args, **kwargs)

        # If set externally we must update it before calling any commands
        if self.refresh_table_asap:
            await self.connection_pool.nodes.initialize()
            self.refresh_table_asap = False

        redirect_addr = None
        asking = False

        try_random_node = False
        slot = self._determine_slot(*args)
        ttl = int(self.RedisClusterRequestTTL)

        while ttl > 0:
            ttl -= 1

            if asking:
                node = self.connection_pool.nodes.nodes[redirect_addr]
                r = self.connection_pool.get_connection_by_node(node)
            elif try_random_node:
                r = self.connection_pool.get_random_connection()
                try_random_node = False
            else:
                if self.refresh_table_asap:
                    # MOVED
                    node = self.connection_pool.get_master_node_by_slot(slot)
                else:
                    node = self.connection_pool.get_node_by_slot(slot)
                r = self.connection_pool.get_connection_by_node(node)

            try:
                if asking:
                    await r.send_command('ASKING')
                    await self.parse_response(r, "ASKING", **kwargs)
                    asking = False

                await r.send_command(*args)
                return await self.parse_response(r, command, **kwargs)
            except (RedisClusterException, BusyLoadingError):
                raise
            except (CancelledError, ConnectionError, TimeoutError):
                try_random_node = True

                if ttl < self.RedisClusterRequestTTL / 2:
                    await asyncio.sleep(0.1)
            except ClusterDownError as e:
                self.connection_pool.disconnect()
                self.connection_pool.reset()
                self.refresh_table_asap = True

                raise e
            except MovedError as e:
                # Reinitialize on ever x number of MovedError.
                # This counter will increase faster when the same client object
                # is shared between multiple threads. To reduce the frequency you
                # can set the variable 'reinitialize_steps' in the constructor.
                self.refresh_table_asap = True
                await self.connection_pool.nodes.increment_reinitialize_counter()

                node = self.connection_pool.nodes.set_node(e.host, e.port, server_type='master')
                self.connection_pool.nodes.slots[e.slot_id][0] = node
            except TryAgainError as e:
                if ttl < self.RedisClusterRequestTTL / 2:
                    await asyncio.sleep(0.05)
            except AskError as e:
                redirect_addr, asking = "{0}:{1}".format(e.host, e.port), True
            finally:
                self.connection_pool.release(r)

        raise ClusterError('TTL exhausted.')

    async def _execute_command_on_nodes(self, nodes, *args, **kwargs):
        command = args[0]
        res = {}

        for node in nodes:
            connection = self.connection_pool.get_connection_by_node(node)

            # copy from redis-py
            try:
                await connection.send_command(*args)
                res[node["name"]] = await self.parse_response(connection, command, **kwargs)
            except CancelledError:
                # do not retry when coroutine is cancelled
                connection.disconnect()
            except (ConnectionError, TimeoutError) as e:
                connection.disconnect()

                if not connection.retry_on_timeout and isinstance(e, TimeoutError):
                    raise

                await connection.send_command(*args)
                res[node["name"]] = await self.parse_response(connection, command, **kwargs)
            finally:
                self.connection_pool.release(connection)
        return self._merge_result(command, res, **kwargs)

    async def pipeline(self, transaction=None, shard_hint=None, watches=None):
        """
        Cluster impl:
            Pipelines do not work in cluster mode the same way they do in normal mode.
            Create a clone of this object so that simulating pipelines will work correctly.
            Each command will be called directly when used and when calling execute() will only return the result stack.
        cluster transaction can only be run with commands in the same node, otherwise error will be raised.
        """
        await self.connection_pool.initialize()
        if shard_hint:
            raise RedisClusterException("shard_hint is deprecated in cluster mode")

        from aredis.pipeline import StrictClusterPipeline
        return StrictClusterPipeline(
            connection_pool=self.connection_pool,
            startup_nodes=self.connection_pool.nodes.startup_nodes,
            result_callbacks=self.result_callbacks,
            response_callbacks=self.response_callbacks,
            transaction=transaction,
            watches=watches
        )
Beispiel #9
0
class ListsCommandMixin:

    RESPONSE_CALLBACKS = dict_merge(
        string_keys_to_dict(
            'BLPOP BRPOP',
            lambda r: r and tuple(r) or None
        ),
        string_keys_to_dict(
            # these return OK, or int if redis-server is >=1.3.4
            'LPUSH RPUSH',
            lambda r: isinstance(r, int) and r or r == b'OK'
        ),
        string_keys_to_dict('LSET LTRIM', bool_ok),
        string_keys_to_dict('LINSERT LLEN LPUSHX RPUSHX', int),
    )

    async def blpop(self, keys, timeout=0):
        """
        LPOP a value off of the first non-empty list
        named in the ``keys`` list.

        If none of the lists in ``keys`` has a value to LPOP, then block
        for ``timeout`` seconds, or until a value gets pushed on to one
        of the lists.

        If timeout is 0, then block indefinitely.
        """
        if timeout is None:
            timeout = 0
        if isinstance(keys, str):
            keys = [keys]
        else:
            keys = list(keys)
        keys.append(timeout)
        return await self.execute_command('BLPOP', *keys)

    async def brpop(self, keys, timeout=0):
        """
        RPOP a value off of the first non-empty list
        named in the ``keys`` list.

        If none of the lists in ``keys`` has a value to LPOP, then block
        for ``timeout`` seconds, or until a value gets pushed on to one
        of the lists.

        If timeout is 0, then block indefinitely.
        """
        if timeout is None:
            timeout = 0
        if isinstance(keys, str):
            keys = [keys]
        else:
            keys = list(keys)
        keys.append(timeout)
        return await self.execute_command('BRPOP', *keys)

    async def brpoplpush(self, src, dst, timeout=0):
        """
        Pop a value off the tail of ``src``, push it on the head of ``dst``
        and then return it.

        This command blocks until a value is in ``src`` or until ``timeout``
        seconds elapse, whichever is first. A ``timeout`` value of 0 blocks
        forever.
        """
        if timeout is None:
            timeout = 0
        return await self.execute_command('BRPOPLPUSH', src, dst, timeout)

    async def lindex(self, name, index):
        """
        Return the item from list ``name`` at position ``index``

        Negative indexes are supported and will return an item at the
        end of the list
        """
        return await self.execute_command('LINDEX', name, index)

    async def linsert(self, name, where, refvalue, value):
        """
        Insert ``value`` in list ``name`` either immediately before or after
        [``where``] ``refvalue``

        Returns the new length of the list on success or -1 if ``refvalue``
        is not in the list.
        """
        return await self.execute_command('LINSERT', name, where, refvalue, value)

    async def llen(self, name):
        "Return the length of the list ``name``"
        return await self.execute_command('LLEN', name)

    async def lpop(self, name):
        "Remove and return the first item of the list ``name``"
        return await self.execute_command('LPOP', name)

    async def lpush(self, name, *values):
        "Push ``values`` onto the head of the list ``name``"
        return await self.execute_command('LPUSH', name, *values)

    async def lpushx(self, name, value):
        "Push ``value`` onto the head of the list ``name`` if ``name`` exists"
        return await self.execute_command('LPUSHX', name, value)

    async def lrange(self, name, start, end):
        """
        Return a slice of the list ``name`` between
        position ``start`` and ``end``

        ``start`` and ``end`` can be negative numbers just like
        Python slicing notation
        """
        return await self.execute_command('LRANGE', name, start, end)

    async def lrem(self, name, count, value):
        """
        Remove the first ``count`` occurrences of elements equal to ``value``
        from the list stored at ``name``.

        The count argument influences the operation in the following ways:
            count > 0: Remove elements equal to value moving from head to tail.
            count < 0: Remove elements equal to value moving from tail to head.
            count = 0: Remove all elements equal to value.
        """
        return await self.execute_command('LREM', name, count, value)

    async def lset(self, name, index, value):
        "Set ``position`` of list ``name`` to ``value``"
        return await self.execute_command('LSET', name, index, value)

    async def ltrim(self, name, start, end):
        """
        Trim the list ``name``, removing all values not within the slice
        between ``start`` and ``end``

        ``start`` and ``end`` can be negative numbers just like
        Python slicing notation
        """
        return await self.execute_command('LTRIM', name, start, end)

    async def rpop(self, name):
        "Remove and return the last item of the list ``name``"
        return await self.execute_command('RPOP', name)

    async def rpoplpush(self, src, dst):
        """
        RPOP a value off of the ``src`` list and atomically LPUSH it
        on to the ``dst`` list.  Returns the value.
        """
        return await self.execute_command('RPOPLPUSH', src, dst)

    async def rpush(self, name, *values):
        "Push ``values`` onto the tail of the list ``name``"
        return await self.execute_command('RPUSH', name, *values)

    async def rpushx(self, name, value):
        "Push ``value`` onto the tail of the list ``name`` if ``name`` exists"
        return await self.execute_command('RPUSHX', name, value)
Beispiel #10
0
class SetsCommandMixin:

    RESPONSE_CALLBACKS = dict_merge(
        string_keys_to_dict(
            'SADD SCARD SDIFFSTORE '
            'SETRANGE SINTERSTORE '
            'SREM SUNIONSTORE', int
        ),
        string_keys_to_dict(
            'SISMEMBER SMOVE', bool
        ),
        string_keys_to_dict(
            'SDIFF SINTER SMEMBERS SUNION',
            lambda r: r and set(r) or set()
        ),
        {
            'SSCAN': parse_sscan,
        }
    )

    async def sadd(self, name, *values):
        "Add ``value(s)`` to set ``name``"
        return await self.execute_command('SADD', name, *values)

    async def scard(self, name):
        "Return the number of elements in set ``name``"
        return await self.execute_command('SCARD', name)

    async def sdiff(self, keys, *args):
        "Return the difference of sets specified by ``keys``"
        args = list_or_args(keys, args)
        return await self.execute_command('SDIFF', *args)

    async def sdiffstore(self, dest, keys, *args):
        """
        Store the difference of sets specified by ``keys`` into a new
        set named ``dest``.  Returns the number of keys in the new set.
        """
        args = list_or_args(keys, args)
        return await self.execute_command('SDIFFSTORE', dest, *args)

    async def sinter(self, keys, *args):
        "Return the intersection of sets specified by ``keys``"
        args = list_or_args(keys, args)
        return await self.execute_command('SINTER', *args)

    async def sinterstore(self, dest, keys, *args):
        """
        Store the intersection of sets specified by ``keys`` into a new
        set named ``dest``.  Returns the number of keys in the new set.
        """
        args = list_or_args(keys, args)
        return await self.execute_command('SINTERSTORE', dest, *args)

    async def sismember(self, name, value):
        "Return a boolean indicating if ``value`` is a member of set ``name``"
        return await self.execute_command('SISMEMBER', name, value)

    async def smembers(self, name):
        "Return all members of the set ``name``"
        return await self.execute_command('SMEMBERS', name)

    async def smove(self, src, dst, value):
        "Move ``value`` from set ``src`` to set ``dst`` atomically"
        return await self.execute_command('SMOVE', src, dst, value)

    async def spop(self, name):
        "Remove and return a random member of set ``name``"
        return await self.execute_command('SPOP', name)

    async def srandmember(self, name, number=None):
        """
        If ``number`` is None, returns a random member of set ``name``.

        If ``number`` is supplied, returns a list of ``number`` random
        memebers of set ``name``. Note this is only available when running
        Redis 2.6+.
        """
        args = number and [number] or []
        return await self.execute_command('SRANDMEMBER', name, *args)

    async def srem(self, name, *values):
        "Remove ``values`` from set ``name``"
        return await self.execute_command('SREM', name, *values)

    async def sunion(self, keys, *args):
        "Return the union of sets specified by ``keys``"
        args = list_or_args(keys, args)
        return await self.execute_command('SUNION', *args)

    async def sunionstore(self, dest, keys, *args):
        """
        Store the union of sets specified by ``keys`` into a new
        set named ``dest``.  Returns the number of keys in the new set.
        """
        args = list_or_args(keys, args)
        return await self.execute_command('SUNIONSTORE', dest, *args)

    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)
Beispiel #11
0
class KeysCommandMixin:

    RESPONSE_CALLBACKS = dict_merge(
        string_keys_to_dict(
            'EXISTS EXPIRE EXPIREAT '
            'MOVE PERSIST RENAMENX', bool
        ),
        {
            'DEL': int,
            'SORT': sort_return_tuples,
            'OBJECT': parse_object,
            'RANDOMKEY': lambda r: r and r or None,
            'SCAN': parse_scan,
            'RENAME': bool_ok,
        }
    )

    async def delete(self, *names):
        """Delete one or more keys specified by ``names``"""
        return await self.execute_command('DEL', *names)

    async def dump(self, name):
        """
        Return a serialized version of the value stored at the specified key.
        If key does not exist a nil bulk reply is returned.
        """
        return await self.execute_command('DUMP', name)

    async def exists(self, name):
        """Returns a boolean indicating whether key ``name`` exists"""
        return await self.execute_command('EXISTS', name)

    async def expire(self, name, time):
        """
        Set an expire flag on key ``name`` for ``time`` seconds. ``time``
        can be represented by an integer or a Python timedelta object.
        """
        if isinstance(time, datetime.timedelta):
            time = time.seconds + time.days * 24 * 3600
        return await self.execute_command('EXPIRE', name, time)

    async def expireat(self, name, when):
        """
        Set an expire flag on key ``name``. ``when`` can be represented
        as an integer indicating unix time or a Python datetime object.
        """
        if isinstance(when, datetime.datetime):
            when = int(mod_time.mktime(when.timetuple()))
        return await self.execute_command('EXPIREAT', name, when)

    async def keys(self, pattern='*'):
        """Returns a list of keys matching ``pattern``"""
        return await self.execute_command('KEYS', pattern)

    async def move(self, name, db):
        """Moves the key ``name`` to a different Redis database ``db``"""
        return await self.execute_command('MOVE', name, db)

    async def object(self, infotype, key):
        """Returns the encoding, idletime, or refcount about the key"""
        return await self.execute_command('OBJECT', infotype, key, infotype=infotype)

    async def persist(self, name):
        """Removes an expiration on ``name``"""
        return await self.execute_command('PERSIST', name)

    async def pexpire(self, name, time):
        """
        Set an expire flag on key ``name`` for ``time`` milliseconds.
        ``time`` can be represented by an integer or a Python timedelta
        object.
        """
        if isinstance(time, datetime.timedelta):
            ms = int(time.microseconds / 1000)
            time = (time.seconds + time.days * 24 * 3600) * 1000 + ms
        return await self.execute_command('PEXPIRE', name, time)

    async def pexpireat(self, name, when):
        """
        Set an expire flag on key ``name``. ``when`` can be represented
        as an integer representing unix time in milliseconds (unix time * 1000)
        or a Python datetime object.
        """
        if isinstance(when, datetime.datetime):
            ms = int(when.microsecond / 1000)
            when = int(mod_time.mktime(when.timetuple())) * 1000 + ms
        return await self.execute_command('PEXPIREAT', name, when)

    async def pttl(self, name):
        """
        Returns the number of milliseconds until the key ``name`` will expire
        """
        return await self.execute_command('PTTL', name)

    async def randomkey(self):
        """Returns the name of a random key"""
        return await self.execute_command('RANDOMKEY')

    async def rename(self, src, dst):
        """
        Renames key ``src`` to ``dst``
        """
        return await self.execute_command('RENAME', src, dst)

    async def renamenx(self, src, dst):
        """Renames key ``src`` to ``dst`` if ``dst`` doesn't already exist"""
        return await self.execute_command('RENAMENX', src, dst)

    async def restore(self, name, ttl, value, replace=False):
        """
        Creates a key using the provided serialized value, previously obtained
        using DUMP.
        """
        params = [name, ttl, value]
        if replace:
            params.append('REPLACE')
        return await self.execute_command('RESTORE', *params)

    async def sort(self, name, start=None, num=None, by=None, get=None,
             desc=False, alpha=False, store=None, groups=False):
        """
        Sorts and returns a list, set or sorted set at ``name``.

        ``start`` and ``num`` are for paginating sorted data

        ``by`` allows using an external key to weight and sort the items.
            Use an "*" to indicate where in the key the item value is located

        ``get`` is for returning items from external keys rather than the
            sorted data itself.  Use an "*" to indicate where int he key
            the item value is located

        ``desc`` is for reversing the sort

        ``alpha`` is for sorting lexicographically rather than numerically

        ``store`` is for storing the result of the sort into
            the key ``store``

        ``groups`` if set to True and if ``get`` contains at least two
            elements, sort will return a list of tuples, each containing the
            values fetched from the arguments to ``get``.

        """
        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 = [name]
        if by is not None:
            pieces.append(b('BY'))
            pieces.append(by)
        if start is not None and num is not None:
            pieces.append(b('LIMIT'))
            pieces.append(start)
            pieces.append(num)
        if get is not None:
            # If get is a string assume we want to get a single value.
            # Otherwise assume it's an interable and we want to get multiple
            # values. We can't just iterate blindly because strings are
            # iterable.
            if isinstance(get, str):
                pieces.append(b('GET'))
                pieces.append(get)
            else:
                for g in get:
                    pieces.append(b('GET'))
                    pieces.append(g)
        if desc:
            pieces.append(b('DESC'))
        if alpha:
            pieces.append(b('ALPHA'))
        if store is not None:
            pieces.append(b('STORE'))
            pieces.append(store)

        if groups:
            if not get or isinstance(get, str) or len(get) < 2:
                raise DataError('when using "groups" the "get" argument '
                                'must be specified and contain at least '
                                'two keys')

        options = {'groups': len(get) if groups else None}
        return await self.execute_command('SORT', *pieces, **options)

    async def touch(self, keys):
        """
        Alters the last access time of a key(s).
        A key is ignored if it does not exist.
        """
        return await self.execute_command('TOUCH', *keys)

    async def ttl(self, name):
        """Returns the number of seconds until the key ``name`` will expire"""
        return await self.execute_command('TTL', name)

    async def type(self, name):
        """Returns the type of key ``name``"""
        return await self.execute_command('TYPE', name)

    async def unlink(self, *keys):
        """Removes the specified keys in a different thread, not blocking"""
        return await self.execute_command('UNLINK', *keys)

    async def wait(self, num_replicas, timeout):
        """
        Redis synchronous replication
        That returns the number of replicas that processed the query when
        we finally have at least ``num_replicas``, or when the ``timeout`` was
        reached.
        """
        return await self.execute_command('WAIT', num_replicas, timeout)

    async def scan(self, cursor=0, match=None, count=None):
        """
        Incrementally return lists of key names. 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 = [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('SCAN', *pieces)
Beispiel #12
0
class ClusterKeysCommandMixin(KeysCommandMixin):

    NODES_FLAGS = dict_merge(
        {
        'MOVE': NodeFlag.BLOCKED,
        'RANDOMKEY': NodeFlag.RANDOM,
        'SCAN': NodeFlag.ALL_MASTERS,
        },
        list_keys_to_dict(
            ['KEYS'],
            NodeFlag.ALL_NODES
        )
    )

    RESULT_CALLBACKS = {
        'KEYS': merge_result,
        'RANDOMKEY': first_key,
        'SCAN': lambda res: res
    }

    async def rename(self, src, dst):
        """
        Rename key ``src`` to ``dst``

        Cluster impl:
            This operation is no longer atomic because each key must be querried
            then set in separate calls because they maybe will change cluster node
        """
        if src == dst:
            raise ResponseError("source and destination objects are the same")

        data = await self.dump(src)

        if data is None:
            raise ResponseError("no such key")

        ttl = await self.pttl(src)

        if ttl is None or ttl < 1:
            ttl = 0

        await self.delete(dst)
        await self.restore(dst, ttl, data)
        await self.delete(src)

        return True

    async def delete(self, *names):
        """
        "Delete one or more keys specified by ``names``"

        Cluster impl:
            Iterate all keys and send DELETE for each key.
            This will go a lot slower than a normal delete call in StrictRedis.

            Operation is no longer atomic.
        """
        count = 0

        for arg in names:
            count += await self.execute_command('DEL', arg)

        return count

    async def renamenx(self, src, dst):
        """
        Rename key ``src`` to ``dst`` if ``dst`` doesn't already exist

        Cluster impl:
            Check if dst key do not exists, then calls rename().

            Operation is no longer atomic.
        """
        if not await self.exists(dst):
            return await self.rename(src, dst)

        return False
Beispiel #13
0
class StringsCommandMixin:
    RESPONSE_CALLBACKS = dict_merge(
        string_keys_to_dict('MSETNX PSETEX SETEX SETNX', bool),
        string_keys_to_dict(
            'BITCOUNT BITPOS DECRBY GETBIT INCRBY '
            'STRLEN SETBIT', int), {
                'INCRBYFLOAT': float,
                'MSET': bool_ok,
                'SET': lambda r: r and r == b'OK',
            })

    async def append(self, key, value):
        """
        Appends the string ``value`` to the value at ``key``. If ``key``
        doesn't already exist, create it with a value of ``value``.
        Returns the new length of the value at ``key``.
        """
        return await self.execute_command('APPEND', key, value)

    async def bitcount(self, key, start=None, end=None):
        """
        Returns the count of set bits in the value of ``key``.  Optional
        ``start`` and ``end`` paramaters indicate which bytes to consider
        """
        params = [key]
        if start is not None and end is not None:
            params.append(start)
            params.append(end)
        elif (start is not None and end is None) or \
                (end is not None and start is None):
            raise RedisError("Both start and end must be specified")
        return await self.execute_command('BITCOUNT', *params)

    async def bitop(self, operation, dest, *keys):
        """
        Perform a bitwise operation using ``operation`` between ``keys`` and
        store the result in ``dest``.
        """
        return await self.execute_command('BITOP', operation, dest, *keys)

    async def bitpos(self, key, bit, start=None, end=None):
        """
        Return the position of the first bit set to 1 or 0 in a string.
        ``start`` and ``end`` difines search range. The range is interpreted
        as a range of bytes and not a range of bits, so start=0 and end=2
        means to look at the first three bytes.
        """
        if bit not in (0, 1):
            raise RedisError('bit must be 0 or 1')
        params = [key, bit]

        start is not None and params.append(start)

        if start is not None and end is not None:
            params.append(end)
        elif start is None and end is not None:
            raise RedisError("start argument is not set, "
                             "when end is specified")
        return await self.execute_command('BITPOS', *params)

    def bitfield(self, key):
        return BitField(self, key)

    async def decr(self, name, amount=1):
        """
        Decrements the value of ``key`` by ``amount``.  If no key exists,
        the value will be initialized as 0 - ``amount``
        """
        return await self.execute_command('DECRBY', name, amount)

    async def get(self, name):
        """
        Return the value at key ``name``, or None if the key doesn't exist
        """
        return await self.execute_command('GET', name)

    async def getbit(self, name, offset):
        "Returns a boolean indicating the value of ``offset`` in ``name``"
        return await self.execute_command('GETBIT', name, offset)

    async def getrange(self, key, start, end):
        """
        Returns the substring of the string value stored at ``key``,
        determined by the offsets ``start`` and ``end`` (both are inclusive)
        """
        return await self.execute_command('GETRANGE', key, start, end)

    async def getset(self, name, value):
        """
        Sets the value at key ``name`` to ``value``
        and returns the old value at key ``name`` atomically.
        """
        return await self.execute_command('GETSET', name, value)

    async def incr(self, name, amount=1):
        """
        Increments the value of ``key`` by ``amount``.  If no key exists,
        the value will be initialized as ``amount``
        """
        return await self.execute_command('INCRBY', name, amount)

    async def incrby(self, name, amount=1):
        """
        Increments the value of ``key`` by ``amount``.  If no key exists,
        the value will be initialized as ``amount``
        """

        # An alias for ``incr()``, because it is already implemented
        # as INCRBY redis command.
        return await self.incr(name, amount)

    async def incrbyfloat(self, name, amount=1.0):
        """
        Increments the value at key ``name`` by floating ``amount``.
        If no key exists, the value will be initialized as ``amount``
        """
        return await self.execute_command('INCRBYFLOAT', name, amount)

    async def mget(self, keys, *args):
        """
        Returns a list of values ordered identically to ``keys``
        """
        args = list_or_args(keys, args)
        return await self.execute_command('MGET', *args)

    async def mset(self, *args, **kwargs):
        """
        Sets key/values based on a mapping. Mapping can be supplied as a single
        dictionary argument or as kwargs.
        """
        if args:
            if len(args) != 1 or not isinstance(args[0], dict):
                raise RedisError('MSET requires **kwargs or a single dict arg')
            kwargs.update(args[0])
        items = []
        for pair in iteritems(kwargs):
            items.extend(pair)
        return await self.execute_command('MSET', *items)

    async def msetnx(self, *args, **kwargs):
        """
        Sets key/values based on a mapping if none of the keys are already set.
        Mapping can be supplied as a single dictionary argument or as kwargs.
        Returns a boolean indicating if the operation was successful.
        """
        if args:
            if len(args) != 1 or not isinstance(args[0], dict):
                raise RedisError('MSETNX requires **kwargs or a single '
                                 'dict arg')
            kwargs.update(args[0])
        items = []
        for pair in iteritems(kwargs):
            items.extend(pair)
        return await self.execute_command('MSETNX', *items)

    async def psetex(self, name, time_ms, value):
        """
        Set the value of key ``name`` to ``value`` that expires in ``time_ms``
        milliseconds. ``time_ms`` can be represented by an integer or a Python
        timedelta object
        """
        if isinstance(time_ms, datetime.timedelta):
            ms = int(time_ms.microseconds / 1000)
            time_ms = (time_ms.seconds + time_ms.days * 24 * 3600) * 1000 + ms
        return await self.execute_command('PSETEX', name, time_ms, value)

    async def set(self, name, value, ex=None, px=None, nx=False, xx=False):
        """
        Set the value at key ``name`` to ``value``

        ``ex`` sets an expire flag on key ``name`` for ``ex`` seconds.

        ``px`` sets an expire flag on key ``name`` for ``px`` milliseconds.

        ``nx`` if set to True, set the value at key ``name`` to ``value`` if it
            does not already exist.

        ``xx`` if set to True, set the value at key ``name`` to ``value`` if it
            already exists.
        """
        pieces = [name, value]
        if ex is not None:
            pieces.append('EX')
            if isinstance(ex, datetime.timedelta):
                ex = ex.seconds + ex.days * 24 * 3600
            pieces.append(ex)
        if px is not None:
            pieces.append('PX')
            if isinstance(px, datetime.timedelta):
                ms = int(px.microseconds / 1000)
                px = (px.seconds + px.days * 24 * 3600) * 1000 + ms
            pieces.append(px)

        if nx:
            pieces.append('NX')
        if xx:
            pieces.append('XX')
        return await self.execute_command('SET', *pieces)

    async def setbit(self, name, offset, value):
        """
        Flag the ``offset`` in ``name`` as ``value``. Returns a boolean
        indicating the previous value of ``offset``.
        """
        value = value and 1 or 0
        return await self.execute_command('SETBIT', name, offset, value)

    async def setex(self, name, time, value):
        """
        Set the value of key ``name`` to ``value`` that expires in ``time``
        seconds. ``time`` can be represented by an integer or a Python
        timedelta object.
        """
        if isinstance(time, datetime.timedelta):
            time = time.seconds + time.days * 24 * 3600
        return await self.execute_command('SETEX', name, time, value)

    async def setnx(self, name, value):
        "Set the value of key ``name`` to ``value`` if key doesn't exist"
        return await self.execute_command('SETNX', name, value)

    async def setrange(self, name, offset, value):
        """
        Overwrite bytes in the value of ``name`` starting at ``offset`` with
        ``value``. If ``offset`` plus the length of ``value`` exceeds the
        length of the original value, the new value will be larger than before.
        If ``offset`` exceeds the length of the original value, null bytes
        will be used to pad between the end of the previous value and the start
        of what's being injected.

        Returns the length of the new string.
        """
        return await self.execute_command('SETRANGE', name, offset, value)

    async def strlen(self, name):
        "Return the number of bytes stored in the value of ``name``"
        return await self.execute_command('STRLEN', name)

    async def substr(self, name, start, end=-1):
        """
        Return a substring of the string at key ``name``. ``start`` and ``end``
        are 0-based integers specifying the portion of the string to return.
        """
        return await self.execute_command('SUBSTR', name, start, end)
Beispiel #14
0
class StreamsCommandMixin:
    RESPONSE_CALLBACKS = dict_merge(
        string_keys_to_dict('XREVRANGE XRANGE', stream_list),
        string_keys_to_dict('XREAD XREADGROUP', multi_stream_list), {
            'XINFO GROUPS': list_of_pairs_to_dict,
            'XINFO STREAM': parse_xinfo_stream,
            'XINFO CONSUMERS': list_of_pairs_to_dict,
            'XGROUP SETID': bool_ok,
            'XGROUP CREATE': bool_ok
        })

    async def xadd(self,
                   name: str,
                   entry: dict,
                   max_len=None,
                   stream_id='*',
                   approximate=True) -> str:
        """
        Appends the specified stream entry to the stream at the specified key.
        If the key does not exist, as a side effect of running
        this command the key is created with a stream value.
        Available since 5.0.0.
        Time complexity: O(log(N)) with N being the number of items already into the stream.

        :param name: name of the stream
        :param entry: key-values to be appended to the stream
        :param max_len: max length of the stream
        length will not be limited max_len is set to None
        notice: max_len should be int greater than 0,
        if set to 0 or negative, the stream length will not be limited

        :param stream_id: id of the options appended to the stream.
        The XADD command will auto-generate a unique id for you
        if the id argument specified is the * character.
        ID are specified by two numbers separated by a "-" character

        :param approximate: whether redis will limit
        the stream with given max length exactly, if set to True,
        there will be a few tens of entries more,
        but never less than 1000 items

        :return: id auto generated or the specified id given.
        notice: specified id without "-" character will be completed like "id-0"
        """
        pieces = []
        if max_len is not None:
            if not isinstance(max_len, int) or max_len < 1:
                raise RedisError("XADD maxlen must be a positive integer")
            pieces.append('MAXLEN')
            if approximate:
                pieces.append('~')
            pieces.append(str(max_len))
        pieces.append(stream_id)
        for kv in entry.items():
            pieces.extend(list(kv))
        return await self.execute_command('XADD', name, *pieces)

    async def xlen(self, name: str) -> int:
        """
        Returns the number of elements in a given stream.
        """
        return await self.execute_command('XLEN', name)

    async def xrange(self, name: str, start='-', end='+', count=None) -> list:
        """
        Read stream values within an interval.

        Available since 5.0.0.
        Time complexity: O(log(N)+M) with N being the number of elements in the stream and M the number
        of elements being returned. If M is constant (e.g. always asking for the first 10 elements with COUNT),
        you can consider it O(log(N)).

        :param name: name of the stream.
        :param start: first stream ID. defaults to '-',
               meaning the earliest available.
        :param end: last stream ID. defaults to '+',
                meaning the latest available.
        :param count: if set, only return this many items, beginning with the
               earliest available.
        :return list of (stream_id, entry(k-v pair))
        """

        pieces = [start, end]
        if count is not None:
            if not isinstance(count, int) or count < 1:
                raise RedisError("XRANGE count must be a positive integer")
            pieces.append("COUNT")
            pieces.append(str(count))
        return await self.execute_command('XRANGE', name, *pieces)

    async def xrevrange(self,
                        name: str,
                        start='+',
                        end='-',
                        count=None) -> list:
        """
        Read stream values within an interval, in reverse order.

        Available since 5.0.0.
        Time complexity: O(log(N)+M) with N being the number of elements in the stream and M the number
        of elements being returned. If M is constant (e.g. always asking for the first 10 elements with COUNT),
        you can consider it O(log(N)).

        :param name: name of the stream
        :param start: first stream ID. defaults to '+',
               meaning the latest available.
        :param end: last stream ID. defaults to '-',
                meaning the earliest available.
        :param count: if set, only return this many items, beginning with the
               latest available.

        """
        pieces = [start, end]
        if count is not None:
            if not isinstance(count, int) or count < 1:
                raise RedisError("XREVRANGE count must be a positive integer")
            pieces.append("COUNT")
            pieces.append(str(count))
        return await self.execute_command('XREVRANGE', name, *pieces)

    async def xread(self, count=None, block=None, **streams) -> dict:
        """
        Available since 5.0.0.

        Time complexity:
        For each stream mentioned: O(log(N)+M) with N being the number
        of elements in the stream and M the number of elements being returned.
        If M is constant (e.g. always asking for the first 10 elements with COUNT),
        you can consider it O(log(N)). On the other side, XADD will pay the O(N)
        time in order to serve the N clients blocked on the stream getting new data.

        Read data from one or multiple streams,
        only returning entries with an ID greater
        than the last received ID reported by the caller.

        :param count: int, if set, only return this many items, beginning with the
               earliest available.
        :param block: int, milliseconds we want to block before timing out,
                if the BLOCK option is not used, the command is synchronous
        :param streams: stream_name - stream_id mapping
        :return dict like {stream_name: [(stream_id: entry), ...]}
        """
        pieces = []
        if block is not None:
            if not isinstance(block, int) or block < 0:
                raise RedisError("XREAD block must be a positive integer")
            pieces.append("BLOCK")
            pieces.append(str(block))
        if count is not None:
            if not isinstance(count, int) or count < 1:
                raise RedisError("XREAD count must be a positive integer")
            pieces.append("COUNT")
            pieces.append(str(count))
        pieces.append("STREAMS")
        ids = []
        for partial_stream in streams.items():
            pieces.append(partial_stream[0])
            ids.append(partial_stream[1])
        pieces.extend(ids)
        return await self.execute_command('XREAD', *pieces)

    async def xreadgroup(self,
                         group: str,
                         consumer_id: str,
                         count=None,
                         block=None,
                         **streams):
        """
        Available since 5.0.0.

        Time complexity:
        For each stream mentioned: O(log(N)+M) with N being the number of elements
        in the stream and M the number of elements being returned.
        If M is constant (e.g. always asking for the first 10 elements with COUNT),
        you can consider it O(log(N)). On the other side,
        XADD will pay the O(N) time in order to serve
        the N clients blocked on the stream getting new data.

        Read data from one or multiple streams via the consumer group,
        only returning entries with an ID greater
        than the last received ID reported by the caller.

        :param group: the name of the consumer group
        :param consumer_id: the name of the consumer that is attempting to read
        :param count: int, if set, only return this many items, beginning with the
               earliest available.
        :param block: int, milliseconds we want to block before timing out,
                if the BLOCK option is not used, the command is synchronous
        :param streams: stream_name - stream_id mapping
        :return dict like {stream_name: [(stream_id: entry), ...]}
        """
        pieces = ['GROUP', group, consumer_id]
        if block is not None:
            if not isinstance(block, int) or block < 1:
                raise RedisError("XREAD block must be a positive integer")
            pieces.append("BLOCK")
            pieces.append(str(block))
        if count is not None:
            if not isinstance(count, int) or count < 1:
                raise RedisError("XREAD count must be a positive integer")
            pieces.append("COUNT")
            pieces.append(str(count))
        pieces.append("STREAMS")
        ids = []
        for partial_stream in streams.items():
            pieces.append(partial_stream[0])
            ids.append(partial_stream[1])
        pieces.extend(ids)
        return await self.execute_command('XREADGROUP', *pieces)

    async def xpending(self,
                       name: str,
                       group: str,
                       start='-',
                       end='+',
                       count=None,
                       consumer=None) -> list:
        """
        Available since 5.0.0.

        Time complexity:
        O(log(N)+M) with N being the number of elements in the consumer
        group pending entries list, and M the number of elements being returned.
        When the command returns just the summary it runs in O(1)
        time assuming the list of consumers is small,
        otherwise there is additional O(N) time needed to iterate every consumer.

        Fetching data from a stream via a consumer group,
        and not acknowledging such data,
        has the effect of creating pending entries.
        The XPENDING command is the interface to inspect the list of pending messages.

        :param name: name of the stream
        :param group: name of the consumer group
        :param start: first stream ID. defaults to '-',
               meaning the earliest available.
        :param end: last stream ID. defaults to '+',
                meaning the latest available.
        :param count: int, number of entries
                [NOTICE] only when count is set to int,
                start & end options will have effect
                and detail of pending entries will be returned
        :param consumer: str, consumer of the stream in the group
                [NOTICE] only when count is set to int,
                this option can be appended to
                query pending entries of given consumer
        """
        pieces = [name, group]
        if count is not None:
            pieces.extend([start, end, count])
            if consumer is not None:
                pieces.append(str(consumer))
        # todo: may there be a parse function
        return await self.execute_command('XPENDING', *pieces)

    async def xtrim(self, name: str, max_len: int, approximate=True) -> int:
        """
        [NOTICE] Not officially released yet

        XTRIM is designed to accept different trimming strategies,
        even if currently only MAXLEN is implemented.

        :param name: name of the stream
        :param max_len: max length of the stream after being trimmed
        :param approximate: whether redis will limit
        the stream with given max length exactly, if set to True,
        there will be a few tens of entries more,
        but never less than 1000 items:

        :return: number of entries trimmed
        """
        pieces = ['MAXLEN']
        if approximate:
            pieces.append('~')
        pieces.append(max_len)
        return await self.execute_command('XTRIM', name, *pieces)

    async def xdel(self, name: str, stream_id: str) -> int:
        """
        [NOTICE] Not officially released yet
        [NOTICE] In the current implementation, memory is not
        really reclaimed until a macro node is completely empty,
        so you should not abuse this feature.

        remove items from the middle of a stream, just by ID.

        :param name: name of the stream
        :param stream_id: id of the options appended to the stream.
        """
        return await self.execute_command('XDEL', name, stream_id)

    async def xinfo_consumers(self, name: str, group: str) -> list:
        """
        [NOTICE] Not officially released yet

        XINFO command is an observability interface that can be used
        with sub-commands in order to get information
        about streams or consumer groups.

        :param name: name of the stream
        :param group: name of the consumer group
        """
        return await self.execute_command('XINFO CONSUMERS', name, group)

    async def xinfo_groups(self, name: str) -> list:
        """
        [NOTICE] Not officially released yet

        XINFO command is an observability interface that can be used
        with sub-commands in order to get information
        about streams or consumer groups.

        :param name: name of the stream
        """
        return await self.execute_command('XINFO GROUPS', name)

    async def xinfo_stream(self, name: str) -> dict:
        """
        [NOTICE] Not officially released yet

        XINFO command is an observability interface that can be used
        with sub-commands in order to get information
        about streams or consumer groups.

        :param name: name of the stream
        """
        return await self.execute_command('XINFO STREAM', name)

    async def xack(self, name: str, group: str, stream_id: str) -> int:
        """
        [NOTICE] Not officially released yet

        XACK is the command that allows a consumer to mark a pending message as correctly processed.

        :param name: name of the stream
        :param group: name of the consumer group
        :param stream_id: id of the entry the consumer wants to mark
        :return: number of entry marked
        """
        return await self.execute_command('XACK', name, group, stream_id)

    async def xclaim(self, name: str, group: str, consumer: str,
                     min_idle_time: int, *stream_ids):
        """
        [NOTICE] Not officially released yet

        Gets ownership of one or multiple messages in the Pending Entries List of a given stream consumer group.

        :param name: name of the stream
        :param group: name of the consumer group
        :param consumer: name of the consumer
        :param min_idle_time: ms
            If the message ID (among the specified ones) exists, and its idle time greater
            or equal to min_idle_time, then the message new owner
            becomes the specified <consumer>. If the minimum idle time specified
            is zero, messages are claimed regardless of their idle time.
        :param stream_ids:
        """
        return await self.execute_command('XCLAIM', name, group, consumer,
                                          min_idle_time, *stream_ids)

    async def xgroup_create(self,
                            name: str,
                            group: str,
                            stream_id='$') -> bool:
        """
        [NOTICE] Not officially released yet
        XGROUP is used in order to create, destroy and manage consumer groups.
        :param name: name of the stream
        :param group: name of the consumer group
        :param stream_id:
            If we provide $ as we did, then only new messages arriving
            in the stream from now on will be provided to the consumers in the group.
            If we specify 0 instead the consumer group will consume all the messages
            in the stream history to start with.
            Of course, you can specify any other valid ID
        """
        return await self.execute_command('XGROUP CREATE', name, group,
                                          stream_id)

    async def xgroup_set_id(self, name: str, group: str,
                            stream_id: str) -> bool:
        """
        [NOTICE] Not officially released yet
        :param name: name of the stream
        :param group: name of the consumer group
        :param stream_id:
            If we provide $ as we did, then only new messages arriving
            in the stream from now on will be provided to the consumers in the group.
            If we specify 0 instead the consumer group will consume all the messages
            in the stream history to start with.
            Of course, you can specify any other valid ID
        """
        return await self.execute_command('XGROUP SETID', name, group,
                                          stream_id)

    async def xgroup_destroy(self, name: str, group: str) -> int:
        """
        [NOTICE] Not officially released yet
        XGROUP is used in order to create, destroy and manage consumer groups.
        :param name: name of the stream
        :param group: name of the consumer group
        """
        return await self.execute_command('XGROUP DESTROY', name, group)

    async def xgroup_del_consumer(self, name: str, group: str,
                                  consumer: str) -> int:
        """
        [NOTICE] Not officially released yet
        XGROUP is used in order to create, destroy and manage consumer groups.
        :param name: name of the stream
        :param group: name of the consumer group
        :param consumer: name of the consumer
        """
        return await self.execute_command('XGROUP DELCONSUMER', name, group,
                                          consumer)
Beispiel #15
0
class ServerCommandMixin:
    RESPONSE_CALLBACKS = dict_merge(
        string_keys_to_dict('BGREWRITEAOF BGSAVE', lambda r: True),
        string_keys_to_dict('FLUSHALL FLUSHDB SAVE '
                            'SHUTDOWN SLAVEOF', bool_ok), {
                                'ROLE': parse_role,
                                'SLOWLOG GET': parse_slowlog_get,
                                'SLOWLOG LEN': int,
                                'SLOWLOG RESET': bool_ok,
                                'CLIENT GETNAME': lambda r: r and nativestr(r),
                                'CLIENT KILL': bool_ok,
                                'CLIENT LIST': parse_client_list,
                                'CLIENT SETNAME': bool_ok,
                                'CLIENT PAUSE': bool_ok,
                                'CONFIG GET': parse_config_get,
                                'CONFIG RESETSTAT': bool_ok,
                                'CONFIG SET': bool_ok,
                                'DEBUG OBJECT': parse_debug_object,
                                'INFO': parse_info,
                                'LASTSAVE': timestamp_to_datetime,
                                'TIME': lambda x: (int(x[0]), int(x[1])),
                            })

    async def bgrewriteaof(self):
        "Tell the Redis server to rewrite the AOF file from data in memory."
        return await self.execute_command('BGREWRITEAOF')

    async def bgsave(self):
        """
        Tell the Redis server to save its data to disk.  Unlike save(),
        this method is asynchronous and returns immediately.
        """
        return await self.execute_command('BGSAVE')

    async def client_kill(self, address):
        "Disconnects the client at ``address`` (ip:port)"
        return await self.execute_command('CLIENT KILL', address)

    async def client_list(self):
        "Returns a list of currently connected clients"
        return await self.execute_command('CLIENT LIST')

    async def client_getname(self):
        "Returns the current connection name"
        return await self.execute_command('CLIENT GETNAME')

    async def client_setname(self, name):
        "Sets the current connection name"
        return await self.execute_command('CLIENT SETNAME', name)

    async def client_pause(self, timeout=0):
        """suspend all the Redis clients for the
        specified amount of time (in milliseconds)."""
        return await self.execute_command('CLIENT PAUSE', timeout)

    async def config_get(self, pattern="*"):
        "Return a dictionary of configuration based on the ``pattern``"
        return await self.execute_command('CONFIG GET', pattern)

    async def config_set(self, name, value):
        "Set config item ``name`` with ``value``"
        return await self.execute_command('CONFIG SET', name, value)

    async def config_resetstat(self):
        "Reset runtime statistics"
        return await self.execute_command('CONFIG RESETSTAT')

    async def config_rewrite(self):
        "Rewrite config file with the minimal change to reflect running config"
        return await self.execute_command('CONFIG REWRITE')

    async def dbsize(self):
        "Returns the number of keys in the current database"
        return await self.execute_command('DBSIZE')

    async def debug_object(self, key):
        "Returns version specific meta information about a given key"
        return await self.execute_command('DEBUG OBJECT', key)

    async def flushall(self):
        "Delete all keys in all databases on the current host"
        return await self.execute_command('FLUSHALL')

    async def flushdb(self):
        "Delete all keys in the current database"
        return await self.execute_command('FLUSHDB')

    async def info(self, section=None):
        """
        Returns a dictionary containing information about the Redis server

        The ``section`` option can be used to select a specific section
        of information

        The section option is not supported by older versions of Redis Server,
        and will generate ResponseError
        """
        if section is None:
            return await self.execute_command('INFO')
        else:
            return await self.execute_command('INFO', section)

    async def lastsave(self):
        """
        Return a Python datetime object representing the last time the
        Redis database was saved to disk
        """
        return await self.execute_command('LASTSAVE')

    async def save(self):
        """
        Tell the Redis server to save its data to disk,
        blocking until the save is complete
        """
        return await self.execute_command('SAVE')

    async def shutdown(self):
        "Shutdown the server"
        try:
            await self.execute_command('SHUTDOWN')
        except ConnectionError:
            # a ConnectionError here is expected
            return
        raise RedisError("SHUTDOWN seems to have failed.")

    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 slowlog_get(self, num=None):
        """
        Get the entries from the slowlog. If ``num`` is specified, get the
        most recent ``num`` items.
        """
        args = ['SLOWLOG GET']
        if num is not None:
            args.append(num)
        return await self.execute_command(*args)

    async def slowlog_len(self):
        "Get the number of items in the slowlog"
        return await self.execute_command('SLOWLOG LEN')

    async def slowlog_reset(self):
        "Remove all items in the slowlog"
        return await self.execute_command('SLOWLOG RESET')

    async def time(self):
        """
        Returns the server time as a 2-item tuple of ints:
        (seconds since epoch, microseconds into this second).
        """
        return await self.execute_command('TIME')

    async def role(self):
        """
        Provide information on the role of a Redis instance in the context of replication,
        by returning if the instance is currently a master, slave, or sentinel.
        The command also returns additional information about the state of the replication
        (if the role is master or slave)
        or the list of monitored master names (if the role is sentinel).
        :return:
        """
        return await self.execute_command('ROLE')
Beispiel #16
0
class SortedSetCommandMixin:

    RESPONSE_CALLBACKS = dict_merge(
        string_keys_to_dict(
            'ZADD ZCARD ZLEXCOUNT '
            'ZREM ZREMRANGEBYLEX '
            'ZREMRANGEBYRANK '
            'ZREMRANGEBYSCORE', int),
        string_keys_to_dict('ZSCORE ZINCRBY', float_or_none),
        string_keys_to_dict('ZRANGE ZRANGEBYSCORE ZREVRANGE ZREVRANGEBYSCORE',
                            zset_score_pairs),
        string_keys_to_dict('ZRANK ZREVRANK', int_or_none), {
            'ZSCAN': parse_zscan,
        })

    async def zadd(self, name, *args, **kwargs):
        """
        Set any number of score, element-name pairs to the key ``name``. Pairs
        can be specified in two ways:

        As *args, in the form of: score1, name1, score2, name2, ...
        or as **kwargs, in the form of: name1=score1, name2=score2, ...

        The following example would add four values to the 'my-key' key:
        redis.zadd('my-key', 1.1, 'name1', 2.2, 'name2', name3=3.3, name4=4.4)
        """
        pieces = []
        if args:
            if len(args) % 2 != 0:
                raise RedisError("ZADD requires an equal number of "
                                 "values and scores")
            pieces.extend(args)
        for pair in iteritems(kwargs):
            pieces.append(pair[1])
            pieces.append(pair[0])
        return await self.execute_command('ZADD', name, *pieces)

    async def zaddoption(self, name, option=None, *args, **kwargs):
        """
        Differs from zadd in that you can set either 'XX' or 'NX' option as
        described here: https://redis.io/commands/zadd. Only for Redis 3.0.2 or
        later.

        The following example would add four values to the 'my-key' key:
        redis.zaddoption('my-key', 'XX', 1.1, 'name1', 2.2, 'name2', name3=3.3, name4=4.4)
        redis.zaddoption('my-key', 'NX CH', name1=2.2)
        """
        if not option:
            raise RedisError("ZADDOPTION must take options")
        options = set(opt.upper() for opt in option.split())
        if options - VALID_ZADD_OPTIONS:
            raise RedisError("ZADD only takes XX, NX, CH, or INCR")
        if 'NX' in options and 'XX' in options:
            raise RedisError("ZADD only takes one of XX or NX")
        pieces = list(options)
        members = []
        if args:
            if len(args) % 2 != 0:
                raise RedisError("ZADD requires an equal number of "
                                 "values and scores")
            members.extend(args)
        for pair in iteritems(kwargs):
            members.append(pair[1])
            members.append(pair[0])
        if 'INCR' in options and len(members) != 2:
            raise RedisError("ZADD with INCR only takes one score-name pair")
        return await self.execute_command('ZADD', name, *pieces, *members)

    async def zcard(self, name):
        """Returns the number of elements in the sorted set ``name``"""
        return await self.execute_command('ZCARD', name)

    async def zcount(self, name, min, max):
        """
        Returns the number of elements in the sorted set at key ``name`` with
        a score between ``min`` and ``max``.
        """
        return await self.execute_command('ZCOUNT', name, min, max)

    async def zincrby(self, name, value, amount=1):
        """
        Increments the score of ``value`` in sorted set ``name`` by ``amount``
        """
        return await self.execute_command('ZINCRBY', name, amount, value)

    async def zinterstore(self, dest, keys, aggregate=None):
        """
        Intersects multiple sorted sets specified by ``keys`` into
        a new sorted set, ``dest``. Scores in the destination will be
        aggregated based on the ``aggregate``, or SUM if none is provided.
        """
        return await self._zaggregate('ZINTERSTORE', dest, keys, aggregate)

    async def zlexcount(self, name, min, max):
        """
        Returns the number of items in the sorted set ``name`` between the
        lexicographical range ``min`` and ``max``.
        """
        return await self.execute_command('ZLEXCOUNT', name, min, max)

    async def zrange(self,
                     name,
                     start,
                     end,
                     desc=False,
                     withscores=False,
                     score_cast_func=float):
        """
        Returns a range of values from sorted set ``name`` between
        ``start`` and ``end`` sorted in ascending order.

        ``start`` and ``end`` can be negative, indicating the end of the range.

        ``desc`` a boolean indicating whether to sort the results descendingly

        ``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 desc:
            return await self.zrevrange(name, start, end, withscores,
                                        score_cast_func)
        pieces = ['ZRANGE', 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 zrangebylex(self, name, min, max, start=None, num=None):
        """
        Returns the lexicographical range of values from sorted set ``name``
        between ``min`` and ``max``.

        If ``start`` and ``num`` are specified, then return a slice of the
        range.
        """
        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 = ['ZRANGEBYLEX', name, min, max]
        if start is not None and num is not None:
            pieces.extend([b('LIMIT'), start, num])
        return await self.execute_command(*pieces)

    async def zrevrangebylex(self, name, max, min, start=None, num=None):
        """
        Returns the reversed lexicographical range of values from sorted set
        ``name`` between ``max`` and ``min``.

        If ``start`` and ``num`` are specified, then return a slice of the
        range.
        """
        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 = ['ZREVRANGEBYLEX', name, max, min]
        if start is not None and num is not None:
            pieces.extend([b('LIMIT'), start, num])
        return await self.execute_command(*pieces)

    async def zrangebyscore(self,
                            name,
                            min,
                            max,
                            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``.

        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 = ['ZRANGEBYSCORE', name, min, max]
        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)

    async def zrank(self, name, value):
        """
        Returns a 0-based value indicating the rank of ``value`` in sorted set
        ``name``
        """
        return await self.execute_command('ZRANK', name, value)

    async def zrem(self, name, *values):
        """Removes member ``values`` from sorted set ``name``"""
        return await self.execute_command('ZREM', name, *values)

    async def zremrangebylex(self, name, min, max):
        """
        Removes all elements in the sorted set ``name`` between the
        lexicographical range specified by ``min`` and ``max``.

        Returns the number of elements removed.
        """
        return await self.execute_command('ZREMRANGEBYLEX', name, min, max)

    async def zremrangebyrank(self, name, min, max):
        """
        Removes all elements in the sorted set ``name`` with ranks between
        ``min`` and ``max``. Values are 0-based, ordered from smallest score
        to largest. Values can be negative indicating the highest scores.
        Returns the number of elements removed
        """
        return await self.execute_command('ZREMRANGEBYRANK', name, min, max)

    async def zremrangebyscore(self, name, min, max):
        """
        Removes all elements in the sorted set ``name`` with scores
        between ``min`` and ``max``. Returns the number of elements removed.
        """
        return await self.execute_command('ZREMRANGEBYSCORE', name, min, max)

    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 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)

    async def zrevrank(self, name, value):
        """
        Returns a 0-based value indicating the descending rank of
        ``value`` in sorted set ``name``
        """
        return await self.execute_command('ZREVRANK', name, value)

    async def zscore(self, name, value):
        "Return the score of element ``value`` in sorted set ``name``"
        return await self.execute_command('ZSCORE', name, value)

    async def zunionstore(self, dest, keys, aggregate=None):
        """
        Performs Union on multiple sorted sets specified by ``keys`` into
        a new sorted set, ``dest``. Scores in the destination will be
        aggregated based on the ``aggregate``, or SUM if none is provided.
        """
        return await self._zaggregate('ZUNIONSTORE', dest, keys, aggregate)

    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 zscan(self,
                    name,
                    cursor=0,
                    match=None,
                    count=None,
                    score_cast_func=float):
        """
        Incrementally returns lists of elements in a sorted set. 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

        ``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)
Beispiel #17
0
class StrictRedis(*mixins):
    """
    Implementation of the Redis protocol.

    This abstract class provides a Python interface to all Redis commands
    and an implementation of the Redis protocol.

    Connection and Pipeline derive from this, implementing how
    the commands are sent and received to the Redis server
    """

    RESPONSE_CALLBACKS = dict_merge(
        *[mixin.RESPONSE_CALLBACKS for mixin in mixins])

    @classmethod
    def from_url(cls, url, db=None, **kwargs):
        """
        Return a Redis client object configured from the given URL, which must
        use either `the ``redis://`` scheme
        <http://www.iana.org/assignments/uri-schemes/prov/redis>`_ for RESP
        connections or the ``unix://`` scheme for Unix domain sockets.

        For example:

        redis://[:password]@localhost:6379/0
        unix://[:password]@/path/to/socket.sock?db=0

        There are several ways to specify a database number. The parse function
        will return the first specified option:
        1. A ``db`` querystring option, e.g. redis://localhost?db=0
        2. If using the redis:// scheme, the path argument of the url, e.g.
        redis://localhost/0
        3. The ``db`` argument to this function.

        If none of these options are specified, db=0 is used.

        Any additional querystring arguments and keyword arguments will be
        passed along to the ConnectionPool class's initializer. In the case
        of conflicting arguments, querystring arguments always win.
        """
        connection_pool = ConnectionPool.from_url(url, db=db, **kwargs)
        return cls(connection_pool=connection_pool)

    def __init__(self,
                 host='localhost',
                 port=6379,
                 db=0,
                 password=None,
                 stream_timeout=None,
                 connect_timeout=None,
                 connection_pool=None,
                 unix_socket_path=None,
                 encoding='utf-8',
                 decode_responses=False,
                 ssl=False,
                 ssl_context=None,
                 ssl_keyfile=None,
                 ssl_certfile=None,
                 ssl_cert_reqs=None,
                 ssl_ca_certs=None,
                 max_connections=None,
                 retry_on_timeout=False,
                 max_idle_time=0,
                 idle_check_interval=1,
                 client_name=None,
                 loop=None,
                 **kwargs):
        if not connection_pool:
            kwargs = {
                'db': db,
                'password': password,
                'encoding': encoding,
                'stream_timeout': stream_timeout,
                'connect_timeout': connect_timeout,
                'max_connections': max_connections,
                'retry_on_timeout': retry_on_timeout,
                'decode_responses': decode_responses,
                'max_idle_time': max_idle_time,
                'idle_check_interval': idle_check_interval,
                'client_name': client_name,
                'loop': loop
            }
            # based on input, setup appropriate connection args
            if unix_socket_path is not None:
                kwargs.update({
                    'path': unix_socket_path,
                    'connection_class': UnixDomainSocketConnection
                })
            else:
                # TCP specific options
                kwargs.update({'host': host, 'port': port})
                if ssl_context is not None:
                    kwargs['ssl_context'] = ssl_context
                elif ssl:
                    ssl_context = RedisSSLContext(ssl_keyfile, ssl_certfile,
                                                  ssl_cert_reqs,
                                                  ssl_ca_certs).get()
                    kwargs['ssl_context'] = ssl_context
            connection_pool = ConnectionPool(**kwargs)
        self.connection_pool = connection_pool
        self._use_lua_lock = None

        self.response_callbacks = self.__class__.RESPONSE_CALLBACKS.copy()

    def __repr__(self):
        return "{}<{}>".format(type(self).__name__, repr(self.connection_pool))

    def set_response_callback(self, command, callback):
        """Sets a custom Response Callback"""
        self.response_callbacks[command] = callback

    # COMMAND EXECUTION AND PROTOCOL PARSING
    async def execute_command(self, *args, **options):
        """Executes a command and returns a parsed response"""
        pool = self.connection_pool
        command_name = args[0]
        connection = pool.get_connection()
        try:
            await connection.send_command(*args)
            return await self.parse_response(connection, command_name,
                                             **options)
        except CancelledError:
            # do not retry when coroutine is cancelled
            connection.disconnect()
            raise
        except (ConnectionError, TimeoutError) as e:
            connection.disconnect()
            if not connection.retry_on_timeout and isinstance(e, TimeoutError):
                raise
            await connection.send_command(*args)
            return await self.parse_response(connection, command_name,
                                             **options)
        finally:
            pool.release(connection)

    async def parse_response(self, connection, command_name, **options):
        """Parses a response from the Redis server"""
        response = await connection.read_response()
        if command_name in self.response_callbacks:
            callback = self.response_callbacks[command_name]
            return callback(response, **options)
        return response

    async def pipeline(self, transaction=True, shard_hint=None):
        """
        Returns a new pipeline object that can queue multiple commands for
        later execution. ``transaction`` indicates whether all commands
        should be executed atomically. Apart from making a group of operations
        atomic, pipelines are useful for reducing the back-and-forth overhead
        between the client and server.
        """
        from aredis.pipeline import StrictPipeline
        pipeline = StrictPipeline(self.connection_pool,
                                  self.response_callbacks, transaction,
                                  shard_hint)
        await pipeline.reset()
        return pipeline
Beispiel #18
0
def test_dict_merge_empty_list():
    assert dict_merge([]) == {}
Beispiel #19
0
def test_dict_merge():
    a = {"a": 1}
    b = {"b": 2}
    c = {"c": 3}
    assert dict_merge(a, b, c) == {"a": 1, "b": 2, "c": 3}
Beispiel #20
0
class ClusterCommandMixin:

    NODES_FLAGS = dict_merge(
        {
            'CLUSTER INFO': NodeFlag.ALL_NODES,
            'CLUSTER COUNTKEYSINSLOT': NodeFlag.SLOT_ID
        },
        list_keys_to_dict(['CLUSTER NODES', 'CLUSTER SLOTS'], NodeFlag.RANDOM))

    RESPONSE_CALLBACKS = {
        'CLUSTER ADDSLOTS': bool_ok,
        'CLUSTER COUNT-FAILURE-REPORTS': lambda x: int(x),
        'CLUSTER COUNTKEYSINSLOT': lambda x: int(x),
        'CLUSTER DELSLOTS': bool_ok,
        'CLUSTER FAILOVER': bool_ok,
        'CLUSTER FORGET': bool_ok,
        'CLUSTER INFO': parse_cluster_info,
        'CLUSTER KEYSLOT': lambda x: int(x),
        'CLUSTER MEET': bool_ok,
        'CLUSTER NODES': parse_cluster_nodes,
        'CLUSTER REPLICATE': bool_ok,
        'CLUSTER RESET': bool_ok,
        'CLUSTER SAVECONFIG': bool_ok,
        'CLUSTER SET-CONFIG-EPOCH': bool_ok,
        'CLUSTER SETSLOT': bool_ok,
        'CLUSTER SLAVES': parse_cluster_nodes,
        'CLUSTER SLOTS': parse_cluster_slots,
        'ASKING': bool_ok,
        'READONLY': bool_ok,
        'READWRITE': bool_ok,
    }

    RESULT_CALLBACKS = dict_merge(
        list_keys_to_dict([
            'CLUSTER INFO', 'CLUSTER ADDSLOTS',
            'CLUSTER COUNT-FAILURE-REPORTS', 'CLUSTER DELSLOTS',
            'CLUSTER FAILOVER', 'CLUSTER FORGET'
        ], lambda res: res))

    def _nodes_slots_to_slots_nodes(self, mapping):
        """
        Converts a mapping of
        {id: <node>, slots: (slot1, slot2)}
        to
        {slot1: <node>, slot2: <node>}

        Operation is expensive so use with caution
        """
        out = {}
        for node in mapping:
            for slot in node['slots']:
                out[str(slot)] = node['id']
        return out

    async def cluster_addslots(self, node_id, *slots):
        """
        Assign new hash slots to receiving node

        Sends to specefied node
        """
        return await self.execute_command('CLUSTER ADDSLOTS',
                                          *slots,
                                          node_id=node_id)

    async def cluster_count_failure_report(self, node_id=''):
        """
        Return the number of failure reports active for a given node

        Sends to specefied node
        """
        return await self.execute_command('CLUSTER COUNT-FAILURE-REPORTS',
                                          node_id=node_id)

    async def cluster_countkeysinslot(self, slot_id):
        """
        Return the number of local keys in the specified hash slot

        Send to node based on specefied slot_id
        """
        return await self.execute_command('CLUSTER COUNTKEYSINSLOT', slot_id)

    async def cluster_delslots(self, *slots):
        """
        Set hash slots as unbound in the cluster.
        It determines by it self what node the slot is in and sends it there

        Returns a list of the results for each processed slot.
        """
        cluster_nodes = self._nodes_slots_to_slots_nodes(await
                                                         self.cluster_nodes())
        res = list()
        for slot in slots:
            res.append(await self.execute_command('CLUSTER DELSLOTS',
                                                  slot,
                                                  node_id=cluster_nodes[slot]))
        return res

    async def cluster_failover(self, node_id, option):
        """
        Forces a slave to perform a manual failover of its master

        Sends to specefied node
        """
        if not isinstance(option, str) or option.upper() not in {
                'FORCE', 'TAKEOVER'
        }:
            raise ClusterError('Wrong option provided')
        return await self.execute_command('CLUSTER FAILOVER',
                                          option,
                                          node_id=node_id)

    async def cluster_forget(self, node_id):
        """
        remove a node via its node ID from the set of known nodes
        of the Redis Cluster node receiving the command

        Sends to all nodes in the cluster
        """
        return await self.execute_command('CLUSTER FORGET', node_id)

    async def cluster_info(self):
        """
        Provides info about Redis Cluster node state

        Sends to random node in the cluster
        """
        return await self.execute_command('CLUSTER INFO')

    async def cluster_keyslot(self, name):
        """
        Returns the hash slot of the specified key

        Sends to random node in the cluster
        """
        return await self.execute_command('CLUSTER KEYSLOT', name)

    async def cluster_meet(self, node_id, host, port):
        """
        Force a node cluster to handshake with another node.

        Sends to specefied node
        """
        return await self.execute_command('CLUSTER MEET',
                                          host,
                                          port,
                                          node_id=node_id)

    async def cluster_nodes(self):
        """
        Force a node cluster to handshake with another node

        Sends to random node in the cluster
        """
        return await self.execute_command('CLUSTER NODES')

    async def cluster_replicate(self, target_node_id):
        """
        Reconfigure a node as a slave of the specified master node

        Sends to specefied node
        """
        return await self.execute_command('CLUSTER REPLICATE', target_node_id)

    async def cluster_reset(self, node_id, soft=True):
        """
        Reset a Redis Cluster node

        If 'soft' is True then it will send 'SOFT' argument
        If 'soft' is False then it will send 'HARD' argument

        Sends to specefied node
        """
        option = 'SOFT' if soft else 'HARD'
        return await self.execute_command('CLUSTER RESET',
                                          option,
                                          node_id=node_id)

    async def cluster_reset_all_nodes(self, soft=True):
        """
        Send CLUSTER RESET to all nodes in the cluster

        If 'soft' is True then it will send 'SOFT' argument
        If 'soft' is False then it will send 'HARD' argument

        Sends to all nodes in the cluster
        """
        option = 'SOFT' if soft else 'HARD'
        res = list()
        for node in await self.cluster_nodes():
            res.append(await self.execute_command('CLUSTER RESET',
                                                  option,
                                                  node_id=node['id']))
        return res

    async def cluster_save_config(self):
        """
        Forces the node to save cluster state on disk

        Sends to all nodes in the cluster
        """
        return await self.execute_command('CLUSTER SAVECONFIG')

    async def cluster_set_config_epoch(self, node_id, epoch):
        """
        Set the configuration epoch in a new node

        Sends to specefied node
        """
        return await self.execute_command('CLUSTER SET-CONFIG-EPOCH',
                                          epoch,
                                          node_id=node_id)

    async def cluster_setslot(self, node_id, slot_id, state):
        """
        Bind an hash slot to a specific node

        Sends to specified node
        """
        if state.upper() in {'IMPORTING', 'MIGRATING', 'NODE'
                             } and node_id is not None:
            return await self.execute_command('CLUSTER SETSLOT', slot_id,
                                              state, node_id)
        elif state.upper() == 'STABLE':
            return await self.execute_command('CLUSTER SETSLOT', slot_id,
                                              'STABLE')
        else:
            raise RedisError('Invalid slot state: {0}'.format(state))

    async def cluster_get_keys_in_slot(self, slot_id, count):
        """
        Return local key names in the specified hash slot
        Sends to specified node
        """
        return await self.execute_command('CLUSTER GETKEYSINSLOT', slot_id,
                                          count)

    async def cluster_slaves(self, target_node_id):
        """
        Force a node cluster to handshake with another node

        Sends to targeted cluster node
        """
        return await self.execute_command('CLUSTER SLAVES', target_node_id)

    async def cluster_slots(self):
        """
        Get array of Cluster slot to node mappings

        Sends to random node in the cluster
        """
        return await self.execute_command('CLUSTER SLOTS')
Beispiel #21
0
class HashCommandMixin:

    RESPONSE_CALLBACKS = dict_merge(
        string_keys_to_dict('HDEL HLEN', int),
        string_keys_to_dict('HEXISTS HMSET', bool), {
            'HGETALL': lambda r: r and pairs_to_dict(r) or {},
            'HINCRBYFLOAT': float,
            'HSCAN': parse_hscan,
        })

    async def hdel(self, name, *keys):
        "Delete ``keys`` from hash ``name``"
        return await self.execute_command('HDEL', name, *keys)

    async def hexists(self, name, key):
        "Returns a boolean indicating if ``key`` exists within hash ``name``"
        return await self.execute_command('HEXISTS', name, key)

    async def hget(self, name, key):
        "Return the value of ``key`` within the hash ``name``"
        return await self.execute_command('HGET', name, key)

    async def hgetall(self, name):
        "Return a Python dict of the hash's name/value pairs"
        return await self.execute_command('HGETALL', name)

    async def hincrby(self, name, key, amount=1):
        "Increment the value of ``key`` in hash ``name`` by ``amount``"
        return await self.execute_command('HINCRBY', name, key, amount)

    async def hincrbyfloat(self, name, key, amount=1.0):
        """
        Increment the value of ``key`` in hash ``name`` by floating ``amount``
        """
        return await self.execute_command('HINCRBYFLOAT', name, key, amount)

    async def hkeys(self, name):
        "Return the list of keys within hash ``name``"
        return await self.execute_command('HKEYS', name)

    async def hlen(self, name):
        "Return the number of elements in hash ``name``"
        return await self.execute_command('HLEN', name)

    async def hset(self, name, key, value):
        """
        Set ``key`` to ``value`` within hash ``name``
        Returns 1 if HSET created a new field, otherwise 0
        """
        return await self.execute_command('HSET', name, key, value)

    async def hsetnx(self, name, key, value):
        """
        Set ``key`` to ``value`` within hash ``name`` if ``key`` does not
        exist.  Returns 1 if HSETNX created a field, otherwise 0.
        """
        return await self.execute_command('HSETNX', name, key, value)

    async def hmset(self, name, mapping):
        """
        Set key to value within hash ``name`` for each corresponding
        key and value from the ``mapping`` dict.
        """
        if not mapping:
            raise DataError("'hmset' with 'mapping' of length 0")
        items = []
        for pair in iteritems(mapping):
            items.extend(pair)
        return await self.execute_command('HMSET', name, *items)

    async def hmget(self, name, keys, *args):
        "Returns a list of values ordered identically to ``keys``"
        args = list_or_args(keys, args)
        return await self.execute_command('HMGET', name, *args)

    async def hvals(self, name):
        "Return the list of values within hash ``name``"
        return await self.execute_command('HVALS', name)

    async def hscan(self, name, cursor=0, match=None, count=None):
        """
        Incrementally return key/value slices in a hash. 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('HSCAN', *pieces)

    async def hstrlen(self, name, key):
        """
        Returns the string length of the value associated
        with field in the hash stored at key.
        If the key or the field do not exist, 0 is returned.
        """
        return await self.execute_command('HSTRLEN', name, key)
Beispiel #22
0
class StrictPipeline(BasePipeline, *pipeline_mixins):
    "Pipeline for the StrictRedis class"
    RESPONSE_CALLBACKS = dict_merge(
        *[mixin.RESPONSE_CALLBACKS for mixin in pipeline_mixins]
    )
    pass