From bab08235f0050aebdc4cf2d85a8fcba08b6dc59b Mon Sep 17 00:00:00 2001
From: laughingman7743 <laughingman7743@gmail.com>
Date: Sat, 3 Oct 2026 18:25:05 +0900
Subject: [PATCH] Check the part limit in AioS3FileSystem transaction writes

#988 added _pipe_file_in_transaction() and _put_file_in_transaction()
to AioS3FileSystem while #968 changed S3FileSystem.pipe_file() and
put_file(), and the two were merged in parallel. Inside a transaction,
pipe() and put() therefore did not reject data that needs more than
10,000 blocks before uploading, and put() sent block_size to the S3 API
instead of open().

Apply the same check and the same block_size handling as the sync
methods.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 pyathena/filesystem/s3_async.py            | 17 ++++++++--
 tests/pyathena/filesystem/test_s3_async.py | 37 ++++++++++++++++++++++
 2 files changed, 52 insertions(+), 2 deletions(-)

diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py
index 7006c6e47..d98e827da 100644
--- a/pyathena/filesystem/s3_async.py
+++ b/pyathena/filesystem/s3_async.py
@@ -187,7 +187,12 @@ def _pipe_file_in_transaction(
         Raises:
             FileExistsError: If the mode is "create" and the path already
                 exists.
+            ValueError: If the data takes more than
+                ``MULTIPART_UPLOAD_MAX_PARTS`` blocks.
         """
+        block_size = kwargs.get("block_size") or self._sync_fs.default_block_size
+        # The size in bytes; the length of a memoryview counts its items.
+        self._sync_fs._check_multipart_upload_size(path, memoryview(value).nbytes, block_size)
         if mode == "create" and self._sync_fs.exists(path):
             raise FileExistsError(path)
         with self.open(path, "wb", **kwargs) as f:
@@ -211,6 +216,11 @@ def _put_file_in_transaction(self, lpath: str, rpath: str, callback, **kwargs) -
             rpath: S3 destination path (s3://bucket/key).
             callback: Progress callback for tracking upload progress.
             **kwargs: Additional S3 parameters (e.g., ContentType, StorageClass).
+                The ``block_size`` parameter of ``open()`` is also accepted.
+
+        Raises:
+            ValueError: If the file takes more than
+                ``MULTIPART_UPLOAD_MAX_PARTS`` blocks.
         """
         if os.path.isdir(lpath):
             return
@@ -218,14 +228,17 @@ def _put_file_in_transaction(self, lpath: str, rpath: str, callback, **kwargs) -
         if not key:
             return
 
-        callback.set_size(os.path.getsize(lpath))
+        size = os.path.getsize(lpath)
+        block_size = kwargs.pop("block_size", None) or self._sync_fs.default_block_size
+        self._sync_fs._check_multipart_upload_size(rpath, size, block_size)
+        callback.set_size(size)
         if "ContentType" not in kwargs:
             content_type, _ = mimetypes.guess_type(lpath)
             if content_type is not None:
                 kwargs["ContentType"] = content_type
 
         with (
-            self.open(rpath, "wb", s3_additional_kwargs=kwargs) as remote,
+            self.open(rpath, "wb", block_size=block_size, s3_additional_kwargs=kwargs) as remote,
             open(lpath, "rb") as local,
         ):
             while data := local.read(remote.blocksize):
diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py
index 9ba1e3cb7..2e6f44c31 100644
--- a/tests/pyathena/filesystem/test_s3_async.py
+++ b/tests/pyathena/filesystem/test_s3_async.py
@@ -236,6 +236,43 @@ def test_transaction_pipe_file_create_existing(self):
             fs.pipe_file("s3://bucket/key", b"data", mode="create")
         fs._sync_fs.exists.assert_called_once_with("s3://bucket/key")
 
+    @pytest.mark.parametrize("kwargs", [{"block_size": 4}, {}])
+    def test_transaction_pipe_put_file_exceeding_max_parts(self, tmp_path, kwargs):
+        # GH-953: in a transaction, as outside one, pipe_file() and put_file()
+        # reject data that does not fit in the maximum number of parts before
+        # opening the file.
+        fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True)
+        fs._sync_fs.MULTIPART_UPLOAD_MAX_PARTS = 3
+        fs._sync_fs.default_block_size = 4
+        fs._sync_fs._call = mock.MagicMock()
+        fs.open = mock.MagicMock()
+        local = tmp_path / "local"
+        local.write_bytes(b"a" * 13)
+
+        with fs.transaction:
+            with pytest.raises(ValueError, match="block_size"):
+                fs.pipe_file("s3://bucket/k1", b"a" * 13, **kwargs)
+            with pytest.raises(ValueError, match="block_size"):
+                fs.put_file(str(local), "s3://bucket/k2", **kwargs)
+        fs.open.assert_not_called()
+        fs._sync_fs._call.assert_not_called()
+
+    def test_transaction_put_file_block_size(self, tmp_path):
+        # In a transaction, put_file() passes block_size to open() instead of
+        # the S3 API, as outside one.
+        fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True)
+        fs.open = mock.MagicMock()
+        fs.open.return_value.__enter__.return_value.blocksize = 8
+        local = tmp_path / "local"
+        local.write_bytes(b"a" * 13)
+
+        with fs.transaction:
+            fs.put_file(str(local), "s3://bucket/key", block_size=8)
+
+        fs.open.assert_called_once_with(
+            "s3://bucket/key", "wb", block_size=8, s3_additional_kwargs={}
+        )
+
     def test_touch_sync_wrapper(self):
         # GH-977: touch() used to be fsspec's open()-based default, which
         # dropped the PutObject parameters and returned None.