Skip to content

Commit fe0f671

Browse files
authored
fix: Fix Event loop is closed error on dynamodb test (feast-dev#5480)
Signed-off-by: ntkathole <[email protected]>
1 parent af99f0a commit fe0f671

File tree

1 file changed

+13
-4
lines changed

1 file changed

+13
-4
lines changed

sdk/python/feast/infra/online_stores/dynamodb.py

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -599,6 +599,7 @@ def _to_client_batch_get_payload(online_config, table_name, batch):
599599

600600
_aioboto_session = None
601601
_aioboto_client = None
602+
_aioboto_context_stack = None
602603

603604

604605
def _get_aioboto_session():
@@ -618,7 +619,7 @@ async def _get_aiodynamodb_client(
618619
total_max_retry_attempts: Union[int, None],
619620
retry_mode: Union[Literal["legacy", "standard", "adaptive"], None],
620621
):
621-
global _aioboto_client
622+
global _aioboto_client, _aioboto_context_stack
622623
if _aioboto_client is None:
623624
logger.debug("initializing the aiobotocore dynamodb client")
624625

@@ -639,15 +640,23 @@ async def _get_aiodynamodb_client(
639640
connector_args={"keepalive_timeout": keepalive_timeout},
640641
),
641642
)
642-
context_stack = contextlib.AsyncExitStack()
643-
_aioboto_client = await context_stack.enter_async_context(client_context)
643+
_aioboto_context_stack = contextlib.AsyncExitStack()
644+
_aioboto_client = await _aioboto_context_stack.enter_async_context(
645+
client_context
646+
)
644647
return _aioboto_client
645648

646649

647650
async def _aiodynamodb_close():
648-
global _aioboto_client
651+
global _aioboto_client, _aioboto_session, _aioboto_context_stack
649652
if _aioboto_client:
650653
await _aioboto_client.close()
654+
_aioboto_client = None
655+
if _aioboto_context_stack:
656+
await _aioboto_context_stack.aclose()
657+
_aioboto_context_stack = None
658+
if _aioboto_session:
659+
_aioboto_session = None
651660

652661

653662
def _initialize_dynamodb_client(

0 commit comments

Comments
 (0)