import types
import asyncio
import socket
import warnings
import sys
from functools import partial
from collections import deque
from contextlib import contextmanager
from .util import (
encode_command,
wait_ok,
_NOTSET,
_set_result,
_set_exception,
coerced_keys_dict,
decode,
parse_url,
get_event_loop,
)
from .parser import Reader
from .stream import open_connection, open_unix_connection
from .errors import (
ConnectionClosedError,
ConnectionForcedCloseError,
RedisError,
ProtocolError,
ReplyError,
WatchVariableError,
ReadOnlyError,
MaxClientsError
)
from .pubsub import Channel
from .abc import AbcChannel
from .abc import AbcConnection
from .log import logger
__all__ = ['create_connection', 'RedisConnection']
MAX_CHUNK_SIZE = 65536
_PUBSUB_COMMANDS = (
'SUBSCRIBE', b'SUBSCRIBE',
'PSUBSCRIBE', b'PSUBSCRIBE',
'UNSUBSCRIBE', b'UNSUBSCRIBE',
'PUNSUBSCRIBE', b'PUNSUBSCRIBE',
)
async def create_connection(address, *, db=None, password=None, ssl=None,
encoding=None, parser=None, loop=None,
timeout=None, connection_cls=None):
"""Creates redis connection.
Opens connection to Redis server specified by address argument.
Address argument can be one of the following:
* A tuple representing (host, port) pair for TCP connections;
* A string representing either Redis URI or unix domain socket path.
SSL argument is passed through to asyncio.create_connection.
By default SSL/TLS is not used.
By default any timeout is applied at the connection stage, however
you can set a limitted time used trying to open a connection via
the `timeout` Kw.
Encoding argument can be used to decode byte-replies to strings.
By default no decoding is done.
Parser parameter can be used to pass custom Redis protocol parser class.
By default hiredis.Reader is used (unless it is missing or platform
is not CPython).
Return value is RedisConnection instance or a connection_cls if it is
given.
This function is a coroutine.
"""
assert isinstance(address, (tuple, list, str)), "tuple or str expected"
if isinstance(address, str):
address, options = parse_url(address)
logger.debug("Parsed Redis URI %r", address)
db = options.setdefault('db', db)
password = options.setdefault('password', password)
encoding = options.setdefault('encoding', encoding)
timeout = options.setdefault('timeout', timeout)
if 'ssl' in options:
assert options['ssl'] or (not options['ssl'] and not ssl), (
"Conflicting ssl options are set", options['ssl'], ssl)
ssl = ssl or options['ssl']
if timeout is not None and timeout <= 0:
raise ValueError("Timeout has to be None or a number greater than 0")
if connection_cls:
assert issubclass(connection_cls, AbcConnection),\
"connection_class does not meet the AbcConnection contract"
cls = connection_cls
else:
cls = RedisConnection
if loop is not None and sys.version_info >= (3, 8, 0):
warnings.warn("The loop argument is deprecated",
DeprecationWarning)
if isinstance(address, (list, tuple)):
host, port = address
logger.debug("Creating tcp connection to %r", address)
reader, writer = await asyncio.wait_for(open_connection(
host, port, limit=MAX_CHUNK_SIZE, ssl=ssl),
timeout)
sock = writer.transport.get_extra_info('socket')
if sock is not None:
sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
address = sock.getpeername()
address = tuple(address[:2])
else:
logger.debug("Creating unix connection to %r", address)
reader, writer = await asyncio.wait_for(open_unix_connection(
address, ssl=ssl, limit=MAX_CHUNK_SIZE),
timeout)
sock = writer.transport.get_extra_info('socket')
if sock is not None:
address = sock.getpeername()
conn = cls(reader, writer, encoding=encoding,
address=address, parser=parser)
try:
if password is not None:
await conn.auth(password)
if db is not None:
await conn.select(db)
except Exception:
conn.close()
await conn.wait_closed()
raise
return conn
class RedisConnection(AbcConnection):
"""Redis connection."""
def __init__(self, reader, writer, *, address, encoding=None,
parser=None, loop=None):
if loop is not None and sys.version_info >= (3, 8):
warnings.warn("The loop argument is deprecated",
DeprecationWarning)
if parser is None:
parser = Reader
assert callable(parser), (
"Parser argument is not callable", parser)
self._reader = reader
self._writer = writer
self._address = address
self._waiters = deque()
self._reader.set_parser(
parser(protocolError=ProtocolError, replyError=ReplyError)
)
self._reader_task = asyncio.ensure_future(self._read_data())
self._close_msg = None
self._db = 0
self._closing = False
self._closed = False
self._close_state = asyncio.Event()
self._reader_task.add_done_callback(lambda x: self._close_state.set())
self._in_transaction = None
self._transaction_error = None # XXX: never used?
self._in_pubsub = 0
self._pubsub_channels = coerced_keys_dict()
self._pubsub_patterns = coerced_keys_dict()
self._encoding = encoding
self._pipeline_buffer = None
def __repr__(self):
return '<RedisConnection [db:{}]>'.format(self._db)
async def _read_data(self):
"""Response reader task."""
last_error = ConnectionClosedError(
"Connection has been closed by server")
while not self._reader.at_eof():
try:
obj = await self._reader.readobj()
except asyncio.CancelledError:
# NOTE: reader can get cancelled from `close()` method only.
last_error = RuntimeError('this is unexpected')
break
except ProtocolError as exc:
# ProtocolError is fatal
# so connection must be closed
if self._in_transaction is not None:
self._transaction_error = exc
last_error = exc
break
except Exception as exc:
# NOTE: for QUIT command connection error can be received
# before response
last_error = exc
break
else:
if (obj == b'' or obj is None) and self._reader.at_eof():
logger.debug("Connection has been closed by server,"
" response: %r", obj)
last_error = ConnectionClosedError("Reader at end of file")
break
if isinstance(obj, MaxClientsError):
last_error = obj
break
if self._in_pubsub:
self._process_pubsub(obj)
else:
self._process_data(obj)
self._closing = True
get_event_loop().call_soon(self._do_close, last_error)
def _process_data(self, obj):
"""Processes command results."""
assert len(self._waiters) > 0, (type(obj), obj)
waiter, encoding, cb = self._waiters.popleft()
if isinstance(obj, RedisError):
if isinstance(obj, ReplyError):
if obj.args[0].startswith('READONLY'):
obj = ReadOnlyError(obj.args[0])
_set_exception(waiter, obj)
if self._in_transaction is not None:
self._transaction_error = obj
else:
if encoding is not None:
try:
obj = decode(obj, encoding)
except Exception as exc:
_set_exception(waiter, exc)
return
if cb is not None:
try:
obj = cb(obj)
except Exception as exc:
_set_exception(waiter, exc)
return
_set_result(waiter, obj)
if self._in_transaction is not None:
self._in_transaction.append((encoding, cb))
def _process_pubsub(self, obj, *, process_waiters=True):
"""Processes pubsub messages."""
kind, *args, data = obj
if kind in (b'subscribe', b'unsubscribe'):
chan, = args
if process_waiters and self._in_pubsub and self._waiters:
self._process_data(obj)
if kind == b'unsubscribe':
ch = self._pubsub_channels.pop(chan, None)
if ch:
ch.close()
self._in_pubsub = data
elif kind in (b'psubscribe', b'punsubscribe'):
chan, = args
if process_waiters and self._in_pubsub and self._waiters:
self._process_data(obj)
if kind == b'punsubscribe':
ch = self._pubsub_patterns.pop(chan, None)
if ch:
ch.close()
self._in_pubsub = data
elif kind == b'message':
chan, = args
self._pubsub_channels[chan].put_nowait(data)
elif kind == b'pmessage':
pattern, chan = args
self._pubsub_patterns[pattern].put_nowait((chan, data))
elif kind == b'pong':
if process_waiters and self._in_pubsub and self._waiters:
self._process_data(data or b'PONG')
else:
logger.warning("Unknown pubsub message received %r", obj)
@contextmanager
def _buffered(self):
# XXX: we must ensure that no await happens
# as long as we buffer commands.
# Probably we can set some error-raising callback on enter
# and remove it on exit
# if some await happens in between -> throw an error.
# This is creepy solution, 'cause some one might want to await
# on some other source except redis.
# So we must only raise error we someone tries to await
# pending aioredis future
# One of solutions is to return coroutine instead of a future
# in `execute` method.
# In a coroutine we can check if buffering is enabled and raise error.
# TODO: describe in docs difference in pipeline mode for
# conn.execute vs pipeline.execute()
if self._pipeline_buffer is None:
self._pipeline_buffer = bytearray()
try:
yield self
buf = self._pipeline_buffer
self._writer.write(buf)
finally:
self._pipeline_buffer = None
else:
yield self
def execute(self, command, *args, encoding=_NOTSET):
"""Executes redis command and returns Future waiting for the answer.
Raises:
* TypeError if any of args can not be encoded as bytes.
* ReplyError on redis '-ERR' responses.
* ProtocolError when response can not be decoded meaning connection
is broken.
* ConnectionClosedError when either client or server has closed the
connection.
"""
if self._reader is None or self._reader.at_eof():
msg = self._close_msg or "Connection closed or corrupted"
raise ConnectionClosedError(msg)
if command is None:
raise TypeError("command must not be None")
if None in args:
raise TypeError("args must not contain None")
command = command.upper().strip()
is_pubsub = command in _PUBSUB_COMMANDS
is_ping = command in ('PING', b'PING')
if self._in_pubsub and not (is_pubsub or is_ping):
raise RedisError("Connection in SUBSCRIBE mode")
elif is_pubsub:
logger.warning("Deprecated. Use `execute_pubsub` method directly")
return self.execute_pubsub(command, *args)
if command in ('SELECT', b'SELECT'):
cb = partial(self._set_db, args=args)
elif command in ('MULTI', b'MULTI'):
cb = self._start_transaction
elif command in ('EXEC', b'EXEC'):
cb = partial(self._end_transaction, discard=False)
encoding = None
elif command in ('DISCARD', b'DISCARD'):
cb = partial(self._end_transaction, discard=True)
else:
cb = None
if encoding is _NOTSET:
encoding = self._encoding
fut = get_event_loop().create_future()
if self._pipeline_buffer is None:
self._writer.write(encode_command(command, *args))
else:
encode_command(command, *args, buf=self._pipeline_buffer)
self._waiters.append((fut, encoding, cb))
return fut
def execute_pubsub(self, command, *channels):
"""Executes redis (p)subscribe/(p)unsubscribe commands.
Returns asyncio.gather coroutine waiting for all channels/patterns
to receive answers.
"""
command = command.upper().strip()
assert command in _PUBSUB_COMMANDS, (
"Pub/Sub command expected", command)
if self._reader is None or self._reader.at_eof():
raise ConnectionClosedError("Connection closed or corrupted")
if None in set(channels):
raise TypeError("args must not contain None")
if not len(channels):
raise TypeError("No channels/patterns supplied")
is_pattern = len(command) in (10, 12)
mkchannel = partial(Channel, is_pattern=is_pattern)
channels = [ch if isinstance(ch, AbcChannel) else mkchannel(ch)
for ch in channels]
if not all(ch.is_pattern == is_pattern for ch in channels):
raise ValueError("Not all channels {} match command {}"
.format(channels, command))
cmd = encode_command(command, *(ch.name for ch in channels))
res = []
for ch in channels:
fut = get_event_loop().create_future()
res.append(fut)
cb = partial(self._update_pubsub, ch=ch)
self._waiters.append((fut, None, cb))
if self._pipeline_buffer is None:
self._writer.write(cmd)
else:
self._pipeline_buffer.extend(cmd)
return asyncio.gather(*res)
def close(self):
"""Close connection."""
self._do_close(ConnectionForcedCloseError())
def _do_close(self, exc):
if self._closed:
return
self._closed = True
self._closing = False
self._writer.transport.close()
self._reader_task.cancel()
self._reader_task = None
self._writer = None
self._reader = None
self._pipeline_buffer = None
if exc is not None:
self._close_msg = str(exc)
while self._waiters:
waiter, *spam = self._waiters.popleft()
logger.debug("Cancelling waiter %r", (waiter, spam))
if exc is None:
_set_exception(waiter, ConnectionForcedCloseError())
else:
_set_exception(waiter, exc)
while self._pubsub_channels:
_, ch = self._pubsub_channels.popitem()
logger.debug("Closing pubsub channel %r", ch)
ch.close(exc)
while self._pubsub_patterns:
_, ch = self._pubsub_patterns.popitem()
logger.debug("Closing pubsub pattern %r", ch)
ch.close(exc)
@property
def closed(self):
"""True if connection is closed."""
closed = self._closing or self._closed
if not closed and self._reader and self._reader.at_eof():
self._closing = closed = True
get_event_loop().call_soon(self._do_close, None)
return closed
async def wait_closed(self):
"""Coroutine waiting until connection is closed."""
await self._close_state.wait()
@property
def db(self):
"""Currently selected db index."""
return self._db
@property
def encoding(self):
"""Current set codec or None."""
return self._encoding
@property
def address(self):
"""Redis server address, either host-port tuple or str."""
return self._address
def select(self, db):
"""Change the selected database for the current connection."""
if not isinstance(db, int):
raise TypeError("DB must be of int type, not {!r}".format(db))
if db < 0:
raise ValueError("DB must be greater or equal 0, got {!r}"
.format(db))
fut = self.execute('SELECT', db)
return wait_ok(fut)
def _set_db(self, ok, args):
assert ok in {b'OK', 'OK'}, ("Unexpected result of SELECT", ok)
self._db = args[0]
return ok
def _start_transaction(self, ok):
assert self._in_transaction is None, (
"Connection is already in transaction", self._in_transaction)
self._in_transaction = deque()
self._transaction_error = None
return ok
def _end_transaction(self, obj, discard):
assert self._in_transaction is not None, (
"Connection is not in transaction", obj)
self._transaction_error = None
recall, self._in_transaction = self._in_transaction, None
recall.popleft() # ignore first (its _start_transaction)
if discard:
return obj
assert isinstance(obj, list) or (obj is None and not discard), (
"Unexpected MULTI/EXEC result", obj, recall)
# TODO: need to be able to re-try transaction
if obj is None:
err = WatchVariableError("WATCH variable has changed")
obj = [err] * len(recall)
assert len(obj) == len(recall), (
"Wrong number of result items in mutli-exec", obj, recall)
res = []
for o, (encoding, cb) in zip(obj, recall):
if not isinstance(o, RedisError):
try:
if encoding:
o = decode(o, encoding)
if cb:
o = cb(o)
except Exception as err:
res.append(err)
continue
res.append(o)
return res
def _update_pubsub(self, obj, *, ch):
kind, *pattern, channel, subscriptions = obj
self._in_pubsub, was_in_pubsub = subscriptions, self._in_pubsub
# XXX: the channels/patterns storage should be refactored.
# if code which supposed to read from channel/pattern
# failed (exception in reader or else) than
# the channel object will still reside in memory
# and leak memory (messages will be put in queue).
if kind == b'subscribe' and channel not in self._pubsub_channels:
self._pubsub_channels[channel] = ch
elif kind == b'psubscribe' and channel not in self._pubsub_patterns:
self._pubsub_patterns[channel] = ch
if not was_in_pubsub:
self._process_pubsub(obj, process_waiters=False)
return obj
@property
def in_transaction(self):
"""Set to True when MULTI command was issued."""
return self._in_transaction is not None
@property
def in_pubsub(self):
"""Indicates that connection is in PUB/SUB mode.
Provides the number of subscribed channels.
"""
return self._in_pubsub
@property
def pubsub_channels(self):
"""Returns read-only channels dict."""
return types.MappingProxyType(self._pubsub_channels)
@property
def pubsub_patterns(self):
"""Returns read-only patterns dict."""
return types.MappingProxyType(self._pubsub_patterns)
def auth(self, password):
"""Authenticate to server."""
fut = self.execute('AUTH', password)
return wait_ok(fut)
