storages: use asynccontextmanager instead of sync to close client (#1521)

Follow-up to #1481, use the asyncontextmanager with `async with` as only
used in async functions (which call run_in_executor)
This commit is contained in:
Ilya Kreymer 2024-02-08 08:28:53 -08:00 committed by GitHub
parent b2a5dbf2cd
commit 65fec64197
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -12,7 +12,7 @@ from typing import (
TYPE_CHECKING,
)
from urllib.parse import urlsplit
from contextlib import asynccontextmanager, contextmanager
from contextlib import asynccontextmanager
import asyncio
import heapq
@ -276,8 +276,10 @@ class StorageOps:
) as client:
yield client, bucket, key
@contextmanager
def get_sync_client(self, org: Organization) -> Iterator[tuple[S3Client, str, str]]:
@asynccontextmanager
async def get_sync_client(
self, org: Organization
) -> AsyncIterator[tuple[S3Client, str, str]]:
"""context manager for s3 client"""
storage = self.get_org_primary_storage(org)
@ -517,7 +519,7 @@ class StorageOps:
contexts: List[str],
) -> Iterator[bytes]:
"""Return filtered stream of logs from specified WACZs sorted by timestamp"""
with self.get_sync_client(org) as (client, bucket, key):
async with self.get_sync_client(org) as (client, bucket, key):
loop = asyncio.get_event_loop()
resp = await loop.run_in_executor(
@ -665,7 +667,7 @@ class StorageOps:
) -> Iterator[bytes]:
"""return an iter for downloading a stream nested wacz file
from list of files"""
with self.get_sync_client(org) as (client, bucket, key):
async with self.get_sync_client(org) as (client, bucket, key):
loop = asyncio.get_event_loop()
resp = await loop.run_in_executor(