Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix: close RedisStore client in case the with_client classmethod is used #3111

Merged
merged 32 commits into from
Feb 27, 2024
Merged
Show file tree
Hide file tree
Changes from 30 commits
Commits
Show all changes
32 commits
Select commit Hold shift + click to select a range
1630f5b
Added store lifetime note
euri10 Feb 13, 2024
c800ff7
Added failing test from issue, will write a "simpler" one after hopef…
euri10 Feb 13, 2024
0523767
Make the RedisStore a context manager and see
euri10 Feb 14, 2024
8650cb6
Implement those good suggstions
euri10 Feb 14, 2024
6f56678
Better naming for the boolean that decides if we should handle stuff …
euri10 Feb 14, 2024
bfe8262
Removed the classmethod arg
euri10 Feb 14, 2024
3cb8b3f
Make the async context manager not abstract
euri10 Feb 14, 2024
3da05e5
No need to implement __aenter__ anymore
euri10 Feb 14, 2024
f732783
Fix from previous changes and rebase
euri10 Feb 14, 2024
18c4aba
Make the bool coherent with the docstring
euri10 Feb 14, 2024
702d5ef
Avoid linter B027 empty-method-without-abstract-decorator on purpose
euri10 Feb 14, 2024
545fe54
Always add the store to lifespan_manager then let the store implement…
euri10 Feb 14, 2024
77ac9c0
Tamed mypy, aclose is definitely a Redis method, so not sure why it i…
euri10 Feb 14, 2024
92d3516
Ruff stuff
euri10 Feb 14, 2024
98ad6c8
Mypy on tests
euri10 Feb 14, 2024
ba6fcc2
Some flakiness? Trying to solve using the redis xdist
euri10 Feb 14, 2024
570ac9b
Cleant the redis_client fixture making it flushall each time, and abo…
euri10 Feb 15, 2024
e85b428
Made the introduced test redis "aware"
euri10 Feb 15, 2024
7b4741b
Just comment the introduced test to see if fxture change doesn't mess…
euri10 Feb 15, 2024
45ee20e
Do we need really a test, or should it uses a mock like all times wit…
euri10 Feb 15, 2024
40f2253
Do we need really a test, or should it uses a mock like all times wit…
euri10 Feb 15, 2024
f7ae116
Removed sync redis instance from fixture
euri10 Feb 15, 2024
e2bc0c8
You cannot yield the redis client weirdly enough...
euri10 Feb 22, 2024
58272d4
Use async context manager in docker stuff as well
euri10 Feb 22, 2024
dc1f3ca
Merge branch 'redis_test_cleaning' into 3083
euri10 Feb 22, 2024
3fc2cf2
We need to explicitely pass this otherwise using with_namespace after…
euri10 Feb 25, 2024
99a814d
Let's keep the fixtures as-is
euri10 Feb 26, 2024
fd2b54b
Add a test for with_client _shutdown
euri10 Feb 26, 2024
3651fe5
Remove the redis_service fixture from the test, note that it won't ru…
euri10 Feb 26, 2024
5e0f521
Merge branch 'main' into 3083
euri10 Feb 26, 2024
2c91e09
Update litestar/stores/redis.py
euri10 Feb 26, 2024
aac3fb5
Merge branch 'main' into 3083
euri10 Feb 27, 2024
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions docs/usage/stores.rst
Original file line number Diff line number Diff line change
Expand Up @@ -259,3 +259,11 @@ with minimal boilerplate:

Without any extra configuration, every call to ``app.stores.get`` with a unique name will return a namespace for this
name only, while re-using the underlying Redis instance.


Store lifetime
++++++++++++++

