This is an automated email from the ASF dual-hosted git repository.
dabla pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new a3cd2a59c23 Add an async upload to S3Hook: aload_bytes (#75242)
a3cd2a59c23 is described below
commit a3cd2a59c234ebd9c5b0e9533b7441b14101b1ad
Author: David Blain <[email protected]>
AuthorDate: Thu Oct 22 14:31:52 2026 +0200
Add an async upload to S3Hook: aload_bytes (#75242)
An async task or trigger that writes to S3 has to hand load_bytes() to a
thread, because S3Hook only has async methods to read and check keys. A
native coroutine on the aiobotocore client of get_async_conn() lets it
upload without blocking its event loop or leaving it.
The new method does not reuse _upload_file_obj(), which adds the per-call
encryption and ACL settings to the hook's own extra_args, so that they do
not leak into later uploads through the same hook.
---
.../src/airflow/providers/amazon/aws/hooks/s3.py | 52 +++++++++++++++
.../amazon/tests/unit/amazon/aws/hooks/test_s3.py | 77 ++++++++++++++++++++++
2 files changed, 129 insertions(+)
diff --git a/providers/amazon/src/airflow/providers/amazon/aws/hooks/s3.py
b/providers/amazon/src/airflow/providers/amazon/aws/hooks/s3.py
index bba4ac7399f..e9089f95bfc 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/hooks/s3.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/hooks/s3.py
@@ -1333,6 +1333,58 @@ class S3Hook(AwsBaseHook):
with BytesIO(bytes_data) as f:
self._upload_file_obj(f, key, bucket_name, replace, encrypt,
acl_policy)
+ @unify_bucket_name_and_key
+ @provide_bucket_name
+ async def aload_bytes(
+ self,
+ bytes_data: bytes,
+ key: str,
+ bucket_name: str | None = None,
+ replace: bool = False,
+ encrypt: bool = False,
+ acl_policy: str | None = None,
+ ) -> None:
+ """
+ Load bytes to S3 asynchronously.
+
+ The async counterpart of :meth:`load_bytes`, on the aiobotocore client
of
+ :meth:`get_async_conn`, so an async task or trigger can upload without
blocking its
+ event loop. It sends the object in a single ``PutObject`` request:
there is no
+ multipart transfer and ``transfer_config`` is not used, so it is meant
for objects
+ below the 5 GB limit of a single upload.
+
+ .. seealso::
+ - :external+boto3:py:meth:`S3.Client.put_object`
+
+ :param bytes_data: bytes to set as content for the key.
+ :param key: S3 key that will point to the file
+ :param bucket_name: Name of the bucket in which to store the file
+ :param replace: A flag to decide whether or not to overwrite the key
+ if it already exists
+ :param encrypt: If True, the file will be encrypted on the server-side
+ by S3 and will be stored in an encrypted form while at rest in S3.
+ :param acl_policy: The string to specify the canned ACL policy for the
+ object to be uploaded
+ """
+ async with await self.get_async_conn() as client:
+ if not replace and await self.get_head_object_async(client, key,
bucket_name):
+ raise ValueError(f"The key {key} already exists.")
+
+ # A copy, so the hook's own extra_args keep no per-call settings.
+ extra_args = dict(self.extra_args)
+ if encrypt:
+ extra_args["ServerSideEncryption"] = "AES256"
+ if acl_policy:
+ extra_args["ACL"] = acl_policy
+ if self._requester_pays:
+ extra_args["RequestPayer"] = "requester"
+
+ await client.put_object(Bucket=bucket_name, Key=key,
Body=bytes_data, **extra_args)
+
+ get_hook_lineage_collector().add_output_asset(
+ context=self, scheme="s3", asset_kwargs={"bucket": bucket_name,
"key": key}
+ )
+
@unify_bucket_name_and_key
@provide_bucket_name
def load_file_obj(
diff --git a/providers/amazon/tests/unit/amazon/aws/hooks/test_s3.py
b/providers/amazon/tests/unit/amazon/aws/hooks/test_s3.py
index 2bc0a03bcfe..c5f373a4e28 100644
--- a/providers/amazon/tests/unit/amazon/aws/hooks/test_s3.py
+++ b/providers/amazon/tests/unit/amazon/aws/hooks/test_s3.py
@@ -1310,6 +1310,83 @@ class TestAwsS3Hook:
assert response["Grants"][1]["Permission"] == "READ"
assert response["Grants"][0]["Permission"] == "FULL_CONTROL"
+ @staticmethod
+ def _mock_async_conn(mock_get_async_conn, head_object=None):
+ """Make ``get_async_conn()`` yield an aiobotocore client mock, as
``async with await`` uses it."""
+
+ class S3Client:
+ """The aiobotocore S3 client methods aload_bytes calls."""
+
+ async def head_object(self, **kwargs): ...
+
+ async def put_object(self, **kwargs): ...
+
+ client = AsyncMock(spec=S3Client)
+ client.head_object.return_value = head_object
+
+ class AsyncConn:
+ async def __aenter__(self):
+ return client
+
+ async def __aexit__(self, *exc_info):
+ return None
+
+ mock_get_async_conn.return_value = AsyncConn()
+ return client
+
+ @pytest.mark.asyncio
+
@async_mock.patch("airflow.providers.amazon.aws.hooks.s3.S3Hook.get_async_conn",
autospec=True)
+ async def test_aload_bytes(self, mock_get_async_conn,
hook_lineage_collector):
+ client = self._mock_async_conn(mock_get_async_conn)
+ hook = S3Hook(aws_conn_id=None)
+
+ await hook.aload_bytes(b"Content", "my_key", "my_bucket")
+
+ client.head_object.assert_awaited_once_with(Bucket="my_bucket",
Key="my_key")
+ client.put_object.assert_awaited_once_with(Bucket="my_bucket",
Key="my_key", Body=b"Content")
+ assert len(hook_lineage_collector.collected_assets.outputs) == 1
+ assert hook_lineage_collector.collected_assets.outputs[0].asset ==
Asset(uri="s3://my_bucket/my_key")
+
+ @pytest.mark.asyncio
+
@async_mock.patch("airflow.providers.amazon.aws.hooks.s3.S3Hook.get_async_conn",
autospec=True)
+ async def test_aload_bytes_takes_the_bucket_from_an_s3_url(self,
mock_get_async_conn):
+ client = self._mock_async_conn(mock_get_async_conn)
+
+ await S3Hook(aws_conn_id=None).aload_bytes(b"Content",
"s3://my_bucket/dir/my_key", replace=True)
+
+ client.head_object.assert_not_called()
+ client.put_object.assert_awaited_once_with(Bucket="my_bucket",
Key="dir/my_key", Body=b"Content")
+
+ @pytest.mark.asyncio
+
@async_mock.patch("airflow.providers.amazon.aws.hooks.s3.S3Hook.get_async_conn",
autospec=True)
+ async def test_aload_bytes_refuses_an_existing_key_without_replace(self,
mock_get_async_conn):
+ client = self._mock_async_conn(mock_get_async_conn,
head_object={"ContentLength": 7})
+
+ with pytest.raises(ValueError, match="The key my_key already exists."):
+ await S3Hook(aws_conn_id=None).aload_bytes(b"Content", "my_key",
"my_bucket")
+ client.put_object.assert_not_called()
+
+ @pytest.mark.asyncio
+
@async_mock.patch("airflow.providers.amazon.aws.hooks.s3.S3Hook.get_async_conn",
autospec=True)
+ async def test_aload_bytes_extra_args_do_not_leak_into_the_hook(self,
mock_get_async_conn):
+ client = self._mock_async_conn(mock_get_async_conn)
+ hook = S3Hook(aws_conn_id=None, extra_args={"ContentType":
"text/plain"}, requester_pays=True)
+
+ await hook.aload_bytes(
+ b"Content", "my_key", "my_bucket", replace=True, encrypt=True,
acl_policy="public-read"
+ )
+
+ client.put_object.assert_awaited_once_with(
+ Bucket="my_bucket",
+ Key="my_key",
+ Body=b"Content",
+ ContentType="text/plain",
+ ServerSideEncryption="AES256",
+ ACL="public-read",
+ RequestPayer="requester",
+ )
+ assert hook.extra_args == {"ContentType": "text/plain"}
+
def test_load_fileobj(self, s3_bucket, tmp_path):
hook = S3Hook()
path = tmp_path / "testfile"
