From 712a48b96516f09d3ebd49e3eab0914d20609d84 Mon Sep 17 00:00:00 2001 From: ntkathole Date: Mon, 30 Jun 2025 20:53:56 +0530 Subject: [PATCH] fix: Fix Event loop is closed error on dynamodb test Signed-off-by: ntkathole --- .../feast/infra/online_stores/dynamodb.py | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/sdk/python/feast/infra/online_stores/dynamodb.py b/sdk/python/feast/infra/online_stores/dynamodb.py index c73a6104495..e4ff4fbe513 100644 --- a/sdk/python/feast/infra/online_stores/dynamodb.py +++ b/sdk/python/feast/infra/online_stores/dynamodb.py @@ -599,6 +599,7 @@ def _to_client_batch_get_payload(online_config, table_name, batch): _aioboto_session = None _aioboto_client = None +_aioboto_context_stack = None def _get_aioboto_session(): @@ -618,7 +619,7 @@ async def _get_aiodynamodb_client( total_max_retry_attempts: Union[int, None], retry_mode: Union[Literal["legacy", "standard", "adaptive"], None], ): - global _aioboto_client + global _aioboto_client, _aioboto_context_stack if _aioboto_client is None: logger.debug("initializing the aiobotocore dynamodb client") @@ -639,15 +640,23 @@ async def _get_aiodynamodb_client( connector_args={"keepalive_timeout": keepalive_timeout}, ), ) - context_stack = contextlib.AsyncExitStack() - _aioboto_client = await context_stack.enter_async_context(client_context) + _aioboto_context_stack = contextlib.AsyncExitStack() + _aioboto_client = await _aioboto_context_stack.enter_async_context( + client_context + ) return _aioboto_client async def _aiodynamodb_close(): - global _aioboto_client + global _aioboto_client, _aioboto_session, _aioboto_context_stack if _aioboto_client: await _aioboto_client.close() + _aioboto_client = None + if _aioboto_context_stack: + await _aioboto_context_stack.aclose() + _aioboto_context_stack = None + if _aioboto_session: + _aioboto_session = None def _initialize_dynamodb_client(