Stores may not be automatically closed when the application is shut down.
This is the case in particular for the RedisStore if you are not using the class method :meth:`RedisStore.with_client <.redis.RedisStore.with_client>` and passing in your own Redis instance.
In this case you're responsible to close the Redis instance yourself.
9 changes: 5 additions & 4 deletions litestar/app.py
Original file line number Diff line number Diff line change
Expand Up @@ -388,7 +388,12 @@ def __init__(

self._openapi_schema: OpenAPI | None = None
self._debug: bool = True
self.stores: StoreRegistry = (
config.stores if isinstance(config.stores, StoreRegistry) else StoreRegistry(config.stores)
)
self._lifespan_managers = config.lifespan
for store in self.stores._stores.values():
self._lifespan_managers.append(store)
self._server_lifespan_managers = [p.server_lifespan for p in config.plugins or [] if isinstance(p, CLIPlugin)]
self.experimental_features = frozenset(config.experimental_features or [])
self.get_logger: GetLogger = get_logger_placeholder
Expand Down Expand Up @@ -471,10 +476,6 @@ def __init__(

self.asgi_handler = self._create_asgi_handler()

self.stores: StoreRegistry = (
config.stores if isinstance(config.stores, StoreRegistry) else StoreRegistry(config.stores)
)

@property
@deprecated(version="2.6.0", kind="property", info="Use create_static_files router instead")
def static_files_config(self) -> list[StaticFilesConfig]:
Expand Down
13 changes: 13 additions & 0 deletions litestar/stores/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,8 @@
from msgspec.msgpack import encode as msgpack_encode

if TYPE_CHECKING:
from types import TracebackType

from typing_extensions import Self


Expand Down Expand Up @@ -76,6 +78,17 @@ async def expires_in(self, key: str) -> int | None:
"""
raise NotImplementedError

async def __aenter__(self) -> None: # noqa: B027
pass

async def __aexit__( # noqa: B027
self,
exc_type: type[BaseException] | None,
exc_val: BaseException | None,
exc_tb: TracebackType | None,
) -> None:
pass


class NamespacedStore(Store):
"""A subclass of :class:`Store`, offering hierarchical namespacing.
Expand Down
36 changes: 32 additions & 4 deletions litestar/stores/redis.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
from __future__ import annotations

from datetime import timedelta
from typing import cast
from typing import TYPE_CHECKING, cast

from redis.asyncio import Redis
from redis.asyncio.connection import ConnectionPool
Expand All @@ -14,23 +14,30 @@

__all__ = ("RedisStore",)

if TYPE_CHECKING:
from types import TracebackType


class RedisStore(NamespacedStore):
"""Redis based, thread and process safe asynchronous key/value store."""

__slots__ = ("_redis",)

def __init__(self, redis: Redis, namespace: str | None | EmptyType = Empty) -> None:
def __init__(
self, redis: Redis, namespace: str | None | EmptyType = Empty, handle_client_shutdown: bool = False
) -> None:
"""Initialize :class:`RedisStore`

Args:
redis: An :class:`redis.asyncio.Redis` instance
namespace: A key prefix to simulate a namespace in redis. If not given,
defaults to ``LITESTAR``. Namespacing can be explicitly disabled by passing
``None``. This will make :meth:`.delete_all` unavailable.
handle_client_shutdown: If ``True``, handle the shutdown of the `redis` instance automatically during the store's lifespan. Should be set to `True` unless the shutdown is handled externally
"""
self._redis = redis
self.namespace: str | None = value_or_default(namespace, "LITESTAR")
self.handle_client_shutdown = handle_client_shutdown

# script to get and renew a key in one atomic step
self._get_and_renew_script = self._redis.register_script(
Expand Down Expand Up @@ -64,6 +71,18 @@ def __init__(self, redis: Redis, namespace: str | None | EmptyType = Empty) -> N
"""
)

async def _shutdown(self) -> None:
if self.handle_client_shutdown:
await self._redis.aclose(close_connection_pool=True) # type: ignore[attr-defined]

async def __aexit__(
self,
exc_type: type[BaseException] | None,
exc_val: BaseException | None,
exc_tb: TracebackType | None,
) -> None:
await self._shutdown()

@classmethod
def with_client(
cls,
Expand All @@ -84,6 +103,7 @@ def with_client(
username: Redis username to use
password: Redis password to use
namespace: Virtual key namespace to use
automatic_store_lifetime: If ``True``, the store lifetime will be handled automatically, it is the default with this method.
euri10 marked this conversation as resolved.
Show resolved Hide resolved
"""
pool = ConnectionPool.from_url(
url=url,
Expand All @@ -93,14 +113,22 @@ def with_client(
username=username,
password=password,
)
return cls(redis=Redis(connection_pool=pool), namespace=namespace)
return cls(
redis=Redis(connection_pool=pool),
namespace=namespace,
handle_client_shutdown=True,
)

def with_namespace(self, namespace: str) -> RedisStore:
"""Return a new :class:`RedisStore` with a nested virtual key namespace.
The current instances namespace will serve as a prefix for the namespace, so it
can be considered the parent namespace.
"""
return type(self)(redis=self._redis, namespace=f"{self.namespace}_{namespace}" if self.namespace else namespace)
euri10 marked this conversation as resolved.
Show resolved Hide resolved
return type(self)(
redis=self._redis,
namespace=f"{self.namespace}_{namespace}" if self.namespace else namespace,
handle_client_shutdown=self.handle_client_shutdown,
)

def _make_key(self, key: str) -> str:
prefix = f"{self.namespace}:" if self.namespace else ""
Expand Down
13 changes: 13 additions & 0 deletions tests/unit/test_stores.py
Original file line number Diff line number Diff line change
Expand Up @@ -366,3 +366,16 @@ async def test_file_store_handle_rename_fail(file_store: FileStore, mocker: Mock
await file_store.set("foo", "bar")
mock_unlink.assert_called_once()
assert Path(mock_unlink.call_args_list[0].args[0]).with_suffix("") == file_store.path.joinpath("foo")


async def test_redis_store_with_client_shutdown() -> None:
redis_store = RedisStore.with_client(url="redis://localhost:6397")
assert await redis_store._redis.ping()
# remove the private shutdown and the assert below fails
# the check on connection is a mimic of https://github.com/redis/redis-py/blob/d529c2ad8d2cf4dcfb41bfd93ea68cfefd81aa66/tests/test_asyncio/test_connection_pool.py#L35-L39
await redis_store._shutdown()
assert not any(
x.is_connected
for x in redis_store._redis.connection_pool._available_connections
+ list(redis_store._redis.connection_pool._in_use_connections)
)
Loading