def __init__(self, host, port, db, password=None, encoding=None): self.host = host self.port = port self.db = db if encoding: self.connection = Connection( host=self.host, port=self.port, db=self.db, password=password, encoding=encoding, decode_responses=True, encoding_errors="replace", ) else: self.connection = Connection( host=self.host, port=self.port, db=self.db, password=password, decode_responses=False, ) # all command upper case self.answer_callbacks = command2callback self.callbacks = self.reder_funcname_mapping()
def __init__(self, *args, **kwargs): TCPClient.__init__(self, kwargs.pop("resolver", None)) Connection.__init__(self, parser_class=AsyncParser, *args, **kwargs) self._stream = None
def __init__(self, *args, **kwargs): TCPClient.__init__(self, kwargs.pop("resolver", None), kwargs.pop("io_loop", None)) Connection.__init__(self, parser_class=AsyncParser, *args, **kwargs) self._stream = None
def test_disconnect(self): conn = Connection() mock_sock = mock.Mock() conn._sock = mock_sock conn.disconnect() mock_sock.shutdown.assert_called_once() mock_sock.close.assert_called_once() assert conn._sock is None
def test_memoryviews_are_not_packed(self): c = Connection() arg = memoryview(b"some_arg") arg_list = ["SOME_COMMAND", arg] cmd = c.pack_command(*arg_list) assert cmd[1] is arg cmds = c.pack_commands([arg_list, arg_list]) assert cmds[1] is arg assert cmds[3] is arg
def test_disconnect__close_OSError(self): """An OSError on socket close will still clear out the socket.""" conn = Connection() mock_sock = mock.Mock() conn._sock = mock_sock conn._sock.close.side_effect = OSError conn.disconnect() mock_sock.shutdown.assert_called_once() mock_sock.close.assert_called_once() assert conn._sock is None
def test_connect_without_retry_on_os_error(self): """Test that the _connect function is not being retried in case of a OSError""" with patch.object(Connection, "_connect") as _connect: _connect.side_effect = OSError("") conn = Connection(retry_on_timeout=True, retry=Retry(NoBackoff(), 2)) with pytest.raises(ConnectionError): conn.connect() assert _connect.call_count == 1 self.clear(conn)
def __init__(self, **kw): args = self.defaults.copy() self.original_args = args args.update(kw) failure_callback = args.pop('failure_callback') BaseCxn.__init__(self, **args) self.attempts = 0 self._depth = count() self.failure_callback = failure_callback is not None and \ failure_callback or self.default_callback
def to_blocking_connection(self, socket_read_size=65536): """ Convert asynchronous connection to blocking socket connection """ conn = Connection( self.host, self.port, self.db, self.password, self.socket_timeout, self.socket_connect_timeout, self.socket_keepalive, self.socket_keepalive_options, self.retry_on_timeout, parser_class=DefaultParser, socket_read_size=socket_read_size) conn._sock = self._stream.socket return conn
def to_blocking_connection(self, socket_read_size=65536): """ Convert asynchronous connection to blocking socket connection """ conn = Connection( self.host, self.port, self.db, self.password, self.socket_timeout, self.socket_connect_timeout, self.socket_keepalive, self.socket_keepalive_options, self.retry_on_timeout, self.encoding, self.encoding_errors, self.decode_responses, DefaultParser, socket_read_size) conn._sock = self._stream.socket return conn
def test_connect_timeout_error_without_retry(self): """Test that the _connect function is not being retried if retry_on_timeout is set to False""" conn = Connection(retry_on_timeout=False) conn._connect = mock.Mock() conn._connect.side_effect = socket.timeout with pytest.raises(TimeoutError) as e: conn.connect() assert conn._connect.call_count == 1 assert str(e.value) == "Timeout connecting to server" self.clear(conn)
def connect(self): self.attempts += 1 try: BaseCxn.connect(self) self.attempts = 0 self.depth = 0 return except ConnectionError, e: out = self.failure_callback(self, e) if out is False: raise return out
def test_close_connection_in_child(self): """ A connection owned by a parent and closed by a child doesn't destroy the file descriptors so a parent can still use it. """ conn = Connection() conn.send_command('ping') assert conn.read_response() == b'PONG' def target(conn): conn.send_command('ping') assert conn.read_response() == b'PONG' conn.disconnect() proc = multiprocessing.Process(target=target, args=(conn, )) proc.start() proc.join(3) assert proc.exitcode is 0 # The connection was created in the parent but disconnected in the # child. The child called socket.close() but did not call # socket.shutdown() because it wasn't the "owning" process. # Therefore the connection still works in the parent. conn.send_command('ping') assert conn.read_response() == b'PONG'
def test_retry_on_timeout_retry(self, retries): retry_on_timeout = retries > 0 c = Connection(retry_on_timeout=retry_on_timeout, retry=Retry(NoBackoff(), retries)) assert c.retry_on_timeout == retry_on_timeout assert isinstance(c.retry, Retry) assert c.retry._retries == retries
def __init__(self, zk_client, max_attempts=3, **kwargs): print "kw args are", kwargs self.zk = zk_client self.host = None self.port = None master = self.zk.get("/redis/master")[0] if master: host, port = json.loads(master)['address'].split(':') else: master = self.elect_master() host, port = master['address'].split(':') kwargs['host'] = host kwargs['port'] = int(port) BaseConnection.__init__(self, **kwargs) self.max_attempts = max_attempts
def connect(self): # get current leader from zookeeper, # if no current leader, elect one. for i in range(self.max_attempts): try: return BaseConnection.connect(self) except ConnectionError, e: #print "connection error", e # if we've failed max_attempts times, if i == self.max_attempts-1: # check if someone else has already updated the master master = self.zk.get("/redis/master")[0] host, port = json.loads(master)['address'].split(':') port = int(port) if host == self.host and port == self.port: # if not remove this master and elect a new one. master = self.elect_master() #print "master is ", master if master is False: raise host, port = master['address'].split(':') port = int(port) # update and recure self.update(host=host, port=port) return self.connect()
def test_retry_connect_on_timeout_error(self): """Test that the _connect function is retried in case of a timeout""" conn = Connection(retry_on_timeout=True, retry=Retry(NoBackoff(), 3)) origin_connect = conn._connect conn._connect = mock.Mock() def mock_connect(): # connect only on the last retry if conn._connect.call_count <= 2: raise socket.timeout else: return origin_connect() conn._connect.side_effect = mock_connect conn.connect() assert conn._connect.call_count == 3 self.clear(conn)
def test_close_connection_in_child(self): """ A connection owned by a parent and closed by a child doesn't destroy the file descriptors so a parent can still use it. """ conn = Connection() conn.send_command('ping') assert conn.read_response() == b'PONG' def target(conn): conn.send_command('ping') assert conn.read_response() == b'PONG' conn.disconnect() proc = multiprocessing.Process(target=target, args=(conn,)) proc.start() proc.join(3) assert proc.exitcode is 0 # The connection was created in the parent but disconnected in the # child. The child called socket.close() but did not call # socket.shutdown() because it wasn't the "owning" process. # Therefore the connection still works in the parent. conn.send_command('ping') assert conn.read_response() == b'PONG'
def test_close_connection_in_parent(self): """ A connection owned by a parent is unusable by a child if the parent (the owning process) closes the connection. """ conn = Connection() conn.send_command('ping') assert conn.read_response() == b'PONG' def target(conn, ev): ev.wait() # the parent closed the connection. because it also created the # connection, the connection is shutdown and the child # cannot use it. with pytest.raises(ConnectionError): conn.send_command('ping') ev = multiprocessing.Event() proc = multiprocessing.Process(target=target, args=(conn, ev)) proc.start() conn.disconnect() ev.set() proc.join(3) assert proc.exitcode is 0
def __init__(self, host, port, db, password=None, encoding=None, get_info=True): self.host = host self.port = port self.db = db if encoding: self.connection = Connection( host=self.host, port=self.port, db=self.db, password=password, encoding=encoding, decode_responses=True, encoding_errors="replace", ) else: self.connection = Connection( host=self.host, port=self.port, db=self.db, password=password, decode_responses=False, ) # all command upper case self.answer_callbacks = command2callback self.callbacks = self.reder_funcname_mapping() self.connection.connect() if get_info: try: self.get_server_info() except Exception as e: logger.warn(f"[After Connection] {str(e)}") config.no_version_reason = str(e) else: config.no_version_reason = "--no-info flag activated"
def handle(self, *args, **options): owner_username = args[0] audience_name = args[1] filename = os.path.abspath(args[2]) redis_key = Profiles.ip_audiences_key # the connection is not necessary, but the object is used to # format redis commands to redis protocol using pack_command redis_connection = Connection(host=settings.profiles_redis_host, port=settings.profiles_redis_port) user = User.objects.get(username=owner_username) audience, _ = Audience.objects.get_or_create(name=audience_name, owner=user.account, is_ip=True) audience_id = audience.public_id with open(filename) as csvfile: reader = csv.reader(csvfile) for row in reader: ip = IPAddress(row[0]).packed sys.stdout.write( redis_connection.pack_command("SADD", redis_key.format(ip), audience_id))
def __init__(self, host, port, db, password=None): self.host = host self.port = port self.db = db if config.decode: self.connection = Connection( host=self.host, port=self.port, db=self.db, password=password, encoding=config.decode, decode_responses=True, encoding_errors="replace", socket_keepalive=config.socket_keepalive, ) else: self.connection = Connection( host=self.host, port=self.port, db=self.db, password=password, decode_responses=False, socket_keepalive=config.socket_keepalive, ) # all command upper case self.answer_callbacks = command2callback self.callbacks = self.reder_funcname_mapping() try: self.connection.connect() except Exception as e: print(str(e), file=sys.stderr) if not config.no_info: try: self.get_server_info() except Exception as e: logger.warning(f"[After Connection] {str(e)}") config.no_version_reason = str(e) else: config.no_version_reason = "--no-info flag activated"
class Client: """ iRedis client, hold a redis-py Client to interact with Redis. """ def reder_funcname_mapping(self): mapping = {} for func_name, func in renders.__dict__.items(): if callable(func): mapping[func_name] = func return mapping def __init__(self, host, port, db, password=None, encoding=None): self.host = host self.port = port self.db = db if encoding: self.connection = Connection( host=self.host, port=self.port, db=self.db, password=password, encoding=encoding, decode_responses=True, encoding_errors="replace", ) else: self.connection = Connection( host=self.host, port=self.port, db=self.db, password=password, decode_responses=False, ) # all command upper case self.answer_callbacks = command2callback self.callbacks = self.reder_funcname_mapping() def __str__(self): if self.db: return f"{self.host}:{self.port}[{self.db}]" return f"{self.host}:{self.port}" def execute_command(self, completer, command_name, *args, **options): "Execute a command and return a parsed response" # === pre hook === try: self.connection.send_command(command_name, *args) resp = self.parse_response(self.connection, completer, command_name, **options) # retry on timeout except (ConnectionError, TimeoutError) as e: self.connection.disconnect() if not (self.connection.retry_on_timeout and isinstance(e, TimeoutError)): raise self.connection.send_command(command_name, *args) resp = self.parse_response(self.connection, completer, command_name, **options) # === After hook === # SELECT db on AUTH if command_name.upper() == "AUTH" and self.db: select_result = self.execute_command(completer, "SELECT", self.db) if nativestr(select_result) != "OK": raise ConnectionError("Invalid Database") if command_name.upper() == "SELECT": logger.debug("[Pre hook] Command is SELECT, change self.db.") self.db = int(args[0]) return resp def parse_response(self, connection, completer, command_name, **options): "Parses a response from the Redis server" response = connection.read_response() logger.info(f"[Redis-Server] Response: {response}") command_upper = command_name.upper() if (command_upper in self.answer_callbacks and self.answer_callbacks[command_upper]): callback_name = self.answer_callbacks[command_upper] callback = self.callbacks[callback_name] rendered = callback(response, completer) else: rendered = response logger.info(f"[rendered] {rendered}") return rendered def send_command(self, command, completer): """ Send command to redis-server, return parsed response. :param command: text command, not parsed :param completer: RedisGrammarCompleter will update completer based on redis response. eg: update key completer after ``keys`` command """ try: input_command, args = split_command_args(command, all_commands) self.patch_completers(command, completer) redis_resp = self.execute_command(completer, input_command, *args) except Exception as e: logger.exception(e) return render_error(str(e)) return redis_resp def patch_completers(self, command, completer): """ Before execute command, patch completers first. Eg: When user run `GET foo`, key completer need to touch foo. Only works when compile-grammar thread is done. """ if not completer: logger.warning( "[Pre patch completer] Complter not ready, not patched...") return redis_grammar = completer.compiled_grammar m = redis_grammar.match(command) if not m: # invalide command! return variables = m.variables() # parse keys keys_token = variables.getall("keys") if keys_token: for key in _strip_quote_args(keys_token): completer.completers['key'].touch(key) key_token = variables.getall("key") if key_token: # NOTE variables.getall always be a list for single_key in _strip_quote_args(key_token): completer.completers['key'].touch(single_key) logger.debug( f"[Complter key] Done: {completer.completers['key'].words}")
class Client: """ iRedis client, hold a redis-py Client to interact with Redis. """ def reder_funcname_mapping(self): mapping = {} for func_name, func in renders.__dict__.items(): if callable(func): mapping[func_name] = func return mapping def __init__(self, host, port, db, encoding=None): self.host = host self.port = port self.db = db if encoding: self.connection = Connection( host=self.host, port=self.port, db=self.db, encoding=encoding, decode_responses=True, encoding_errors="replace", ) else: self.connection = Connection( host=self.host, port=self.port, db=self.db, decode_responses=False ) # all command upper case self.answer_callbacks = command2callback self.callbacks = self.reder_funcname_mapping() def __str__(self): if self.db: return f"{self.host}:{self.port}[{self.db}]" return f"{self.host}:{self.port}" def execute_command(self, completer, command_name, *args, **options): "Execute a command and return a parsed response" # === pre hook === if command_name.upper() == "SELECT": logger.debug("[Pre hook] Command is SELECT, change self.db.") self.db = int(args[0]) try: self.connection.send_command(command_name, *args) resp = self.parse_response( self.connection, completer, command_name, **options ) # retry on timeout except (ConnectionError, TimeoutError) as e: self.connection.disconnect() if not (self.connection.retry_on_timeout and isinstance(e, TimeoutError)): raise self.connection.send_command(command_name, *args) resp = self.parse_response( self.connection, completer, command_name, **options ) # === After hook === # SELECT db on AUTH if command_name.upper() == "AUTH" and self.db: select_result = self.execute_command(completer, "SELECT", self.db) if nativestr(select_result) != "OK": raise ConnectionError("Invalid Database") return resp def parse_response(self, connection, completer, command_name, **options): "Parses a response from the Redis server" try: response = connection.read_response() logger.info(f"[Redis-Server] Response: {response}") except ResponseError as e: logger.warn(f"[Redis-Server] ERROR: {str(e)}") response = str(e) command_upper = command_name.upper() if command_upper in self.answer_callbacks and self.answer_callbacks[command_upper]: callback_name = self.answer_callbacks[command_upper] callback = self.callbacks[callback_name] rendered = callback(response, completer) else: rendered = response return rendered def parse_input(self, input_command): """ parse input command to command and args. convert input to upper case, we use upper case command internally; strip quotes in args """ return None def _valide_token(self, words): token = "".join(words).strip() if token: yield token def _strip_quote_args(self, s): """ Given string s, split it into args.(Like bash paring) Handle with all quote cases. Raise ``InvalidArguments`` if quotes not match :return: args list. """ sperator = re.compile(r"\s") word = [] in_quote = None pre_back_slash = False for char in s: if in_quote: # close quote if char == in_quote: if not pre_back_slash: yield from self._valide_token(word) word = [] in_quote = None else: # previous char is \ , merge with current " word[-1] = char else: word.append(char) # not in quote else: # sperator if sperator.match(char): if word: yield from self._valide_token(word) word = [] else: word.append(char) # open quotes elif char in ["'", '"']: in_quote = char else: word.append(char) if char == "\\" and not pre_back_slash: pre_back_slash = True else: pre_back_slash = False if word: yield from self._valide_token(word) # quote not close if in_quote: raise InvalidArguments() def send_command(self, command, completer): """ Send command to redis-server, return parsed response. :param command: text command, not parsed :param completer: RedisGrammarCompleter will update completer based on redis response. eg: update key completer after ``keys`` command """ # Parse command-name and args upper_raw_command = command.upper() for command_name in all_commands: if upper_raw_command.startswith(command_name): l = len(command_name) input_command = command[:l] input_args = command[l:] break else: raise InvalidArguments(r"`{command} is not a valide Redis Command") args = list(self._strip_quote_args(input_args)) logger.debug(f"[Parsed comamnd name] {input_command}") logger.debug(f"[Parsed comamnd args] {args}") redis_resp = self.execute_command(completer, input_command, *args) return redis_resp
def test_retry_on_timeout_boolean(self, retry_on_timeout): c = Connection(retry_on_timeout=retry_on_timeout) assert c.retry_on_timeout == retry_on_timeout assert isinstance(c.retry, Retry) assert c.retry._retries == (1 if retry_on_timeout else 0)
def __init__(self, *args, **kwargs): Connection.__init__(self, *args, **kwargs) # 上次心跳包时间 self.last_beat = 0 # 初始化的时候立刻链接 self.connect()
class Client: """ iRedis client, hold a redis-py Client to interact with Redis. """ def reder_funcname_mapping(self): mapping = {} for func_name, func in renders.__dict__.items(): if callable(func): mapping[func_name] = func return mapping def __init__(self, host, port, db, password=None, encoding=None): self.host = host self.port = port self.db = db if encoding: self.connection = Connection( host=self.host, port=self.port, db=self.db, password=password, encoding=encoding, decode_responses=True, encoding_errors="replace", ) else: self.connection = Connection( host=self.host, port=self.port, db=self.db, password=password, decode_responses=False, ) # all command upper case self.answer_callbacks = command2callback self.callbacks = self.reder_funcname_mapping() def __str__(self): if self.db: return f"{self.host}:{self.port}[{self.db}]" return f"{self.host}:{self.port}" def execute_command_and_read_response( self, completer, command_name, *args, **options ): "Execute a command and return a parsed response" # === pre hook === # TRANSATION state chage if command_name.upper() in ["EXEC", "DISCARD"]: logger.debug(f"[After hook] Command is {command_name}, unset transaction.") config.transaction = False if command_name.upper() in ["ZSCAN", "ZPOPMAX", "ZPOPMIN"]: config.withscores = True try: self.connection.send_command(command_name, *args) resp = self.parse_response( self.connection, completer, command_name, **options ) # retry on timeout except (ConnectionError, TimeoutError) as e: self.connection.disconnect() if not (self.connection.retry_on_timeout and isinstance(e, TimeoutError)): raise self.connection.send_command(command_name, *args) resp = self.parse_response( self.connection, completer, command_name, **options ) except redis.exceptions.ExecAbortError: config.transaction = False raise # === After hook === # SELECT db on AUTH if command_name.upper() == "AUTH" and self.db: select_result = self.execute_command_and_read_response( completer, "SELECT", self.db ) if nativestr(select_result) != "OK": raise ConnectionError("Invalid Database") elif command_name.upper() == "SELECT": logger.debug("[After hook] Command is SELECT, change self.db.") self.db = int(args[0]) if command_name.upper() == "MULTI": logger.debug("[After hook] Command is MULTI, start transaction.") config.transaction = True return resp def render_command_result(self, command_name, response, completer): """ Render command result using callback :param command_name: command name, (will be converted to UPPER case; :param completer: completers to be patched; """ command_upper = command_name.upper() # else, use defined callback if ( command_upper in self.answer_callbacks and self.answer_callbacks[command_upper] ): callback_name = self.answer_callbacks[command_upper] callback = self.callbacks[callback_name] rendered = callback(response, completer) # FIXME # not implemented command, use no transaction # this `else` should be deleted finally else: rendered = response logger.info(f"[rendered] {rendered}") return rendered def parse_response(self, connection, completer, command_name, **options): "Parses a response from the Redis server" response = connection.read_response() logger.info(f"[Redis-Server] Response: {response}") # if in transaction, use queue render first if config.transaction: callback = renders.render_transaction_queue rendered = callback(response, completer) else: rendered = self.render_command_result(command_name, response, completer) return rendered def send_command(self, command, completer): """ Send command to redis-server, return parsed response. :param command: text command, not parsed :param completer: RedisGrammarCompleter will update completer based on redis response. eg: update key completer after ``keys`` command """ input_command = "" try: input_command, args = split_command_args(command, all_commands) self.pre_hook(command, completer) redis_resp = self.execute_command_and_read_response( completer, input_command, *args ) except Exception as e: logger.exception(e) return render_error(str(e)) finally: config.withscores = False return redis_resp def pre_hook(self, command, completer): """ Before execute command, patch completers first. Eg: When user run `GET foo`, key completer need to touch foo. Only works when compile-grammar thread is done. """ if not completer: logger.warning("[Pre patch completer] Complter not ready, not patched...") return redis_grammar = completer.compiled_grammar m = redis_grammar.match(command) if not m: # invalide command! return variables = m.variables() # zset withscores withscores = variables.get("withscores") logger.debug(f"[PRE HOOK] withscores: {withscores}") if withscores: config.withscores = True # auto update LatestUsedFirstWordCompleter for _token, _completer in completer.completers.items(): if not isinstance(_completer, LatestUsedFirstWordCompleter): continue # getall always returns a [] tokens_in_command = variables.getall(_token) for tokens_in_command in tokens_in_command: # prompt_toolkit didn't support multi tokens # like DEL key1 key2 key3 # so we have to split them manualy for single_token in _strip_quote_args(tokens_in_command): _completer.touch(single_token) logger.debug(f"[Complter {_token} updated] Done: {_completer.words}")
def tearDown(self): self._redis.redis.flushdb() Connection(self._redis.redis).disconnect() super(TestRedisListenStore, self).tearDown()
class Client: """ iRedis client, hold a redis-py Client to interact with Redis. """ def __init__(self, host, port, db, password=None): self.host = host self.port = port self.db = db if config.decode: self.connection = Connection( host=self.host, port=self.port, db=self.db, password=password, encoding=config.decode, decode_responses=True, encoding_errors="replace", socket_keepalive=config.socket_keepalive, ) else: self.connection = Connection( host=self.host, port=self.port, db=self.db, password=password, decode_responses=False, socket_keepalive=config.socket_keepalive, ) # all command upper case self.answer_callbacks = command2callback try: self.connection.connect() except Exception as e: print(str(e), file=sys.stderr) if not config.no_info: try: self.get_server_info() except Exception as e: logger.warning(f"[After Connection] {str(e)}") config.no_version_reason = str(e) else: config.no_version_reason = "--no-info flag activated" def get_server_info(self): # safe to decode Redis's INFO response info_resp = nativestr(self.execute("INFO")) version = re.findall(r"^redis_version:([\d\.]+)\r\n", info_resp, re.MULTILINE)[0] logger.debug(f"[Redis Version] {version}") config.version = version def __str__(self): if self.db: return f"{self.host}:{self.port}[{self.db}]" return f"{self.host}:{self.port}" def client_execute_command(self, command_name, *args): command = command_name.upper() if command == "HELP": yield self.do_help(*args) if command == "PEEK": yield from self.do_peek(*args) if command == "CLEAR": clear() if command == "EXIT": exit() def execute(self, command_name, *args, **options): """Execute a command and return a parsed response Here we retry once for ConnectionError. """ retry_times = config.retry_times # FIXME configureable last_error = None need_refresh_connection = False while retry_times >= 0: try: if need_refresh_connection: print( f"{str(last_error)} retrying... retry left: {retry_times+1}", file=sys.stderr, ) self.connection.disconnect() self.connection.connect() logger.info( f"New connection created, retry on {self.connection}.") self.connection.send_command(command_name, *args) response = self.connection.read_response() except AuthenticationError: raise except (ConnectionError, TimeoutError) as e: logger.warning(f"Connection Error, got {e}, retrying...") last_error = e retry_times -= 1 need_refresh_connection = True except redis.exceptions.ExecAbortError: config.transaction = False raise else: return response raise last_error def _dynamic_render(self, command_name, response): """ Render command result using callback :param command_name: command name, (will be converted to UPPER case; """ return OutputRender.dynamic_render(command_name=command_name, response=response) def render_response(self, response, command_name): "Parses a response from the Redis server" logger.info(f"[Redis-Server] Response: {response}") # if in transaction, use queue render first if config.transaction: callback = renders.OutputRender.render_transaction_queue rendered = callback(response) else: rendered = self._dynamic_render(command_name, response) return rendered def monitor(self): """Redis' MONITOR command: https://redis.io/commands/monitor This command need to read from a stream resp, so it's different """ while 1: response = self.connection.read_response() yield OutputRender.render_bulk_string_decode(response) def subscribing(self): while 1: response = self.connection.read_response() yield OutputRender.render_subscribe(response) def unsubscribing(self): "unsubscribe from all channels" response = self.execute("UNSUBSCRIBE") yield OutputRender.render_subscribe(response) def split_command_and_pipeline(self, rawinput, completer: IRedisCompleter): """ split user raw input to redis command and shell pipeline. eg: GET json | jq .key return: GET json, jq . key """ grammar = completer.get_completer(input_text=rawinput).compiled_grammar matched = grammar.match(rawinput) if not matched: # invalide command! return rawinput, None variables = matched.variables() shell_command = variables.get("shellcommand") if shell_command: redis_command = rawinput.replace(shell_command, "") shell_command = shell_command.lstrip("| ") return redis_command, shell_command return rawinput, None def send_command(self, raw_command, completer=None): # noqa """ Send raw_command to redis-server, return parsed response. :param raw_command: text raw_command, not parsed :param completer: RedisGrammarCompleter will update completer based on redis response. eg: update key completer after ``keys`` raw_command """ if completer is None: # not in a tty redis_command, shell_command = raw_command, None else: redis_command, shell_command = self.split_command_and_pipeline( raw_command, completer) logger.info( f"[Prepare command] Redis: {redis_command}, Shell: {shell_command}" ) try: command_name, args = split_command_args(redis_command, all_commands) logger.info( f"[Split command] command: {command_name}, args: {args}") input_command_upper = command_name.upper() # Confirm for dangerous command if config.warning: confirm = confirm_dangerous_command(input_command_upper) # if we can prompt to user, it's always a tty # so we always yield FormattedText here. if confirm is False: yield FormattedText([("class:warning", "Canceled!")]) return if confirm is True: yield FormattedText([("class:warning", "Your Call!!")]) self.pre_hook(raw_command, command_name, args, completer) # if raw_command is not supposed to send to server if input_command_upper in CLIENT_COMMANDS: logger.info(f"{input_command_upper} is an iredis command.") yield from self.client_execute_command(command_name, *args) return redis_resp = self.execute(command_name, *args) # if shell, do not render, just run in shell pipe and show the # subcommand's stdout/stderr if shell_command: # pass the raw response of redis to shell command if isinstance(redis_resp, list): stdin = b"\n".join(redis_resp) else: stdin = redis_resp run(shell_command, input=stdin, stdout=sys.stdout, shell=True) return self.after_hook(raw_command, command_name, args, completer, redis_resp) yield self.render_response(redis_resp, command_name) # FIXME generator response do not support pipeline if input_command_upper == "MONITOR": # TODO special render for monitor try: yield from self.monitor() except KeyboardInterrupt: pass elif input_command_upper in [ "SUBSCRIBE", "PSUBSCRIBE", ]: # enter subscribe mode try: yield from self.subscribing() except KeyboardInterrupt: yield from self.unsubscribing() except Exception as e: logger.exception(e) yield OutputRender.render_error(str(e)) finally: config.withscores = False def after_hook(self, command, command_name, args, completer, response): # === After hook === # SELECT db on AUTH if command_name.upper() == "AUTH": if self.db: select_result = self.execute("SELECT", self.db) if nativestr(select_result) != "OK": raise ConnectionError("Invalid Database") # When the connection is TimeoutError or ConnectionError, reconnect the connection will use it self.connection.password = args[0] elif command_name.upper() == "SELECT": logger.debug("[After hook] Command is SELECT, change self.db.") self.db = int(args[0]) # When the connection is TimeoutError or ConnectionError, reconnect the connection will use it self.connection.db = self.db elif command_name.upper() == "MULTI": logger.debug("[After hook] Command is MULTI, start transaction.") config.transaction = True if completer: completer.update_completer_for_response(command_name, response) def pre_hook(self, command, command_name, args, completer: IRedisCompleter): """ Before execute command, patch completers first. Eg: When user run `GET foo`, key completer need to touch foo. Only works when compile-grammar thread is done. """ # TRANSATION state chage if command_name.upper() in ["EXEC", "DISCARD"]: logger.debug( f"[After hook] Command is {command_name}, unset transaction.") config.transaction = False # score display for sorted set if command_name.upper() in ["ZSCAN", "ZPOPMAX", "ZPOPMIN"]: config.withscores = True # not a tty if not completer: logger.warning( "[Pre patch completer] Complter is None, not a tty, " "not patch completers, not set withscores") return completer.update_completer_for_input(command) redis_grammar = completer.get_completer(command).compiled_grammar m = redis_grammar.match(command) if not m: # invalide command! return variables = m.variables() # zset withscores withscores = variables.get("withscores") logger.debug(f"[PRE HOOK] withscores: {withscores}") if withscores: config.withscores = True def do_help(self, *args): command_docs_name = "-".join(args).lower() command_summary_name = " ".join(args).upper() try: doc_file = open(project_data / "commands" / f"{command_docs_name}.md") except FileNotFoundError: raise NotRedisCommand( f"{command_summary_name} is not a valide Redis command.") with doc_file as doc_file: doc = doc_file.read() rendered_detail = markdown.render(doc) summary_dict = commands_summary[command_summary_name] avaiable_version = summary_dict.get("since", "?") server_version = config.version # FIXME anything strange with single quotes? logger.debug(f"[--version--] '{server_version}'") try: is_avaiable = StrictVersion(server_version) > StrictVersion( avaiable_version) except Exception as e: logger.exception(e) is_avaiable = None if is_avaiable: avaiable_text = f"(Avaiable on your redis-server: {server_version})" elif is_avaiable is False: avaiable_text = f"(Not avaiable on your redis-server: {server_version})" else: avaiable_text = "" since_text = f"{avaiable_version} {avaiable_text}" summary = [ ("", "\n"), ("class:doccommand", " " + command_summary_name), ("", "\n"), ("class:dockey", " summary: "), ("", summary_dict.get("summary", "No summary")), ("", "\n"), ("class:dockey", " complexity: "), ("", summary_dict.get("complexity", "?")), ("", "\n"), ("class:dockey", " since: "), ("", since_text), ("", "\n"), ("class:dockey", " group: "), ("", summary_dict.get("group", "?")), ("", "\n"), ("class:dockey", " syntax: "), ("", command_summary_name), # command *compose_command_syntax(summary_dict, style_class=""), # command args ("", "\n\n"), ] return FormattedText(summary + rendered_detail) def do_peek(self, key): """ PEEK command implementation. It's a generator, will run different redis commands based on the key's type, yields FormattedText once a command reached result. Redis current supported types: string, list, set, zset, hash and stream. """ def _string(key): strlen = self.execute("strlen", key) yield FormattedText([("class:dockey", "strlen: "), ("", str(strlen))]) value = self.execute("GET", key) yield FormattedText([ ("class:dockey", "value: "), ("", renders.OutputRender.render_bulk_string(value)), ]) def _list(key): llen = self.execute("llen", key) yield FormattedText([("class:dockey", "llen: "), ("", str(llen))]) if llen <= 20: contents = self.execute(f"LRANGE {key} 0 -1") else: first_10 = self.execute(f"LRANGE {key} 0 9") last_10 = self.execute(f"LRANGE {key} -10 -1") contents = first_10 + [f"{llen-20} elements was omitted ..." ] + last_10 yield FormattedText([("class:dockey", "elements: ")]) yield renders.OutputRender.render_list(contents) def _set(key): cardinality = self.execute("scard", key) yield FormattedText([("class:dockey", "cardinality: "), ("", str(cardinality))]) if cardinality <= 20: contents = self.execute("smembers", key) yield FormattedText([("class:dockey", "members: ")]) yield renders.OutputRender.render_list(contents) else: _, contents = self.execute(f"sscan {key} 0 count 20") first_n = len(contents) yield FormattedText([("class:dockey", f"members (first {first_n}): ")]) yield renders.OutputRender.render_members(contents) # TODO update completers def _zset(key): count = self.execute(f"zcount {key} -inf +inf") yield FormattedText([("class:dockey", "zcount: "), ("", str(count))]) if count <= 20: contents = self.execute(f"zrange {key} 0 -1 withscores") yield FormattedText([("class:dockey", "members: ")]) yield renders.OutputRender.render_members(contents) else: _, contents = self.execute(f"zscan {key} 0 count 20") first_n = len(contents) // 2 yield FormattedText([("class:dockey", f"members (first {first_n}): ")]) config.withscores = True output = renders.OutputRender.render_members(contents) config.withscores = False yield output def _hash(key): hlen = self.execute(f"hlen {key}") yield FormattedText([("class:dockey", "hlen: "), ("", str(hlen))]) if hlen <= 20: contents = self.execute(f"hgetall {key}") yield FormattedText([("class:dockey", "fields: ")]) else: _, contents = self.execute(f"hscan {key} 0 count 20") first_n = len(contents) // 2 yield FormattedText([("class:dockey", f"fields (first {first_n}): ")]) yield renders.OutputRender.render_hash_pairs(contents) def _stream(key): xinfo = self.execute("xinfo stream", key) yield FormattedText([("class:dockey", "XINFO: ")]) yield renders.OutputRender.render_list(xinfo) def _none(key): yield f"Key {key} doesn't exist." resp = nativestr(self.execute("type", key)) # FIXME raw write_result parse FormattedText yield FormattedText([("class:dockey", "type: "), ("", resp)]) if resp == "none": return encoding = nativestr(self.execute("object encoding", key)) yield FormattedText([("class:dockey", "object encoding: "), ("", encoding)]) memory_usage = str(self.execute("memory usage", key)) yield FormattedText([("class:dockey", "memory usage(bytes): "), ("", memory_usage)]) ttl = str(self.execute("ttl", key)) yield FormattedText([("class:dockey", "ttl: "), ("", ttl)]) yield from { "string": _string, "list": _list, "set": _set, "zset": _zset, "hash": _hash, "stream": _stream, "none": _none, }[resp](key)
def __init__(self, resolver, host='localhost', port=6379, **kwargs): host, port = resolver.resolve_host_port(host, port) _Connection.__init__(self, host=host, port=port, **kwargs)
class Client: """ iRedis client, hold a redis-py Client to interact with Redis. """ def reder_funcname_mapping(self): mapping = {} for func_name, func in renders.__dict__.items(): if callable(func): mapping[func_name] = func return mapping def __init__(self, host, port, db, password=None, encoding=None, get_info=True): self.host = host self.port = port self.db = db if encoding: self.connection = Connection( host=self.host, port=self.port, db=self.db, password=password, encoding=encoding, decode_responses=True, encoding_errors="replace", ) else: self.connection = Connection( host=self.host, port=self.port, db=self.db, password=password, decode_responses=False, ) # all command upper case self.answer_callbacks = command2callback self.callbacks = self.reder_funcname_mapping() self.connection.connect() if get_info: try: self.get_server_info() except Exception as e: logger.warn(f"[After Connection] {str(e)}") config.no_version_reason = str(e) else: config.no_version_reason = "--no-info flag activated" def get_server_info(self): self.connection.send_command("INFO") # safe to decode Redis's INFO response info_resp = utils.ensure_str(self.connection.read_response()) version = re.findall(r"^redis_version:([\d\.]+)\r\n", info_resp, re.MULTILINE)[0] logger.debug(f"[Redis Version] {version}") config.version = version def __str__(self): if self.db: return f"{self.host}:{self.port}[{self.db}]" return f"{self.host}:{self.port}" def client_execute_command(self, command_name, *args): command = command_name.upper() if command == "HELP": return self.do_help(*args) def execute_command_and_read_response(self, completer, command_name, *args, **options): "Execute a command and return a parsed response" try: self.connection.send_command(command_name, *args) response = self.connection.read_response() # retry on timeout except (ConnectionError, TimeoutError) as e: self.connection.disconnect() if not (self.connection.retry_on_timeout and isinstance(e, TimeoutError)): raise self.connection.send_command(command_name, *args) response = self.connection.read_response() except redis.exceptions.ExecAbortError: config.transaction = False raise return self.render_response(response, completer, command_name, **options) def render_command_result(self, command_name, response, completer): """ Render command result using callback :param command_name: command name, (will be converted to UPPER case; :param completer: completers to be patched; """ command_upper = command_name.upper() # else, use defined callback if (command_upper in self.answer_callbacks and self.answer_callbacks[command_upper]): callback_name = self.answer_callbacks[command_upper] callback = self.callbacks[callback_name] rendered = callback(response, completer) # FIXME # not implemented command, use no conversion # this `else` should be deleted finally else: rendered = response logger.info(f"[rendered] {rendered}") return rendered def render_response(self, response, completer, command_name, **options): "Parses a response from the Redis server" logger.info(f"[Redis-Server] Response: {response}") # if in transaction, use queue render first if config.transaction: callback = renders.render_transaction_queue rendered = callback(response, completer) else: rendered = self.render_command_result(command_name, response, completer) return rendered def monitor(self): """Redis' MONITOR command: https://redis.io/commands/monitor This command need to read from a stream resp, so it's different """ while 1: response = self.connection.read_response() yield render_bulk_string_decode(response) def subscribing(self): while 1: response = self.connection.read_response() yield render_subscribe(response) def unsubscribing(self): "unsubscribe from all channels" self.connection.send_command("UNSUBSCRIBE") response = self.connection.read_response() yield render_subscribe(response) def send_command(self, raw_command, completer=None): """ Send raw_command to redis-server, return parsed response. :param raw_command: text raw_command, not parsed :param completer: RedisGrammarCompleter will update completer based on redis response. eg: update key completer after ``keys`` raw_command """ command_name = "" try: command_name, args = split_command_args(raw_command, all_commands) # if raw_command is not supposed to send to server if command_name.upper() in CLIENT_COMMANDS: redis_resp = self.client_execute_command(command_name, *args) yield redis_resp return self.pre_hook(raw_command, command_name, args, completer) redis_resp = self.execute_command_and_read_response( completer, command_name, *args) self.after_hook(raw_command, command_name, args, completer) yield redis_resp if command_name.upper() == "MONITOR": # TODO special render for monitor try: yield from self.monitor() except KeyboardInterrupt: pass elif command_name.upper() in [ "SUBSCRIBE", "PSUBSCRIBE", ]: # enter subscribe mode try: yield from self.subscribing() except KeyboardInterrupt: yield from self.unsubscribing() except Exception as e: logger.exception(e) yield render_error(str(e)) finally: config.withscores = False def after_hook(self, command, command_name, args, completer): # === After hook === # SELECT db on AUTH if command_name.upper() == "AUTH" and self.db: select_result = self.execute_command_and_read_response( completer, "SELECT", self.db) if nativestr(select_result) != "OK": raise ConnectionError("Invalid Database") elif command_name.upper() == "SELECT": logger.debug("[After hook] Command is SELECT, change self.db.") self.db = int(args[0]) if command_name.upper() == "MULTI": logger.debug("[After hook] Command is MULTI, start transaction.") config.transaction = True def pre_hook(self, command, command_name, args, completer): """ Before execute command, patch completers first. Eg: When user run `GET foo`, key completer need to touch foo. Only works when compile-grammar thread is done. """ # TRANSATION state chage if command_name.upper() in ["EXEC", "DISCARD"]: logger.debug( f"[After hook] Command is {command_name}, unset transaction.") config.transaction = False # score display for sorted set if command_name.upper() in ["ZSCAN", "ZPOPMAX", "ZPOPMIN"]: config.withscores = True # patch completers if not completer: logger.warning( "[Pre patch completer] Complter not ready, not patched...") return redis_grammar = completer.compiled_grammar m = redis_grammar.match(command) if not m: # invalide command! return variables = m.variables() # zset withscores withscores = variables.get("withscores") logger.debug(f"[PRE HOOK] withscores: {withscores}") if withscores: config.withscores = True # auto update LatestUsedFirstWordCompleter for _token, _completer in completer.completers.items(): if not isinstance(_completer, LatestUsedFirstWordCompleter): continue # getall always returns a [] tokens_in_command = variables.getall(_token) for tokens_in_command in tokens_in_command: # prompt_toolkit didn't support multi tokens # like DEL key1 key2 key3 # so we have to split them manualy for single_token in _strip_quote_args(tokens_in_command): _completer.touch(single_token) logger.debug( f"[Complter {_token} updated] Done: {_completer.words}") def do_help(self, *args): command_docs_name = "-".join(args).lower() command_summary_name = " ".join(args).upper() try: doc_file = open(project_path / "redis-doc" / "commands" / f"{command_docs_name}.md") except FileNotFoundError: raise NotRedisCommand( f"{command_summary_name} is not a valide Redis command.") with doc_file as doc_file: doc = doc_file.read() rendered_detail = markdown.render(doc) summary_dict = commands_summary[command_summary_name] avaiable_version = summary_dict.get("since", "?") server_version = config.version # FIXME anything strange with single quotes? logger.debug(f"[--version--] '{server_version}'") try: is_avaiable = StrictVersion(server_version) > StrictVersion( avaiable_version) except Exception as e: logger.exception(e) is_avaiable = None if is_avaiable: avaiable_text = f"(Avaiable on your redis-server: {server_version})" elif is_avaiable is False: avaiable_text = f"(Not avaiable on your redis-server: {server_version})" else: avaiable_text = "" since_text = f"{avaiable_version} {avaiable_text}" summary = [ ("", "\n"), ("class:doccommand", " " + command_summary_name), ("", "\n"), ("class:dockey", " summary: "), ("", summary_dict.get("summary", "No summary")), ("", "\n"), ("class:dockey", " complexity: "), ("", summary_dict.get("complexity", "?")), ("", "\n"), ("class:dockey", " since: "), ("", since_text), ("", "\n"), ("class:dockey", " group: "), ("", summary_dict.get("group", "?")), ("", "\n"), ("class:dockey", " syntax: "), ("", command_summary_name), # command *compose_command_syntax(summary_dict, style_class=""), # command args ("", "\n\n"), ] return FormattedText(summary + rendered_detail)
def __init__(self): self.conn = Connection()