class ClusterParser(DefaultParser): EXCEPTION_CLASSES = dict_merge( DefaultParser.EXCEPTION_CLASSES, { 'ASK': AskError, 'TRYAGAIN': TryAgainError, 'MOVED': MovedError, 'CLUSTERDOWN': ClusterDownError, 'CROSSSLOT': ClusterCrossSlotError, })
def __init__(self, *args, serialize_fn, deserialize_fn, **kwargs): super().__init__(*args, **kwargs) self.serialize_fn = serialize_fn self.deserialize_fn = deserialize_fn # Chain response callbacks to deserialize output FROM_SERIALIZED_CALLBACKS = dict_merge( string_keys_to_dict('KEYS TYPE SCAN HKEYS', self.decode), string_keys_to_dict( 'MGET HVALS HMGET LRANGE SRANDMEMBER GET GETSET HGET LPOP ' 'RPOPLPUSH BRPOPLPUSH LINDEX SPOP', self.parse_list), string_keys_to_dict('SMEMBERS SDIFF SINTER SUNION', self.parse_set), string_keys_to_dict('HGETALL', self.parse_hgetall), string_keys_to_dict('HSCAN', self.parse_hscan), string_keys_to_dict('SSCAN', self.parse_sscan), string_keys_to_dict( 'ZRANGE ZRANGEBYSCORE ZREVRANGE ZREVRANGEBYSCORE', self.parse_zrange), string_keys_to_dict('ZSCAN', self.parse_zscan), string_keys_to_dict('BLPOP BRPOP', self.parse_bpop), { 'PUBSUB CHANNELS': self.decode, 'PUBSUB NUMSUB': self.decode, }, ) for cmd in FROM_SERIALIZED_CALLBACKS: if cmd in self.response_callbacks: self.response_callbacks[cmd] = chain_functions( self.response_callbacks[cmd], FROM_SERIALIZED_CALLBACKS[cmd]) else: self.response_callbacks[cmd] = FROM_SERIALIZED_CALLBACKS[cmd] # For the following we call first our callback as the redis-py callback returns nativestr, which may no be what we want for cmd in 'GEORADIUS', 'GEORADIUSBYMEMBER': self.response_callbacks[cmd] = chain_functions( self.parse_georadius, self.response_callbacks[cmd])
class Yedis(StrictRedis): """Provides access to timeseries via the Yedis API of YugabyteDB.""" # Overridden callbacks RESPONSE_CALLBACKS = dict_merge( StrictRedis.RESPONSE_CALLBACKS, string_keys_to_dict( 'TSADD TSREM', bool_ok ), string_keys_to_dict( 'TSCARD', int ), string_keys_to_dict( 'TSLASTN TSRANGEBYTIME TSREVRANGEBYTIME', timeseries_time_value_pairs ) ) def tsadd(self, name, *times_and_values, **options): """ Add any number of ``time``, ``value`` pairs to the time series that is specified by the given key ``name``. Pairs can be specified as ``times_and_values``, in the form of: ``time1``, ``value1``, ``time2``, ``value2``,... This is useful in storing time series like data where the ``name`` could define a metric, the ``time`` is the time when the metric was generated and ``value`` is the value of the metric at the given time. ``times`` can be of different types. They are converted to timestamps - 64 bit signed integers - using a callable ``time_cast_func`` within options. ``times`` should be python datetimes if ``time_cast_func`` isn't set. A naive datetime is accounted an UTC datetime. ``times`` could be the time in seconds since the epoch (unix time) as a floating point number in case ``time_cast_func`` is set to ``unixtime_to_timestamp`` for instance. Note: The unix time is naive, that is, it doesn't know its timezone. ``times`` MUST be in UTC! To provide raw timestamps set ``time_cast_func`` to None. """ pieces = ['TSADD', name] if times_and_values: if len(times_and_values) % 2 != 0: raise RedisError("TSADD requires an equal number of " "times and values") time_cast_func = options.get('time_cast_func', DatetimeToTimestamp(pytz.utc)) if time_cast_func is None: pieces.extend(times_and_values) else: pieces.extend( chain.from_iterable(izip(imap(time_cast_func, times_and_values[0::2]), times_and_values[1::2]))) return self.execute_command(*pieces) def tsrem(self, name, *times, **options): """ Removes one or more specified ``times`` from the time series that is specified by the given key ``name``. ``times`` can be of different types. They are converted to timestamps - 64 bit signed integers - using a callable ``time_cast_func`` within options. ``times`` should be python datetimes if ``time_cast_func`` isn't set. A naive datetime is accounted an UTC datetime. ``times`` could be the time in seconds since the epoch (unix time) as a floating point number in case ``time_cast_func`` is set to ``unixtime_to_timestamp`` for instance. Note: The unix time is naive, that is, it doesn't know its timezone. ``times`` MUST be in UTC! To provide raw timestamps set ``time_cast_func`` to None. """ pieces = ['TSREM', name] if times: time_cast_func = options.get('time_cast_func', DatetimeToTimestamp(pytz.utc)) if time_cast_func is None: pieces.extend(times) else: pieces.extend(imap(time_cast_func, times)) return self.execute_command(*pieces) def tscard(self, name): """ Return the number of entires in the time series that is specified by the given key ``name``. """ return self.execute_command('TSCARD', name) def tsget(self, name, tm, time_cast_func=DatetimeToTimestamp(pytz.utc)): """ Return the value for the given time ``tm`` in the time series that is specified by the given key ``name``. ``time_cast_func`` a callable used to cast the time ``tm`` to a timestamp - 64 bit signed integer (cf. ``tsadd``). """ if time_cast_func is None: return self.execute_command('TSGET', name, tm) return self.execute_command('TSGET', name, time_cast_func(tm)) def tslastn(self, name, num, timestamp_cast_func=TimestampToDatetime(pytz.utc)): """ Returns the latest ``num`` entries in the time series that is specified by the given key ``name``. The entries are returned in ascending order of the times. ``timestamp_cast_func`` a callable used to cast the timestamp return values. It should reflect how timestamp were inserted (cf. ``time_cast_func``). """ if timestamp_cast_func is None: return self.execute_command('TSLASTN', name, num, timestamp_cast_func=int) return self.execute_command('TSLASTN', name, num, timestamp_cast_func=timestamp_cast_func) def tsrangebytime(self, name, tm_low=None, tm_high=None, time_cast_func=DatetimeToTimestamp(pytz.utc), timestamp_cast_func=TimestampToDatetime(pytz.utc)): """ Returns the entries for the given time range from ``tm_low`` to ``tm_high`` in the time series that is specified by the given key ``name``. The entries are returned in ascending order of the times. Special bounds -inf (``tm_low`` is None or -inf) and +inf (``tm_high`` is None or inf) are also supported to retrieve an entire range. ``time_cast_func`` a callable used to cast the time ``tm`` to a timestamp - 64 bit signed integer (cf. ``tsadd``). ``timestamp_cast_func`` a callable used to cast the timestamp return values. It should reflect how timestamp were inserted (cf. ``time_cast_func``). """ pieces = ['TSRANGEBYTIME', name] if tm_low is None or (isinstance(tm_low, float) and isinf(tm_low)): pieces.append(m_inf) elif time_cast_func is None: pieces.append(tm_low) else: pieces.append(time_cast_func(tm_low)) if tm_high is None or (isinstance(tm_high, float) and isinf(tm_high)): pieces.append(p_inf) elif time_cast_func is None: pieces.append(tm_high) else: pieces.append(time_cast_func(tm_high)) if timestamp_cast_func is None: return self.execute_command(*pieces, timestamp_cast_func=int) return self.execute_command(*pieces, timestamp_cast_func=timestamp_cast_func) def tsrevrangebytime(self, name, tm_low=None, tm_high=None, num=None, time_cast_func=DatetimeToTimestamp(pytz.utc), timestamp_cast_func=TimestampToDatetime(pytz.utc)): """ Returns the entries for the given time range from ``tm_low`` to ``tm_high`` in the time series that is specified by the given key ``name``. The entries are returned in descending order of the times. Special bounds -inf (``tm_low`` is None or -inf) and +inf (``tm_high`` is None or inf) are also supported to retrieve an entire range. If ``num`` is specified, then at most ``num`` entries will be fetched. ``time_cast_func`` a callable used to cast the time ``tm`` to a timestamp - 64 bit signed integer (cf. ``tsadd``). ``timestamp_cast_func`` a callable used to cast the timestamp return values. It should reflect how timestamp were inserted (cf. ``time_cast_func``). """ pieces = ['TSREVRANGEBYTIME', name] if tm_low is None or (isinstance(tm_low, float) and isinf(tm_low)): pieces.append(m_inf) elif time_cast_func is None: pieces.append(tm_low) else: pieces.append(time_cast_func(tm_low)) if tm_high is None or (isinstance(tm_high, float) and isinf(tm_high)): pieces.append(p_inf) elif time_cast_func is None: pieces.append(tm_high) else: pieces.append(time_cast_func(tm_high)) if num is not None: pieces.extend([Token.get_token('LIMIT'), num]) if timestamp_cast_func is None: return self.execute_command(*pieces, timestamp_cast_func=int) return self.execute_command(*pieces, timestamp_cast_func=timestamp_cast_func)
class DisqueAlpha(object): """ 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 """ _job_score = None RESPONSE_CALLBACKS = dict_merge( string_keys_to_dict('GETJOB', parse_job_resp), string_keys_to_dict('QLEN ACKJOB FASTACK', int), string_keys_to_dict( 'ADDJOB', lambda r: six.text_type(six.binary_type(r).decode())), { 'INFO': parse_info, 'CLIENT GETNAME': lambda r: r and six.text_type(r), 'CLIENT KILL': bool_ok, 'CLIENT LIST': parse_client_list, 'CLIENT SETNAME': bool_ok, 'CONFIG GET': parse_config_get, 'CONFIG RESETSTAT': bool_ok, 'CONFIG SET': bool_ok, 'CLUSTER NODES': parse_cluster_nodes, 'HELLO': parse_hello, 'TIME': parse_time, }, string_keys_to_dict('BGREWRITEAOF', lambda r: True), ) @classmethod def from_url(cls, url, **kwargs): """ Return a Disque client object configured from the given URL. For example:: disque://[:password]@localhost:6379 unix://[:password]@/path/to/socket.sock 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, **kwargs) return cls(connection_pool=connection_pool) def __init__(self, host='localhost', port=7711, password=None, socket_timeout=None, socket_connect_timeout=None, socket_keepalive=None, socket_keepalive_options=None, connection_pool=None, unix_socket_path=None, encoding='utf-8', encoding_errors='strict', decode_responses=False, retry_on_timeout=False, job_origin_ttl_secs=5, record_job_origin=False): """ job_origin_ttl_secs is the number of seconds to store counts of incoming jobs. The higher the throughput you're expecting, the lower this number should be. """ self.record_job_origin = record_job_origin kwargs = { 'password': password, 'socket_timeout': socket_timeout, 'encoding': encoding, 'encoding_errors': encoding_errors, 'decode_responses': decode_responses, 'retry_on_timeout': retry_on_timeout, 'db': 0, } # 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, 'socket_connect_timeout': socket_connect_timeout, 'socket_keepalive': socket_keepalive, 'socket_keepalive_options': socket_keepalive_options, }) if not connection_pool: connection_pool = ConnectionPool(**kwargs) self.response_callbacks = self.__class__.RESPONSE_CALLBACKS.copy() self.connection_pool = {'default': connection_pool} self.default_node = 'default' self._job_score = RollingCounter(ttl_secs=job_origin_ttl_secs) self.__connect_cluster(kwargs) def __connect_cluster(self, connection_kwargs): hi = self.hello() self.default_node = bin_to_str(hi['id'][:8]) self.connection_pool.pop('default') for node, ip, port, version in hi['nodes']: connection_kwargs.update(dict(host=ip, port=port)) self.connection_pool[bin_to_str( node[:8])] = ConnectionPool(**connection_kwargs) def __repr__(self): return "%s<%s>" % (type(self).__name__, repr(self.connection_pool)) def set_response_callback(self, command, callback): "Set a custom Response Callback" self.response_callbacks[command] = callback __read_cmds = {'GETJOB': 0, 'ACKJOB': 0, 'FASTACK': 0} def _get_connection(self, command_name, **options): node = self.default_node if self.record_job_origin and command_name in self.__read_cmds: node = self._job_score.max(node) pool = self.connection_pool.get(node) if pool is None: pool = self.connection_pool[self.default_node] node = self.default_node return pool.get_connection(command_name, **options), node def _release_connection(self, connection, node): return self.connection_pool[node].release(connection) def execute_command(self, *args, **options): "Execute a command and return a parsed response" command_name = args[0] connection, node = self._get_connection(command_name, **options) try: connection.send_command(*args) return 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 connection.send_command(*args) return self.parse_response(connection, command_name, **options) finally: self._release_connection(connection, node) def parse_response(self, connection, command_name, **options): "Parses a response from the Redis server" response = connection.read_response() if command_name in self.response_callbacks: return self.response_callbacks[command_name](response, **options) return response # SERVER INFORMATION def bgrewriteaof(self): "Tell the Redis server to rewrite the AOF file from data in memory." return self.execute_command('BGREWRITEAOF') def client_kill(self, address): "Disconnects the client at ``address`` (ip:port)" return self.execute_command('CLIENT KILL', address) def client_list(self): "Returns a list of currently connected clients" return self.execute_command('CLIENT LIST') def client_getname(self): "Returns the current connection name" return self.execute_command('CLIENT GETNAME') def client_setname(self, name): "Sets the current connection name" return self.execute_command('CLIENT SETNAME', name) def client_pause(self, pause_msec): return self.execute_command('CLIENT PAUSE', pause_msec) def config_get(self, pattern="*"): "Return a dictionary of configuration based on the ``pattern``" return self.execute_command('CONFIG GET', pattern) def config_set(self, name, value): "Set config item ``name`` with ``value``" return self.execute_command('CONFIG SET', name, value) def config_resetstat(self): "Reset runtime statistics" return self.execute_command('CONFIG RESETSTAT') def config_rewrite(self): "Rewrite config file with the minimal change to reflect running config" return self.execute_command('CONFIG REWRITE') # Danger: debug commands ahead def debug_segfault(self): """ Danger: will segfault connected Disque instance""" return self.execute_command('DEBUG SEGFAULT') def debug_oom(self): """ Danger: will OOM connected Disque instance""" return self.execute_command('DEBUG OOM') def debug_flushall(self): return self.execute_command('DEBUG FLUSHALL') def debug_loadaof(self): return self.execute_command('DEBUG LOADAOF') def debug_sleep(self, sleep_secs): return self.execute_command('DEBUG SLEEP', sleep_secs) def debug_error(self, message): return self.execute_command('DEBUG ERROR', message) def debug_structsize(self): return self.execute_command('DEBUG STRUCTSIZE') # Cluster admin commands def cluster_meet(self, ip, port): return self.execute_command('CLUSTER MEET', ip, port) def cluster_nodes(self): return self.execute_command('CLUSTER NODES') def cluster_saveconfig(self): return self.execute_command('CLUSTER SAVECONFIG') def cluster_forget(self, node): return self.execute_command('CLUSTER FORGET', node) def _cluster_reset(self, reset): return self.execute_command('CLUSTER RESET', reset) def cluster_reset_hard(self): return self._cluster_reset(Token('HARD')) def cluster_reset_soft(self): return self._cluster_reset(Token('SOFT')) def cluster_info(self): return self.execute_command('CLUSTER INFO') def hello(self): return self.execute_command('HELLO') def info(self, section=None): """ Returns a dictionary containing information about the Disque server The ``section`` option can be used to select a specific section of information Valid section names are: SERVER, CLIENTS, MEMORY, JOBS, QUEUES, PERSISTENCE, STATS, CPU """ if section is None: return self.execute_command('INFO') else: return self.execute_command('INFO', section) def ping(self): "Ping the Redis server" return self.execute_command('PING') def shutdown(self): "Shutdown the server" try: self.execute_command('SHUTDOWN') except ConnectionError: # a ConnectionError here is expected return raise DisqueError("SHUTDOWN seems to have failed.") 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 self.execute_command(*args) def slowlog_len(self): "Get the number of items in the slowlog" return self.execute_command('SLOWLOG LEN') def slowlog_reset(self): "Remove all items in the slowlog" return self.execute_command('SLOWLOG RESET') def time(self): """ Returns the server time as a 2-item tuple of ints: (seconds since epoch, microseconds into this second). """ return self.execute_command('TIME') # BASIC JOB COMMANDS def addjob(self, queue, body, timeout_ms=0, replicate=0, delay_secs=0, retry_secs=-1, ttl_secs=0, maxlen=0, async=False): args = ['ADDJOB', queue, body, timeout_ms] if replicate > 0: args += [Token('REPLICATE'), replicate] if delay_secs > 0: args += [Token('DELAY'), delay_secs] if retry_secs >= 0: args += [Token('RETRY'), retry_secs] if ttl_secs > 0: args += [Token('TTL'), ttl_secs] if maxlen > 0: args += [Token('MAXLEN'), maxlen] if async: args += [Token('ASYNC')] return self.execute_command(*args)