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 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_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
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)
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)
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}")
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
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}")