hiredis.py 8.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230
  1. import asyncio
  2. import socket
  3. import sys
  4. from typing import Callable, List, Optional, Union
  5. if sys.version_info.major >= 3 and sys.version_info.minor >= 11:
  6. from asyncio import timeout as async_timeout
  7. else:
  8. from async_timeout import timeout as async_timeout
  9. from redis.compat import TypedDict
  10. from ..exceptions import ConnectionError, InvalidResponse, RedisError
  11. from ..typing import EncodableT
  12. from ..utils import HIREDIS_AVAILABLE
  13. from .base import AsyncBaseParser, BaseParser
  14. from .socket import (
  15. NONBLOCKING_EXCEPTION_ERROR_NUMBERS,
  16. NONBLOCKING_EXCEPTIONS,
  17. SENTINEL,
  18. SERVER_CLOSED_CONNECTION_ERROR,
  19. )
  20. # Used to signal that hiredis-py does not have enough data to parse.
  21. # Using `False` or `None` is not reliable, given that the parser can
  22. # return `False` or `None` for legitimate reasons from RESP payloads.
  23. NOT_ENOUGH_DATA = object()
  24. class _HiredisReaderArgs(TypedDict, total=False):
  25. protocolError: Callable[[str], Exception]
  26. replyError: Callable[[str], Exception]
  27. encoding: Optional[str]
  28. errors: Optional[str]
  29. class _HiredisParser(BaseParser):
  30. "Parser class for connections using Hiredis"
  31. def __init__(self, socket_read_size):
  32. if not HIREDIS_AVAILABLE:
  33. raise RedisError("Hiredis is not installed")
  34. self.socket_read_size = socket_read_size
  35. self._buffer = bytearray(socket_read_size)
  36. def __del__(self):
  37. try:
  38. self.on_disconnect()
  39. except Exception:
  40. pass
  41. def on_connect(self, connection, **kwargs):
  42. import hiredis
  43. self._sock = connection._sock
  44. self._socket_timeout = connection.socket_timeout
  45. kwargs = {
  46. "protocolError": InvalidResponse,
  47. "replyError": self.parse_error,
  48. "errors": connection.encoder.encoding_errors,
  49. "notEnoughData": NOT_ENOUGH_DATA,
  50. }
  51. if connection.encoder.decode_responses:
  52. kwargs["encoding"] = connection.encoder.encoding
  53. self._reader = hiredis.Reader(**kwargs)
  54. self._next_response = NOT_ENOUGH_DATA
  55. def on_disconnect(self):
  56. self._sock = None
  57. self._reader = None
  58. self._next_response = NOT_ENOUGH_DATA
  59. def can_read(self, timeout):
  60. if not self._reader:
  61. raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR)
  62. if self._next_response is NOT_ENOUGH_DATA:
  63. self._next_response = self._reader.gets()
  64. if self._next_response is NOT_ENOUGH_DATA:
  65. return self.read_from_socket(timeout=timeout, raise_on_timeout=False)
  66. return True
  67. def read_from_socket(self, timeout=SENTINEL, raise_on_timeout=True):
  68. sock = self._sock
  69. custom_timeout = timeout is not SENTINEL
  70. try:
  71. if custom_timeout:
  72. sock.settimeout(timeout)
  73. bufflen = self._sock.recv_into(self._buffer)
  74. if bufflen == 0:
  75. raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR)
  76. self._reader.feed(self._buffer, 0, bufflen)
  77. # data was read from the socket and added to the buffer.
  78. # return True to indicate that data was read.
  79. return True
  80. except socket.timeout:
  81. if raise_on_timeout:
  82. raise TimeoutError("Timeout reading from socket")
  83. return False
  84. except NONBLOCKING_EXCEPTIONS as ex:
  85. # if we're in nonblocking mode and the recv raises a
  86. # blocking error, simply return False indicating that
  87. # there's no data to be read. otherwise raise the
  88. # original exception.
  89. allowed = NONBLOCKING_EXCEPTION_ERROR_NUMBERS.get(ex.__class__, -1)
  90. if not raise_on_timeout and ex.errno == allowed:
  91. return False
  92. raise ConnectionError(f"Error while reading from socket: {ex.args}")
  93. finally:
  94. if custom_timeout:
  95. sock.settimeout(self._socket_timeout)
  96. def read_response(self, disable_decoding=False):
  97. if not self._reader:
  98. raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR)
  99. # _next_response might be cached from a can_read() call
  100. if self._next_response is not NOT_ENOUGH_DATA:
  101. response = self._next_response
  102. self._next_response = NOT_ENOUGH_DATA
  103. return response
  104. if disable_decoding:
  105. response = self._reader.gets(False)
  106. else:
  107. response = self._reader.gets()
  108. while response is NOT_ENOUGH_DATA:
  109. self.read_from_socket()
  110. if disable_decoding:
  111. response = self._reader.gets(False)
  112. else:
  113. response = self._reader.gets()
  114. # if the response is a ConnectionError or the response is a list and
  115. # the first item is a ConnectionError, raise it as something bad
  116. # happened
  117. if isinstance(response, ConnectionError):
  118. raise response
  119. elif (
  120. isinstance(response, list)
  121. and response
  122. and isinstance(response[0], ConnectionError)
  123. ):
  124. raise response[0]
  125. return response
  126. class _AsyncHiredisParser(AsyncBaseParser):
  127. """Async implementation of parser class for connections using Hiredis"""
  128. __slots__ = ("_reader",)
  129. def __init__(self, socket_read_size: int):
  130. if not HIREDIS_AVAILABLE:
  131. raise RedisError("Hiredis is not available.")
  132. super().__init__(socket_read_size=socket_read_size)
  133. self._reader = None
  134. def on_connect(self, connection):
  135. import hiredis
  136. self._stream = connection._reader
  137. kwargs: _HiredisReaderArgs = {
  138. "protocolError": InvalidResponse,
  139. "replyError": self.parse_error,
  140. "notEnoughData": NOT_ENOUGH_DATA,
  141. }
  142. if connection.encoder.decode_responses:
  143. kwargs["encoding"] = connection.encoder.encoding
  144. kwargs["errors"] = connection.encoder.encoding_errors
  145. self._reader = hiredis.Reader(**kwargs)
  146. self._connected = True
  147. def on_disconnect(self):
  148. self._connected = False
  149. async def can_read_destructive(self):
  150. if not self._connected:
  151. raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR)
  152. if self._reader.gets() is not NOT_ENOUGH_DATA:
  153. return True
  154. try:
  155. async with async_timeout(0):
  156. return await self.read_from_socket()
  157. except asyncio.TimeoutError:
  158. return False
  159. async def read_from_socket(self):
  160. buffer = await self._stream.read(self._read_size)
  161. if not buffer or not isinstance(buffer, bytes):
  162. raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR) from None
  163. self._reader.feed(buffer)
  164. # data was read from the socket and added to the buffer.
  165. # return True to indicate that data was read.
  166. return True
  167. async def read_response(
  168. self, disable_decoding: bool = False
  169. ) -> Union[EncodableT, List[EncodableT]]:
  170. # If `on_disconnect()` has been called, prohibit any more reads
  171. # even if they could happen because data might be present.
  172. # We still allow reads in progress to finish
  173. if not self._connected:
  174. raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR) from None
  175. if disable_decoding:
  176. response = self._reader.gets(False)
  177. else:
  178. response = self._reader.gets()
  179. while response is NOT_ENOUGH_DATA:
  180. await self.read_from_socket()
  181. if disable_decoding:
  182. response = self._reader.gets(False)
  183. else:
  184. response = self._reader.gets()
  185. # if the response is a ConnectionError or the response is a list and
  186. # the first item is a ConnectionError, raise it as something bad
  187. # happened
  188. if isinstance(response, ConnectionError):
  189. raise response
  190. elif (
  191. isinstance(response, list)
  192. and response
  193. and isinstance(response[0], ConnectionError)
  194. ):
  195. raise response[0]
  196. return response