You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 
 

1264 lines
47 KiB

  1. from __future__ import unicode_literals
  2. from distutils.version import StrictVersion
  3. from itertools import chain
  4. from time import time
  5. import errno
  6. import io
  7. import os
  8. import socket
  9. import sys
  10. import threading
  11. import warnings
  12. from redis._compat import (xrange, imap, byte_to_chr, unicode, long,
  13. nativestr, basestring, iteritems,
  14. LifoQueue, Empty, Full, urlparse, parse_qs,
  15. recv, recv_into, unquote, BlockingIOError,
  16. sendall, shutdown, ssl_wrap_socket)
  17. from redis.exceptions import (
  18. AuthenticationError,
  19. BusyLoadingError,
  20. ConnectionError,
  21. DataError,
  22. ExecAbortError,
  23. InvalidResponse,
  24. NoScriptError,
  25. ReadOnlyError,
  26. RedisError,
  27. ResponseError,
  28. TimeoutError,
  29. )
  30. from redis.utils import HIREDIS_AVAILABLE
  31. try:
  32. import ssl
  33. ssl_available = True
  34. except ImportError:
  35. ssl_available = False
  36. NONBLOCKING_EXCEPTION_ERROR_NUMBERS = {
  37. BlockingIOError: errno.EWOULDBLOCK,
  38. }
  39. if ssl_available:
  40. if hasattr(ssl, 'SSLWantReadError'):
  41. NONBLOCKING_EXCEPTION_ERROR_NUMBERS[ssl.SSLWantReadError] = 2
  42. NONBLOCKING_EXCEPTION_ERROR_NUMBERS[ssl.SSLWantWriteError] = 2
  43. else:
  44. NONBLOCKING_EXCEPTION_ERROR_NUMBERS[ssl.SSLError] = 2
  45. # In Python 2.7 a socket.error is raised for a nonblocking read.
  46. # The _compat module aliases BlockingIOError to socket.error to be
  47. # Python 2/3 compatible.
  48. # However this means that all socket.error exceptions need to be handled
  49. # properly within these exception handlers.
  50. # We need to make sure socket.error is included in these handlers and
  51. # provide a dummy error number that will never match a real exception.
  52. if socket.error not in NONBLOCKING_EXCEPTION_ERROR_NUMBERS:
  53. NONBLOCKING_EXCEPTION_ERROR_NUMBERS[socket.error] = -999999
  54. NONBLOCKING_EXCEPTIONS = tuple(NONBLOCKING_EXCEPTION_ERROR_NUMBERS.keys())
  55. if HIREDIS_AVAILABLE:
  56. import hiredis
  57. hiredis_version = StrictVersion(hiredis.__version__)
  58. HIREDIS_SUPPORTS_CALLABLE_ERRORS = \
  59. hiredis_version >= StrictVersion('0.1.3')
  60. HIREDIS_SUPPORTS_BYTE_BUFFER = \
  61. hiredis_version >= StrictVersion('0.1.4')
  62. HIREDIS_SUPPORTS_ENCODING_ERRORS = \
  63. hiredis_version >= StrictVersion('1.0.0')
  64. if not HIREDIS_SUPPORTS_BYTE_BUFFER:
  65. msg = ("redis-py works best with hiredis >= 0.1.4. You're running "
  66. "hiredis %s. Please consider upgrading." % hiredis.__version__)
  67. warnings.warn(msg)
  68. HIREDIS_USE_BYTE_BUFFER = True
  69. # only use byte buffer if hiredis supports it
  70. if not HIREDIS_SUPPORTS_BYTE_BUFFER:
  71. HIREDIS_USE_BYTE_BUFFER = False
  72. SYM_STAR = b'*'
  73. SYM_DOLLAR = b'$'
  74. SYM_CRLF = b'\r\n'
  75. SYM_EMPTY = b''
  76. SERVER_CLOSED_CONNECTION_ERROR = "Connection closed by server."
  77. SENTINEL = object()
  78. class Encoder(object):
  79. "Encode strings to bytes and decode bytes to strings"
  80. def __init__(self, encoding, encoding_errors, decode_responses):
  81. self.encoding = encoding
  82. self.encoding_errors = encoding_errors
  83. self.decode_responses = decode_responses
  84. def encode(self, value):
  85. "Return a bytestring representation of the value"
  86. if isinstance(value, bytes):
  87. return value
  88. elif isinstance(value, bool):
  89. # special case bool since it is a subclass of int
  90. raise DataError("Invalid input of type: 'bool'. Convert to a "
  91. "byte, string or number first.")
  92. elif isinstance(value, float):
  93. value = repr(value).encode()
  94. elif isinstance(value, (int, long)):
  95. # python 2 repr() on longs is '123L', so use str() instead
  96. value = str(value).encode()
  97. elif not isinstance(value, basestring):
  98. # a value we don't know how to deal with. throw an error
  99. typename = type(value).__name__
  100. raise DataError("Invalid input of type: '%s'. Convert to a "
  101. "byte, string or number first." % typename)
  102. if isinstance(value, unicode):
  103. value = value.encode(self.encoding, self.encoding_errors)
  104. return value
  105. def decode(self, value, force=False):
  106. "Return a unicode string from the byte representation"
  107. if (self.decode_responses or force) and isinstance(value, bytes):
  108. value = value.decode(self.encoding, self.encoding_errors)
  109. return value
  110. class BaseParser(object):
  111. EXCEPTION_CLASSES = {
  112. 'ERR': {
  113. 'max number of clients reached': ConnectionError,
  114. 'Client sent AUTH, but no password is set': AuthenticationError,
  115. 'invalid password': AuthenticationError,
  116. },
  117. 'EXECABORT': ExecAbortError,
  118. 'LOADING': BusyLoadingError,
  119. 'NOSCRIPT': NoScriptError,
  120. 'READONLY': ReadOnlyError,
  121. 'NOAUTH': AuthenticationError,
  122. }
  123. def parse_error(self, response):
  124. "Parse an error response"
  125. error_code = response.split(' ')[0]
  126. if error_code in self.EXCEPTION_CLASSES:
  127. response = response[len(error_code) + 1:]
  128. exception_class = self.EXCEPTION_CLASSES[error_code]
  129. if isinstance(exception_class, dict):
  130. exception_class = exception_class.get(response, ResponseError)
  131. return exception_class(response)
  132. return ResponseError(response)
  133. class SocketBuffer(object):
  134. def __init__(self, socket, socket_read_size, socket_timeout):
  135. self._sock = socket
  136. self.socket_read_size = socket_read_size
  137. self.socket_timeout = socket_timeout
  138. self._buffer = io.BytesIO()
  139. # number of bytes written to the buffer from the socket
  140. self.bytes_written = 0
  141. # number of bytes read from the buffer
  142. self.bytes_read = 0
  143. @property
  144. def length(self):
  145. return self.bytes_written - self.bytes_read
  146. def _read_from_socket(self, length=None, timeout=SENTINEL,
  147. raise_on_timeout=True):
  148. sock = self._sock
  149. socket_read_size = self.socket_read_size
  150. buf = self._buffer
  151. buf.seek(self.bytes_written)
  152. marker = 0
  153. custom_timeout = timeout is not SENTINEL
  154. try:
  155. if custom_timeout:
  156. sock.settimeout(timeout)
  157. while True:
  158. data = recv(self._sock, socket_read_size)
  159. # an empty string indicates the server shutdown the socket
  160. if isinstance(data, bytes) and len(data) == 0:
  161. raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR)
  162. buf.write(data)
  163. data_length = len(data)
  164. self.bytes_written += data_length
  165. marker += data_length
  166. if length is not None and length > marker:
  167. continue
  168. return True
  169. except socket.timeout:
  170. if raise_on_timeout:
  171. raise TimeoutError("Timeout reading from socket")
  172. return False
  173. except NONBLOCKING_EXCEPTIONS as ex:
  174. # if we're in nonblocking mode and the recv raises a
  175. # blocking error, simply return False indicating that
  176. # there's no data to be read. otherwise raise the
  177. # original exception.
  178. allowed = NONBLOCKING_EXCEPTION_ERROR_NUMBERS.get(ex.__class__, -1)
  179. if not raise_on_timeout and ex.errno == allowed:
  180. return False
  181. raise ConnectionError("Error while reading from socket: %s" %
  182. (ex.args,))
  183. finally:
  184. if custom_timeout:
  185. sock.settimeout(self.socket_timeout)
  186. def can_read(self, timeout):
  187. return bool(self.length) or \
  188. self._read_from_socket(timeout=timeout,
  189. raise_on_timeout=False)
  190. def read(self, length):
  191. length = length + 2 # make sure to read the \r\n terminator
  192. # make sure we've read enough data from the socket
  193. if length > self.length:
  194. self._read_from_socket(length - self.length)
  195. self._buffer.seek(self.bytes_read)
  196. data = self._buffer.read(length)
  197. self.bytes_read += len(data)
  198. # purge the buffer when we've consumed it all so it doesn't
  199. # grow forever
  200. if self.bytes_read == self.bytes_written:
  201. self.purge()
  202. return data[:-2]
  203. def readline(self):
  204. buf = self._buffer
  205. buf.seek(self.bytes_read)
  206. data = buf.readline()
  207. while not data.endswith(SYM_CRLF):
  208. # there's more data in the socket that we need
  209. self._read_from_socket()
  210. buf.seek(self.bytes_read)
  211. data = buf.readline()
  212. self.bytes_read += len(data)
  213. # purge the buffer when we've consumed it all so it doesn't
  214. # grow forever
  215. if self.bytes_read == self.bytes_written:
  216. self.purge()
  217. return data[:-2]
  218. def purge(self):
  219. self._buffer.seek(0)
  220. self._buffer.truncate()
  221. self.bytes_written = 0
  222. self.bytes_read = 0
  223. def close(self):
  224. try:
  225. self.purge()
  226. self._buffer.close()
  227. except Exception:
  228. # issue #633 suggests the purge/close somehow raised a
  229. # BadFileDescriptor error. Perhaps the client ran out of
  230. # memory or something else? It's probably OK to ignore
  231. # any error being raised from purge/close since we're
  232. # removing the reference to the instance below.
  233. pass
  234. self._buffer = None
  235. self._sock = None
  236. class PythonParser(BaseParser):
  237. "Plain Python parsing class"
  238. def __init__(self, socket_read_size):
  239. self.socket_read_size = socket_read_size
  240. self.encoder = None
  241. self._sock = None
  242. self._buffer = None
  243. def __del__(self):
  244. try:
  245. self.on_disconnect()
  246. except Exception:
  247. pass
  248. def on_connect(self, connection):
  249. "Called when the socket connects"
  250. self._sock = connection._sock
  251. self._buffer = SocketBuffer(self._sock,
  252. self.socket_read_size,
  253. connection.socket_timeout)
  254. self.encoder = connection.encoder
  255. def on_disconnect(self):
  256. "Called when the socket disconnects"
  257. self._sock = None
  258. if self._buffer is not None:
  259. self._buffer.close()
  260. self._buffer = None
  261. self.encoder = None
  262. def can_read(self, timeout):
  263. return self._buffer and self._buffer.can_read(timeout)
  264. def read_response(self):
  265. response = self._buffer.readline()
  266. if not response:
  267. raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR)
  268. byte, response = byte_to_chr(response[0]), response[1:]
  269. if byte not in ('-', '+', ':', '$', '*'):
  270. raise InvalidResponse("Protocol Error: %s, %s" %
  271. (str(byte), str(response)))
  272. # server returned an error
  273. if byte == '-':
  274. response = nativestr(response)
  275. error = self.parse_error(response)
  276. # if the error is a ConnectionError, raise immediately so the user
  277. # is notified
  278. if isinstance(error, ConnectionError):
  279. raise error
  280. # otherwise, we're dealing with a ResponseError that might belong
  281. # inside a pipeline response. the connection's read_response()
  282. # and/or the pipeline's execute() will raise this error if
  283. # necessary, so just return the exception instance here.
  284. return error
  285. # single value
  286. elif byte == '+':
  287. pass
  288. # int value
  289. elif byte == ':':
  290. response = long(response)
  291. # bulk response
  292. elif byte == '$':
  293. length = int(response)
  294. if length == -1:
  295. return None
  296. response = self._buffer.read(length)
  297. # multi-bulk response
  298. elif byte == '*':
  299. length = int(response)
  300. if length == -1:
  301. return None
  302. response = [self.read_response() for i in xrange(length)]
  303. if isinstance(response, bytes):
  304. response = self.encoder.decode(response)
  305. return response
  306. class HiredisParser(BaseParser):
  307. "Parser class for connections using Hiredis"
  308. def __init__(self, socket_read_size):
  309. if not HIREDIS_AVAILABLE:
  310. raise RedisError("Hiredis is not installed")
  311. self.socket_read_size = socket_read_size
  312. if HIREDIS_USE_BYTE_BUFFER:
  313. self._buffer = bytearray(socket_read_size)
  314. def __del__(self):
  315. try:
  316. self.on_disconnect()
  317. except Exception:
  318. pass
  319. def on_connect(self, connection):
  320. self._sock = connection._sock
  321. self._socket_timeout = connection.socket_timeout
  322. kwargs = {
  323. 'protocolError': InvalidResponse,
  324. 'replyError': self.parse_error,
  325. }
  326. # hiredis < 0.1.3 doesn't support functions that create exceptions
  327. if not HIREDIS_SUPPORTS_CALLABLE_ERRORS:
  328. kwargs['replyError'] = ResponseError
  329. if connection.encoder.decode_responses:
  330. kwargs['encoding'] = connection.encoder.encoding
  331. if HIREDIS_SUPPORTS_ENCODING_ERRORS:
  332. kwargs['errors'] = connection.encoder.encoding_errors
  333. self._reader = hiredis.Reader(**kwargs)
  334. self._next_response = False
  335. def on_disconnect(self):
  336. self._sock = None
  337. self._reader = None
  338. self._next_response = False
  339. def can_read(self, timeout):
  340. if not self._reader:
  341. raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR)
  342. if self._next_response is False:
  343. self._next_response = self._reader.gets()
  344. if self._next_response is False:
  345. return self.read_from_socket(timeout=timeout,
  346. raise_on_timeout=False)
  347. return True
  348. def read_from_socket(self, timeout=SENTINEL, raise_on_timeout=True):
  349. sock = self._sock
  350. custom_timeout = timeout is not SENTINEL
  351. try:
  352. if custom_timeout:
  353. sock.settimeout(timeout)
  354. if HIREDIS_USE_BYTE_BUFFER:
  355. bufflen = recv_into(self._sock, self._buffer)
  356. if bufflen == 0:
  357. raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR)
  358. self._reader.feed(self._buffer, 0, bufflen)
  359. else:
  360. buffer = recv(self._sock, self.socket_read_size)
  361. # an empty string indicates the server shutdown the socket
  362. if not isinstance(buffer, bytes) or len(buffer) == 0:
  363. raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR)
  364. self._reader.feed(buffer)
  365. # data was read from the socket and added to the buffer.
  366. # return True to indicate that data was read.
  367. return True
  368. except socket.timeout:
  369. if raise_on_timeout:
  370. raise TimeoutError("Timeout reading from socket")
  371. return False
  372. except NONBLOCKING_EXCEPTIONS as ex:
  373. # if we're in nonblocking mode and the recv raises a
  374. # blocking error, simply return False indicating that
  375. # there's no data to be read. otherwise raise the
  376. # original exception.
  377. allowed = NONBLOCKING_EXCEPTION_ERROR_NUMBERS.get(ex.__class__, -1)
  378. if not raise_on_timeout and ex.errno == allowed:
  379. return False
  380. raise ConnectionError("Error while reading from socket: %s" %
  381. (ex.args,))
  382. finally:
  383. if custom_timeout:
  384. sock.settimeout(self._socket_timeout)
  385. def read_response(self):
  386. if not self._reader:
  387. raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR)
  388. # _next_response might be cached from a can_read() call
  389. if self._next_response is not False:
  390. response = self._next_response
  391. self._next_response = False
  392. return response
  393. response = self._reader.gets()
  394. while response is False:
  395. self.read_from_socket()
  396. response = self._reader.gets()
  397. # if an older version of hiredis is installed, we need to attempt
  398. # to convert ResponseErrors to their appropriate types.
  399. if not HIREDIS_SUPPORTS_CALLABLE_ERRORS:
  400. if isinstance(response, ResponseError):
  401. response = self.parse_error(response.args[0])
  402. elif isinstance(response, list) and response and \
  403. isinstance(response[0], ResponseError):
  404. response[0] = self.parse_error(response[0].args[0])
  405. # if the response is a ConnectionError or the response is a list and
  406. # the first item is a ConnectionError, raise it as something bad
  407. # happened
  408. if isinstance(response, ConnectionError):
  409. raise response
  410. elif isinstance(response, list) and response and \
  411. isinstance(response[0], ConnectionError):
  412. raise response[0]
  413. return response
  414. if HIREDIS_AVAILABLE:
  415. DefaultParser = HiredisParser
  416. else:
  417. DefaultParser = PythonParser
  418. class Connection(object):
  419. "Manages TCP communication to and from a Redis server"
  420. description_format = "Connection<host=%(host)s,port=%(port)s,db=%(db)s>"
  421. def __init__(self, host='localhost', port=6379, db=0, password=None,
  422. socket_timeout=None, socket_connect_timeout=None,
  423. socket_keepalive=False, socket_keepalive_options=None,
  424. socket_type=0, retry_on_timeout=False, encoding='utf-8',
  425. encoding_errors='strict', decode_responses=False,
  426. parser_class=DefaultParser, socket_read_size=65536,
  427. health_check_interval=0):
  428. self.pid = os.getpid()
  429. self.host = host
  430. self.port = int(port)
  431. self.db = db
  432. self.password = password
  433. self.socket_timeout = socket_timeout
  434. self.socket_connect_timeout = socket_connect_timeout or socket_timeout
  435. self.socket_keepalive = socket_keepalive
  436. self.socket_keepalive_options = socket_keepalive_options or {}
  437. self.socket_type = socket_type
  438. self.retry_on_timeout = retry_on_timeout
  439. self.health_check_interval = health_check_interval
  440. self.next_health_check = 0
  441. self.encoder = Encoder(encoding, encoding_errors, decode_responses)
  442. self._sock = None
  443. self._parser = parser_class(socket_read_size=socket_read_size)
  444. self._description_args = {
  445. 'host': self.host,
  446. 'port': self.port,
  447. 'db': self.db,
  448. }
  449. self._connect_callbacks = []
  450. self._buffer_cutoff = 6000
  451. def __repr__(self):
  452. return self.description_format % self._description_args
  453. def __del__(self):
  454. try:
  455. self.disconnect()
  456. except Exception:
  457. pass
  458. def register_connect_callback(self, callback):
  459. self._connect_callbacks.append(callback)
  460. def clear_connect_callbacks(self):
  461. self._connect_callbacks = []
  462. def connect(self):
  463. "Connects to the Redis server if not already connected"
  464. if self._sock:
  465. return
  466. try:
  467. sock = self._connect()
  468. except socket.timeout:
  469. raise TimeoutError("Timeout connecting to server")
  470. except socket.error:
  471. e = sys.exc_info()[1]
  472. raise ConnectionError(self._error_message(e))
  473. self._sock = sock
  474. try:
  475. self.on_connect()
  476. except RedisError:
  477. # clean up after any error in on_connect
  478. self.disconnect()
  479. raise
  480. # run any user callbacks. right now the only internal callback
  481. # is for pubsub channel/pattern resubscription
  482. for callback in self._connect_callbacks:
  483. callback(self)
  484. def _connect(self):
  485. "Create a TCP socket connection"
  486. # we want to mimic what socket.create_connection does to support
  487. # ipv4/ipv6, but we want to set options prior to calling
  488. # socket.connect()
  489. err = None
  490. for res in socket.getaddrinfo(self.host, self.port, self.socket_type,
  491. socket.SOCK_STREAM):
  492. family, socktype, proto, canonname, socket_address = res
  493. sock = None
  494. try:
  495. sock = socket.socket(family, socktype, proto)
  496. # TCP_NODELAY
  497. sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
  498. # TCP_KEEPALIVE
  499. if self.socket_keepalive:
  500. sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
  501. for k, v in iteritems(self.socket_keepalive_options):
  502. sock.setsockopt(socket.IPPROTO_TCP, k, v)
  503. # set the socket_connect_timeout before we connect
  504. sock.settimeout(self.socket_connect_timeout)
  505. # connect
  506. sock.connect(socket_address)
  507. # set the socket_timeout now that we're connected
  508. sock.settimeout(self.socket_timeout)
  509. return sock
  510. except socket.error as _:
  511. err = _
  512. if sock is not None:
  513. sock.close()
  514. if err is not None:
  515. raise err
  516. raise socket.error("socket.getaddrinfo returned an empty list")
  517. def _error_message(self, exception):
  518. # args for socket.error can either be (errno, "message")
  519. # or just "message"
  520. if len(exception.args) == 1:
  521. return "Error connecting to %s:%s. %s." % \
  522. (self.host, self.port, exception.args[0])
  523. else:
  524. return "Error %s connecting to %s:%s. %s." % \
  525. (exception.args[0], self.host, self.port, exception.args[1])
  526. def on_connect(self):
  527. "Initialize the connection, authenticate and select a database"
  528. self._parser.on_connect(self)
  529. # if a password is specified, authenticate
  530. if self.password:
  531. # avoid checking health here -- PING will fail if we try
  532. # to check the health prior to the AUTH
  533. self.send_command('AUTH', self.password, check_health=False)
  534. if nativestr(self.read_response()) != 'OK':
  535. raise AuthenticationError('Invalid Password')
  536. # if a database is specified, switch to it
  537. if self.db:
  538. self.send_command('SELECT', self.db)
  539. if nativestr(self.read_response()) != 'OK':
  540. raise ConnectionError('Invalid Database')
  541. def disconnect(self):
  542. "Disconnects from the Redis server"
  543. self._parser.on_disconnect()
  544. if self._sock is None:
  545. return
  546. try:
  547. if os.getpid() == self.pid:
  548. shutdown(self._sock, socket.SHUT_RDWR)
  549. self._sock.close()
  550. except socket.error:
  551. pass
  552. self._sock = None
  553. def check_health(self):
  554. "Check the health of the connection with a PING/PONG"
  555. if self.health_check_interval and time() > self.next_health_check:
  556. try:
  557. self.send_command('PING', check_health=False)
  558. if nativestr(self.read_response()) != 'PONG':
  559. raise ConnectionError(
  560. 'Bad response from PING health check')
  561. except (ConnectionError, TimeoutError) as ex:
  562. self.disconnect()
  563. self.send_command('PING', check_health=False)
  564. if nativestr(self.read_response()) != 'PONG':
  565. raise ConnectionError(
  566. 'Bad response from PING health check')
  567. def send_packed_command(self, command, check_health=True):
  568. "Send an already packed command to the Redis server"
  569. if not self._sock:
  570. self.connect()
  571. # guard against health check recurrsion
  572. if check_health:
  573. self.check_health()
  574. try:
  575. if isinstance(command, str):
  576. command = [command]
  577. for item in command:
  578. sendall(self._sock, item)
  579. except socket.timeout:
  580. self.disconnect()
  581. raise TimeoutError("Timeout writing to socket")
  582. except socket.error:
  583. e = sys.exc_info()[1]
  584. self.disconnect()
  585. if len(e.args) == 1:
  586. errno, errmsg = 'UNKNOWN', e.args[0]
  587. else:
  588. errno = e.args[0]
  589. errmsg = e.args[1]
  590. raise ConnectionError("Error %s while writing to socket. %s." %
  591. (errno, errmsg))
  592. except: # noqa: E722
  593. self.disconnect()
  594. raise
  595. def send_command(self, *args, **kwargs):
  596. "Pack and send a command to the Redis server"
  597. self.send_packed_command(self.pack_command(*args),
  598. check_health=kwargs.get('check_health', True))
  599. def can_read(self, timeout=0):
  600. "Poll the socket to see if there's data that can be read."
  601. sock = self._sock
  602. if not sock:
  603. self.connect()
  604. sock = self._sock
  605. return self._parser.can_read(timeout)
  606. def read_response(self):
  607. "Read the response from a previously sent command"
  608. try:
  609. response = self._parser.read_response()
  610. except socket.timeout:
  611. self.disconnect()
  612. raise TimeoutError("Timeout reading from %s:%s" %
  613. (self.host, self.port))
  614. except socket.error:
  615. self.disconnect()
  616. e = sys.exc_info()[1]
  617. raise ConnectionError("Error while reading from %s:%s : %s" %
  618. (self.host, self.port, e.args))
  619. except: # noqa: E722
  620. self.disconnect()
  621. raise
  622. if self.health_check_interval:
  623. self.next_health_check = time() + self.health_check_interval
  624. if isinstance(response, ResponseError):
  625. raise response
  626. return response
  627. def pack_command(self, *args):
  628. "Pack a series of arguments into the Redis protocol"
  629. output = []
  630. # the client might have included 1 or more literal arguments in
  631. # the command name, e.g., 'CONFIG GET'. The Redis server expects these
  632. # arguments to be sent separately, so split the first argument
  633. # manually. These arguments should be bytestrings so that they are
  634. # not encoded.
  635. if isinstance(args[0], unicode):
  636. args = tuple(args[0].encode().split()) + args[1:]
  637. elif b' ' in args[0]:
  638. args = tuple(args[0].split()) + args[1:]
  639. buff = SYM_EMPTY.join((SYM_STAR, str(len(args)).encode(), SYM_CRLF))
  640. buffer_cutoff = self._buffer_cutoff
  641. for arg in imap(self.encoder.encode, args):
  642. # to avoid large string mallocs, chunk the command into the
  643. # output list if we're sending large values
  644. if len(buff) > buffer_cutoff or len(arg) > buffer_cutoff:
  645. buff = SYM_EMPTY.join(
  646. (buff, SYM_DOLLAR, str(len(arg)).encode(), SYM_CRLF))
  647. output.append(buff)
  648. output.append(arg)
  649. buff = SYM_CRLF
  650. else:
  651. buff = SYM_EMPTY.join(
  652. (buff, SYM_DOLLAR, str(len(arg)).encode(),
  653. SYM_CRLF, arg, SYM_CRLF))
  654. output.append(buff)
  655. return output
  656. def pack_commands(self, commands):
  657. "Pack multiple commands into the Redis protocol"
  658. output = []
  659. pieces = []
  660. buffer_length = 0
  661. buffer_cutoff = self._buffer_cutoff
  662. for cmd in commands:
  663. for chunk in self.pack_command(*cmd):
  664. chunklen = len(chunk)
  665. if buffer_length > buffer_cutoff or chunklen > buffer_cutoff:
  666. output.append(SYM_EMPTY.join(pieces))
  667. buffer_length = 0
  668. pieces = []
  669. if chunklen > self._buffer_cutoff:
  670. output.append(chunk)
  671. else:
  672. pieces.append(chunk)
  673. buffer_length += chunklen
  674. if pieces:
  675. output.append(SYM_EMPTY.join(pieces))
  676. return output
  677. class SSLConnection(Connection):
  678. description_format = "SSLConnection<host=%(host)s,port=%(port)s,db=%(db)s>"
  679. def __init__(self, ssl_keyfile=None, ssl_certfile=None,
  680. ssl_cert_reqs='required', ssl_ca_certs=None, **kwargs):
  681. if not ssl_available:
  682. raise RedisError("Python wasn't built with SSL support")
  683. super(SSLConnection, self).__init__(**kwargs)
  684. self.keyfile = ssl_keyfile
  685. self.certfile = ssl_certfile
  686. if ssl_cert_reqs is None:
  687. ssl_cert_reqs = ssl.CERT_NONE
  688. elif isinstance(ssl_cert_reqs, basestring):
  689. CERT_REQS = {
  690. 'none': ssl.CERT_NONE,
  691. 'optional': ssl.CERT_OPTIONAL,
  692. 'required': ssl.CERT_REQUIRED
  693. }
  694. if ssl_cert_reqs not in CERT_REQS:
  695. raise RedisError(
  696. "Invalid SSL Certificate Requirements Flag: %s" %
  697. ssl_cert_reqs)
  698. ssl_cert_reqs = CERT_REQS[ssl_cert_reqs]
  699. self.cert_reqs = ssl_cert_reqs
  700. self.ca_certs = ssl_ca_certs
  701. def _connect(self):
  702. "Wrap the socket with SSL support"
  703. sock = super(SSLConnection, self)._connect()
  704. if hasattr(ssl, "create_default_context"):
  705. context = ssl.create_default_context()
  706. context.check_hostname = False
  707. context.verify_mode = self.cert_reqs
  708. if self.certfile and self.keyfile:
  709. context.load_cert_chain(certfile=self.certfile,
  710. keyfile=self.keyfile)
  711. if self.ca_certs:
  712. context.load_verify_locations(self.ca_certs)
  713. sock = ssl_wrap_socket(context, sock, server_hostname=self.host)
  714. else:
  715. # In case this code runs in a version which is older than 2.7.9,
  716. # we want to fall back to old code
  717. sock = ssl_wrap_socket(ssl,
  718. sock,
  719. cert_reqs=self.cert_reqs,
  720. keyfile=self.keyfile,
  721. certfile=self.certfile,
  722. ca_certs=self.ca_certs)
  723. return sock
  724. class UnixDomainSocketConnection(Connection):
  725. description_format = "UnixDomainSocketConnection<path=%(path)s,db=%(db)s>"
  726. def __init__(self, path='', db=0, password=None,
  727. socket_timeout=None, encoding='utf-8',
  728. encoding_errors='strict', decode_responses=False,
  729. retry_on_timeout=False,
  730. parser_class=DefaultParser, socket_read_size=65536,
  731. health_check_interval=0):
  732. self.pid = os.getpid()
  733. self.path = path
  734. self.db = db
  735. self.password = password
  736. self.socket_timeout = socket_timeout
  737. self.retry_on_timeout = retry_on_timeout
  738. self.health_check_interval = health_check_interval
  739. self.next_health_check = 0
  740. self.encoder = Encoder(encoding, encoding_errors, decode_responses)
  741. self._sock = None
  742. self._parser = parser_class(socket_read_size=socket_read_size)
  743. self._description_args = {
  744. 'path': self.path,
  745. 'db': self.db,
  746. }
  747. self._connect_callbacks = []
  748. self._buffer_cutoff = 6000
  749. def _connect(self):
  750. "Create a Unix domain socket connection"
  751. sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
  752. sock.settimeout(self.socket_timeout)
  753. sock.connect(self.path)
  754. return sock
  755. def _error_message(self, exception):
  756. # args for socket.error can either be (errno, "message")
  757. # or just "message"
  758. if len(exception.args) == 1:
  759. return "Error connecting to unix socket: %s. %s." % \
  760. (self.path, exception.args[0])
  761. else:
  762. return "Error %s connecting to unix socket: %s. %s." % \
  763. (exception.args[0], self.path, exception.args[1])
  764. FALSE_STRINGS = ('0', 'F', 'FALSE', 'N', 'NO')
  765. def to_bool(value):
  766. if value is None or value == '':
  767. return None
  768. if isinstance(value, basestring) and value.upper() in FALSE_STRINGS:
  769. return False
  770. return bool(value)
  771. URL_QUERY_ARGUMENT_PARSERS = {
  772. 'socket_timeout': float,
  773. 'socket_connect_timeout': float,
  774. 'socket_keepalive': to_bool,
  775. 'retry_on_timeout': to_bool,
  776. 'max_connections': int,
  777. 'health_check_interval': int,
  778. }
  779. class ConnectionPool(object):
  780. "Generic connection pool"
  781. @classmethod
  782. def from_url(cls, url, db=None, decode_components=False, **kwargs):
  783. """
  784. Return a connection pool configured from the given URL.
  785. For example::
  786. redis://[:password]@localhost:6379/0
  787. rediss://[:password]@localhost:6379/0
  788. unix://[:password]@/path/to/socket.sock?db=0
  789. Three URL schemes are supported:
  790. - ```redis://``
  791. <https://www.iana.org/assignments/uri-schemes/prov/redis>`_ creates a
  792. normal TCP socket connection
  793. - ```rediss://``
  794. <https://www.iana.org/assignments/uri-schemes/prov/rediss>`_ creates
  795. a SSL wrapped TCP socket connection
  796. - ``unix://`` creates a Unix Domain Socket connection
  797. There are several ways to specify a database number. The parse function
  798. will return the first specified option:
  799. 1. A ``db`` querystring option, e.g. redis://localhost?db=0
  800. 2. If using the redis:// scheme, the path argument of the url, e.g.
  801. redis://localhost/0
  802. 3. The ``db`` argument to this function.
  803. If none of these options are specified, db=0 is used.
  804. The ``decode_components`` argument allows this function to work with
  805. percent-encoded URLs. If this argument is set to ``True`` all ``%xx``
  806. escapes will be replaced by their single-character equivalents after
  807. the URL has been parsed. This only applies to the ``hostname``,
  808. ``path``, and ``password`` components.
  809. Any additional querystring arguments and keyword arguments will be
  810. passed along to the ConnectionPool class's initializer. The querystring
  811. arguments ``socket_connect_timeout`` and ``socket_timeout`` if supplied
  812. are parsed as float values. The arguments ``socket_keepalive`` and
  813. ``retry_on_timeout`` are parsed to boolean values that accept
  814. True/False, Yes/No values to indicate state. Invalid types cause a
  815. ``UserWarning`` to be raised. In the case of conflicting arguments,
  816. querystring arguments always win.
  817. """
  818. url = urlparse(url)
  819. url_options = {}
  820. for name, value in iteritems(parse_qs(url.query)):
  821. if value and len(value) > 0:
  822. parser = URL_QUERY_ARGUMENT_PARSERS.get(name)
  823. if parser:
  824. try:
  825. url_options[name] = parser(value[0])
  826. except (TypeError, ValueError):
  827. warnings.warn(UserWarning(
  828. "Invalid value for `%s` in connection URL." % name
  829. ))
  830. else:
  831. url_options[name] = value[0]
  832. if decode_components:
  833. password = unquote(url.password) if url.password else None
  834. path = unquote(url.path) if url.path else None
  835. hostname = unquote(url.hostname) if url.hostname else None
  836. else:
  837. password = url.password
  838. path = url.path
  839. hostname = url.hostname
  840. # We only support redis://, rediss:// and unix:// schemes.
  841. if url.scheme == 'unix':
  842. url_options.update({
  843. 'password': password,
  844. 'path': path,
  845. 'connection_class': UnixDomainSocketConnection,
  846. })
  847. elif url.scheme in ('redis', 'rediss'):
  848. url_options.update({
  849. 'host': hostname,
  850. 'port': int(url.port or 6379),
  851. 'password': password,
  852. })
  853. # If there's a path argument, use it as the db argument if a
  854. # querystring value wasn't specified
  855. if 'db' not in url_options and path:
  856. try:
  857. url_options['db'] = int(path.replace('/', ''))
  858. except (AttributeError, ValueError):
  859. pass
  860. if url.scheme == 'rediss':
  861. url_options['connection_class'] = SSLConnection
  862. else:
  863. valid_schemes = ', '.join(('redis://', 'rediss://', 'unix://'))
  864. raise ValueError('Redis URL must specify one of the following'
  865. 'schemes (%s)' % valid_schemes)
  866. # last shot at the db value
  867. url_options['db'] = int(url_options.get('db', db or 0))
  868. # update the arguments from the URL values
  869. kwargs.update(url_options)
  870. # backwards compatability
  871. if 'charset' in kwargs:
  872. warnings.warn(DeprecationWarning(
  873. '"charset" is deprecated. Use "encoding" instead'))
  874. kwargs['encoding'] = kwargs.pop('charset')
  875. if 'errors' in kwargs:
  876. warnings.warn(DeprecationWarning(
  877. '"errors" is deprecated. Use "encoding_errors" instead'))
  878. kwargs['encoding_errors'] = kwargs.pop('errors')
  879. return cls(**kwargs)
  880. def __init__(self, connection_class=Connection, max_connections=None,
  881. **connection_kwargs):
  882. """
  883. Create a connection pool. If max_connections is set, then this
  884. object raises redis.ConnectionError when the pool's limit is reached.
  885. By default, TCP connections are created unless connection_class is
  886. specified. Use redis.UnixDomainSocketConnection for unix sockets.
  887. Any additional keyword arguments are passed to the constructor of
  888. connection_class.
  889. """
  890. max_connections = max_connections or 2 ** 31
  891. if not isinstance(max_connections, (int, long)) or max_connections < 0:
  892. raise ValueError('"max_connections" must be a positive integer')
  893. self.connection_class = connection_class
  894. self.connection_kwargs = connection_kwargs
  895. self.max_connections = max_connections
  896. self.reset()
  897. def __repr__(self):
  898. return "%s<%s>" % (
  899. type(self).__name__,
  900. repr(self.connection_class(**self.connection_kwargs)),
  901. )
  902. def reset(self):
  903. self.pid = os.getpid()
  904. self._created_connections = 0
  905. self._available_connections = []
  906. self._in_use_connections = set()
  907. self._check_lock = threading.Lock()
  908. def _checkpid(self):
  909. if self.pid != os.getpid():
  910. with self._check_lock:
  911. if self.pid == os.getpid():
  912. # another thread already did the work while we waited
  913. # on the lock.
  914. return
  915. self.reset()
  916. def get_connection(self, command_name, *keys, **options):
  917. "Get a connection from the pool"
  918. self._checkpid()
  919. try:
  920. connection = self._available_connections.pop()
  921. except IndexError:
  922. connection = self.make_connection()
  923. self._in_use_connections.add(connection)
  924. try:
  925. # ensure this connection is connected to Redis
  926. connection.connect()
  927. # connections that the pool provides should be ready to send
  928. # a command. if not, the connection was either returned to the
  929. # pool before all data has been read or the socket has been
  930. # closed. either way, reconnect and verify everything is good.
  931. try:
  932. if connection.can_read():
  933. raise ConnectionError('Connection has data')
  934. except ConnectionError:
  935. connection.disconnect()
  936. connection.connect()
  937. if connection.can_read():
  938. raise ConnectionError('Connection not ready')
  939. except: # noqa: E722
  940. # release the connection back to the pool so that we don't leak it
  941. self.release(connection)
  942. raise
  943. return connection
  944. def get_encoder(self):
  945. "Return an encoder based on encoding settings"
  946. kwargs = self.connection_kwargs
  947. return Encoder(
  948. encoding=kwargs.get('encoding', 'utf-8'),
  949. encoding_errors=kwargs.get('encoding_errors', 'strict'),
  950. decode_responses=kwargs.get('decode_responses', False)
  951. )
  952. def make_connection(self):
  953. "Create a new connection"
  954. if self._created_connections >= self.max_connections:
  955. raise ConnectionError("Too many connections")
  956. self._created_connections += 1
  957. return self.connection_class(**self.connection_kwargs)
  958. def release(self, connection):
  959. "Releases the connection back to the pool"
  960. self._checkpid()
  961. if connection.pid != self.pid:
  962. return
  963. self._in_use_connections.remove(connection)
  964. self._available_connections.append(connection)
  965. def disconnect(self):
  966. "Disconnects all connections in the pool"
  967. self._checkpid()
  968. all_conns = chain(self._available_connections,
  969. self._in_use_connections)
  970. for connection in all_conns:
  971. connection.disconnect()
  972. class BlockingConnectionPool(ConnectionPool):
  973. """
  974. Thread-safe blocking connection pool::
  975. >>> from redis.client import Redis
  976. >>> client = Redis(connection_pool=BlockingConnectionPool())
  977. It performs the same function as the default
  978. ``:py:class: ~redis.connection.ConnectionPool`` implementation, in that,
  979. it maintains a pool of reusable connections that can be shared by
  980. multiple redis clients (safely across threads if required).
  981. The difference is that, in the event that a client tries to get a
  982. connection from the pool when all of connections are in use, rather than
  983. raising a ``:py:class: ~redis.exceptions.ConnectionError`` (as the default
  984. ``:py:class: ~redis.connection.ConnectionPool`` implementation does), it
  985. makes the client wait ("blocks") for a specified number of seconds until
  986. a connection becomes available.
  987. Use ``max_connections`` to increase / decrease the pool size::
  988. >>> pool = BlockingConnectionPool(max_connections=10)
  989. Use ``timeout`` to tell it either how many seconds to wait for a connection
  990. to become available, or to block forever:
  991. # Block forever.
  992. >>> pool = BlockingConnectionPool(timeout=None)
  993. # Raise a ``ConnectionError`` after five seconds if a connection is
  994. # not available.
  995. >>> pool = BlockingConnectionPool(timeout=5)
  996. """
  997. def __init__(self, max_connections=50, timeout=20,
  998. connection_class=Connection, queue_class=LifoQueue,
  999. **connection_kwargs):
  1000. self.queue_class = queue_class
  1001. self.timeout = timeout
  1002. super(BlockingConnectionPool, self).__init__(
  1003. connection_class=connection_class,
  1004. max_connections=max_connections,
  1005. **connection_kwargs)
  1006. def reset(self):
  1007. self.pid = os.getpid()
  1008. self._check_lock = threading.Lock()
  1009. # Create and fill up a thread safe queue with ``None`` values.
  1010. self.pool = self.queue_class(self.max_connections)
  1011. while True:
  1012. try:
  1013. self.pool.put_nowait(None)
  1014. except Full:
  1015. break
  1016. # Keep a list of actual connection instances so that we can
  1017. # disconnect them later.
  1018. self._connections = []
  1019. def make_connection(self):
  1020. "Make a fresh connection."
  1021. connection = self.connection_class(**self.connection_kwargs)
  1022. self._connections.append(connection)
  1023. return connection
  1024. def get_connection(self, command_name, *keys, **options):
  1025. """
  1026. Get a connection, blocking for ``self.timeout`` until a connection
  1027. is available from the pool.
  1028. If the connection returned is ``None`` then creates a new connection.
  1029. Because we use a last-in first-out queue, the existing connections
  1030. (having been returned to the pool after the initial ``None`` values
  1031. were added) will be returned before ``None`` values. This means we only
  1032. create new connections when we need to, i.e.: the actual number of
  1033. connections will only increase in response to demand.
  1034. """
  1035. # Make sure we haven't changed process.
  1036. self._checkpid()
  1037. # Try and get a connection from the pool. If one isn't available within
  1038. # self.timeout then raise a ``ConnectionError``.
  1039. connection = None
  1040. try:
  1041. connection = self.pool.get(block=True, timeout=self.timeout)
  1042. except Empty:
  1043. # Note that this is not caught by the redis client and will be
  1044. # raised unless handled by application code. If you want never to
  1045. raise ConnectionError("No connection available.")
  1046. # If the ``connection`` is actually ``None`` then that's a cue to make
  1047. # a new connection to add to the pool.
  1048. if connection is None:
  1049. connection = self.make_connection()
  1050. try:
  1051. # ensure this connection is connected to Redis
  1052. connection.connect()
  1053. # connections that the pool provides should be ready to send
  1054. # a command. if not, the connection was either returned to the
  1055. # pool before all data has been read or the socket has been
  1056. # closed. either way, reconnect and verify everything is good.
  1057. try:
  1058. if connection.can_read():
  1059. raise ConnectionError('Connection has data')
  1060. except ConnectionError:
  1061. connection.disconnect()
  1062. connection.connect()
  1063. if connection.can_read():
  1064. raise ConnectionError('Connection not ready')
  1065. except: # noqa: E722
  1066. # release the connection back to the pool so that we don't leak it
  1067. self.release(connection)
  1068. raise
  1069. return connection
  1070. def release(self, connection):
  1071. "Releases the connection back to the pool."
  1072. # Make sure we haven't changed process.
  1073. self._checkpid()
  1074. if connection.pid != self.pid:
  1075. return
  1076. # Put the connection back into the pool.
  1077. try:
  1078. self.pool.put_nowait(connection)
  1079. except Full:
  1080. # perhaps the pool has been reset() after a fork? regardless,
  1081. # we don't want this connection
  1082. pass
  1083. def disconnect(self):
  1084. "Disconnects all connections in the pool."
  1085. self._checkpid()
  1086. for connection in self._connections:
  1087. connection.disconnect()