feat(bot): 🎉 Create a redis client from scratch.
This commit is contained in:
Nyeki.py
@@ -0,0 +1,826 @@
|
||||
from typing import TypeAlias, Awaitable, AsyncGenerator
|
||||
|
||||
from urllib import parse
|
||||
|
||||
from datetime import datetime, timedelta
|
||||
|
||||
import asyncio
|
||||
|
||||
import hiredis
|
||||
|
||||
|
||||
Resp3Repr: TypeAlias = str | int | bytes | bytearray | float | bool | None
|
||||
|
||||
|
||||
class _Ellipsis:
|
||||
"""Represents an incomplete response of Redis."""
|
||||
|
||||
|
||||
def _encode(
|
||||
arg: (
|
||||
Resp3Repr
|
||||
| list[Resp3Repr]
|
||||
| tuple[Resp3Repr]
|
||||
| dict[Resp3Repr, Resp3Repr]
|
||||
| set[Resp3Repr]
|
||||
)
|
||||
) -> bytes:
|
||||
"""Encode a value into its RESP3 representation.
|
||||
|
||||
This function supports str (converted to bytes), int, bytes, bytearray, float, bool, and None
|
||||
as basic types. It also supports lists (arrays), dicts (maps), and sets.
|
||||
|
||||
Args:
|
||||
arg (Resp3Repr | list[Resp3Repr] | tuple[Resp3Repr] | dict[Resp3Repr, Resp3Repr] | set[Resp3Repr]):
|
||||
The argument to encode.
|
||||
|
||||
Raises:
|
||||
ValueError: The value provided is invalid or not encodable.
|
||||
|
||||
Returns:
|
||||
bytes: The encoded representation of the value.
|
||||
"""
|
||||
|
||||
if isinstance(arg, (str, int)):
|
||||
arg = str(arg).encode()
|
||||
|
||||
if isinstance(arg, (bytes, bytearray)):
|
||||
return f"${len(arg)}\r\n".encode() + bytes(arg) + b"\r\n"
|
||||
|
||||
elif isinstance(arg, int):
|
||||
return f":{arg}\r\n".encode()
|
||||
|
||||
elif isinstance(arg, type(None)):
|
||||
return b"_\r\n"
|
||||
|
||||
elif isinstance(arg, float):
|
||||
result = f",{arg}\r\n".encode()
|
||||
|
||||
elif isinstance(arg, bool):
|
||||
if arg:
|
||||
return b"#t\r\n"
|
||||
|
||||
return b"#f\r\n"
|
||||
|
||||
elif isinstance(arg, (list, tuple)):
|
||||
return f"*{len(arg)}\r\n".encode() + b"".join([_encode(i) for i in arg])
|
||||
|
||||
elif isinstance(arg, dict):
|
||||
result = f"%{len(arg.keys())}\r\n".encode()
|
||||
|
||||
for key, value in arg.items():
|
||||
result += _encode(key) + _encode(value)
|
||||
|
||||
return result
|
||||
|
||||
elif isinstance(arg, set):
|
||||
result = f"~{len(arg)}\r\n".encode()
|
||||
|
||||
for item in arg:
|
||||
result += _encode(item)
|
||||
|
||||
return result
|
||||
|
||||
raise ValueError(f"Value of type {type(arg)} is not encodable into RESP3.")
|
||||
|
||||
|
||||
def _get_secs(time: int | timedelta) -> int:
|
||||
"""Converts an integer or timedelta into an amount of seconds.
|
||||
|
||||
Args:
|
||||
time (int | timedelta): The number or timedelta to convert.
|
||||
|
||||
Returns:
|
||||
int: The amount of seconds.
|
||||
"""
|
||||
|
||||
if isinstance(time, timedelta):
|
||||
return int(time.total_seconds())
|
||||
|
||||
return time
|
||||
|
||||
|
||||
def _get_millis(time: int | timedelta) -> int:
|
||||
"""Converts an integer or timedelta into an amount of milliseconds.
|
||||
|
||||
Args:
|
||||
time (int | timedelta): The number or timedelta to convert.
|
||||
|
||||
Returns:
|
||||
int: The amount of milliseconds.
|
||||
"""
|
||||
|
||||
if isinstance(time, timedelta):
|
||||
return int(time.total_seconds() * 1000)
|
||||
|
||||
return time
|
||||
|
||||
|
||||
def _get_timestamp_secs(dtime: datetime | int) -> int:
|
||||
"""Converts a datetime or timestamp into it's seconds timestamp.
|
||||
|
||||
Args:
|
||||
dtime (datetime | int): The datetime object to convert or an integer value which will be
|
||||
returned as-is.
|
||||
|
||||
Returns:
|
||||
int: The seconds timestamp.
|
||||
"""
|
||||
|
||||
if isinstance(dtime, datetime):
|
||||
return int(dtime.timestamp())
|
||||
|
||||
return dtime
|
||||
|
||||
|
||||
def _get_timestamp_millis(dtime: datetime | int) -> int:
|
||||
"""Converts a datetime or timestamp into it's milliseconds timestamp.
|
||||
|
||||
Args:
|
||||
dtime (datetime | int): The datetime object to convert or an integer value which will be
|
||||
returned as-is.
|
||||
|
||||
Returns:
|
||||
int: The milliseconds timestamp.
|
||||
"""
|
||||
|
||||
if isinstance(dtime, datetime):
|
||||
return int(dtime.timestamp() * 1000)
|
||||
|
||||
return dtime
|
||||
|
||||
|
||||
class RedisError(Exception):
|
||||
"""Raised whenever an error is raised from this library."""
|
||||
|
||||
|
||||
class FatalRedisError(RedisError):
|
||||
"""Raised whenever an error occurs that does not allow the connection to stay alive."""
|
||||
|
||||
|
||||
class Runnable:
|
||||
"""The base class for all Redis commands.
|
||||
|
||||
This class is inherited by Connection, Pipeline, and Transaction to implement their own
|
||||
functionality.
|
||||
"""
|
||||
|
||||
_client: "Redis"
|
||||
|
||||
def __init__(self, client: "Redis") -> None:
|
||||
"""Initializes the runner with the client's configuration.
|
||||
|
||||
Args:
|
||||
client (Redis): The client that initialized the runner.
|
||||
"""
|
||||
|
||||
self._client = client
|
||||
|
||||
def command(self, *args: tuple[Resp3Repr]) -> Resp3Repr | Awaitable[Resp3Repr]:
|
||||
"""Runs a command taking it as a list of arguments.
|
||||
|
||||
For example:
|
||||
```
|
||||
await Runnable.command("SET", "key", "value")
|
||||
```
|
||||
|
||||
This method is not implemented by default, so you need a Connection, a Pipeline, or a
|
||||
Transaction to use it.
|
||||
|
||||
Returns:
|
||||
Resp3Repr | Awaitable[Resp3Repr]: The value(s) returned by the execution of the
|
||||
command.
|
||||
"""
|
||||
|
||||
raise NotImplementedError(
|
||||
"This method needs to be implemented by a subclass of Runner."
|
||||
)
|
||||
|
||||
def ping(self, arg: str = "PONG") -> Resp3Repr:
|
||||
"""Returns PONG if no argument is provided, otherwise return a copy of the argument as a
|
||||
bulk.
|
||||
|
||||
Args:
|
||||
arg (str): The argument to return by the command.
|
||||
|
||||
Returns:
|
||||
Resp3Repr: The response of the command.
|
||||
"""
|
||||
|
||||
return self.command("PING", str(arg))
|
||||
|
||||
|
||||
class RunnableGenericMixin:
|
||||
"""Mixin containing all Redis generic commands."""
|
||||
|
||||
def delete(self, *keys: tuple[str]) -> None:
|
||||
"""Deletes one or more keys."""
|
||||
|
||||
return self.command("DEL", *[str(key) for key in keys])
|
||||
|
||||
def exists(self, *keys: tuple[str]) -> int:
|
||||
"""Returns if key exists.
|
||||
|
||||
The user should be aware that if the same existing key is mentioned in the arguments
|
||||
multiple times, it will be counted multiple times. So if somekey exists, `EXISTS somekey
|
||||
somekey` will return 2.
|
||||
|
||||
Args:
|
||||
*keys (tuple[str]): The keys to check for existance.
|
||||
|
||||
Returns:
|
||||
int: The amount of keys that exist.
|
||||
"""
|
||||
|
||||
return self.command("EXISTS", *[str(key) for key in keys])
|
||||
|
||||
|
||||
class RunnableHashMixin:
|
||||
"""Mixin containing all Redis hash commands."""
|
||||
|
||||
def hget(self, key: str, field: str) -> str | None:
|
||||
"""Returns the value associated with field in the hash stored at key.
|
||||
|
||||
Args:
|
||||
key (str): The key of the hash.
|
||||
field (str): The field of the hash to get.
|
||||
|
||||
Returns:
|
||||
str | None: The value of the field, or None if it does not exist.
|
||||
"""
|
||||
|
||||
return self.command("HGET", str(key), str(field))
|
||||
|
||||
def hgetall(self, key: str) -> dict[Resp3Repr, Resp3Repr]:
|
||||
"""Returns all fields and values of the hash stored at key.
|
||||
|
||||
Args:
|
||||
key (str): The key to get the fields and values from.
|
||||
|
||||
Returns:
|
||||
dict[Resp3Repr, Resp3Repr]: A dict containing all the fields mapped to their respective
|
||||
values.
|
||||
"""
|
||||
|
||||
return self.command("HGETALL", str(key))
|
||||
|
||||
def hincrby(self, key: str, field: str, increment: int) -> int:
|
||||
"""Increments the number stored at field in the hash stored at key by increment.
|
||||
|
||||
If key does not exist, a new key holding a hash is created. If field does not exist the
|
||||
value is set to 0 before the operation is performed. The range of values supported by
|
||||
HINCRBY is limited to 64 bit signed integers.
|
||||
|
||||
Args:
|
||||
key (str): The key of the hash whose field to increment.
|
||||
field (str): The field to increment.
|
||||
increment (int): The amount to increment by.
|
||||
|
||||
Returns:
|
||||
int: The amount returned
|
||||
"""
|
||||
|
||||
return self.command("HINCRBY", str(key), str(field), str(increment))
|
||||
|
||||
def hset(self, key: str, mapping: dict[Resp3Repr, Resp3Repr]) -> int:
|
||||
"""Sets the specified fields to their respective values in the hash stored at key.
|
||||
|
||||
This command overwrites the values of specified fields that exist in the hash. If key
|
||||
doesn't exist, a new key holding a hash is created.
|
||||
|
||||
Args:
|
||||
key (str): The key to set the values for.
|
||||
mapping (dict[Resp3Repr, Resp3Repr]): The fields mapped to their values.
|
||||
|
||||
Returns:
|
||||
int: The number of fields that were added.
|
||||
"""
|
||||
|
||||
items = []
|
||||
|
||||
for k, v in mapping.items():
|
||||
items.append(k)
|
||||
items.append(v)
|
||||
|
||||
return self.command("HSET", str(key), *items)
|
||||
|
||||
|
||||
class RunnableSetMixin:
|
||||
"""Mixin containing all Redis set commands."""
|
||||
|
||||
def sadd(self, key: str, *values: tuple[str]) -> int:
|
||||
"""Adds the specified members to the set stored at key.
|
||||
|
||||
Specified members that are already a member of this set are ignored. If key does not exist,
|
||||
a new set is created before adding the specified members. An error is raised when the value
|
||||
stored at key is not a set.
|
||||
|
||||
Args:
|
||||
key (str): The key to add the values to.
|
||||
*values (tuple[str]): The values to add to the set.
|
||||
|
||||
Returns:
|
||||
int: The amount of new values that were added.
|
||||
"""
|
||||
|
||||
return self.command("SADD", str(key), *values)
|
||||
|
||||
def scard(self, key: str) -> int:
|
||||
"""Returns the set cardinality (number of elements) of the set stored at key.
|
||||
|
||||
Args:
|
||||
key (str): The key of the set to get the cardinality from.
|
||||
|
||||
Returns:
|
||||
int: The number of elements, or 0 if it does not exist.
|
||||
"""
|
||||
|
||||
return self.command("SCARD", str(key))
|
||||
|
||||
def smembers(self, key: str) -> set[str]:
|
||||
"""Returns all the members of the set value stored at key.
|
||||
|
||||
Args:
|
||||
key (str): The key of the set to get the members from.
|
||||
|
||||
Returns:
|
||||
set[str]: The members of the set.
|
||||
"""
|
||||
|
||||
return self.command("SMEMBERS", str(key))
|
||||
|
||||
def sismember(self, key: str, value: str) -> bool:
|
||||
"""Returns whether the value is a member of the set.
|
||||
|
||||
Args:
|
||||
key (str): The key of the set.
|
||||
value (str): The value to check whether it is a member of the set.
|
||||
|
||||
Returns:
|
||||
bool: Whether the value is a member of the set.
|
||||
"""
|
||||
|
||||
return self.command("SISMEMBER", str(key), str(value))
|
||||
|
||||
def spop(self, key: str, count: int = 1) -> str | list[str]:
|
||||
"""Removes and returns one or more random members from the set value store at key."""
|
||||
|
||||
return self.command("SPOP", str(key), count)
|
||||
|
||||
def srem(self, key: str, *values: tuple[str]) -> int:
|
||||
"""Removes the specified members from the set stored at key.
|
||||
|
||||
Specified members that are not a member of this set are ignored. If key does not exist,
|
||||
it is treated as an empty set and this command returns 0. An error is returned when the
|
||||
value stored at key is not a set.
|
||||
|
||||
Args:
|
||||
key (str): The key of the set.
|
||||
*values (tuple[str]): The values to remove from the set.
|
||||
|
||||
Returns:
|
||||
int: The amount of values that were removed from the set.
|
||||
"""
|
||||
|
||||
return self.command("SREM", str(key), *values)
|
||||
|
||||
|
||||
class RunnableStringMixin:
|
||||
"""Mixin containing all Redis string commands."""
|
||||
|
||||
def append(self, key: str, value: str) -> int:
|
||||
"""Appends a string to an existing key, or creates a new one if it does not exist.
|
||||
|
||||
If key already exists and is a string, this command appends the value at the end of the
|
||||
string. If key does not exist it is created and set as an empty string, so APPEND will be
|
||||
similar to SET in this special case.
|
||||
|
||||
Args:
|
||||
key (str): The key to append to.
|
||||
value (str): The string to append.
|
||||
|
||||
Returns:
|
||||
int: The new length of the value.
|
||||
"""
|
||||
|
||||
return self.command("APPEND", str(key), str(value))
|
||||
|
||||
def decr(self, key: str) -> int:
|
||||
"""Decrements the number stored at key by one.
|
||||
|
||||
If the key does not exist, it is set to 0 before performing the operation. An error is
|
||||
returned if the key contains a value of the wrong type or contains a string that can not be
|
||||
represented as integer. This operation is limited to 64 bit signed integers.
|
||||
|
||||
Args:
|
||||
key (str): The key to decrement.
|
||||
|
||||
Returns:
|
||||
int: The decremented value.
|
||||
"""
|
||||
|
||||
return self.command("DECR", str(key))
|
||||
|
||||
def decrby(self, key: str, decrement: int) -> int:
|
||||
"""Reduces the value stored at the specified key by the specified decrement.
|
||||
|
||||
If the key does not exist, it is initialized with a value of 0 before performing the
|
||||
operation. If the key's value is not of the correct type or cannot be represented as an
|
||||
integer, an error is returned. This operation is limited to 64-bit signed integers.
|
||||
|
||||
Args:
|
||||
key (str): The key to decrement.
|
||||
decrement (int): The amount to decrement by.
|
||||
|
||||
Returns:
|
||||
int: The decremented value.
|
||||
"""
|
||||
|
||||
return self.command("DECRBY", str(key), decrement)
|
||||
|
||||
def get(self, key: str) -> str | None:
|
||||
"""Returns the string value of a key.
|
||||
|
||||
If the key does not exist None is returned. An error is returned if the value stored at key
|
||||
is not a string, because GET only handles string values.
|
||||
|
||||
Args:
|
||||
key (str): The key to get the value from.
|
||||
|
||||
Returns:
|
||||
str | None: The value of the key, of None if it does not exist.
|
||||
"""
|
||||
|
||||
return self.command("GET", str(key))
|
||||
|
||||
def getdel(self, key: str) -> str | None:
|
||||
"""Returns the value of key and delete the key.
|
||||
|
||||
This command is similar to GET, except for the fact that it also deletes the key on success
|
||||
(if and only if the key's value type is a string).
|
||||
|
||||
Args:
|
||||
key (str): The key to get and delete.
|
||||
|
||||
Returns:
|
||||
str | None: The value of the key.
|
||||
"""
|
||||
|
||||
return self.command("GETDEL", str(key))
|
||||
|
||||
def set(
|
||||
self,
|
||||
key: str,
|
||||
value: str,
|
||||
nx: bool = False,
|
||||
xx: bool = False,
|
||||
get: bool = False,
|
||||
ex: timedelta | int | None = None,
|
||||
px: timedelta | int | None = None,
|
||||
exat: datetime | int | None = None,
|
||||
pxat: datetime | int | None = None,
|
||||
keepttl: bool = False,
|
||||
) -> str | None:
|
||||
"""Set key to hold the string value.
|
||||
|
||||
If key already holds a value, it is overwritten, regardless of its type. Any previous time
|
||||
to live associated with the key is discarded on successful SET operation.
|
||||
|
||||
Args:
|
||||
key (str): The key to set.
|
||||
value (str): The value to set to the specified key.
|
||||
nx (bool, optional): Only set the key if it does not already exist. Defaults to False.
|
||||
xx (bool, optional): Only set the key if it already exists. Defaults to False.
|
||||
get (bool, optional): Return the old string stored at key, or None if key did not
|
||||
exist. An error is returned and SET aborted if the value stored at key is not a
|
||||
string. Defaults to False.
|
||||
ex (timedelta | int | None, optional): Set the specified expire time, in seconds (a positive
|
||||
integer). Defaults to None.
|
||||
px (timedelta | int | None, optional): Set the specified expire time, in milliseconds (a positive
|
||||
integer). Defaults to None.
|
||||
exat (datetime | int | None, optional): Set the specified time at which the key will
|
||||
expire, in seconds (a positive integer) or a datetime object. Defaults to None.
|
||||
pxat (datetime | int | None, optional): Set the specified time at which the key will
|
||||
expire, in milliseconds (a positive integer) or a datetime object. Defaults to None.
|
||||
keepttl (bool, optional): Retain the time to live associated with the key. Defaults to
|
||||
False.
|
||||
|
||||
Returns:
|
||||
str | None: The old key in case `get` was set, or None.
|
||||
"""
|
||||
|
||||
if nx and xx:
|
||||
raise ValueError("Only one of XX or NX can be set, but not both.")
|
||||
|
||||
if len([i for i in (ex, px, exat, pxat, keepttl) if i]) > 1:
|
||||
raise ValueError(
|
||||
"Only one of EX, PX, EXAT, PXAT, or KEEPTTL can be set, but not multiple of them."
|
||||
)
|
||||
|
||||
command = ["SET", str(key), str(value)]
|
||||
|
||||
if nx:
|
||||
command += ["NX"]
|
||||
|
||||
elif xx:
|
||||
command += ["XX"]
|
||||
|
||||
if get:
|
||||
command += ["GET"]
|
||||
|
||||
if ex:
|
||||
command += ["EX", str(_get_secs(ex))]
|
||||
|
||||
elif px:
|
||||
command += ["PX", str(_get_millis(px))]
|
||||
|
||||
elif exat:
|
||||
command += ["EXAT", str(_get_timestamp_secs(exat))]
|
||||
|
||||
elif pxat:
|
||||
command += ["PXAT", str(_get_timestamp_millis(pxat))]
|
||||
|
||||
elif keepttl:
|
||||
command += ["KEEPTTL"]
|
||||
|
||||
return self.command(*command)
|
||||
|
||||
|
||||
class FullRunnable(
|
||||
Runnable,
|
||||
RunnableGenericMixin,
|
||||
RunnableHashMixin,
|
||||
RunnableSetMixin,
|
||||
RunnableStringMixin,
|
||||
):
|
||||
"""A runnable containing all the command mixins."""
|
||||
|
||||
|
||||
class Connection(FullRunnable):
|
||||
_reader: asyncio.StreamReader
|
||||
_writer: asyncio.StreamWriter
|
||||
_parser: hiredis.Reader
|
||||
_lock: asyncio.Lock
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
client: "Redis",
|
||||
reader: asyncio.StreamReader,
|
||||
writer: asyncio.StreamWriter,
|
||||
parser: hiredis.Reader = hiredis.Reader(notEnoughData=_Ellipsis),
|
||||
) -> None:
|
||||
"""Initializes the connection object.
|
||||
|
||||
Args:
|
||||
client (Redis): The client that initialized the connection.
|
||||
reader (asyncio.StreamReader): The socket stream reader.
|
||||
writer (asyncio.StreamWriter): The socket stream writer.
|
||||
parser (hiredis.Reader, optional): The RESP3 parser. Defaults to a basic implementation
|
||||
if not set.
|
||||
"""
|
||||
|
||||
self._client = client
|
||||
self._reader = reader
|
||||
self._writer = writer
|
||||
self._parser = parser
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
super().__init__(client)
|
||||
|
||||
def release(self) -> None:
|
||||
"""Returns the connection back to the pool."""
|
||||
|
||||
self._client._pool.append(self)
|
||||
|
||||
def destroy(self) -> None:
|
||||
"""Closes the connection and frees it's space in the client's pool."""
|
||||
|
||||
self._client._open_connections -= 1
|
||||
|
||||
self._writer.close()
|
||||
|
||||
async def __aenter__(self) -> "Connection":
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *args) -> None:
|
||||
self.release()
|
||||
|
||||
async def command(self, *args: list[Resp3Repr]) -> Resp3Repr:
|
||||
"""Executes a command and returns it's response.
|
||||
|
||||
Raises:
|
||||
hiredis.ReplyError: Any errors returned by Redis.
|
||||
"""
|
||||
|
||||
args = list(args)
|
||||
|
||||
command = _encode(args)
|
||||
|
||||
async def run():
|
||||
async with self._lock:
|
||||
self._writer.write(command)
|
||||
await self._writer.drain()
|
||||
|
||||
while True:
|
||||
self._parser.feed(await self._reader.readuntil(b"\r\n"))
|
||||
|
||||
result = self._parser.gets()
|
||||
|
||||
if result == _Ellipsis:
|
||||
continue
|
||||
|
||||
elif isinstance(result, hiredis.ReplyError):
|
||||
raise result
|
||||
|
||||
return result
|
||||
|
||||
try:
|
||||
return await asyncio.wait_for(run(), self._client._command_timeout)
|
||||
|
||||
except (
|
||||
hiredis.ProtocolError,
|
||||
asyncio.TimeoutError,
|
||||
ConnectionResetError,
|
||||
) as exc:
|
||||
self.destroy()
|
||||
|
||||
raise FatalRedisError() from exc
|
||||
|
||||
async def hello(self, username: str = None, password: str = None) -> dict:
|
||||
"""Sends the initial HELLO command to the server.
|
||||
|
||||
Args:
|
||||
username (str, optional): The username to authenticate with. Defaults to None, or
|
||||
"default" if `password` is set.
|
||||
password (str, optional): The password to authenticate with. Defaults to None.
|
||||
|
||||
Returns:
|
||||
dict: The metadata map returned by the server.
|
||||
"""
|
||||
|
||||
command = ["HELLO", "3"]
|
||||
|
||||
if password is not None:
|
||||
command += ["AUTH", username or "default", password]
|
||||
|
||||
return await self.command(*command)
|
||||
|
||||
|
||||
class Pipeline(FullRunnable):
|
||||
"""A Redis pipeline."""
|
||||
|
||||
_commands: list[tuple[Resp3Repr]]
|
||||
|
||||
def __init__(self, client: "Redis") -> None:
|
||||
self._commands = []
|
||||
|
||||
super().__init__(client)
|
||||
|
||||
def command(self, *args: tuple[Resp3Repr]) -> None:
|
||||
self._commands.append(args)
|
||||
|
||||
async def execute(self):
|
||||
"""Executes the pipeline.
|
||||
|
||||
Returns:
|
||||
tuple[Resp3Repr]: The results returned by the execution of the commands.
|
||||
"""
|
||||
|
||||
command = b"".join([_encode(comm) for comm in self._commands])
|
||||
|
||||
async with await self._client.connection() as conn:
|
||||
|
||||
async def run():
|
||||
async with conn._lock:
|
||||
conn._writer.write(command)
|
||||
await conn._writer.drain()
|
||||
|
||||
results = []
|
||||
|
||||
while True:
|
||||
conn._parser.feed(await conn._reader.readuntil(b"\r\n"))
|
||||
|
||||
result = conn._parser.gets()
|
||||
|
||||
if result == _Ellipsis:
|
||||
continue
|
||||
|
||||
elif isinstance(result, hiredis.ReplyError):
|
||||
raise RedisError(
|
||||
"An error was returned from the Redis server."
|
||||
) from result
|
||||
|
||||
results.append(result)
|
||||
|
||||
if len(results) == len(self._commands):
|
||||
return results
|
||||
|
||||
try:
|
||||
return tuple(
|
||||
await asyncio.wait_for(run(), self._client._command_timeout)
|
||||
)
|
||||
|
||||
except (
|
||||
hiredis.ProtocolError,
|
||||
asyncio.TimeoutError,
|
||||
ConnectionResetError,
|
||||
) as exc:
|
||||
conn.destroy()
|
||||
|
||||
raise FatalRedisError() from exc
|
||||
|
||||
|
||||
class Redis(FullRunnable):
|
||||
_pool: list[Connection] = []
|
||||
_open_connections: int = 0
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
url: str = "",
|
||||
max_connections: int = 100,
|
||||
connection_timeout: int = 3,
|
||||
command_timeout: int = 3,
|
||||
encoding: str = "utf-8",
|
||||
) -> None:
|
||||
"""Initializes the client class.
|
||||
|
||||
Args:
|
||||
url (str): The redis URL to connect to. Defaults to "redis://localhost:6379"
|
||||
max_connections (int, optional): The maximum amount of connections that can be opened.
|
||||
Defaults to 100.
|
||||
connection_timeout (int, optional): The amount of seconds to wait for a connection to
|
||||
establish before raising a TimeoutError. Defaults to 3.
|
||||
command_timeout (int, optional): The amount of seconds to wait for a command to execute
|
||||
before raising a TimeoutError. Defaults to 3.
|
||||
encoding (str, optional): The encoding with which to decode Redis response values.
|
||||
"""
|
||||
|
||||
self._client = self
|
||||
|
||||
self._lock = asyncio.Lock()
|
||||
self._max_connections = max_connections
|
||||
self._connection_timeout = connection_timeout
|
||||
self._command_timeout = command_timeout
|
||||
self._encoding = encoding
|
||||
|
||||
parsed_url = parse.urlparse(url)
|
||||
|
||||
self._host = parsed_url.hostname or "localhost"
|
||||
self._port = parsed_url.port or 6379
|
||||
self._username = parsed_url.username
|
||||
self._password = parsed_url.password
|
||||
|
||||
async def connection(self) -> Connection:
|
||||
"""Returns an available connection or creates a new one if none are available and the pool
|
||||
has space remaining.
|
||||
|
||||
Raises:
|
||||
ConnectionError: If the pool has reached it's max size and no connections are available.
|
||||
"""
|
||||
|
||||
async with self._lock:
|
||||
try:
|
||||
connection = self._pool.pop()
|
||||
return connection
|
||||
|
||||
except IndexError as exc:
|
||||
if self._open_connections >= self._max_connections:
|
||||
raise ConnectionError(
|
||||
"The pool has reached it's limit of connections and no connections are "
|
||||
"available at the moment."
|
||||
) from exc
|
||||
|
||||
reader, writer = await asyncio.wait_for(
|
||||
asyncio.open_connection(host=self._host, port=self._port),
|
||||
self._connection_timeout,
|
||||
)
|
||||
|
||||
connection = Connection(
|
||||
client=self,
|
||||
reader=reader,
|
||||
writer=writer,
|
||||
parser=hiredis.Reader(
|
||||
notEnoughData=_Ellipsis, encoding=self._encoding
|
||||
),
|
||||
)
|
||||
|
||||
await connection.hello(username=self._username, password=self._password)
|
||||
|
||||
self._open_connections += 1
|
||||
|
||||
return connection
|
||||
|
||||
async def command(self, *args: tuple[Resp3Repr]) -> Awaitable[Resp3Repr]:
|
||||
"""Executes a command with a free connection from the pool.
|
||||
|
||||
Returns:
|
||||
Awaitable[Resp3Repr]: The command's response.
|
||||
"""
|
||||
|
||||
async with await self.connection() as conn:
|
||||
return await conn.command(*args)
|
||||
|
||||
def pipeline(self) -> Pipeline:
|
||||
"""Creates a new pipeline."""
|
||||
|
||||
return Pipeline(self)
|
||||
Reference in New Issue
Block a user