|
| 1 | +"""S3 storage driver client abstraction for the S3 storage driver. |
| 2 | +
|
| 3 | +.. warning:: |
| 4 | + This API is experimental. |
| 5 | +""" |
| 6 | + |
| 7 | +from __future__ import annotations |
| 8 | + |
| 9 | +import io |
| 10 | +from abc import ABC, abstractmethod |
| 11 | + |
| 12 | +from botocore.exceptions import ClientError |
| 13 | +from types_aiobotocore_s3.client import S3Client |
| 14 | + |
| 15 | + |
| 16 | +class S3StorageDriverClient(ABC): |
| 17 | + """Abstract base class for S3 object operations. |
| 18 | +
|
| 19 | + Implementations must support ``put_object`` and ``get_object``. Multipart |
| 20 | + upload handling (if needed) is an internal concern of each implementation. |
| 21 | +
|
| 22 | + .. warning:: |
| 23 | + This API is experimental. |
| 24 | + """ |
| 25 | + |
| 26 | + @abstractmethod |
| 27 | + async def put_object(self, *, bucket: str, key: str, data: bytes) -> None: |
| 28 | + """Upload *data* to the given S3 *bucket* and *key*.""" |
| 29 | + |
| 30 | + @abstractmethod |
| 31 | + async def object_exists(self, *, bucket: str, key: str) -> bool: |
| 32 | + """Return ``True`` if an object exists at the given *bucket* and *key*.""" |
| 33 | + |
| 34 | + @abstractmethod |
| 35 | + async def get_object(self, *, bucket: str, key: str) -> bytes: |
| 36 | + """Download and return the bytes stored at the given S3 *bucket* and *key*.""" |
| 37 | + |
| 38 | + |
| 39 | +class Aioboto3StorageDriverClient(S3StorageDriverClient): |
| 40 | + """Adapter that wraps an aioboto3 S3 client as an :class:`S3StorageDriverClient`. |
| 41 | +
|
| 42 | + Internally delegates to ``upload_fileobj`` for uploads (which handles |
| 43 | + multipart automatically for objects above the multipart threshold) and |
| 44 | + ``get_object`` for downloads. |
| 45 | +
|
| 46 | + .. warning:: |
| 47 | + This API is experimental. |
| 48 | + """ |
| 49 | + |
| 50 | + def __init__(self, client: S3Client) -> None: |
| 51 | + """Wrap an aioboto3 S3 client. |
| 52 | +
|
| 53 | + Args: |
| 54 | + client: An aioboto3 S3 client, typically obtained from |
| 55 | + ``aioboto3.Session().client("s3")``. |
| 56 | + """ |
| 57 | + self._client = client |
| 58 | + |
| 59 | + async def object_exists(self, *, bucket: str, key: str) -> bool: |
| 60 | + """Check existence via aioboto3's ``head_object``.""" |
| 61 | + try: |
| 62 | + await self._client.head_object(Bucket=bucket, Key=key) |
| 63 | + return True |
| 64 | + except ClientError as e: |
| 65 | + # head_object returns 404 as a ClientError when the key doesn't exist. |
| 66 | + if e.response.get("Error", {}).get("Code") == "404": |
| 67 | + return False |
| 68 | + raise |
| 69 | + |
| 70 | + async def put_object(self, *, bucket: str, key: str, data: bytes) -> None: |
| 71 | + """Upload *data* via aioboto3's ``upload_fileobj``.""" |
| 72 | + # upload_fileobj is an aioboto3-specific method not in the |
| 73 | + # types_aiobotocore_s3 stubs; it handles multipart automatically. |
| 74 | + await self._client.upload_fileobj(io.BytesIO(data), bucket, key) # type: ignore[arg-type] |
| 75 | + |
| 76 | + async def get_object(self, *, bucket: str, key: str) -> bytes: |
| 77 | + """Download bytes via aioboto3's ``get_object``.""" |
| 78 | + response = await self._client.get_object(Bucket=bucket, Key=key) |
| 79 | + # StreamingBody.read() is untyped in aiobotocore, returns bytes at runtime. |
| 80 | + return await response["Body"].read() # type: ignore[no-any-return] |
| 81 | + |
| 82 | + |
| 83 | +def new_aioboto3_client(client: S3Client) -> Aioboto3StorageDriverClient: |
| 84 | + """Create an :class:`S3StorageDriverClient` from an aioboto3 S3 client. |
| 85 | +
|
| 86 | + This is a convenience factory. Equivalent to ``Aioboto3StorageDriverClient(client)``. |
| 87 | +
|
| 88 | + Args: |
| 89 | + client: An aioboto3 S3 client, typically obtained from |
| 90 | + ``aioboto3.Session().client("s3")``. |
| 91 | + """ |
| 92 | + return Aioboto3StorageDriverClient(client) |
0 commit comments