From a59e533699bfe78427eccaaada5af247cebda77d Mon Sep 17 00:00:00 2001 From: Malcolm Jones Date: Thu, 1 Aug 2024 14:47:26 -0400 Subject: [PATCH] chg: goobredis fixes --- .prompts.d/annotations.txt | 18 +++++----- src/goob_ai/backend/cache/goobredis.py | 39 +++++++++----------- tests/backend/cache/test_goobredis.py | 50 ++++++++++++++++++++++++-- 3 files changed, 74 insertions(+), 33 deletions(-) diff --git a/.prompts.d/annotations.txt b/.prompts.d/annotations.txt index ebcf38db..ec56eb64 100644 --- a/.prompts.d/annotations.txt +++ b/.prompts.d/annotations.txt @@ -9,23 +9,23 @@ using @src/goob_ai/backend/__init__.py as context, I want you to make sure to add typing annotations to each function or class in side of @src/goob_ai/backend/__init__.py. Be sure to include return types when necessary. add descriptive docstrings to all functions and classes inside of @src/goob_ai/backend/__init__.py. Please use pep257 convention, in the style of 'google'. Update existing docstrings if they do not conform to is. Don't forget to include any necessary typing modules at the top of the file. - using @src/goob_ai/backend/cache/__init__.py as context, I want you to make sure to add typing annotations to each function or class in side of @src/goob_ai/backend/cache/__init__.py. Be sure to include return types when necessary. add descriptive docstrings to all functions and classes inside of @src/goob_ai/backend/cache/__init__.py. Please use pep257 convention, in the style of 'google'. Update existing docstrings if they do not conform to is. Don't forget to include any necessary typing modules at the top of the file. + using @__init__.py as context, I want you to make sure to add typing annotations to each function or class in side of @__init__.py. Be sure to include return types when necessary. add descriptive docstrings to all functions and classes inside of @__init__.py. Please use pep257 convention, in the style of 'google'. Update existing docstrings if they do not conform to is. Don't forget to include any necessary typing modules at the top of the file. - using @src/goob_ai/backend/cache/goobredis.py as context, I want you to make sure to add typing annotations to each function or class in side of @src/goob_ai/backend/cache/goobredis.py. Be sure to include return types when necessary. add descriptive docstrings to all functions and classes inside of @src/goob_ai/backend/cache/goobredis.py. Please use pep257 convention, in the style of 'google'. Update existing docstrings if they do not conform to is. Don't forget to include any necessary typing modules at the top of the file. + using @goobredis.py as context, I want you to make sure to add typing annotations to each function or class in side of @goobredis.py. Be sure to include return types when necessary. add descriptive docstrings to all functions and classes inside of @goobredis.py. Please use pep257 convention, in the style of 'google'. Update existing docstrings if they do not conform to is. Don't forget to include any necessary typing modules at the top of the file. - using @src/goob_ai/backend/cache/goobredis_session.py as context, I want you to make sure to add typing annotations to each function or class in side of @src/goob_ai/backend/cache/goobredis_session.py. Be sure to include return types when necessary. add descriptive docstrings to all functions and classes inside of @src/goob_ai/backend/cache/goobredis_session.py. Please use pep257 convention, in the style of 'google'. Update existing docstrings if they do not conform to is. Don't forget to include any necessary typing modules at the top of the file. + using @goobredis_session.py as context, I want you to make sure to add typing annotations to each function or class in side of @goobredis_session.py. Be sure to include return types when necessary. add descriptive docstrings to all functions and classes inside of @goobredis_session.py. Please use pep257 convention, in the style of 'google'. Update existing docstrings if they do not conform to is. Don't forget to include any necessary typing modules at the top of the file. - using @src/goob_ai/backend/discord/discord_client.py as context, I want you to make sure to add typing annotations to each function or class in side of @src/goob_ai/backend/discord/discord_client.py. Be sure to include return types when necessary. add descriptive docstrings to all functions and classes inside of @src/goob_ai/backend/discord/discord_client.py. Please use pep257 convention, in the style of 'google'. Update existing docstrings if they do not conform to is. Don't forget to include any necessary typing modules at the top of the file. + using @discord_client.py as context, I want you to make sure to add typing annotations to each function or class in side of @discord_client.py. Be sure to include return types when necessary. add descriptive docstrings to all functions and classes inside of @discord_client.py. Please use pep257 convention, in the style of 'google'. Update existing docstrings if they do not conform to is. Don't forget to include any necessary typing modules at the top of the file. using @src/goob_ai/backend/ml/__init__.py as context, I want you to make sure to add typing annotations to each function or class in side of @src/goob_ai/backend/ml/__init__.py. Be sure to include return types when necessary. add descriptive docstrings to all functions and classes inside of @src/goob_ai/backend/ml/__init__.py. Please use pep257 convention, in the style of 'google'. Update existing docstrings if they do not conform to is. Don't forget to include any necessary typing modules at the top of the file. - using @src/goob_ai/backend/ml/devices.py as context, I want you to make sure to add typing annotations to each function or class in side of @src/goob_ai/backend/ml/devices.py. Be sure to include return types when necessary. add descriptive docstrings to all functions and classes inside of @src/goob_ai/backend/ml/devices.py. Please use pep257 convention, in the style of 'google'. Update existing docstrings if they do not conform to is. Don't forget to include any necessary typing modules at the top of the file. + using @devices.py as context, I want you to make sure to add typing annotations to each function or class in side of @src/goob_ai/backend/ml/devices.py. Be sure to include return types when necessary. add descriptive docstrings to all functions and classes inside of @src/goob_ai/backend/ml/devices.py. Please use pep257 convention, in the style of 'google'. Update existing docstrings if they do not conform to is. Don't forget to include any necessary typing modules at the top of the file. - using @src/goob_ai/base.py as context, I want you to make sure to add typing annotations to each function or class in side of @src/goob_ai/base.py. Be sure to include return types when necessary. add descriptive docstrings to all functions and classes inside of @src/goob_ai/base.py. Please use pep257 convention, in the style of 'google'. Update existing docstrings if they do not conform to is. Don't forget to include any necessary typing modules at the top of the file. + using @base.py as context, I want you to make sure to add typing annotations to each function or class in side of @src/goob_ai/base.py. Be sure to include return types when necessary. add descriptive docstrings to all functions and classes inside of @src/goob_ai/base.py. Please use pep257 convention, in the style of 'google'. Update existing docstrings if they do not conform to is. Don't forget to include any necessary typing modules at the top of the file. using @src/goob_ai/bot_logger/__init__.py as context, I want you to make sure to add typing annotations to each function or class in side of @src/goob_ai/bot_logger/__init__.py. Be sure to include return types when necessary. add descriptive docstrings to all functions and classes inside of @src/goob_ai/bot_logger/__init__.py. Please use pep257 convention, in the style of 'google'. Update existing docstrings if they do not conform to is. Don't forget to include any necessary typing modules at the top of the file. - using @src/goob_ai/chromadb_manager.py as context, I want you to make sure to add typing annotations to each function or class in side of @src/goob_ai/chromadb_manager.py. Be sure to include return types when necessary. add descriptive docstrings to all functions and classes inside of @src/goob_ai/chromadb_manager.py. Please use pep257 convention, in the style of 'google'. Update existing docstrings if they do not conform to is. Don't forget to include any necessary typing modules at the top of the file. + using @chromadb_manager.py as context, I want you to make sure to add typing annotations to each function or class in side of @src/goob_ai/chromadb_manager.py. Be sure to include return types when necessary. add descriptive docstrings to all functions and classes inside of @src/goob_ai/chromadb_manager.py. Please use pep257 convention, in the style of 'google'. Update existing docstrings if they do not conform to is. Don't forget to include any necessary typing modules at the top of the file. using @src/goob_ai/cli.py as context, I want you to make sure to add typing annotations to each function or class in side of @src/goob_ai/cli.py. Be sure to include return types when necessary. add descriptive docstrings to all functions and classes inside of @src/goob_ai/cli.py. Please use pep257 convention, in the style of 'google'. Update existing docstrings if they do not conform to is. Don't forget to include any necessary typing modules at the top of the file. @@ -149,13 +149,13 @@ using @src/goob_ai/services/__init__.py as context, I want you to make sure to add typing annotations to each function or class in side of @src/goob_ai/services/__init__.py. Be sure to include return types when necessary. add descriptive docstrings to all functions and classes inside of @src/goob_ai/services/__init__.py. Please use pep257 convention, in the style of 'google'. Update existing docstrings if they do not conform to is. Don't forget to include any necessary typing modules at the top of the file. - using @src/goob_ai/services/chroma_service.py as context, I want you to make sure to add typing annotations to each function or class in side of @src/goob_ai/services/chroma_service.py. Be sure to include return types when necessary. add descriptive docstrings to all functions and classes inside of @src/goob_ai/services/chroma_service.py. Please use pep257 convention, in the style of 'google'. Update existing docstrings if they do not conform to is. Don't forget to include any necessary typing modules at the top of the file. + using @chroma_service.py as context, I want you to make sure to add typing annotations to each function or class in side of @src/goob_ai/services/chroma_service.py. Be sure to include return types when necessary. add descriptive docstrings to all functions and classes inside of @src/goob_ai/services/chroma_service.py. Please use pep257 convention, in the style of 'google'. Update existing docstrings if they do not conform to is. Don't forget to include any necessary typing modules at the top of the file. using @src/goob_ai/services/cromadb_service.py as context, I want you to make sure to add typing annotations to each function or class in side of @src/goob_ai/services/cromadb_service.py. Be sure to include return types when necessary. add descriptive docstrings to all functions and classes inside of @src/goob_ai/services/cromadb_service.py. Please use pep257 convention, in the style of 'google'. Update existing docstrings if they do not conform to is. Don't forget to include any necessary typing modules at the top of the file. using @src/goob_ai/services/image_service.py as context, I want you to make sure to add typing annotations to each function or class in side of @src/goob_ai/services/image_service.py. Be sure to include return types when necessary. add descriptive docstrings to all functions and classes inside of @src/goob_ai/services/image_service.py. Please use pep257 convention, in the style of 'google'. Update existing docstrings if they do not conform to is. Don't forget to include any necessary typing modules at the top of the file. - using @src/goob_ai/services/pinecone_service.py as context, I want you to make sure to add typing annotations to each function or class in side of @src/goob_ai/services/pinecone_service.py. Be sure to include return types when necessary. add descriptive docstrings to all functions and classes inside of @src/goob_ai/services/pinecone_service.py. Please use pep257 convention, in the style of 'google'. Update existing docstrings if they do not conform to is. Don't forget to include any necessary typing modules at the top of the file. + using @pinecone_service.py as context, I want you to make sure to add typing annotations to each function or class in side of @src/goob_ai/services/pinecone_service.py. Be sure to include return types when necessary. add descriptive docstrings to all functions and classes inside of @src/goob_ai/services/pinecone_service.py. Please use pep257 convention, in the style of 'google'. Update existing docstrings if they do not conform to is. Don't forget to include any necessary typing modules at the top of the file. using @src/goob_ai/services/screencrop_service.py as context, I want you to make sure to add typing annotations to each function or class in side of @src/goob_ai/services/screencrop_service.py. Be sure to include return types when necessary. add descriptive docstrings to all functions and classes inside of @src/goob_ai/services/screencrop_service.py. Please use pep257 convention, in the style of 'google'. Update existing docstrings if they do not conform to is. Don't forget to include any necessary typing modules at the top of the file. diff --git a/src/goob_ai/backend/cache/goobredis.py b/src/goob_ai/backend/cache/goobredis.py index b4072834..74cfb641 100644 --- a/src/goob_ai/backend/cache/goobredis.py +++ b/src/goob_ai/backend/cache/goobredis.py @@ -106,6 +106,7 @@ async def _connect(self) -> None: ) self._pool = redis.asyncio.Redis(connection_pool=self._conn_pool) self._pubsub_channels: dict[str, PubSub] = {} + self.auto_close_connection_pool: bool = self._pool.auto_close_connection_pool async def finalize(self) -> None: """Finalize the Redis client.""" @@ -293,28 +294,6 @@ async def unsubscribe(self, channel_name: str) -> None: finally: await p.__aexit__(None, None, None) - # async def subscribe(self, channel_name: str) -> AsyncGenerator[Any, None]: - # """ - # Subscribe to a Redis channel. - - # Args: - # channel_name: The name of the channel to subscribe to. - - # Yields: - # Messages received from the subscribed channel. - - # Raises: - # NoRedisConfigured: If no Redis is configured. - # """ - # if self._pool is None: - # raise NoRedisConfigured() - - # p: PubSub = await self._pool.pubsub().__aenter__() - # self._pubsub_channels[channel_name] = p - # await p.subscribe(channel_name) - # async for message in self._listener(p): - # yield message - async def subscribe(self, channel_name: str) -> AsyncGenerator[Any, None]: """ Subscribe to a Redis channel. @@ -351,6 +330,22 @@ async def _listener(self, p: PubSub) -> AsyncGenerator[Any, None]: if message is not None: yield message["data"] + async def aclose(self, close_connection_pool: Optional[bool] = None) -> None: + """ + Closes Redis client connection + + :param close_connection_pool: decides whether to close the connection pool used + by this Redis client, overriding Redis.auto_close_connection_pool. By default, + let Redis.auto_close_connection_pool decide whether to close the connection + pool. + """ + conn = self._pool + if conn: + self._pool = None + await self._pool.release(conn) + if close_connection_pool or (close_connection_pool is None and self.auto_close_connection_pool): + await self._pool.disconnect() + # def set(self, key, value, expiration=3600): # try: diff --git a/tests/backend/cache/test_goobredis.py b/tests/backend/cache/test_goobredis.py index b1eb3214..8c532526 100644 --- a/tests/backend/cache/test_goobredis.py +++ b/tests/backend/cache/test_goobredis.py @@ -8,6 +8,8 @@ from typing import TYPE_CHECKING +import redis.asyncio as redis + from goob_ai.backend.cache.goobredis import get_driver from loguru import logger as LOGGER @@ -22,10 +24,54 @@ from pytest_mock.plugin import MockerFixture +########################################################################## + + +@pytest.fixture(name="create_redis") +async def create_redis(request): + """Wrapper around redis.create_redis.""" + # (single_connection,) = request.param + + teardown_clients = [] + + async def client_factory( + **kwargs, + ) -> GoobRedisClient: + client = await get_driver() + + async def teardown(): + try: + await client.aclose() + # await client.connection_pool.disconnect() + except redis.ConnectionError: + pass + + teardown_clients.append(teardown) + return client + + yield client_factory + + for teardown in teardown_clients: + await teardown() + + +@pytest.fixture(name="r") +async def r(create_redis): + return await create_redis() + + +@pytest.fixture(name="decoded_r") +async def decoded_r(create_redis): + return await create_redis(decode_responses=True) + + +########################################################################## + # @pytest.mark.app_settings({"applications": ["guillotina", "guillotina.contrib.redis"]}) -async def test_redis_ops(caplog: LogCaptureFixture): - driver = await get_driver() +async def test_redis_ops(caplog: LogCaptureFixture, create_redis, **kwargs): + driver: GoobRedisClient = await get_driver() + r = await create_redis(**kwargs) assert driver.initialized assert driver.pool is not None