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