Skip to content

Commit

Permalink
hi
Browse files Browse the repository at this point in the history
  • Loading branch information
dycw committed Sep 16, 2024
1 parent 764df13 commit 5375a91
Showing 1 changed file with 30 additions and 22 deletions.
52 changes: 30 additions & 22 deletions src/utilities/redis.py
Original file line number Diff line number Diff line change
Expand Up @@ -1587,6 +1587,7 @@ def __str__(self) -> str:
@contextmanager
def yield_client(
*,
client: redis.Redis | None = None,
host: str = _HOST,
port: int = _PORT,
db: int = 0,
Expand All @@ -1596,24 +1597,28 @@ def yield_client(
**kwargs: Any,
) -> Iterator[redis.Redis]:
"""Yield a synchronous client."""
client = redis.Redis(
host=host,
port=port,
db=db,
password=password,
connection_pool=connection_pool,
decode_responses=decode_responses,
**kwargs,
)
if client is None:
client_use = redis.Redis(
host=host,
port=port,
db=db,
password=password,
connection_pool=connection_pool,
decode_responses=decode_responses,
**kwargs,
)
else:
client_use = client
try:
yield client
yield client_use
finally:
client.close()
client_use.close()


@asynccontextmanager
async def yield_client_async(
*,
client: redis.asyncio.Redis | None = None,
host: str = _HOST,
port: int = _PORT,
db: str | int = 0,
Expand All @@ -1623,19 +1628,22 @@ async def yield_client_async(
**kwargs: Any,
) -> AsyncIterator[redis.asyncio.Redis]:
"""Yield an asynchronous client."""
client = redis.asyncio.Redis(
host=host,
port=port,
db=db,
password=password,
connection_pool=connection_pool,
decode_responses=decode_responses,
**kwargs,
)
if client is None:
client_use = redis.asyncio.Redis(
host=host,
port=port,
db=db,
password=password,
connection_pool=connection_pool,
decode_responses=decode_responses,
**kwargs,
)
else:
client_use = client
try:
yield client
yield client_use
finally:
match client.connection_pool:
match client_use.connection_pool:
case redis.ConnectionPool() as pool:
pool.disconnect(inuse_connections=False) # pragma: no cover
case redis.asyncio.ConnectionPool() as pool:
Expand Down

0 comments on commit 5375a91

Please sign in to comment.