123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122 |
- from __future__ import unicode_literals
- from distutils.version import StrictVersion
- from itertools import chain
- import io
- import os
- import socket
- import sys
- import threading
- import warnings
- try:
- import ssl
- ssl_available = True
- except ImportError:
- ssl_available = False
- from redis._compat import (xrange, imap, byte_to_chr, unicode, bytes, long,
- nativestr, basestring, iteritems,
- LifoQueue, Empty, Full, urlparse, parse_qs,
- recv, recv_into, select, unquote)
- from redis.exceptions import (
- DataError,
- RedisError,
- ConnectionError,
- TimeoutError,
- BusyLoadingError,
- ResponseError,
- InvalidResponse,
- AuthenticationError,
- NoScriptError,
- ExecAbortError,
- ReadOnlyError
- )
- from redis.utils import HIREDIS_AVAILABLE
- if HIREDIS_AVAILABLE:
- import hiredis
- hiredis_version = StrictVersion(hiredis.__version__)
- HIREDIS_SUPPORTS_CALLABLE_ERRORS = \
- hiredis_version >= StrictVersion('0.1.3')
- HIREDIS_SUPPORTS_BYTE_BUFFER = \
- hiredis_version >= StrictVersion('0.1.4')
- if not HIREDIS_SUPPORTS_BYTE_BUFFER:
- msg = ("redis-py works best with hiredis >= 0.1.4. You're running "
- "hiredis %s. Please consider upgrading." % hiredis.__version__)
- warnings.warn(msg)
- HIREDIS_USE_BYTE_BUFFER = True
- # only use byte buffer if hiredis supports it
- if not HIREDIS_SUPPORTS_BYTE_BUFFER:
- HIREDIS_USE_BYTE_BUFFER = False
- SYM_STAR = b'*'
- SYM_DOLLAR = b'$'
- SYM_CRLF = b'\r\n'
- SYM_EMPTY = b''
- SERVER_CLOSED_CONNECTION_ERROR = "Connection closed by server."
- class Token(object):
- """
- Literal strings in Redis commands, such as the command names and any
- hard-coded arguments are wrapped in this class so we know not to apply
- and encoding rules on them.
- """
- _cache = {}
- @classmethod
- def get_token(cls, value):
- "Gets a cached token object or creates a new one if not already cached"
- # Use try/except because after running for a short time most tokens
- # should already be cached
- try:
- return cls._cache[value]
- except KeyError:
- token = Token(value)
- cls._cache[value] = token
- return token
- def __init__(self, value):
- if isinstance(value, Token):
- value = value.value
- self.value = value
- self.encoded_value = value.encode()
- def __repr__(self):
- return self.value
- def __str__(self):
- return self.value
- class Encoder(object):
- "Encode strings to bytes and decode bytes to strings"
- def __init__(self, encoding, encoding_errors, decode_responses):
- self.encoding = encoding
- self.encoding_errors = encoding_errors
- self.decode_responses = decode_responses
- def encode(self, value):
- "Return a bytestring representation of the value"
- if isinstance(value, Token):
- return value.encoded_value
- elif isinstance(value, bytes):
- return value
- elif isinstance(value, bool):
- # special case bool since it is a subclass of int
- raise DataError("Invalid input of type: 'bool'. Convert to a "
- "byte, string or number first.")
- elif isinstance(value, float):
- value = repr(value).encode()
- elif isinstance(value, (int, long)):
- # python 2 repr() on longs is '123L', so use str() instead
- value = str(value).encode()
- elif not isinstance(value, basestring):
- # a value we don't know how to deal with. throw an error
- typename = type(value).__name__
- raise DataError("Invalid input of type: '%s'. Convert to a "
- "byte, string or number first." % typename)
- if isinstance(value, unicode):
- value = value.encode(self.encoding, self.encoding_errors)
- return value
- def decode(self, value, force=False):
- "Return a unicode string from the byte representation"
- if (self.decode_responses or force) and isinstance(value, bytes):
- value = value.decode(self.encoding, self.encoding_errors)
- return value
- class BaseParser(object):
- EXCEPTION_CLASSES = {
- 'ERR': {
- 'max number of clients reached': ConnectionError
- },
- 'EXECABORT': ExecAbortError,
- 'LOADING': BusyLoadingError,
- 'NOSCRIPT': NoScriptError,
- 'READONLY': ReadOnlyError,
- }
- def parse_error(self, response):
- "Parse an error response"
- error_code = response.split(' ')[0]
- if error_code in self.EXCEPTION_CLASSES:
- response = response[len(error_code) + 1:]
- exception_class = self.EXCEPTION_CLASSES[error_code]
- if isinstance(exception_class, dict):
- exception_class = exception_class.get(response, ResponseError)
- return exception_class(response)
- return ResponseError(response)
- class SocketBuffer(object):
- def __init__(self, socket, socket_read_size):
- self._sock = socket
- self.socket_read_size = socket_read_size
- self._buffer = io.BytesIO()
- # number of bytes written to the buffer from the socket
- self.bytes_written = 0
- # number of bytes read from the buffer
- self.bytes_read = 0
- @property
- def length(self):
- return self.bytes_written - self.bytes_read
- def _read_from_socket(self, length=None):
- socket_read_size = self.socket_read_size
- buf = self._buffer
- buf.seek(self.bytes_written)
- marker = 0
- try:
- while True:
- data = recv(self._sock, socket_read_size)
- # an empty string indicates the server shutdown the socket
- if isinstance(data, bytes) and len(data) == 0:
- raise socket.error(SERVER_CLOSED_CONNECTION_ERROR)
- buf.write(data)
- data_length = len(data)
- self.bytes_written += data_length
- marker += data_length
- if length is not None and length > marker:
- continue
- break
- except socket.timeout:
- raise TimeoutError("Timeout reading from socket")
- except socket.error:
- e = sys.exc_info()[1]
- raise ConnectionError("Error while reading from socket: %s" %
- (e.args,))
- def read(self, length):
- length = length + 2 # make sure to read the \r\n terminator
- # make sure we've read enough data from the socket
- if length > self.length:
- self._read_from_socket(length - self.length)
- self._buffer.seek(self.bytes_read)
- data = self._buffer.read(length)
- self.bytes_read += len(data)
- # purge the buffer when we've consumed it all so it doesn't
- # grow forever
- if self.bytes_read == self.bytes_written:
- self.purge()
- return data[:-2]
- def readline(self):
- buf = self._buffer
- buf.seek(self.bytes_read)
- data = buf.readline()
- while not data.endswith(SYM_CRLF):
- # there's more data in the socket that we need
- self._read_from_socket()
- buf.seek(self.bytes_read)
- data = buf.readline()
- self.bytes_read += len(data)
- # purge the buffer when we've consumed it all so it doesn't
- # grow forever
- if self.bytes_read == self.bytes_written:
- self.purge()
- return data[:-2]
- def purge(self):
- self._buffer.seek(0)
- self._buffer.truncate()
- self.bytes_written = 0
- self.bytes_read = 0
- def close(self):
- try:
- self.purge()
- self._buffer.close()
- except Exception:
- # issue #633 suggests the purge/close somehow raised a
- # BadFileDescriptor error. Perhaps the client ran out of
- # memory or something else? It's probably OK to ignore
- # any error being raised from purge/close since we're
- # removing the reference to the instance below.
- pass
- self._buffer = None
- self._sock = None
- class PythonParser(BaseParser):
- "Plain Python parsing class"
- def __init__(self, socket_read_size):
- self.socket_read_size = socket_read_size
- self.encoder = None
- self._sock = None
- self._buffer = None
- def __del__(self):
- try:
- self.on_disconnect()
- except Exception:
- pass
- def on_connect(self, connection):
- "Called when the socket connects"
- self._sock = connection._sock
- self._buffer = SocketBuffer(self._sock, self.socket_read_size)
- self.encoder = connection.encoder
- def on_disconnect(self):
- "Called when the socket disconnects"
- if self._sock is not None:
- self._sock.close()
- self._sock = None
- if self._buffer is not None:
- self._buffer.close()
- self._buffer = None
- self.encoder = None
- def can_read(self):
- return self._buffer and bool(self._buffer.length)
- def read_response(self):
- response = self._buffer.readline()
- if not response:
- raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR)
- byte, response = byte_to_chr(response[0]), response[1:]
- if byte not in ('-', '+', ':', '$', '*'):
- raise InvalidResponse("Protocol Error: %s, %s" %
- (str(byte), str(response)))
- # server returned an error
- if byte == '-':
- response = nativestr(response)
- error = self.parse_error(response)
- # if the error is a ConnectionError, raise immediately so the user
- # is notified
- if isinstance(error, ConnectionError):
- raise error
- # otherwise, we're dealing with a ResponseError that might belong
- # inside a pipeline response. the connection's read_response()
- # and/or the pipeline's execute() will raise this error if
- # necessary, so just return the exception instance here.
- return error
- # single value
- elif byte == '+':
- pass
- # int value
- elif byte == ':':
- response = long(response)
- # bulk response
- elif byte == '$':
- length = int(response)
- if length == -1:
- return None
- response = self._buffer.read(length)
- # multi-bulk response
- elif byte == '*':
- length = int(response)
- if length == -1:
- return None
- response = [self.read_response() for i in xrange(length)]
- if isinstance(response, bytes):
- response = self.encoder.decode(response)
- return response
- class HiredisParser(BaseParser):
- "Parser class for connections using Hiredis"
- def __init__(self, socket_read_size):
- if not HIREDIS_AVAILABLE:
- raise RedisError("Hiredis is not installed")
- self.socket_read_size = socket_read_size
- if HIREDIS_USE_BYTE_BUFFER:
- self._buffer = bytearray(socket_read_size)
- def __del__(self):
- try:
- self.on_disconnect()
- except Exception:
- pass
- def on_connect(self, connection):
- self._sock = connection._sock
- kwargs = {
- 'protocolError': InvalidResponse,
- 'replyError': self.parse_error,
- }
- # hiredis < 0.1.3 doesn't support functions that create exceptions
- if not HIREDIS_SUPPORTS_CALLABLE_ERRORS:
- kwargs['replyError'] = ResponseError
- if connection.encoder.decode_responses:
- kwargs['encoding'] = connection.encoder.encoding
- self._reader = hiredis.Reader(**kwargs)
- self._next_response = False
- def on_disconnect(self):
- self._sock = None
- self._reader = None
- self._next_response = False
- def can_read(self):
- if not self._reader:
- raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR)
- if self._next_response is False:
- self._next_response = self._reader.gets()
- return self._next_response is not False
- def read_response(self):
- if not self._reader:
- raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR)
- # _next_response might be cached from a can_read() call
- if self._next_response is not False:
- response = self._next_response
- self._next_response = False
- return response
- response = self._reader.gets()
- socket_read_size = self.socket_read_size
- while response is False:
- try:
- if HIREDIS_USE_BYTE_BUFFER:
- bufflen = recv_into(self._sock, self._buffer)
- if bufflen == 0:
- raise socket.error(SERVER_CLOSED_CONNECTION_ERROR)
- else:
- buffer = recv(self._sock, socket_read_size)
- # an empty string indicates the server shutdown the socket
- if not isinstance(buffer, bytes) or len(buffer) == 0:
- raise socket.error(SERVER_CLOSED_CONNECTION_ERROR)
- except socket.timeout:
- raise TimeoutError("Timeout reading from socket")
- except socket.error:
- e = sys.exc_info()[1]
- raise ConnectionError("Error while reading from socket: %s" %
- (e.args,))
- if HIREDIS_USE_BYTE_BUFFER:
- self._reader.feed(self._buffer, 0, bufflen)
- else:
- self._reader.feed(buffer)
- response = self._reader.gets()
- # if an older version of hiredis is installed, we need to attempt
- # to convert ResponseErrors to their appropriate types.
- if not HIREDIS_SUPPORTS_CALLABLE_ERRORS:
- if isinstance(response, ResponseError):
- response = self.parse_error(response.args[0])
- elif isinstance(response, list) and response and \
- isinstance(response[0], ResponseError):
- response[0] = self.parse_error(response[0].args[0])
- # if the response is a ConnectionError or the response is a list and
- # the first item is a ConnectionError, raise it as something bad
- # happened
- if isinstance(response, ConnectionError):
- raise response
- elif isinstance(response, list) and response and \
- isinstance(response[0], ConnectionError):
- raise response[0]
- return response
- if HIREDIS_AVAILABLE:
- DefaultParser = HiredisParser
- else:
- DefaultParser = PythonParser
- class Connection(object):
- "Manages TCP communication to and from a Redis server"
- description_format = "Connection<host=%(host)s,port=%(port)s,db=%(db)s>"
- def __init__(self, host='localhost', port=6379, db=0, password=None,
- socket_timeout=None, socket_connect_timeout=None,
- socket_keepalive=False, socket_keepalive_options=None,
- socket_type=0, retry_on_timeout=False, encoding='utf-8',
- encoding_errors='strict', decode_responses=False,
- parser_class=DefaultParser, socket_read_size=65536):
- self.pid = os.getpid()
- self.host = host
- self.port = int(port)
- self.db = db
- self.password = password
- self.socket_timeout = socket_timeout
- self.socket_connect_timeout = socket_connect_timeout or socket_timeout
- self.socket_keepalive = socket_keepalive
- self.socket_keepalive_options = socket_keepalive_options or {}
- self.socket_type = socket_type
- self.retry_on_timeout = retry_on_timeout
- self.encoder = Encoder(encoding, encoding_errors, decode_responses)
- self._sock = None
- self._parser = parser_class(socket_read_size=socket_read_size)
- self._description_args = {
- 'host': self.host,
- 'port': self.port,
- 'db': self.db,
- }
- self._connect_callbacks = []
- self._buffer_cutoff = 6000
- def __repr__(self):
- return self.description_format % self._description_args
- def __del__(self):
- try:
- self.disconnect()
- except Exception:
- pass
- def register_connect_callback(self, callback):
- self._connect_callbacks.append(callback)
- def clear_connect_callbacks(self):
- self._connect_callbacks = []
- def connect(self):
- "Connects to the Redis server if not already connected"
- if self._sock:
- return
- try:
- sock = self._connect()
- except socket.timeout:
- raise TimeoutError("Timeout connecting to server")
- except socket.error:
- e = sys.exc_info()[1]
- raise ConnectionError(self._error_message(e))
- self._sock = sock
- try:
- self.on_connect()
- except RedisError:
- # clean up after any error in on_connect
- self.disconnect()
- raise
- # run any user callbacks. right now the only internal callback
- # is for pubsub channel/pattern resubscription
- for callback in self._connect_callbacks:
- callback(self)
- def _connect(self):
- "Create a TCP socket connection"
- # we want to mimic what socket.create_connection does to support
- # ipv4/ipv6, but we want to set options prior to calling
- # socket.connect()
- err = None
- for res in socket.getaddrinfo(self.host, self.port, self.socket_type,
- socket.SOCK_STREAM):
- family, socktype, proto, canonname, socket_address = res
- sock = None
- try:
- sock = socket.socket(family, socktype, proto)
- # TCP_NODELAY
- sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
- # TCP_KEEPALIVE
- if self.socket_keepalive:
- sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
- for k, v in iteritems(self.socket_keepalive_options):
- sock.setsockopt(socket.SOL_TCP, k, v)
- # set the socket_connect_timeout before we connect
- sock.settimeout(self.socket_connect_timeout)
- # connect
- sock.connect(socket_address)
- # set the socket_timeout now that we're connected
- sock.settimeout(self.socket_timeout)
- return sock
- except socket.error as _:
- err = _
- if sock is not None:
- sock.close()
- if err is not None:
- raise err
- raise socket.error("socket.getaddrinfo returned an empty list")
- def _error_message(self, exception):
- # args for socket.error can either be (errno, "message")
- # or just "message"
- if len(exception.args) == 1:
- return "Error connecting to %s:%s. %s." % \
- (self.host, self.port, exception.args[0])
- else:
- return "Error %s connecting to %s:%s. %s." % \
- (exception.args[0], self.host, self.port, exception.args[1])
- def on_connect(self):
- "Initialize the connection, authenticate and select a database"
- self._parser.on_connect(self)
- # if a password is specified, authenticate
- if self.password:
- self.send_command('AUTH', self.password)
- if nativestr(self.read_response()) != 'OK':
- raise AuthenticationError('Invalid Password')
- # if a database is specified, switch to it
- if self.db:
- self.send_command('SELECT', self.db)
- if nativestr(self.read_response()) != 'OK':
- raise ConnectionError('Invalid Database')
- def disconnect(self):
- "Disconnects from the Redis server"
- self._parser.on_disconnect()
- if self._sock is None:
- return
- try:
- self._sock.shutdown(socket.SHUT_RDWR)
- self._sock.close()
- except socket.error:
- pass
- self._sock = None
- def send_packed_command(self, command):
- "Send an already packed command to the Redis server"
- if not self._sock:
- self.connect()
- try:
- if isinstance(command, str):
- command = [command]
- for item in command:
- self._sock.sendall(item)
- except socket.timeout:
- self.disconnect()
- raise TimeoutError("Timeout writing to socket")
- except socket.error:
- e = sys.exc_info()[1]
- self.disconnect()
- if len(e.args) == 1:
- errno, errmsg = 'UNKNOWN', e.args[0]
- else:
- errno = e.args[0]
- errmsg = e.args[1]
- raise ConnectionError("Error %s while writing to socket. %s." %
- (errno, errmsg))
- except Exception as e:
- self.disconnect()
- raise e
- def send_command(self, *args):
- "Pack and send a command to the Redis server"
- self.send_packed_command(self.pack_command(*args))
- def can_read(self, timeout=0):
- "Poll the socket to see if there's data that can be read."
- sock = self._sock
- if not sock:
- self.connect()
- sock = self._sock
- return self._parser.can_read() or \
- bool(select([sock], [], [], timeout)[0])
- def read_response(self):
- "Read the response from a previously sent command"
- try:
- response = self._parser.read_response()
- except Exception as e:
- self.disconnect()
- raise e
- if isinstance(response, ResponseError):
- raise response
- return response
- def pack_command(self, *args):
- "Pack a series of arguments into the Redis protocol"
- output = []
- # the client might have included 1 or more literal arguments in
- # the command name, e.g., 'CONFIG GET'. The Redis server expects these
- # arguments to be sent separately, so split the first argument
- # manually. All of these arguements get wrapped in the Token class
- # to prevent them from being encoded.
- command = args[0]
- if ' ' in command:
- args = tuple(Token.get_token(s)
- for s in command.split()) + args[1:]
- else:
- args = (Token.get_token(command),) + args[1:]
- buff = SYM_EMPTY.join((SYM_STAR, str(len(args)).encode(), SYM_CRLF))
- buffer_cutoff = self._buffer_cutoff
- for arg in imap(self.encoder.encode, args):
- # to avoid large string mallocs, chunk the command into the
- # output list if we're sending large values
- if len(buff) > buffer_cutoff or len(arg) > buffer_cutoff:
- buff = SYM_EMPTY.join(
- (buff, SYM_DOLLAR, str(len(arg)).encode(), SYM_CRLF))
- output.append(buff)
- output.append(arg)
- buff = SYM_CRLF
- else:
- buff = SYM_EMPTY.join(
- (buff, SYM_DOLLAR, str(len(arg)).encode(),
- SYM_CRLF, arg, SYM_CRLF))
- output.append(buff)
- return output
- def pack_commands(self, commands):
- "Pack multiple commands into the Redis protocol"
- output = []
- pieces = []
- buffer_length = 0
- buffer_cutoff = self._buffer_cutoff
- for cmd in commands:
- for chunk in self.pack_command(*cmd):
- chunklen = len(chunk)
- if buffer_length > buffer_cutoff or chunklen > buffer_cutoff:
- output.append(SYM_EMPTY.join(pieces))
- buffer_length = 0
- pieces = []
- if chunklen > self._buffer_cutoff:
- output.append(chunk)
- else:
- pieces.append(chunk)
- buffer_length += chunklen
- if pieces:
- output.append(SYM_EMPTY.join(pieces))
- return output
- class SSLConnection(Connection):
- description_format = "SSLConnection<host=%(host)s,port=%(port)s,db=%(db)s>"
- def __init__(self, ssl_keyfile=None, ssl_certfile=None,
- ssl_cert_reqs='required', ssl_ca_certs=None, **kwargs):
- if not ssl_available:
- raise RedisError("Python wasn't built with SSL support")
- super(SSLConnection, self).__init__(**kwargs)
- self.keyfile = ssl_keyfile
- self.certfile = ssl_certfile
- if ssl_cert_reqs is None:
- ssl_cert_reqs = ssl.CERT_NONE
- elif isinstance(ssl_cert_reqs, basestring):
- CERT_REQS = {
- 'none': ssl.CERT_NONE,
- 'optional': ssl.CERT_OPTIONAL,
- 'required': ssl.CERT_REQUIRED
- }
- if ssl_cert_reqs not in CERT_REQS:
- raise RedisError(
- "Invalid SSL Certificate Requirements Flag: %s" %
- ssl_cert_reqs)
- ssl_cert_reqs = CERT_REQS[ssl_cert_reqs]
- self.cert_reqs = ssl_cert_reqs
- self.ca_certs = ssl_ca_certs
- def _connect(self):
- "Wrap the socket with SSL support"
- sock = super(SSLConnection, self)._connect()
- sock = ssl.wrap_socket(sock,
- cert_reqs=self.cert_reqs,
- keyfile=self.keyfile,
- certfile=self.certfile,
- ca_certs=self.ca_certs)
- return sock
- class UnixDomainSocketConnection(Connection):
- description_format = "UnixDomainSocketConnection<path=%(path)s,db=%(db)s>"
- def __init__(self, path='', db=0, password=None,
- socket_timeout=None, encoding='utf-8',
- encoding_errors='strict', decode_responses=False,
- retry_on_timeout=False,
- parser_class=DefaultParser, socket_read_size=65536):
- self.pid = os.getpid()
- self.path = path
- self.db = db
- self.password = password
- self.socket_timeout = socket_timeout
- self.retry_on_timeout = retry_on_timeout
- self.encoder = Encoder(encoding, encoding_errors, decode_responses)
- self._sock = None
- self._parser = parser_class(socket_read_size=socket_read_size)
- self._description_args = {
- 'path': self.path,
- 'db': self.db,
- }
- self._connect_callbacks = []
- self._buffer_cutoff = 6000
- def _connect(self):
- "Create a Unix domain socket connection"
- sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
- sock.settimeout(self.socket_timeout)
- sock.connect(self.path)
- return sock
- def _error_message(self, exception):
- # args for socket.error can either be (errno, "message")
- # or just "message"
- if len(exception.args) == 1:
- return "Error connecting to unix socket: %s. %s." % \
- (self.path, exception.args[0])
- else:
- return "Error %s connecting to unix socket: %s. %s." % \
- (exception.args[0], self.path, exception.args[1])
- FALSE_STRINGS = ('0', 'F', 'FALSE', 'N', 'NO')
- def to_bool(value):
- if value is None or value == '':
- return None
- if isinstance(value, basestring) and value.upper() in FALSE_STRINGS:
- return False
- return bool(value)
- URL_QUERY_ARGUMENT_PARSERS = {
- 'socket_timeout': float,
- 'socket_connect_timeout': float,
- 'socket_keepalive': to_bool,
- 'retry_on_timeout': to_bool,
- 'max_connections': int,
- }
- class ConnectionPool(object):
- "Generic connection pool"
- @classmethod
- def from_url(cls, url, db=None, decode_components=False, **kwargs):
- """
- Return a connection pool configured from the given URL.
- For example::
- redis://[:password]@localhost:6379/0
- rediss://[:password]@localhost:6379/0
- unix://[:password]@/path/to/socket.sock?db=0
- Three URL schemes are supported:
- - ```redis://``
- <https://www.iana.org/assignments/uri-schemes/prov/redis>`_ creates a
- normal TCP socket connection
- - ```rediss://``
- <https://www.iana.org/assignments/uri-schemes/prov/rediss>`_ creates
- a SSL wrapped TCP socket connection
- - ``unix://`` creates a Unix Domain Socket connection
- There are several ways to specify a database number. The parse function
- will return the first specified option:
- 1. A ``db`` querystring option, e.g. redis://localhost?db=0
- 2. If using the redis:// scheme, the path argument of the url, e.g.
- redis://localhost/0
- 3. The ``db`` argument to this function.
- If none of these options are specified, db=0 is used.
- The ``decode_components`` argument allows this function to work with
- percent-encoded URLs. If this argument is set to ``True`` all ``%xx``
- escapes will be replaced by their single-character equivalents after
- the URL has been parsed. This only applies to the ``hostname``,
- ``path``, and ``password`` components.
- Any additional querystring arguments and keyword arguments will be
- passed along to the ConnectionPool class's initializer. The querystring
- arguments ``socket_connect_timeout`` and ``socket_timeout`` if supplied
- are parsed as float values. The arguments ``socket_keepalive`` and
- ``retry_on_timeout`` are parsed to boolean values that accept
- True/False, Yes/No values to indicate state. Invalid types cause a
- ``UserWarning`` to be raised. In the case of conflicting arguments,
- querystring arguments always win.
- """
- url = urlparse(url)
- url_options = {}
- for name, value in iteritems(parse_qs(url.query)):
- if value and len(value) > 0:
- parser = URL_QUERY_ARGUMENT_PARSERS.get(name)
- if parser:
- try:
- url_options[name] = parser(value[0])
- except (TypeError, ValueError):
- warnings.warn(UserWarning(
- "Invalid value for `%s` in connection URL." % name
- ))
- else:
- url_options[name] = value[0]
- if decode_components:
- password = unquote(url.password) if url.password else None
- path = unquote(url.path) if url.path else None
- hostname = unquote(url.hostname) if url.hostname else None
- else:
- password = url.password
- path = url.path
- hostname = url.hostname
- # We only support redis:// and unix:// schemes.
- if url.scheme == 'unix':
- url_options.update({
- 'password': password,
- 'path': path,
- 'connection_class': UnixDomainSocketConnection,
- })
- else:
- url_options.update({
- 'host': hostname,
- 'port': int(url.port or 6379),
- 'password': password,
- })
- # If there's a path argument, use it as the db argument if a
- # querystring value wasn't specified
- if 'db' not in url_options and path:
- try:
- url_options['db'] = int(path.replace('/', ''))
- except (AttributeError, ValueError):
- pass
- if url.scheme == 'rediss':
- url_options['connection_class'] = SSLConnection
- # last shot at the db value
- url_options['db'] = int(url_options.get('db', db or 0))
- # update the arguments from the URL values
- kwargs.update(url_options)
- # backwards compatability
- if 'charset' in kwargs:
- warnings.warn(DeprecationWarning(
- '"charset" is deprecated. Use "encoding" instead'))
- kwargs['encoding'] = kwargs.pop('charset')
- if 'errors' in kwargs:
- warnings.warn(DeprecationWarning(
- '"errors" is deprecated. Use "encoding_errors" instead'))
- kwargs['encoding_errors'] = kwargs.pop('errors')
- return cls(**kwargs)
- def __init__(self, connection_class=Connection, max_connections=None,
- **connection_kwargs):
- """
- Create a connection pool. If max_connections is set, then this
- object raises redis.ConnectionError when the pool's limit is reached.
- By default, TCP connections are created unless connection_class is
- specified. Use redis.UnixDomainSocketConnection for unix sockets.
- Any additional keyword arguments are passed to the constructor of
- connection_class.
- """
- max_connections = max_connections or 2 ** 31
- if not isinstance(max_connections, (int, long)) or max_connections < 0:
- raise ValueError('"max_connections" must be a positive integer')
- self.connection_class = connection_class
- self.connection_kwargs = connection_kwargs
- self.max_connections = max_connections
- self.reset()
- def __repr__(self):
- return "%s<%s>" % (
- type(self).__name__,
- self.connection_class.description_format % self.connection_kwargs,
- )
- def reset(self):
- self.pid = os.getpid()
- self._created_connections = 0
- self._available_connections = []
- self._in_use_connections = set()
- self._check_lock = threading.Lock()
- def _checkpid(self):
- if self.pid != os.getpid():
- with self._check_lock:
- if self.pid == os.getpid():
- # another thread already did the work while we waited
- # on the lock.
- return
- self.disconnect()
- self.reset()
- def get_connection(self, command_name, *keys, **options):
- "Get a connection from the pool"
- self._checkpid()
- try:
- connection = self._available_connections.pop()
- except IndexError:
- connection = self.make_connection()
- self._in_use_connections.add(connection)
- return connection
- def get_encoder(self):
- "Return an encoder based on encoding settings"
- kwargs = self.connection_kwargs
- return Encoder(
- encoding=kwargs.get('encoding', 'utf-8'),
- encoding_errors=kwargs.get('encoding_errors', 'strict'),
- decode_responses=kwargs.get('decode_responses', False)
- )
- def make_connection(self):
- "Create a new connection"
- if self._created_connections >= self.max_connections:
- raise ConnectionError("Too many connections")
- self._created_connections += 1
- return self.connection_class(**self.connection_kwargs)
- def release(self, connection):
- "Releases the connection back to the pool"
- self._checkpid()
- if connection.pid != self.pid:
- return
- self._in_use_connections.remove(connection)
- self._available_connections.append(connection)
- def disconnect(self):
- "Disconnects all connections in the pool"
- all_conns = chain(self._available_connections,
- self._in_use_connections)
- for connection in all_conns:
- connection.disconnect()
- class BlockingConnectionPool(ConnectionPool):
- """
- Thread-safe blocking connection pool::
- >>> from redis.client import Redis
- >>> client = Redis(connection_pool=BlockingConnectionPool())
- It performs the same function as the default
- ``:py:class: ~redis.connection.ConnectionPool`` implementation, in that,
- it maintains a pool of reusable connections that can be shared by
- multiple redis clients (safely across threads if required).
- The difference is that, in the event that a client tries to get a
- connection from the pool when all of connections are in use, rather than
- raising a ``:py:class: ~redis.exceptions.ConnectionError`` (as the default
- ``:py:class: ~redis.connection.ConnectionPool`` implementation does), it
- makes the client wait ("blocks") for a specified number of seconds until
- a connection becomes available.
- Use ``max_connections`` to increase / decrease the pool size::
- >>> pool = BlockingConnectionPool(max_connections=10)
- Use ``timeout`` to tell it either how many seconds to wait for a connection
- to become available, or to block forever:
- # Block forever.
- >>> pool = BlockingConnectionPool(timeout=None)
- # Raise a ``ConnectionError`` after five seconds if a connection is
- # not available.
- >>> pool = BlockingConnectionPool(timeout=5)
- """
- def __init__(self, max_connections=50, timeout=20,
- connection_class=Connection, queue_class=LifoQueue,
- **connection_kwargs):
- self.queue_class = queue_class
- self.timeout = timeout
- super(BlockingConnectionPool, self).__init__(
- connection_class=connection_class,
- max_connections=max_connections,
- **connection_kwargs)
- def reset(self):
- self.pid = os.getpid()
- self._check_lock = threading.Lock()
- # Create and fill up a thread safe queue with ``None`` values.
- self.pool = self.queue_class(self.max_connections)
- while True:
- try:
- self.pool.put_nowait(None)
- except Full:
- break
- # Keep a list of actual connection instances so that we can
- # disconnect them later.
- self._connections = []
- def make_connection(self):
- "Make a fresh connection."
- connection = self.connection_class(**self.connection_kwargs)
- self._connections.append(connection)
- return connection
- def get_connection(self, command_name, *keys, **options):
- """
- Get a connection, blocking for ``self.timeout`` until a connection
- is available from the pool.
- If the connection returned is ``None`` then creates a new connection.
- Because we use a last-in first-out queue, the existing connections
- (having been returned to the pool after the initial ``None`` values
- were added) will be returned before ``None`` values. This means we only
- create new connections when we need to, i.e.: the actual number of
- connections will only increase in response to demand.
- """
- # Make sure we haven't changed process.
- self._checkpid()
- # Try and get a connection from the pool. If one isn't available within
- # self.timeout then raise a ``ConnectionError``.
- connection = None
- try:
- connection = self.pool.get(block=True, timeout=self.timeout)
- except Empty:
- # Note that this is not caught by the redis client and will be
- # raised unless handled by application code. If you want never to
- raise ConnectionError("No connection available.")
- # If the ``connection`` is actually ``None`` then that's a cue to make
- # a new connection to add to the pool.
- if connection is None:
- connection = self.make_connection()
- return connection
- def release(self, connection):
- "Releases the connection back to the pool."
- # Make sure we haven't changed process.
- self._checkpid()
- if connection.pid != self.pid:
- return
- # Put the connection back into the pool.
- try:
- self.pool.put_nowait(connection)
- except Full:
- # perhaps the pool has been reset() after a fork? regardless,
- # we don't want this connection
- pass
- def disconnect(self):
- "Disconnects all connections in the pool."
- for connection in self._connections:
- connection.disconnect()
|