From 8cf11bf0ce2c578c8ee85221bec835dfd514e7d9 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Thu, 15 Oct 2026 00:37:23 +0900 Subject: [PATCH 1/4] Stop an append and a sync multipart copy at a failed part copy When an S3File write waits for one of the max_workers uploaded parts in flight, it now also checks the part copies of an append that have finished, and raises the error of a failed one, as it does for a failed uploaded part. The part copies are tracked in their own set, so they still do not count toward max_workers. S3FileSystem._finish_multipart_upload() now waits with FIRST_EXCEPTION before collecting the parts in part-number order, so that a failed part copy of a sync multipart copy is handled without waiting for the parts before it, and the parts queued after it are cancelled. Closes #1178 Co-Authored-By: Claude Opus 5.5 --- docs/filesystem.md | 10 ++-- pyathena/filesystem/s3.py | 42 ++++++++++----- tests/pyathena/filesystem/test_s3.py | 79 ++++++++++++++++++++++++++-- 3 files changed, 110 insertions(+), 21 deletions(-) diff --git a/docs/filesystem.md b/docs/filesystem.md index a1d127454..a814b1a72 100644 --- a/docs/filesystem.md +++ b/docs/filesystem.md @@ -91,11 +91,11 @@ this section apply to the compressed data. A multipart write keeps at most `max_workers` parts of the written data in flight and waits for one of them to finish before it uploads another. The parts take about `max_workers` × the block size of memory, in addition to the buffered data of the file. -If one of these parts has failed when the write waits, the write raises its error -instead of uploading the rest of the data, and aborts the multipart upload after the -parts being uploaded finish. A part that fails after the last wait raises its error -when the file is closed, or when the transaction commits, and the multipart upload is -aborted. +If one of these parts, or a part copied from the existing object in an append, has +failed when the write waits, the write raises its error instead of uploading the rest +of the data, and aborts the multipart upload after the parts being uploaded or copied +finish. A part that fails after the last wait raises its error when the file is +closed, or when the transaction commits, and the multipart upload is aborted. `max_workers` is an argument of `S3FileSystem`, `AioS3FileSystem`, `open`, `put`, and `pipe`. diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index f5ac8fb1e..3c759b2cf 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -9,7 +9,7 @@ import os.path import time from collections.abc import Callable, Mapping -from concurrent.futures import FIRST_COMPLETED, Future, as_completed, wait +from concurrent.futures import FIRST_COMPLETED, FIRST_EXCEPTION, Future, as_completed, wait from copy import deepcopy from datetime import datetime from io import BytesIO @@ -1831,7 +1831,8 @@ def _finish_multipart_upload( interrupted, the parts that have not started are cancelled, the running ones are waited for, and the multipart upload is aborted so that no incomplete upload or part is left behind, unless ``abort`` is - false. The original error is then re-raised. + false. The original error is then re-raised. A failed part is + handled as soon as it fails, without waiting for the parts before it. Args: upload: The multipart upload returned by creation. @@ -1847,6 +1848,14 @@ def _finish_multipart_upload( """ request_kwargs = request_kwargs or {} try: + _, pending = wait(futures, return_when=FIRST_EXCEPTION) + if pending: + # A part failed before the others finished. Its error is + # raised without waiting for the parts before it, so that the + # parts queued after it are cancelled. + for future in futures: + if future.done() and not future.cancelled() and future.exception() is not None: + future.result() # The futures are in part-number order. parts = [future.result() for future in futures] return self.core.complete_multipart_upload( @@ -2898,6 +2907,8 @@ def __init__( # The parts uploaded from the buffer that may not have finished. The # part copies of an append hold no data and are not counted. self._uploads_in_flight: set[Future[S3MultipartUploadPart]] = set() + # The part copies of an append that may not have finished. + self._copies_in_flight: set[Future[S3MultipartUploadPart]] = set() if append_info is not None: if append_data is not None: self.write(append_data) @@ -3056,6 +3067,7 @@ def _initiate_upload(self) -> None: # cancel it or wait for it before the abort. part: Future[S3MultipartUploadPart] = Future() self.multipart_upload_parts.append(part) + self._copies_in_flight.add(part) self._executor.submit_to( part, writer.upload_part_copy, @@ -3078,10 +3090,11 @@ def _upload_chunk(self, final: bool = False) -> bool: which uploads it with PutObject. Otherwise, the buffered data is uploaded as parts, and a part is submitted only while fewer than ``max_workers`` uploaded parts are in flight; the write waits for - one of them to finish first. A part found failed by that wait stops - the write with its error; one that fails afterwards is reported by - :meth:`commit`. If the upload of the chunk fails or is interrupted, - the file is closed without committing it. + one of them to finish first. A part, or a part copy of an append, + found failed by that wait stops the write with its error; one that + fails afterwards is reported by :meth:`commit`. If the upload of the + chunk fails or is interrupted, the file is closed without committing + it. Args: final: Whether this is the last flush of the file, after which @@ -3092,8 +3105,9 @@ def _upload_chunk(self, final: bool = False) -> bool: it does for any value other than False. Raises: - BaseException: The error of an uploaded part that the write found - failed while it waited for a part in flight. + BaseException: The error of an uploaded part, or of a part copy + of an append, that the write found failed while it waited for + a part in flight. RuntimeError: If no multipart upload is initialized. ValueError: If the data needs more parts than a multipart upload allows. @@ -3126,10 +3140,13 @@ def _upload_chunk(self, final: bool = False) -> bool: done, self._uploads_in_flight = wait( self._uploads_in_flight, return_when=FIRST_COMPLETED ) - for upload in done: - # A failed part would otherwise keep its data until - # commit() while the rest of the data is uploaded. - upload.result() + copied = {copy for copy in self._copies_in_flight if copy.done()} + self._copies_in_flight -= copied + for finished in done | copied: + # Otherwise the rest of the data is uploaded after a + # failed part or part copy, and a failed part keeps + # its data until commit(). + finished.result() # Tracked before it is scheduled, so that an interrupt while # it is being scheduled still lets discard() cancel it or wait # for it before the abort. @@ -3225,6 +3242,7 @@ def discard(self) -> None: self.multipart_upload = None self.multipart_upload_parts = [] self._uploads_in_flight = set() + self._copies_in_flight = set() def url(self, expiration: int = 3600, **kwargs) -> str: """Generate a presigned HTTP URL to read this file (if it already exists). diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index d112a618c..1af8d40cd 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -3532,9 +3532,11 @@ def upload_part(): release.wait(5) events.append("part 2 stored") - def wait_parts(futures): - release.set() - return wait(futures) + def wait_parts(futures, **kwargs): + if not kwargs: + # The cleanup waits for the running parts. + release.set() + return wait(futures, **kwargs) with ( ThreadPoolExecutor(max_workers=1) as executor, @@ -3550,8 +3552,51 @@ def wait_parts(futures): ) assert events == ["part 2 stored", "abort"] - waited.assert_called_once_with([failed, running]) + assert waited.call_args_list[-1] == mock.call([failed, running]) + assert pending.cancelled() + + def test_finish_multipart_upload_stops_at_failed_part(self): + # GH-1178: a failed part is handled without waiting for the parts + # before it, so that the parts queued after it are cancelled. The + # parts used to be collected in part-number order. + fs = self._make_fs() + upload = S3MultipartUpload({"Bucket": "bucket", "Key": "key", "UploadId": "uploadid"}) + fs.core.complete_multipart_upload = mock.MagicMock() + started = threading.Event() + release = threading.Event() + cleanup = [] + + def copy_part(): + started.set() + # Copying until the cleanup waits for it. + release.wait(5) + + def wait_parts(futures, **kwargs): + if not kwargs: + # Whether the earlier part was still copying when the cleanup + # started. + cleanup.append(running.done()) + release.set() + return wait(futures, **kwargs) + + failed: Future[SimpleNamespace] = Future() + failed.set_exception(PermissionError("Access Denied")) + with ( + ThreadPoolExecutor(max_workers=1) as executor, + mock.patch("pyathena.filesystem.s3.wait", side_effect=wait_parts), + ): + running = executor.submit(copy_part) + pending = executor.submit(SimpleNamespace) + assert started.wait(5) + with pytest.raises(PermissionError, match="Access Denied"): + fs._finish_multipart_upload(upload=upload, futures=[running, failed, pending]) + + assert cleanup == [False] assert pending.cancelled() + fs.core.complete_multipart_upload.assert_not_called() + fs._call.assert_called_once_with( + fs._client.abort_multipart_upload, Bucket="bucket", Key="key", UploadId="uploadid" + ) def test_finish_multipart_upload_does_not_wait_for_cancelled_parts(self): # GH-976: a cancelled part is not waited for, as nothing may @@ -7190,6 +7235,32 @@ def upload_part(**kwargs): fs.core.complete_multipart_upload.assert_not_called() assert f.multipart_upload is None + def test_append_failed_part_copy_while_waiting(self): + # GH-1178: a failed part copy of an append that a write finds while + # it waits for a part stops the write, as a failed uploaded part + # does. The part copies are not counted toward max_workers, and the + # write used to upload the rest of the data. + fs = self._make_append_fs(b"a" * 6) + fs.core.upload_part_copy.side_effect = PermissionError("Access Denied") + # One thread, so that the part copy finishes before the first part. + executor = mock.MagicMock(wraps=S3ThreadPoolExecutor(max_workers=1)) + f = S3File( + fs, "s3://bucket/key.txt", mode="ab", block_size=4, max_workers=1, executor=executor + ) + f.write(b"b" * 4) + with pytest.raises(PermissionError, match="Access Denied"): + f.write(b"b" * 4) + + assert [ + c.kwargs["part_number"] for c in executor.submit_to.call_args_list if "body" in c.kwargs + ] == [2] + fs._call.assert_called_once_with( + S3_CLIENT.abort_multipart_upload, Bucket="bucket", Key="key.txt", UploadId="uploadid" + ) + fs.core.complete_multipart_upload.assert_not_called() + assert f.closed + assert f.multipart_upload_parts == [] + def test_multipart_write_request_parameters(self): # GH-946: the parts receive the parameters of the file that they # accept, such as RequestPayer and SSE-C, as does the completion. From f135e16b99c518a08bd17a489e03c5fa8105c6c8 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Thu, 15 Oct 2026 01:01:27 +0900 Subject: [PATCH 2/4] Check part copies before each part and stop commit() at a failed part The part copies of an append are now checked before each part is submitted, not only when the write waits for a part, so that a failed part copy stops the write with any max_workers. S3File.commit() waits with FIRST_EXCEPTION before collecting the parts, as _finish_multipart_upload() does, which now iterates the futures that the wait returns as done. The tests that replace wait() tell the waits apart by return_when, and the write-file helper sets _copies_in_flight. Co-Authored-By: Claude Opus 5.5 --- docs/filesystem.md | 11 +-- pyathena/filesystem/s3.py | 65 +++++++++-------- tests/pyathena/filesystem/test_s3.py | 83 ++++++++++++++++++---- tests/pyathena/filesystem/test_s3_async.py | 5 +- 4 files changed, 118 insertions(+), 46 deletions(-) diff --git a/docs/filesystem.md b/docs/filesystem.md index a814b1a72..acc96f81f 100644 --- a/docs/filesystem.md +++ b/docs/filesystem.md @@ -91,11 +91,12 @@ this section apply to the compressed data. A multipart write keeps at most `max_workers` parts of the written data in flight and waits for one of them to finish before it uploads another. The parts take about `max_workers` × the block size of memory, in addition to the buffered data of the file. -If one of these parts, or a part copied from the existing object in an append, has -failed when the write waits, the write raises its error instead of uploading the rest -of the data, and aborts the multipart upload after the parts being uploaded or copied -finish. A part that fails after the last wait raises its error when the file is -closed, or when the transaction commits, and the multipart upload is aborted. +If one of these parts has failed when the write waits, or a part copied from the +existing object in an append has failed before the write uploads another part, the +write raises its error instead of uploading the rest of the data, and aborts the +multipart upload after the parts being uploaded or copied finish. A part that fails +later raises its error when the file is closed, or when the transaction commits, and +the multipart upload is aborted. `max_workers` is an argument of `S3FileSystem`, `AioS3FileSystem`, `open`, `put`, and `pipe`. diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index 3c759b2cf..ab3121f4c 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -1848,14 +1848,12 @@ def _finish_multipart_upload( """ request_kwargs = request_kwargs or {} try: - _, pending = wait(futures, return_when=FIRST_EXCEPTION) - if pending: - # A part failed before the others finished. Its error is - # raised without waiting for the parts before it, so that the - # parts queued after it are cancelled. - for future in futures: - if future.done() and not future.cancelled() and future.exception() is not None: - future.result() + done, _ = wait(futures, return_when=FIRST_EXCEPTION) + for future in done: + if not future.cancelled() and future.exception() is not None: + # Raised without waiting for the parts before it, so that + # the parts queued after it are cancelled. + future.result() # The futures are in part-number order. parts = [future.result() for future in futures] return self.core.complete_multipart_upload( @@ -3090,11 +3088,11 @@ def _upload_chunk(self, final: bool = False) -> bool: which uploads it with PutObject. Otherwise, the buffered data is uploaded as parts, and a part is submitted only while fewer than ``max_workers`` uploaded parts are in flight; the write waits for - one of them to finish first. A part, or a part copy of an append, - found failed by that wait stops the write with its error; one that - fails afterwards is reported by :meth:`commit`. If the upload of the - chunk fails or is interrupted, the file is closed without committing - it. + one of them to finish first. A part found failed by that wait, or a + part copy of an append found failed before a part is submitted, stops + the write with its error; one that fails afterwards is reported by + :meth:`commit`. If the upload of the chunk fails or is interrupted, + the file is closed without committing it. Args: final: Whether this is the last flush of the file, after which @@ -3105,9 +3103,9 @@ def _upload_chunk(self, final: bool = False) -> bool: it does for any value other than False. Raises: - BaseException: The error of an uploaded part, or of a part copy - of an append, that the write found failed while it waited for - a part in flight. + BaseException: The error of an uploaded part that the write found + failed while it waited for a part in flight, or of a part copy + of an append found failed before a part is submitted. RuntimeError: If no multipart upload is initialized. ValueError: If the data needs more parts than a multipart upload allows. @@ -3135,18 +3133,22 @@ def _upload_chunk(self, final: bool = False) -> bool: for part_number, data in writer.iter_parts( buffer, first_part_number=len(self.multipart_upload_parts) + 1 ): + if self._copies_in_flight: + copied = {copy for copy in self._copies_in_flight if copy.done()} + self._copies_in_flight -= copied + for copy in copied: + # Otherwise the rest of the data is uploaded after a + # failed part copy. + copy.result() if len(self._uploads_in_flight) >= self.max_workers: # Otherwise the parts waiting for a worker pile up in memory. done, self._uploads_in_flight = wait( self._uploads_in_flight, return_when=FIRST_COMPLETED ) - copied = {copy for copy in self._copies_in_flight if copy.done()} - self._copies_in_flight -= copied - for finished in done | copied: - # Otherwise the rest of the data is uploaded after a - # failed part or part copy, and a failed part keeps - # its data until commit(). - finished.result() + for upload in done: + # A failed part would otherwise keep its data until + # commit() while the rest of the data is uploaded. + upload.result() # Tracked before it is scheduled, so that an interrupt while # it is being scheduled still lets discard() cancel it or wait # for it before the abort. @@ -3173,11 +3175,12 @@ def commit(self) -> None: Creates an empty object if nothing was written, uploads the buffered data with PutObject if no multipart upload part was submitted, and - otherwise completes the multipart upload, which is aborted if the - completion fails or is interrupted. If the abort also fails, the - upload is kept so that :meth:`discard` can abort it. Invalidates the - cache of the path afterwards. Does nothing for a file whose failed - write dropped the written data. + otherwise completes the multipart upload, which is aborted if a part + or the completion fails or is interrupted. A failed part is handled + as soon as it fails, without waiting for the parts before it. If the + abort also fails, the upload is kept so that :meth:`discard` can + abort it. Invalidates the cache of the path afterwards. Does nothing + for a file whose failed write dropped the written data. Raises: FileExistsError: If an object was created at the path after the @@ -3206,6 +3209,12 @@ def commit(self) -> None: upload_id = cast(str, self.multipart_upload.upload_id) try: + done, _ = wait(self.multipart_upload_parts, return_when=FIRST_EXCEPTION) + for future in done: + if not future.cancelled() and future.exception() is not None: + # Raised without waiting for the parts before it, so + # that the parts queued after it are cancelled. + future.result() parts = [future.result() for future in self.multipart_upload_parts] self._get_multipart_writer().complete( parts, **self._get_request_kwargs("complete_multipart_upload") diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index 1af8d40cd..196215641 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -18,7 +18,7 @@ import uuid import weakref from base64 import b64encode -from concurrent.futures import Future, ThreadPoolExecutor, wait +from concurrent.futures import FIRST_COMPLETED, FIRST_EXCEPTION, Future, ThreadPoolExecutor, wait from datetime import UTC, datetime from itertools import chain from pathlib import Path @@ -3552,7 +3552,10 @@ def wait_parts(futures, **kwargs): ) assert events == ["part 2 stored", "abort"] - assert waited.call_args_list[-1] == mock.call([failed, running]) + assert waited.call_args_list == [ + mock.call([failed, running, pending], return_when=FIRST_EXCEPTION), + mock.call([failed, running]), + ] assert pending.cancelled() def test_finish_multipart_upload_stops_at_failed_part(self): @@ -6713,6 +6716,7 @@ def _make_write_file(data: bytes, autocommit: bool): # a part. file.max_workers = S3Core.MULTIPART_UPLOAD_MAX_PARTS file._uploads_in_flight = set() + file._copies_in_flight = set() file._executor = mock.MagicMock() return file @@ -7072,6 +7076,9 @@ def record_upload(future, fn, *args, **kwargs): executor.submit_to = record_upload # type: ignore[method-assign] def release_newest(futures, **kwargs): + if kwargs.get("return_when") != FIRST_COMPLETED: + # The wait of commit() for the parts. + return wait(futures, **kwargs) # Only the newest part in flight finishes, so a write that waits # for the oldest one would time out. gates[max(n for n, upload in uploads.items() if not upload.done())].set() @@ -7095,7 +7102,9 @@ def release_newest(futures, **kwargs): assert self._uploaded_object(fs, existing) == existing + b"".join(writes) # Each part was submitted once fewer than max_workers were in flight. assert in_flight == [min(n, max_workers) for n in range(1, len(uploads) + 1)] - assert waited.call_count == len(uploads) - max_workers + assert [c.kwargs.get("return_when") for c in waited.call_args_list].count( + FIRST_COMPLETED + ) == len(uploads) - max_workers def test_write_interrupted_while_waiting_for_part(self): # GH-1158: an interrupt while a write waits for a part closes the file @@ -7235,19 +7244,26 @@ def upload_part(**kwargs): fs.core.complete_multipart_upload.assert_not_called() assert f.multipart_upload is None - def test_append_failed_part_copy_while_waiting(self): - # GH-1178: a failed part copy of an append that a write finds while - # it waits for a part stops the write, as a failed uploaded part - # does. The part copies are not counted toward max_workers, and the - # write used to upload the rest of the data. + @pytest.mark.parametrize("max_workers", [1, 4]) + def test_append_failed_part_copy(self, max_workers): + # GH-1178: a failed part copy of an append stops the write before it + # uploads another part, whether or not the write waits for a part. + # The part copies are not counted toward max_workers, and the write + # used to upload the rest of the data. fs = self._make_append_fs(b"a" * 6) fs.core.upload_part_copy.side_effect = PermissionError("Access Denied") - # One thread, so that the part copy finishes before the first part. - executor = mock.MagicMock(wraps=S3ThreadPoolExecutor(max_workers=1)) + executor = mock.MagicMock(wraps=S3ThreadPoolExecutor(max_workers=max_workers)) f = S3File( - fs, "s3://bucket/key.txt", mode="ab", block_size=4, max_workers=1, executor=executor + fs, + "s3://bucket/key.txt", + mode="ab", + block_size=4, + max_workers=max_workers, + executor=executor, ) f.write(b"b" * 4) + # The part copy has failed before the next write. + assert wait(f.multipart_upload_parts[:1], timeout=5).done with pytest.raises(PermissionError, match="Access Denied"): f.write(b"b" * 4) @@ -7259,7 +7275,6 @@ def test_append_failed_part_copy_while_waiting(self): ) fs.core.complete_multipart_upload.assert_not_called() assert f.closed - assert f.multipart_upload_parts == [] def test_multipart_write_request_parameters(self): # GH-946: the parts receive the parameters of the file that they @@ -7495,6 +7510,50 @@ def test_commit_failure_and_discard(self, caplog, error, abort_fails): assert file.multipart_upload is None assert file.multipart_upload_parts == [] + def test_commit_stops_at_failed_part(self): + # GH-1178: commit() handles a failed part without waiting for the + # parts before it, so that the parts queued after it, such as the + # part copies of a large append, are cancelled. The parts used to be + # collected in part-number order. + file = self._make_multipart_write_file(b"x" * 16, autocommit=False) + started = threading.Event() + release = threading.Event() + cleanup = [] + + def copy_part(): + started.set() + # Copying until the cleanup waits for it. + release.wait(5) + + def wait_parts(futures, **kwargs): + if not kwargs: + # Whether the earlier part was still copying when the cleanup + # started. + cleanup.append(running.done()) + release.set() + return wait(futures, **kwargs) + + failed: Future[SimpleNamespace] = Future() + failed.set_exception(PermissionError("Access Denied")) + with ( + ThreadPoolExecutor(max_workers=1) as executor, + mock.patch("pyathena.filesystem.s3.wait", side_effect=wait_parts), + ): + running = executor.submit(copy_part) + pending = executor.submit(SimpleNamespace) + assert started.wait(5) + file.multipart_upload_parts = [running, failed, pending] + with pytest.raises(PermissionError, match="Access Denied"): + file.commit() + + assert cleanup == [False] + assert pending.cancelled() + file.fs.core.complete_multipart_upload.assert_not_called() + file.fs._call.assert_called_once_with( + S3_CLIENT.abort_multipart_upload, Bucket="bucket", Key="key.txt", UploadId="uploadid" + ) + assert file.multipart_upload is None + def test_commit_failure_and_interrupted_abort(self): # GH-945: if the abort after a failed completion is interrupted, the # interrupt propagates and the upload is kept so that discard() diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index 7ce28961e..0440d35e0 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -8,7 +8,7 @@ import urllib.parse import urllib.request import uuid -from concurrent.futures import wait +from concurrent.futures import FIRST_COMPLETED, wait from datetime import UTC, datetime from itertools import chain from pathlib import Path @@ -2405,6 +2405,9 @@ def record_upload(executor, future, fn, *args, **kwargs): submit_to(executor, future, fn, *args, **kwargs) def release_newest(futures, **kwargs): + if kwargs.get("return_when") != FIRST_COMPLETED: + # The wait of commit() for the parts. + return wait(futures, **kwargs) # Only the newest part in flight finishes, so a write that waits # for the oldest one would time out. gates[max(n for n, upload in uploads.items() if not upload.done())].set() From cfd6d3b97cea64cc33c116fd89fb778ddf35beda Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Thu, 15 Oct 2026 01:10:42 +0900 Subject: [PATCH 3/4] Hold the part copy until the first write returns in the append test The part copy of test_append_failed_part_copy could fail before the first write checked the part copies, so that the first write raised instead of the second. The copy now waits for an event that the test sets after the first write. Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/filesystem/test_s3.py | 21 ++++++++++++++++----- 1 file changed, 16 insertions(+), 5 deletions(-) diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index 196215641..25d2d4da2 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -7251,7 +7251,14 @@ def test_append_failed_part_copy(self, max_workers): # The part copies are not counted toward max_workers, and the write # used to upload the rest of the data. fs = self._make_append_fs(b"a" * 6) - fs.core.upload_part_copy.side_effect = PermissionError("Access Denied") + release = threading.Event() + + def upload_part_copy(**kwargs): + # Copying until the first write has returned. + assert release.wait(5) + raise PermissionError("Access Denied") + + fs.core.upload_part_copy.side_effect = upload_part_copy executor = mock.MagicMock(wraps=S3ThreadPoolExecutor(max_workers=max_workers)) f = S3File( fs, @@ -7261,11 +7268,15 @@ def test_append_failed_part_copy(self, max_workers): max_workers=max_workers, executor=executor, ) - f.write(b"b" * 4) - # The part copy has failed before the next write. - assert wait(f.multipart_upload_parts[:1], timeout=5).done - with pytest.raises(PermissionError, match="Access Denied"): + try: f.write(b"b" * 4) + release.set() + # The part copy has failed before the next write. + assert wait(f.multipart_upload_parts[:1], timeout=5).done + with pytest.raises(PermissionError, match="Access Denied"): + f.write(b"b" * 4) + finally: + release.set() assert [ c.kwargs["part_number"] for c in executor.submit_to.call_args_list if "body" in c.kwargs From 506372903b04b4c853ce788e09702559d61281d8 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Thu, 15 Oct 2026 01:17:31 +0900 Subject: [PATCH 4/4] Wait for the test's event without a timeout and shut down its executor The part copy of test_append_failed_part_copy waited for its event with a 5-second timeout, whose expiry would fail the first write. The test sets the event in a finally block, so the copy now waits without a timeout, and the block also shuts down the executor when an assertion fails. Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/filesystem/test_s3.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index 25d2d4da2..7eaa1608d 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -7254,8 +7254,9 @@ def test_append_failed_part_copy(self, max_workers): release = threading.Event() def upload_part_copy(**kwargs): - # Copying until the first write has returned. - assert release.wait(5) + # Copying until the first write has returned. The test always + # sets the event, so the wait needs no timeout. + release.wait() raise PermissionError("Access Denied") fs.core.upload_part_copy.side_effect = upload_part_copy @@ -7277,6 +7278,7 @@ def upload_part_copy(**kwargs): f.write(b"b" * 4) finally: release.set() + executor.shutdown() assert [ c.kwargs["part_number"] for c in executor.submit_to.call_args_list if "body" in c.kwargs