diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index 7868fc201..cc603331a 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -3,18 +3,15 @@ import contextlib import gc import gzip +import io import os import tempfile import threading import time -import urllib.parse -import urllib.request +import traceback import uuid import weakref from concurrent.futures import Future, ThreadPoolExecutor, wait -from datetime import UTC, datetime -from itertools import chain -from pathlib import Path from types import SimpleNamespace from unittest import mock @@ -22,22 +19,23 @@ import fsspec import pytest from botocore.stub import Stubber -from fsspec import Callback +from fsspec.callbacks import Callback -from pyathena.filesystem.s3 import S3File, S3FileSystem +from pyathena.filesystem.s3 import S3FileSystem from pyathena.filesystem.s3_async import AioS3File, AioS3FileSystem, AioS3StreamedFile from pyathena.filesystem.s3_core import S3Core, S3MultipartCopyPlan from pyathena.filesystem.s3_executor import S3AioThreadPool, S3Executor from pyathena.filesystem.s3_object import ( + S3Metadata, S3MultipartUpload, S3MultipartUploadPart, S3Object, S3ObjectType, - S3StorageClass, + S3ObjectVersion, ) from pyathena.filesystem.s3_path import S3Path from pyathena.filesystem.s3_writer import S3MultipartStateError -from pyathena.util import RetryConfig, stop_retries_on +from pyathena.util import RetryConfig from tests import ENV from tests.pyathena.conftest import connect from tests.pyathena.util import ( @@ -45,12 +43,28 @@ MULTIPART_COPY_BLOCK_SIZE, MULTIPART_COPY_KWARGS, ThrottledOperation, + await_bounded, overlapping_calls, retried, stub_multipart_copy, + wait_until, ) +class _RecordingCallback(Callback): + """A progress callback that records its updates.""" + + def __init__(self): + super().__init__() + self.events = [] + + def set_size(self, size): + self.events.append(("size", size)) + + def relative_update(self, inc=1): + self.events.append(("update", inc)) + + @pytest.fixture(scope="class") def register_async_filesystem(): fsspec.register_implementation( @@ -63,49 +77,6 @@ def register_async_filesystem(): @pytest.mark.usefixtures("register_async_filesystem") class TestAioS3FileSystem: - @pytest.mark.parametrize("max_workers", [1, 4]) - @pytest.mark.asyncio - async def test_copy_object_with_multipart_upload_part_sizes(self, max_workers): - # GH-951: the parts are within the S3 part size limits whatever the - # number of workers; a single worker used to copy the whole object - # as one part larger than 5 GiB. - fs = AioS3FileSystem( - connection=mock.MagicMock(), max_workers=max_workers, skip_instance_cache=True - ) - sync_fs = fs._sync_fs - sync_fs.core.create_multipart_upload = mock.MagicMock( - return_value=S3MultipartUpload( - {"Bucket": "bucket", "Key": "dst", "UploadId": "uploadid"} - ) - ) - sync_fs.core.upload_part_copy = mock.MagicMock( - side_effect=lambda **kw: SimpleNamespace(etag='"e"', part_number=kw["part_number"]) - ) - sync_fs.core.complete_multipart_upload = mock.MagicMock() - # The HeadObject of the source. - sync_fs._call = sync_fs._core.call = mock.MagicMock( - return_value={"ContentLength": 5 * 2**30 + 2**20} - ) - - await fs._copy_object_with_multipart_upload( - S3Path("bucket", "src"), - S3Path("bucket", "dst"), - # Copy without reading the metadata, tags and annotations of the - # source (GH-973). - MetadataDirective="REPLACE", - TaggingDirective="REPLACE", - AnnotationDirective="EXCLUDE", - ) - - parts = sorted( - (c.kwargs["part_number"], c.kwargs["range_"]) - for c in sync_fs.core.upload_part_copy.call_args_list - ) - assert parts == [ - (1, (0, 5 * 2**29 + 2**19)), - (2, (5 * 2**29 + 2**19, 5 * 2**30 + 2**20)), - ] - @staticmethod async def _multipart_copy(fs=None, **kwargs): # max_workers=1 runs the stubbed requests in a deterministic order. @@ -119,11 +90,13 @@ async def _multipart_copy(fs=None, **kwargs): with Stubber(fs._sync_fs._client) as stubber: stub_multipart_copy(stubber, **kwargs) try: - await fs._copy_object_with_multipart_upload( - S3Path("bucket", "src"), - S3Path("bucket", "dst"), - block_size=MULTIPART_COPY_BLOCK_SIZE, - **MULTIPART_COPY_KWARGS, + await await_bounded( + fs._copy_object_with_multipart_upload( + S3Path("bucket", "src"), + S3Path("bucket", "dst"), + block_size=MULTIPART_COPY_BLOCK_SIZE, + **MULTIPART_COPY_KWARGS, + ) ) finally: stubber.assert_no_pending_responses() @@ -133,30 +106,31 @@ async def _multipart_copy(fs=None, **kwargs): async def test_copy_object_with_multipart_upload_small_head_object_size(self, size): # GH-973: see # TestS3FileSystem.test_copy_object_with_multipart_upload_small_head_object_size. - fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) - sync_fs = fs._sync_fs - sync_fs._call = sync_fs._core.call = mock.MagicMock( - return_value={"ContentLength": size, "VersionId": "v1"} - ) - sync_fs.core.copy_object = mock.MagicMock() - sync_fs.core.create_multipart_upload = mock.MagicMock() - - await fs._copy_object_with_multipart_upload( - S3Path("bucket", "src"), - S3Path("bucket", "dst"), - ContentType="text/csv", - RequestPayer="requester", - ) - - sync_fs.core.copy_object.assert_called_once_with( - S3Path("bucket", "src", "v1"), - S3Path("bucket", "dst"), - ContentType="text/csv", - RequestPayer="requester", + # The async copy has its own single-request branch: the reported + # version is copied with one CopyObject and nothing else is sent. + fs = AioS3FileSystem( + key="dummy", secret="dummy", region_name="us-east-1", skip_instance_cache=True ) - sync_fs.core.create_multipart_upload.assert_not_called() - # Only HeadObject; the tags are not read for the multipart upload. - assert sync_fs._call.call_count == 1 + with Stubber(fs._sync_fs._client) as stubber: + stubber.add_response("head_object", {"ContentLength": size, "VersionId": "v1"}, None) + stubber.add_response( + "copy_object", + {}, + { + "CopySource": {"Bucket": "bucket", "Key": "src", "VersionId": "v1"}, + "Bucket": "bucket", + "Key": "dst", + "ContentType": "text/csv", + "RequestPayer": "requester", + }, + ) + await fs._copy_object_with_multipart_upload( + S3Path("bucket", "src"), + S3Path("bucket", "dst"), + ContentType="text/csv", + RequestPayer="requester", + ) + stubber.assert_no_pending_responses() @pytest.mark.asyncio async def test_copy_object_with_multipart_upload_copies_source(self): @@ -164,12 +138,6 @@ async def test_copy_object_with_multipart_upload_copies_source(self): # TestS3FileSystem.test_copy_object_with_multipart_upload_copies_source. await self._multipart_copy() - @pytest.mark.asyncio - async def test_copy_object_with_multipart_upload_failed_listing(self): - # GH-973: nothing is written when the annotations cannot be listed. - with pytest.raises(PermissionError): - await self._multipart_copy(fail_list=True) - @pytest.mark.asyncio async def test_copy_object_with_multipart_upload_failed_part(self): # GH-973: a failed part copy aborts the upload, as in S3FileSystem. @@ -229,17 +197,24 @@ async def test_copy_object_with_multipart_upload_waits_for_running_parts(self): ) ) events = [] - failed = threading.Event() + # Part 1 retries until the failure of part 2 stops it, so that it + # runs until the copy handles the failure. + operation = ThrottledOperation(operation_name="UploadPartCopy") + throttled = retried(operation, RetryConfig(attempt=2, multiplier=EVENT_TIMEOUT)) def upload_part_copy(**kw): part_number = kw["part_number"] events.append(f"start {part_number}") if part_number == 1: - failed.wait(5) - time.sleep(0.05) - events.append("end 1") - return SimpleNamespace(etag='"e"', part_number=part_number) - failed.set() + try: + return throttled() + finally: + # Ends a little after the stop, so that an abort that + # does not wait for it is recorded first. + time.sleep(0.05) + events.append("end 1") + # Fails once part 1 is running, which the copy cannot cancel. + assert operation.started.wait(EVENT_TIMEOUT) raise OSError("part failed") sync_fs.core.upload_part_copy = mock.MagicMock(side_effect=upload_part_copy) @@ -252,17 +227,20 @@ def upload_part_copy(**kw): ) with pytest.raises(OSError, match="part failed"): - await fs._copy_object_with_multipart_upload( - S3Path("bucket", "src"), - S3Path("bucket", "dst"), - MetadataDirective="REPLACE", - TaggingDirective="REPLACE", - AnnotationDirective="EXCLUDE", + await await_bounded( + fs._copy_object_with_multipart_upload( + S3Path("bucket", "src"), + S3Path("bucket", "dst"), + MetadataDirective="REPLACE", + TaggingDirective="REPLACE", + AnnotationDirective="EXCLUDE", + ) ) # Part 3 waits for a worker and is not started after the failure. assert sorted(events[:2]) == ["start 1", "start 2"] assert events[2:] == ["end 1", "abort"] + assert operation.calls == 1 sync_fs.core.complete_multipart_upload.assert_not_called() @pytest.mark.asyncio @@ -278,10 +256,28 @@ async def test_copy_object_with_multipart_upload_max_workers(self): {"Bucket": "bucket", "Key": "dst", "UploadId": "uploadid"} ) ) - call, state = overlapping_calls(3) + condition = threading.Condition() + state = {"calls": 0, "active": 0, "peak": 0} + # The first three part copies run at once. + first_three = threading.Barrier(3, timeout=EVENT_TIMEOUT) def upload_part_copy(**kw): - call() + with condition: + state["calls"] += 1 + number = state["calls"] + state["active"] += 1 + state["peak"] = max(state["peak"], state["active"]) + condition.notify_all() + try: + if number <= 3: + first_three.wait() + # Waits briefly for a fourth overlapping call, which only a + # copy that exceeds max_workers runs. + with condition: + condition.wait_for(lambda: state["active"] > 3, timeout=0.5) + finally: + with condition: + state["active"] -= 1 return SimpleNamespace(etag='"e"', part_number=kw["part_number"]) sync_fs.core.upload_part_copy = mock.MagicMock(side_effect=upload_part_copy) @@ -364,6 +360,193 @@ async def test_blocking_calls_bypass_default_executor(self): assert info == {"name": "bucket/key"} + @staticmethod + def _serve_keys(fs, keys): + # Answer the S3 requests of the filesystem from a set of keys of + # "bucket", which is returned, and record them by operation name. The + # requests may come from several threads at once. + store = set(keys) + requests = [] + lock = threading.Lock() + + def call(method, **request): + with lock: + return answer(method if isinstance(method, str) else method.__name__, request) + + def answer(name, request): + requests.append((name, request)) + key = request.get("Key") + if name == "head_object": + if key not in store: + raise FileNotFoundError(key) + return {"ContentLength": 2, "ETag": '"e"'} + if name == "get_object": + return {"Body": io.BytesIO(b"aa"), "ContentLength": 2, "ETag": '"e"'} + if name == "list_objects_v2": + prefix, delimiter = request.get("Prefix", ""), request.get("Delimiter") + listed = sorted(k for k in list(store) if k.startswith(prefix)) + directories = sorted( + { + prefix + k[len(prefix) :].split(delimiter)[0] + delimiter + for k in listed + if delimiter and delimiter in k[len(prefix) :] + } + ) + files = [k for k in listed if not delimiter or delimiter not in k[len(prefix) :]] + return { + "Contents": [{"Key": k, "Size": 2, "ETag": '"e"'} for k in files], + "CommonPrefixes": [{"Prefix": d} for d in directories], + "KeyCount": len(files) + len(directories), + } + if name == "copy_object" and request["CopySource"]["Key"] not in store: + raise FileNotFoundError(request["CopySource"]["Key"]) + if name in {"put_object", "copy_object"}: + store.add(key) + return {"ETag": '"e"'} + if name == "delete_object": + store.discard(key) + if name == "delete_objects": + deleted = [o["Key"] for o in request["Delete"]["Objects"]] + store.difference_update(deleted) + return {"Deleted": [{"Key": k} for k in deleted]} + if name == "get_object_tagging": + return {"TagSet": [{"Key": "a", "Value": "1"}]} + if name == "list_object_versions": + return { + "Versions": [{"Key": "key", "VersionId": "v1", "IsLatest": False}], + "DeleteMarkers": [{"Key": "key", "VersionId": "m1", "IsLatest": True}], + } + return {} + + fs._sync_fs._call = fs._sync_fs._core.call = call + return store, requests + + @pytest.mark.parametrize( + ("method", "sync_method", "args", "kwargs"), + [ + ( + "_info", + "info", + ("s3://bucket/key?versionId=v1",), + {"ExpectedBucketOwner": "111122223333"}, + ), + ("_ls", "ls", ("s3://bucket/dir",), {"detail": True}), + ( + "_cat_file", + "cat_file", + ("s3://bucket/key",), + {"start": 1, "end": 2, "ExpectedBucketOwner": "111122223333"}, + ), + ( + "_exists", + "exists", + ("s3://bucket/key?versionId=v1",), + {"ExpectedBucketOwner": "111122223333"}, + ), + ("_rm_file", "rm_file", ("s3://bucket/key?versionId=v1",), {}), + ( + "_pipe_file", + "pipe_file", + ("s3://bucket/new", b"a"), + {"mode": "create", "ContentType": "text/csv"}, + ), + ( + "_put_file", + "put_file", + ("local-source", "s3://bucket/new"), + {"mode": "create", "callback": "recording", "ContentType": "text/csv"}, + ), + ( + "_get_file", + "get_file", + ("s3://bucket/key?versionId=v1", "local-destination"), + {"callback": "recording", "ExpectedBucketOwner": "111122223333"}, + ), + ("_touch", "touch", ("s3://bucket/new",), {"ContentType": "text/csv"}), + # An existing object is not replaced. + ("_touch", "touch", ("s3://bucket/key",), {"truncate": False}), + ("_find", "find", ("s3://bucket/dir",), {"maxdepth": 1, "withdirs": True}), + ( + "_mv", + "mv", + ("s3://bucket/dir", "s3://bucket/out"), + {"recursive": True, "maxdepth": 1, "MetadataDirective": "REPLACE"}, + ), + ( + "metadata", + "metadata", + ("s3://bucket/key?versionId=v1",), + {"ExpectedBucketOwner": "111122223333"}, + ), + ("get_tags", "get_tags", ("s3://bucket/key?versionId=v1",), {}), + ("put_tags", "put_tags", ("s3://bucket/key", {"b": "2"}), {"mode": "m"}), + ("chmod", "chmod", ("s3://bucket/dir", "private"), {"recursive": True}), + ( + "object_version_info", + "object_version_info", + ("s3://bucket/key",), + {"delete_markers": True}, + ), + ], + ) + @pytest.mark.asyncio + async def test_delegations_send_the_requests_of_the_sync_methods( + self, tmp_path, method, sync_method, args, kwargs + ): + # The methods that run the synchronous filesystem's methods pass their + # arguments through: each sends the requests, and returns the result, + # of the synchronous method called with the same arguments. + keys = {"key", "dir/a", "dir/sub/b"} + source = tmp_path / "source" + source.write_bytes(b"a") + + def comparable(value): + # The metadata and the listed versions compare by their public + # properties. + if isinstance(value, list): + return [comparable(v) for v in value] + if isinstance(value, (S3Metadata, S3ObjectVersion)): + cls = type(value) + properties = [n for n in dir(cls) if isinstance(getattr(cls, n), property)] + return cls, {n: getattr(value, n) for n in properties} + return value + + outcomes = [] + for run, name in (("sync", sync_method), ("async", method)): + fs = AioS3FileSystem( + key="dummy", secret="dummy", region_name="us-east-1", skip_instance_cache=True + ) + store, requests = self._serve_keys(fs, keys) + (tmp_path / run).mkdir() + local = { + "local-source": str(source), + "local-destination": str(tmp_path / run / "downloaded"), + } + call_args = [local.get(a, a) if isinstance(a, str) else a for a in args] + callback = _RecordingCallback() + call_kwargs = {k: callback if v == "recording" else v for k, v in kwargs.items()} + try: + if run == "sync": + result = getattr(fs._sync_fs, name)(*call_args, **call_kwargs) + else: + result = getattr(fs, name)(*call_args, **call_kwargs) + if asyncio.iscoroutine(result): + result = await await_bounded(result) + except Exception as e: + result = type(e) + downloaded = {p.name: p.read_bytes() for p in (tmp_path / run).iterdir()} + outcomes.append( + ( + comparable(result), + sorted(map(repr, requests)), + sorted(store), + downloaded, + callback.events, + ) + ) + assert outcomes[0][1], "no request was sent" + assert outcomes[1] == outcomes[0] + @pytest.mark.asyncio async def test_max_workers_none(self): # None takes the default of ThreadPoolExecutor for the threads of the @@ -372,12 +555,41 @@ async def test_max_workers_none(self): fs = AioS3FileSystem( connection=mock.MagicMock(), max_workers=None, skip_instance_cache=True ) - fs._sync_fs.info = mock.MagicMock(return_value={"name": "bucket/key"}) + sync_fs = fs._sync_fs + sync_fs.info = mock.MagicMock(return_value={"name": "bucket/key"}) assert await fs._info("s3://bucket/key") == {"name": "bucket/key"} - # The files and copies take the same default. - cpus = getattr(os, "process_cpu_count", os.cpu_count)() - assert fs._sync_fs.max_workers == fs._thread_pool.max_workers == min(32, (cpus or 1) + 4) + # A multipart copy and a file without max_workers of their own. + sync_fs.core.create_multipart_upload = mock.MagicMock( + return_value=S3MultipartUpload( + {"Bucket": "bucket", "Key": "dst", "UploadId": "uploadid"} + ) + ) + sync_fs.core.upload_part_copy = mock.MagicMock( + side_effect=lambda **kw: SimpleNamespace(etag='"e"', part_number=kw["part_number"]) + ) + sync_fs.core.complete_multipart_upload = mock.MagicMock() + sync_fs.core.put_object = mock.MagicMock() + # The HeadObject of the source, with 2 parts of the default block size. + sync_fs._call = sync_fs._core.call = mock.MagicMock( + return_value={"ContentLength": 2 * S3Core.MULTIPART_UPLOAD_MAX_PART_SIZE} + ) + await fs._copy_object_with_multipart_upload( + S3Path("bucket", "src"), + S3Path("bucket", "dst"), + MetadataDirective="REPLACE", + TaggingDirective="REPLACE", + AnnotationDirective="EXCLUDE", + ) + assert sync_fs.core.upload_part_copy.call_count == 2 + sync_fs.core.complete_multipart_upload.assert_called_once() + + def write(): + with fs.open("s3://bucket/key", "wb") as f: + f.write(b"a") + + await asyncio.to_thread(write) + sync_fs.core.put_object.assert_called_once() @pytest.mark.parametrize("cancellations", [1, 2]) @pytest.mark.asyncio @@ -681,28 +893,6 @@ def record_completion(fn, *args, **kwargs): assert events == ["abort"] - @pytest.mark.asyncio - async def test_cancelled_call_stops_retries(self): - # GH-1203: cancelling a call stops the retries of its request in the - # thread of the filesystem, whose backoff wait would take at least - # EVENT_TIMEOUT seconds. - fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) - operation = ThrottledOperation(block=True, operation_name="HeadObject") - fs._sync_fs.info = retried(operation, RetryConfig(attempt=2, multiplier=EVENT_TIMEOUT)) - task = asyncio.ensure_future(fs._info("s3://bucket/key")) - try: - assert await asyncio.to_thread(operation.started.wait, EVENT_TIMEOUT) - task.cancel() - with pytest.raises(asyncio.CancelledError): - await task - finally: - operation.release.set() - started = time.monotonic() - # Waits for the thread of the cancelled call. - await asyncio.to_thread(fs._thread_pool.shutdown, True) - assert time.monotonic() - started < EVENT_TIMEOUT - assert operation.calls == 1 - @pytest.mark.asyncio async def test_copy_object_with_multipart_upload_failed_part_stops_running_parts(self): # GH-1203: a failed part copy stops the retries of the running ones @@ -733,12 +923,11 @@ def upload_part_copy(part_number, **kwargs): return SimpleNamespace(etag='"e"', part_number=part_number) sync_fs.core.upload_part_copy = mock.MagicMock(side_effect=upload_part_copy) - started = time.monotonic() + # The retries stop before the bound of the wait. with pytest.raises(PermissionError, match="Access Denied"): - await fs._copy_object_with_multipart_upload( - plan.source, plan.destination, max_workers=2 + await await_bounded( + fs._copy_object_with_multipart_upload(plan.source, plan.destination, max_workers=2) ) - assert time.monotonic() - started < EVENT_TIMEOUT assert operation.calls == 1 sync_fs._abort_multipart_upload.assert_called_once() @@ -814,10 +1003,11 @@ def copy_object_annotation(name, *args, **kwargs): return throttled() sync_fs.core.copy_object_annotation = mock.MagicMock(side_effect=copy_object_annotation) - started = time.monotonic() + # The retries stop before the bound of the wait. with pytest.raises(PermissionError, match="Access Denied"): - await fs._copy_object_with_multipart_upload(plan.source, plan.destination) - assert time.monotonic() - started < EVENT_TIMEOUT + await await_bounded( + fs._copy_object_with_multipart_upload(plan.source, plan.destination) + ) assert operation.calls == 1 @pytest.mark.asyncio @@ -844,18 +1034,20 @@ async def test_copy_object_with_multipart_upload_failed_annotation_skips_queued_ return_value=SimpleNamespace(version_id=None, etag='"e"') ) + loop = asyncio.get_running_loop() + def copy_object_annotation(name, *args, **kwargs): - # Fails once the other copies wait for the only thread. - deadline = time.monotonic() + EVENT_TIMEOUT - while fs._thread_pool._get_executor()._work_queue.qsize() < 2: - assert time.monotonic() < deadline - time.sleep(0.01) + # The tasks of the three copies start in one pass of the event + # loop, in which each submits its copy; a callback scheduled now + # runs after that pass, so the other copies wait for the only + # thread when this one fails. + asyncio.run_coroutine_threadsafe(asyncio.sleep(0), loop).result(EVENT_TIMEOUT) raise PermissionError("Access Denied") sync_fs.core.copy_object_annotation = mock.MagicMock(side_effect=copy_object_annotation) with pytest.raises(PermissionError, match="Access Denied"): - await fs._copy_object_with_multipart_upload( - plan.source, plan.destination, max_workers=3 + await await_bounded( + fs._copy_object_with_multipart_upload(plan.source, plan.destination, max_workers=3) ) sync_fs.core.copy_object_annotation.assert_called_once() @@ -1010,28 +1202,6 @@ async def main(): assert events == expected - @pytest.mark.parametrize( - "block_size", - [ - S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE - 1, - S3Core.MULTIPART_UPLOAD_MAX_PART_SIZE + 1, - ], - ) - @pytest.mark.asyncio - async def test_copy_object_with_multipart_upload_invalid_block_size(self, block_size): - # GH-926: the message states the accepted range. - fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) - fs._sync_fs._call = fs._sync_fs._core.call = mock.MagicMock() - - with pytest.raises( - ValueError, - match=r"between 5 MiB \(5242880 bytes\) and 5 GiB \(5368709120 bytes\), inclusive", - ): - await fs._copy_object_with_multipart_upload( - S3Path("bucket", "src"), S3Path("bucket", "dst"), block_size=block_size - ) - fs._sync_fs._call.assert_not_called() - @pytest.mark.parametrize("commit", [True, False]) def test_transaction_pipe_put_file(self, tmp_path, commit): # GH-977: pipe_file() and put_file() join the transaction of this @@ -1212,25 +1382,6 @@ def test_transaction_pipe_put_file_exceeding_max_parts(self, tmp_path, 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 and max_workers to - # open() instead of the S3 API, as outside one. - fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) - fs.open = mock.MagicMock() - 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, - max_workers=fs._sync_fs.max_workers, - 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. @@ -1252,8 +1403,8 @@ def test_touch_sync_wrapper(self): @pytest.mark.asyncio async def test_rm_requests(self): # GH-971: _rm() sent the keys of every bucket to the bucket of the - # first path, dropped the DeleteObjects parameters, ignored per-key - # errors and emptied a bucket path. + # first path and dropped the DeleteObjects parameters. The async + # filesystem sends the request of each bucket in parallel. fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) sync_fs = fs._sync_fs sync_fs._call = sync_fs._core.call = mock.MagicMock(return_value={}) @@ -1267,23 +1418,6 @@ async def test_rm_requests(self): ("b2", [{"Key": "b"}], "111122223333"), ] - sync_fs._call.reset_mock() - with pytest.raises(ValueError, match="Cannot delete the bucket"): - await fs._rm("s3://bucket", recursive=True) - sync_fs._call.assert_not_called() - - sync_fs._call.return_value = { - "Errors": [{"Key": "locked", "Code": "AccessDenied", "Message": "Access Denied"}] - } - with pytest.raises(OSError, match=r"bucket/locked \(AccessDenied: Access Denied\)"): - await fs._rm("s3://bucket/locked") - - sync_fs._call.side_effect = PermissionError("Access Denied") - sync_fs.dircache["bucket/a"] = [] - with pytest.raises(PermissionError, match="Access Denied"): - await fs._rm("s3://bucket/a") - assert "bucket/a" not in sync_fs.dircache - @pytest.mark.asyncio async def test_rm_request_error_keeps_errors(self): fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) @@ -1345,46 +1479,75 @@ async def wait_invalidated(path): @pytest.mark.asyncio async def test_rm_maxdepth(self): - # GH-962: _rm() did not pass maxdepth when expanding the path. - fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) - sync_fs = fs._sync_fs - sync_fs._call = sync_fs._core.call = mock.MagicMock(return_value={}) - sync_fs.find = mock.MagicMock(return_value=["bucket/dir/a"]) - sync_fs.exists = mock.MagicMock(return_value=True) - - # batch_size is part of fsspec's async _rm() signature. - await fs._rm("s3://bucket/dir", recursive=True, maxdepth=1, batch_size=10) - sync_fs.find.assert_called_once_with("bucket/dir", maxdepth=1, withdirs=True, detail=False) - (call,) = sync_fs._call.call_args_list - assert call.kwargs["Delete"]["Objects"] == [{"Key": "dir"}, {"Key": "dir/a"}] + # GH-962: _rm() did not pass maxdepth when expanding the path, so a + # key below maxdepth was deleted too. + fs = AioS3FileSystem( + key="dummy", secret="dummy", region_name="us-east-1", skip_instance_cache=True + ) + with Stubber(fs._sync_fs._client) as stubber: + # One level of the directory: dir/sub/deep is not listed. + stubber.add_response( + "list_objects_v2", + { + "Contents": [{"Key": "dir/a"}], + "CommonPrefixes": [{"Prefix": "dir/sub/"}], + "KeyCount": 2, + }, + {"Bucket": "bucket", "Prefix": "dir/", "Delimiter": "/"}, + ) + stubber.add_response( + "delete_objects", + {}, + { + "Bucket": "bucket", + "Delete": { + "Objects": [{"Key": "dir"}, {"Key": "dir/a"}, {"Key": "dir/sub"}], + "Quiet": True, + }, + }, + ) + # batch_size is part of fsspec's async _rm() signature. + await fs._rm("s3://bucket/dir", recursive=True, maxdepth=1, batch_size=10) + stubber.assert_no_pending_responses() @pytest.mark.parametrize(("mode", "open_mode"), [("overwrite", "wb"), ("create", "xb")]) def test_put_file_in_transaction_open_parameters(self, tmp_path, mode, open_mode): - # GH-969: the open() parameters of put_file() go to open(), and the - # other parameters, also in s3_additional_kwargs, to S3. - # GH-972: fsspec's mode argument selects the mode of the file. + # GH-969: in a transaction, the open() parameters of put_file() go to + # open(), and the other parameters, also in s3_additional_kwargs, to + # S3. GH-972: fsspec's mode argument selects the mode of the file. fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) fs.open = mock.MagicMock() lpath = tmp_path / "data.csv" lpath.write_bytes(b"a") - fs._put_file_in_transaction( - str(lpath), - "s3://bucket/key", - Callback(), - mode, - block_size=S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE, - max_workers=2, - StorageClass="STANDARD_IA", - ) - - fs.open.assert_called_once_with( - "s3://bucket/key", - open_mode, - block_size=S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE, - max_workers=2, - s3_additional_kwargs={"StorageClass": "STANDARD_IA", "ContentType": "text/csv"}, - ) + with fs.transaction: + fs.put_file( + str(lpath), + "s3://bucket/key", + mode=mode, + block_size=S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE, + max_workers=2, + StorageClass="STANDARD_IA", + ) + # Without them, the file gets the defaults of the filesystem. + fs.put_file(str(lpath), "s3://bucket/key", mode=mode) + + assert fs.open.call_args_list == [ + mock.call( + "s3://bucket/key", + open_mode, + block_size=S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE, + max_workers=2, + s3_additional_kwargs={"StorageClass": "STANDARD_IA", "ContentType": "text/csv"}, + ), + mock.call( + "s3://bucket/key", + open_mode, + block_size=fs._sync_fs.default_block_size, + max_workers=fs._sync_fs.max_workers, + s3_additional_kwargs={"ContentType": "text/csv"}, + ), + ] @pytest.mark.asyncio async def test_cp_file_directory(self): @@ -1410,42 +1573,48 @@ async def test_copy_version(self): "bucket/b?versionId=v1", "s3://bucket/d/b", RequestPayer="requester" ) - @pytest.mark.asyncio - async def test_get_version_path_destination(self, tmp_path): - # A Path destination is paired too: the version is downloaded to the - # file, not into a directory of that name. - fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) - fs._sync_fs.isdir = mock.MagicMock(return_value=False) - fs._get_file = mock.AsyncMock() - - await fs._get("s3://bucket/b?versionId=v1", tmp_path / "out.bin") - assert fs._get_file.await_args.args[:2] == ( - "bucket/b?versionId=v1", - (tmp_path / "out.bin").as_posix(), - ) - @pytest.mark.asyncio async def test_expand_path_glob_lists_stem_prefix(self): # Unversioned globs keep fsspec's async expansion, which lists only # the keys that start with the stem before the first wildcard. fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) - fs._sync_fs._find = mock.MagicMock(return_value=[]) + sync_fs = fs._sync_fs + sync_fs._call = sync_fs._core.call = mock.MagicMock( + side_effect=lambda method, **request: {"Contents": [], "KeyCount": 0} + ) with pytest.raises(FileNotFoundError): await fs._expand_path("s3://bucket/reports/2026-*.csv") - assert fs._sync_fs._find.call_args.kwargs["prefix"] == "2026-" - + listings = [ + c.kwargs + for c in sync_fs._call.call_args_list + if c.args[0] is sync_fs._client.list_objects_v2 + ] + assert listings[0]["Prefix"] == "reports/2026-" + + @pytest.mark.parametrize( + ("destination", "expected"), + [ + # GH-979: a version is downloaded to a local path named after its + # key. + ("directory", "b"), + # A Path destination is paired too: the version is downloaded to + # the file, not into a directory of that name. + ("path", "out.bin"), + ], + ids=["directory", "path"], + ) @pytest.mark.asyncio - async def test_get_version(self, tmp_path): - # GH-979: a version is downloaded to a local path named after its key. + async def test_get_version(self, tmp_path, destination, expected): fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) fs._sync_fs.isdir = mock.MagicMock(return_value=False) fs._get_file = mock.AsyncMock() - await fs._get("s3://bucket/b?versionId=v1", f"{tmp_path}/") + lpath = f"{tmp_path}/" if destination == "directory" else tmp_path / "out.bin" + await fs._get("s3://bucket/b?versionId=v1", lpath) assert fs._get_file.await_args.args[:2] == ( "bucket/b?versionId=v1", - f"{tmp_path.as_posix()}/b", + f"{tmp_path.as_posix()}/{expected}", ) @pytest.mark.asyncio @@ -1506,98 +1675,6 @@ async def copy_file(path1, path2, **kwargs): assert finished == ["s3://bucket/b"] fs._sync_fs._call.assert_not_called() - @pytest.mark.parametrize("status", [None, "Suspended", "Enabled"]) - @pytest.mark.asyncio - async def test_mv_null_version_onto_key(self, status): - fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) - fs._sync_fs._call = fs._sync_fs._core.call = mock.MagicMock( - return_value={} if status is None else {"Status": status} - ) - fs._copy_file = mock.AsyncMock(return_value=True) - fs._delete_objects = mock.AsyncMock() - sources = ["s3://bucket/a?versionId=null", "s3a://bucket/b?version_id=null"] - destinations = ["s3://bucket/a", "s3://bucket/b"] - - await fs._mv(sources, destinations, MetadataDirective="COPY") - - fs._sync_fs._call.assert_called_once_with( - fs._sync_fs._client.get_bucket_versioning, Bucket="bucket" - ) - if status == "Enabled": - assert fs._copy_file.await_args_list == [ - mock.call(source, dest, MetadataDirective="COPY") - for source, dest in zip(sources, destinations, strict=True) - ] - fs._delete_objects.assert_awaited_once_with(sources) - else: - fs._copy_file.assert_not_awaited() - fs._delete_objects.assert_awaited_once_with([]) - - @pytest.mark.parametrize("status", [None, "Suspended", "Enabled"]) - @pytest.mark.asyncio - async def test_mv_null_version_conflicts_depend_on_bucket_state(self, status): - fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) - fs._sync_fs._call = fs._sync_fs._core.call = mock.MagicMock( - return_value={} if status is None else {"Status": status} - ) - fs._copy_file = mock.AsyncMock(return_value=True) - fs._delete_objects = mock.AsyncMock() - sources = ["s3://bucket/a", "s3://bucket/b?versionId=null"] - destinations = ["s3a://bucket/b", "s3://bucket/out"] - - with ( - contextlib.nullcontext() - if status == "Enabled" - else pytest.raises(ValueError, match="another path that is moved") - ): - await fs._mv(sources, destinations) - - if status == "Enabled": - assert fs._copy_file.await_args_list == [ - mock.call(source, dest) for source, dest in zip(sources, destinations, strict=True) - ] - fs._delete_objects.assert_awaited_once_with(sources) - else: - fs._copy_file.assert_not_awaited() - fs._delete_objects.assert_not_awaited() - - @pytest.mark.asyncio - async def test_mv_null_version_directory_bucket_does_not_read_bucket_state(self): - fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) - fs._sync_fs._call = fs._sync_fs._core.call = mock.MagicMock() - fs._copy_file = mock.AsyncMock(return_value=True) - fs._delete_objects = mock.AsyncMock() - key = "s3://example--usw2-az1--x-s3/key" - - await fs._mv([f"{key}?versionId=null"], [key]) - - fs._sync_fs._call.assert_not_called() - fs._copy_file.assert_not_awaited() - fs._delete_objects.assert_awaited_once_with([]) - - @pytest.mark.parametrize("stage", ["lookup", "copy"]) - @pytest.mark.asyncio - async def test_mv_null_version_failure_does_not_delete(self, stage): - fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) - fs._sync_fs._call = fs._sync_fs._core.call = mock.MagicMock( - return_value={"Status": "Enabled"} - ) - fs._copy_file = mock.AsyncMock(return_value=True) - fs._delete_objects = mock.AsyncMock() - if stage == "lookup": - fs._sync_fs._call.side_effect = PermissionError("Access Denied") - else: - fs._copy_file.side_effect = PermissionError("Access Denied") - - with pytest.raises(PermissionError, match="Access Denied"): - await fs._mv( - ["s3://bucket/other", "s3://bucket/key?versionId=null"], - ["s3://bucket/out", "s3://bucket/key"], - ) - fs._delete_objects.assert_not_awaited() - if stage == "lookup": - fs._copy_file.assert_not_awaited() - @pytest.mark.parametrize("size", [10, 5 * 2**30 + 1]) @pytest.mark.asyncio async def test_cp_file_multipart_parameters(self, size): @@ -1641,16 +1718,18 @@ def upload_part_copy(**kw): "AnnotationDirective": "EXCLUDE", } - await fs._cp_file( - "s3://bucket/src", - "s3://bucket/dst", - block_size=S3Core.MULTIPART_UPLOAD_MAX_PART_SIZE // 2, - max_workers=1, - RequestPayer="requester", - ContentType="text/csv", - # Copy without reading the metadata, tags and annotations of the - # source (GH-973). - **directives, + await await_bounded( + fs._cp_file( + "s3://bucket/src", + "s3://bucket/dst", + block_size=S3Core.MULTIPART_UPLOAD_MAX_PART_SIZE // 2, + max_workers=1, + RequestPayer="requester", + ContentType="text/csv", + # Copy without reading the metadata, tags and annotations of the + # source (GH-973). + **directives, + ) ) if size <= S3Core.MULTIPART_UPLOAD_MAX_PART_SIZE: @@ -1693,79 +1772,6 @@ def test_internal_file_system_not_cached(self): for fs in (fs1, fs2, fs3): assert fs._sync_fs not in S3FileSystem._cache.values() - @pytest.mark.parametrize( - ("code", "exception"), - [ - (None, None), - ("NoSuchUpload", None), - ("NoSuchBucket", FileNotFoundError), - ("AccessDenied", PermissionError), - ("InternalError", OSError), - ], - ) - def test_clear_multipart_uploads_race(self, code, exception): - fs = AioS3FileSystem( - key="dummy", - secret="dummy", - region_name="us-east-1", - max_workers=1, - retry_config=RetryConfig(attempt=1), - skip_instance_cache=True, - ) - with Stubber(fs.core.client) as stubber: - stubber.add_response( - "list_multipart_uploads", - { - "Uploads": [ - {"Key": "prefix/gone", "UploadId": "gone"}, - { - "Key": "prefix/pending", - "UploadId": "pending", - "ChecksumAlgorithm": "CRC32", - "ChecksumType": "FULL_OBJECT", - }, - ], - "IsTruncated": False, - }, - {"Bucket": "bucket", "Prefix": "prefix/"}, - ) - request = {"Bucket": "bucket", "Key": "prefix/gone", "UploadId": "gone"} - if code: - stubber.add_client_error( - "abort_multipart_upload", - service_error_code=code, - http_status_code=404 if code.startswith("NoSuch") else 500, - expected_params=request, - ) - else: - stubber.add_response("abort_multipart_upload", {}, request) - stubber.add_response( - "abort_multipart_upload", - {}, - {"Bucket": "bucket", "Key": "prefix/pending", "UploadId": "pending"}, - ) - expected = pytest.raises(exception) if exception else contextlib.nullcontext() - with expected: - fs.clear_multipart_uploads("s3://bucket/prefix/") - stubber.assert_no_pending_responses() - - def test_clear_multipart_uploads_empty(self): - fs = AioS3FileSystem( - key="dummy", - secret="dummy", - region_name="us-east-1", - max_workers=1, - skip_instance_cache=True, - ) - with Stubber(fs.core.client) as stubber: - stubber.add_response( - "list_multipart_uploads", - {"Uploads": [], "IsTruncated": False}, - {"Bucket": "bucket", "Prefix": "prefix/"}, - ) - fs.clear_multipart_uploads("s3://bucket/prefix/") - stubber.assert_no_pending_responses() - @pytest.mark.parametrize( ("algorithm", "checksum_type"), [(None, None), ("SHA256", None), ("CRC32", None), ("CRC32", "FULL_OBJECT")], @@ -1838,206 +1844,23 @@ async def test_multipart_copy_uses_creation_algorithm(self, algorithm, checksum_ "AnnotationDirective": "EXCLUDE", **checksum_kwargs, } - await fs._copy_object_with_multipart_upload(**kwargs) + await await_bounded(fs._copy_object_with_multipart_upload(**kwargs)) stubber.assert_no_pending_responses() @pytest.fixture(scope="class") - def fs(self, request): - if not hasattr(request, "param"): - request.param = {} - return AioS3FileSystem(connection=connect(), **request.param) - - @pytest.mark.parametrize( - ("algorithm", "checksum_type"), - [(None, None), ("SHA256", None), ("CRC32", None), ("CRC32", "FULL_OBJECT")], - ) - def test_open_multipart_with_checksum(self, fs, algorithm, checksum_type): - block_size = 5 * 2**20 - data = b"x" * (block_size + 1) - path = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_open_multipart_with_checksum/{uuid.uuid4()}" - ) - kwargs = {"ChecksumAlgorithm": algorithm} if algorithm else {} - if checksum_type: - kwargs["ChecksumType"] = checksum_type - try: - with fs.open(path, "wb", block_size=block_size, s3_additional_kwargs=kwargs) as file: - file.write(data) - assert fs.cat_file(path) == data - assert fs.list_multipart_uploads(path) == [] - finally: - fs.clear_multipart_uploads(path) - if fs.exists(path): - fs.rm(path) - - @pytest.mark.parametrize( - ("algorithm", "checksum_type"), - [(None, None), ("SHA256", None), ("CRC32", None), ("CRC32", "FULL_OBJECT")], - ) - def test_put_file_multipart_with_checksum(self, fs, tmp_path, algorithm, checksum_type): - block_size = 5 * 2**20 - data = b"x" * (block_size + 1) - path = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_put_file_multipart_with_checksum/{uuid.uuid4()}" - ) - kwargs = {"ChecksumAlgorithm": algorithm} if algorithm else {} - if checksum_type: - kwargs["ChecksumType"] = checksum_type - try: - lpath = tmp_path / "data" - lpath.write_bytes(data) - fs.put_file(str(lpath), path, block_size=block_size, s3_additional_kwargs=kwargs) - assert fs.cat_file(path) == data - assert fs.list_multipart_uploads(path) == [] - finally: - fs.clear_multipart_uploads(path) - if fs.exists(path): - fs.rm(path) - - @pytest.mark.parametrize( - ("algorithm", "checksum_type"), - [(None, None), ("SHA256", None), ("CRC32", None), ("CRC32", "FULL_OBJECT")], - ) - def test_pipe_file_multipart_with_checksum(self, fs, algorithm, checksum_type): - block_size = 5 * 2**20 - data = b"x" * (block_size + 1) - path = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_pipe_file_multipart_with_checksum/{uuid.uuid4()}" - ) - kwargs = {"ChecksumAlgorithm": algorithm} if algorithm else {} - if checksum_type: - kwargs["ChecksumType"] = checksum_type - try: - fs.pipe_file(path, data, block_size=block_size, s3_additional_kwargs=kwargs) - assert fs.cat_file(path) == data - assert fs.list_multipart_uploads(path) == [] - finally: - fs.clear_multipart_uploads(path) - if fs.exists(path): - fs.rm(path) - - @pytest.mark.parametrize( - ("algorithm", "checksum_type"), - [(None, None), ("SHA256", None), ("CRC32", None), ("CRC32", "FULL_OBJECT")], - ) - def test_append_multipart_with_checksum(self, fs, algorithm, checksum_type): - block_size = 5 * 2**20 - data = b"x" * (block_size + 1) - path = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_append_multipart_with_checksum/{uuid.uuid4()}" - ) - kwargs = {"ChecksumAlgorithm": algorithm} if algorithm else {} - if checksum_type: - kwargs["ChecksumType"] = checksum_type - try: - original = b"y" * block_size - fs.pipe_file(path, original) - with fs.open(path, "ab", block_size=block_size, s3_additional_kwargs=kwargs) as file: - file.write(data) - assert fs.cat_file(path) == original + data - assert fs.list_multipart_uploads(path) == [] - finally: - fs.clear_multipart_uploads(path) - if fs.exists(path): - fs.rm(path) - - def test_clear_multipart_uploads_after_listed_upload_is_aborted(self, fs): - sync_fs = fs._sync_fs - prefix = ( - f"{ENV.s3_staging_key}{ENV.schema}/filesystem/test_async_clear_multipart_race/" - f"{uuid.uuid4()}/" - ) - path = f"s3://{ENV.s3_staging_bucket}/{prefix}" - gone_path = S3Path(ENV.s3_staging_bucket, f"{prefix}gone") - gone = sync_fs.core.create_multipart_upload(gone_path) - try: - sync_fs.core.create_multipart_upload(S3Path(ENV.s3_staging_bucket, f"{prefix}pending")) - list_uploads = sync_fs.list_multipart_uploads - - def list_then_abort(path): - uploads = list_uploads(path) - assert len(uploads) == 2 - assert any(upload.upload_id == gone.upload_id for upload in uploads) - sync_fs.core.abort_multipart_upload(gone) - return uploads - - with mock.patch.object( - sync_fs, "list_multipart_uploads", side_effect=list_then_abort - ) as list_mock: - fs.clear_multipart_uploads(path) - list_mock.assert_called_once_with(path) - assert fs.list_multipart_uploads(path) == [] - finally: - fs.clear_multipart_uploads(path) - - @pytest.mark.parametrize( - ("fs", "start", "end", "target_data"), - list( - chain( - *[ - [ - ({"default_block_size": x}, 0, 5, b"01234"), - ({"default_block_size": x}, 2, 7, b"23456"), - ({"default_block_size": x}, 0, 10, b"0123456789"), - ] - for x in (S3FileSystem.DEFAULT_BLOCK_SIZE, 3) - ] - ) - ), - indirect=["fs"], - ) - def test_read(self, fs, start, end, target_data): - # lowest level access: use the core - data = fs._sync_fs.core.get_object( - S3Path(ENV.s3_staging_bucket, ENV.s3_filesystem_test_file_key), (start, end) - ) - assert data == target_data, data - with fs.open( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_filesystem_test_file_key}", "rb" - ) as file: - # mid-level access: use _fetch_range - data = file._fetch_range(start, end) - assert data == target_data, data - # high-level: use fileobj seek and read - file.seek(start) - data = file.read(end - start) - assert data == target_data, data - - @pytest.mark.parametrize( - ("base", "exp"), - [ - (1, 2**10), - (1, 2**20), - ], - ) - def test_write(self, fs, base, exp): - data = b"a" * (base * exp) - path = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_write/{uuid.uuid4()}" - ) - with fs.open(path, "wb") as f: - f.write(data) - with fs.open(path, "rb") as f: - actual = f.read() - assert len(actual) == len(data) - assert actual == data + def fs(self): + return AioS3FileSystem(connection=connect()) @pytest.mark.parametrize( "size", [ - 2**10, # < block size: one-shot PutObject path (the GH-719 regression) 10 * 2**20, # > block size (5 MiB): multipart path via the async executor ], ) def test_write_transaction(self, fs, size): - # GH-719 regression for the async filesystem: AioS3File inherits - # _upload_chunk/commit from S3File, so the transaction fix must hold here - # too, including the multipart path driven by the async executor. + # GH-719 regression for the async filesystem: the deferred commit of a + # multipart write whose parts the async executor uploads. The + # single-request path is S3File's, as in the sync test. data = b"a" * size path = ( f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" @@ -2050,50 +1873,6 @@ def test_write_transaction(self, fs, size): assert len(actual) == len(data) assert actual == data - def test_write_transaction_rollback(self, fs): - # Kept small-only on purpose: the multipart discard()/abort path is - # already exercised by the sync test (AioS3File inherits discard()), - # so there is no need to pay for another large upload here. - path = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_write_transaction_rollback/{uuid.uuid4()}" - ) - - def write_then_fail(): - with fs.transaction: - f = fs.open(path, "wb") - f.write(b"hello world") - f.close() - raise RuntimeError("rollback") - - with pytest.raises(RuntimeError): - write_then_fail() - fs.invalidate_cache(path) - assert not fs.exists(path) - - @pytest.mark.parametrize( - ("base", "exp"), - [ - (1, 2**10), - (1, 2**20), - ], - ) - def test_append(self, fs, base, exp): - data = b"a" * (base * exp) - path = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_append/{uuid.uuid4()}" - ) - with fs.open(path, "ab") as f: - f.write(data) - extra = b"extra" - with fs.open(path, "ab") as f: - f.write(extra) - with fs.open(path, "rb") as f: - actual = f.read() - assert len(actual) == len(data + extra) - assert actual == data + extra - @pytest.mark.parametrize( "size", [ @@ -2131,141 +1910,6 @@ async def test_open_async_append_exclusive(self, fs): async with await fs.open_async(path, "rb") as f: assert await f.read() == b"hello world" - @pytest.mark.asyncio - async def test_ls_buckets(self, fs): - fs.invalidate_cache() - actual = await fs._ls("s3://") - assert ENV.s3_staging_bucket in actual, actual - - fs.invalidate_cache() - actual = await fs._ls("s3:///") - assert ENV.s3_staging_bucket in actual, actual - - fs.invalidate_cache() - actual = await fs._ls("s3://", detail=True) - found = next(filter(lambda x: x.name == ENV.s3_staging_bucket, actual), None) - assert found - assert found.name == ENV.s3_staging_bucket - - fs.invalidate_cache() - actual = await fs._ls("s3:///", detail=True) - found = next(filter(lambda x: x.name == ENV.s3_staging_bucket, actual), None) - assert found - assert found.name == ENV.s3_staging_bucket - - @pytest.mark.asyncio - async def test_ls_dirs(self, fs): - dir_ = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_ls_dirs" - ) - for i in range(5): - await fs._pipe_file(f"{dir_}/prefix/test_{i}", bytes(i)) - await fs._touch(f"{dir_}/prefix2") - - assert len(await fs._ls(f"{dir_}/prefix")) == 5 - assert len(await fs._ls(f"{dir_}/prefix/")) == 5 - assert len(await fs._ls(f"{dir_}/prefix/test_")) == 0 - assert len(await fs._ls(f"{dir_}/prefix2")) == 1 - - test_1 = await fs._ls(f"{dir_}/prefix/test_1") - assert len(test_1) == 1 - assert test_1[0] == fs._strip_protocol(f"{dir_}/prefix/test_1") - - test_1_detail = await fs._ls(f"{dir_}/prefix/test_1", detail=True) - assert len(test_1_detail) == 1 - assert test_1_detail[0].name == fs._strip_protocol(f"{dir_}/prefix/test_1") - assert test_1_detail[0].size == 1 - - @pytest.mark.asyncio - async def test_info_bucket(self, fs): - dir_ = f"s3://{ENV.s3_staging_bucket}" - s3_path = S3Path.parse(dir_) - info = await fs._info(dir_) - - assert info.name == fs._strip_protocol(dir_) - assert info.bucket == s3_path.bucket - assert info.key is None - assert info.last_modified is None - assert info.size == 0 - assert info.etag is None - assert info.type == S3ObjectType.S3_OBJECT_TYPE_DIRECTORY - assert info.storage_class == S3StorageClass.S3_STORAGE_CLASS_BUCKET - assert info.version_id == s3_path.version_id - - dir_ = f"s3://{ENV.s3_staging_bucket}/" - s3_path = S3Path.parse(dir_) - info = await fs._info(dir_) - - assert info.name == fs._strip_protocol(dir_) - assert info.bucket == s3_path.bucket - assert info.key is None - assert info.last_modified is None - assert info.size == 0 - assert info.etag is None - assert info.type == S3ObjectType.S3_OBJECT_TYPE_DIRECTORY - assert info.storage_class == S3StorageClass.S3_STORAGE_CLASS_BUCKET - assert info.version_id == s3_path.version_id - - @pytest.mark.asyncio - async def test_info_dir(self, fs): - dir_ = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_info_dir" - ) - file = f"{dir_}/{uuid.uuid4()}" - - fs.invalidate_cache() - with pytest.raises(FileNotFoundError): - await fs._info(f"s3://{uuid.uuid4()}") - - await fs._pipe_file(file, b"a") - s3_path = S3Path.parse(dir_) - fs.invalidate_cache() - info = await fs._info(dir_) - fs.invalidate_cache() - - assert info.name == fs._strip_protocol(dir_) - assert info.bucket == s3_path.bucket - assert info.key == s3_path.key.rstrip("/") - assert info.last_modified is None - assert info.size == 0 - assert info.etag is None - assert info.type == S3ObjectType.S3_OBJECT_TYPE_DIRECTORY - assert info.storage_class == S3StorageClass.S3_STORAGE_CLASS_DIRECTORY - assert info.version_id == s3_path.version_id - - @pytest.mark.asyncio - async def test_info_file(self, fs): - dir_ = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_info_file" - ) - file = f"{dir_}/{uuid.uuid4()}" - - fs.invalidate_cache() - with pytest.raises(FileNotFoundError): - await fs._info(file) - - now = datetime.now(UTC) - await fs._pipe_file(file, b"a") - s3_path = S3Path.parse(file) - fs.invalidate_cache() - info = await fs._info(file) - fs.invalidate_cache() - ls_info = (await fs._ls(file, detail=True))[0] - - assert info == ls_info - assert info.name == fs._strip_protocol(file) - assert info.bucket == s3_path.bucket - assert info.key == s3_path.key - assert info.last_modified >= now - assert info.size == 1 - assert info.etag is not None - assert info.type == S3ObjectType.S3_OBJECT_TYPE_FILE - assert info.storage_class == S3StorageClass.S3_STORAGE_CLASS_STANDARD - assert info.version_id == s3_path.version_id - @pytest.mark.asyncio async def test_find(self, fs): dir_ = ( @@ -2288,6 +1932,10 @@ async def test_find(self, fs): result = await fs._find(f"{dir_}/prefix/test_") assert len(result) == 0 + # maxdepth reaches the listing: only prefix2 is one level down. + result = await fs._find(dir_, maxdepth=1) + assert result == [fs._strip_protocol(f"{dir_}/prefix2")] + result = await fs._find(f"{dir_}/prefix", prefix="test_") assert len(result) == 5 @@ -2305,95 +1953,6 @@ async def test_find(self, fs): ].name == fs._strip_protocol(f"{dir_}/prefix/test_1") assert test_1_detail[fs._strip_protocol(f"{dir_}/prefix/test_1")].size == 1 - @pytest.mark.asyncio - async def test_find_maxdepth(self, fs): - dir_ = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_find_maxdepth" - ) - # Create files at different depths - await fs._touch(f"{dir_}/file0.txt") - await fs._touch(f"{dir_}/level1/file1.txt") - await fs._touch(f"{dir_}/level1/level2/file2.txt") - await fs._touch(f"{dir_}/level1/level2/level3/file3.txt") - - # maxdepth must be at least 1, as in fsspec - with pytest.raises(ValueError, match="maxdepth must be at least 1"): - await fs._find(dir_, maxdepth=0) - - # Test maxdepth=1 (only files in the root) - result = await fs._find(dir_, maxdepth=1) - assert len(result) == 1 - assert fs._strip_protocol(f"{dir_}/file0.txt") in result - - # Test maxdepth=2 (files in root and level1) - result = await fs._find(dir_, maxdepth=2) - assert len(result) == 2 - assert fs._strip_protocol(f"{dir_}/file0.txt") in result - assert fs._strip_protocol(f"{dir_}/level1/file1.txt") in result - - # Test maxdepth=3 (files in root, level1, and level2) - result = await fs._find(dir_, maxdepth=3) - assert len(result) == 3 - assert fs._strip_protocol(f"{dir_}/level1/level2/file2.txt") in result - - # Test no maxdepth (all files) - result = await fs._find(dir_) - assert len(result) == 4 - - @pytest.mark.asyncio - async def test_find_withdirs(self, fs): - dir_ = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_find_withdirs" - ) - # Create directory structure with files - await fs._touch(f"{dir_}/file1.txt") - await fs._touch(f"{dir_}/subdir1/file2.txt") - await fs._touch(f"{dir_}/subdir1/subdir2/file3.txt") - await fs._touch(f"{dir_}/subdir3/file4.txt") - - # Test default behavior (withdirs=False) - result = await fs._find(dir_) - assert len(result) == 4 # Only files - for r in result: - assert r.endswith(".txt") - - # Test withdirs=True - result = await fs._find(dir_, withdirs=True) - assert len(result) > 4 # Files and directories - - # Verify directories are included - dirs = [r for r in result if not r.endswith(".txt")] - assert len(dirs) > 0 - assert any("subdir1" in d for d in dirs) - assert any("subdir2" in d for d in dirs) - assert any("subdir3" in d for d in dirs) - - # Test withdirs=False explicitly - result = await fs._find(dir_, withdirs=False) - assert len(result) == 4 # Only files - - @pytest.mark.asyncio - async def test_du(self, fs): - """Disk usage reports file sizes, their total, and the requested depth.""" - directory = f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/filesystem/test_async_du" - first = f"{directory}/first" - second = f"{directory}/nested/second" - try: - await fs._pipe_file(first, b"abc") - await fs._pipe_file(second, b"12345") - assert await fs._du(directory) == 8 - assert await fs._du(directory, total=False) == { - fs._strip_protocol(first): 3, - fs._strip_protocol(second): 5, - } - assert await fs._du(directory, maxdepth=1) == 3 - assert await fs._du(first) == 3 - finally: - with contextlib.suppress(FileNotFoundError): - await fs._rm(directory, recursive=True) - @pytest.mark.asyncio async def test_glob(self, fs): dir_ = ( @@ -2415,54 +1974,6 @@ async def test_glob(self, fs): with pytest.raises(ValueError): # noqa: PT011 fs.glob("*") - @pytest.mark.asyncio - async def test_exists_bucket(self, fs): - assert await fs._exists("s3://") - assert await fs._exists("s3:///") - - path = f"s3://{ENV.s3_staging_bucket}" - assert await fs._exists(path) - - not_exists_path = f"s3://{uuid.uuid4()}" - assert not await fs._exists(not_exists_path) - - @pytest.mark.asyncio - async def test_exists_object(self, fs): - path = f"s3://{ENV.s3_staging_bucket}/{ENV.s3_filesystem_test_file_key}" - assert await fs._exists(path) - - not_exists_path = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_exists/{uuid.uuid4()}" - ) - assert not await fs._exists(not_exists_path) - - @pytest.mark.asyncio - async def test_rm_file(self, fs): - dir_ = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_rm_file" - ) - file = f"{dir_}/{uuid.uuid4()}" - await fs._touch(file) - await fs._rm_file(file) - - assert not await fs._exists(file) - assert not await fs._exists(dir_) - - @pytest.mark.asyncio - async def test_rm(self, fs): - dir_ = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_rm" - ) - file = f"{dir_}/{uuid.uuid4()}" - await fs._touch(file) - await fs._rm(file) - - assert not await fs._exists(file) - assert not await fs._exists(dir_) - @pytest.mark.asyncio async def test_rm_recursive(self, fs): dir_ = ( @@ -2470,7 +1981,7 @@ async def test_rm_recursive(self, fs): f"filesystem/test_async_rm_recursive" ) - files = [f"{dir_}/{uuid.uuid4()}" for _ in range(10)] + files = [f"{dir_}/{uuid.uuid4()}" for _ in range(2)] for f in files: await fs._touch(f) @@ -2484,120 +1995,12 @@ async def test_rm_recursive(self, fs): assert not await fs._exists(f) assert not await fs._exists(dir_) - @pytest.mark.asyncio - async def test_touch(self, fs): - path = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_touch/{uuid.uuid4()}" - ) - assert not await fs._exists(path) - await fs._touch(path) - assert await fs._exists(path) - info = await fs._info(path) - assert info.size == 0 - - with fs.open(path, "wb") as f: - f.write(b"data") - info = await fs._info(path, refresh=True) - assert info.size == 4 - await fs._touch(path, truncate=True) - info = await fs._info(path, refresh=True) - assert info.size == 0 - - with fs.open(path, "wb") as f: - f.write(b"data") - info = await fs._info(path, refresh=True) - assert info.size == 4 - with pytest.raises(ValueError, match="Cannot touch"): - await fs._touch(path, truncate=False) - info = await fs._info(path, refresh=True) - assert info.size == 4 - - @pytest.mark.asyncio - @pytest.mark.parametrize( - ("base", "exp"), - [ - (1, 2**10), - (1, 2**20), - ], - ) - async def test_pipe_cat(self, fs, base, exp): - data = b"a" * (base * exp) - path = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_pipe_cat/{uuid.uuid4()}" - ) - await fs._pipe_file(path, data) - assert await fs._cat_file(path) == data - - @pytest.mark.asyncio - async def test_cat_ranges(self, fs): - data = b"1234567890abcdefghijklmnopqrstuvwxyz" - path = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_cat_ranges/{uuid.uuid4()}" - ) - await fs._pipe_file(path, data) - - assert await fs._cat_file(path) == data - assert await fs._cat_file(path, start=5) == data[5:] - assert await fs._cat_file(path, end=5) == data[:5] - assert await fs._cat_file(path, start=1, end=-1) == data[1:-1] - assert await fs._cat_file(path, start=-5) == data[-5:] - - @pytest.mark.asyncio - @pytest.mark.parametrize( - ("base", "exp"), - [ - (1, 2**10), - (1, 2**20), - ], - ) - async def test_put(self, fs, base, exp): - with tempfile.NamedTemporaryFile(delete=False) as tmp: - data = b"a" * (base * exp) - tmp.write(data) - tmp.flush() - - rpath = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_put/{uuid.uuid4()}" - ) - await fs._put_file(lpath=tmp.name, rpath=rpath) - tmp.seek(0) - assert await fs._cat_file(rpath) == tmp.read() - @pytest.mark.asyncio @pytest.mark.parametrize( ("base", "exp"), [ + # Below the maximum part size, a single CopyObject request. (1, 2**10), - (1, 2**20), - ], - ) - async def test_put_with_callback(self, fs, base, exp): - with tempfile.NamedTemporaryFile(delete=False) as tmp: - data = b"a" * (base * exp) - tmp.write(data) - tmp.flush() - - rpath = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_put_with_callback/{uuid.uuid4()}" - ) - callback = Callback() - await fs._put_file(lpath=tmp.name, rpath=rpath, callback=callback) - tmp.seek(0) - assert await fs._cat_file(rpath) == tmp.read() - assert callback.size == os.stat(tmp.name).st_size - assert callback.value == callback.size - - @pytest.mark.asyncio - @pytest.mark.parametrize( - ("base", "exp"), - [ - (1, 2**10), - (1, 2**20), ], ) async def test_upload_cp_file(self, fs, base, exp): @@ -2665,65 +2068,6 @@ async def test_move_null_version_onto_key(self, fs, versioning_buckets, status): else: assert after == before - @pytest.mark.skipif( - os.getenv("AWS_ATHENA_S3_VERSIONING_TESTS") != "1", - reason="Set AWS_ATHENA_S3_VERSIONING_TESTS=1 to create versioning test buckets.", - ) - @pytest.mark.parametrize( - "status", - [ - pytest.param(None, id="async-wrapper-None"), - pytest.param("Enabled", id="async-wrapper-Enabled"), - pytest.param("Suspended", id="async-wrapper-Suspended"), - ], - ) - def test_move_null_version_onto_key_sync_wrapper(self, fs, versioning_buckets, status): - client, buckets = versioning_buckets - bucket = buckets[status] - key = "async-wrapper" - path = f"s3://{bucket}/{key}" - before = [ - v - for v in client.list_object_versions(Bucket=bucket, Prefix=key)["Versions"] - if v["Key"] == key - ] - assert any(v["VersionId"] == "null" for v in before) - if status: - assert not next(v for v in before if v["VersionId"] == "null")["IsLatest"] - - fs.mv(f"{path}?versionId=null", path) - - with client.get_object(Bucket=bucket, Key=key)["Body"] as body: - assert body.read() == (b"original" if status != "Suspended" else b"current") - after = [ - v - for v in client.list_object_versions(Bucket=bucket, Prefix=key)["Versions"] - if v["Key"] == key - ] - if status == "Enabled": - assert not any(v["VersionId"] == "null" for v in after) - assert len(after) == len(before) - latest = next(v for v in after if v["IsLatest"]) - assert latest["VersionId"] not in {v["VersionId"] for v in before} - else: - assert after == before - - @pytest.mark.asyncio - async def test_move(self, fs): - path1 = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_move/{uuid.uuid4()}" - ) - path2 = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_move/{uuid.uuid4()}" - ) - data = b"a" - await fs._pipe_file(path1, data) - fs.mv(path1, path2) - assert await fs._cat_file(path2) == data - assert not await fs._exists(path1) - @pytest.mark.asyncio async def test_move_recursive(self, fs): # GH-974: the directory entries used to make mv() fail, after copying @@ -2740,98 +2084,6 @@ async def test_move_recursive(self, fs): assert not await fs._exists(f"{base}/src/a") assert not await fs._exists(f"{base}/src/sub/b") - @pytest.mark.asyncio - async def test_get_file(self, fs): - with tempfile.TemporaryDirectory() as tmp: - rpath = f"s3://{ENV.s3_staging_bucket}/{ENV.s3_filesystem_test_file_key}" - lpath = Path(f"{tmp}/{uuid.uuid4()}") - callback = Callback() - await fs._get_file(rpath=rpath, lpath=str(lpath), callback=callback) - - assert lpath.open("rb").read() == await fs._cat_file(rpath) - assert callback.size == os.stat(lpath).st_size - assert callback.value == callback.size - - def test_open_returns_aio_s3_file(self, fs): - path = f"s3://{ENV.s3_staging_bucket}/{ENV.s3_filesystem_test_file_key}" - with fs.open(path, "rb") as f: - assert isinstance(f, AioS3File) - data = f.read() - assert data == b"0123456789" - - def test_checksum(self, fs): - path = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_checksum/{uuid.uuid4()}" - ) - s3_path = S3Path.parse(path) - - fs.pipe_file(path, b"foo") - checksum = fs.checksum(path) - fs.ls(path) # caching - fs._sync_fs.core.put_object(s3_path, b"bar") - assert checksum == fs.checksum(path) - assert checksum != fs.checksum(path, refresh=True) - - fs.pipe_file(path, b"foo") - checksum = fs.checksum(path) - fs.ls(path) # caching - fs.core.delete_object(s3_path) - assert checksum == fs.checksum(path) - with pytest.raises(FileNotFoundError): - fs.checksum(path, refresh=True) - - def test_sign(self, fs): - path = f"s3://{ENV.s3_staging_bucket}/{ENV.s3_filesystem_test_file_key}" - requested = time.time() - time.sleep(1) - url = fs.sign(path, expiration=100) - parsed = urllib.parse.urlparse(url) - query = urllib.parse.parse_qs(parsed.query) - expires = int(query["Expires"][0]) - with urllib.request.urlopen(url) as r: - data = r.read() - - assert "https" in url - assert requested + 100 < expires - assert data == b"0123456789" - - @pytest.mark.asyncio - async def test_mkdir_and_rmdir(self, fs): - # The bucket already exists. - with pytest.raises(FileExistsError): - await fs._mkdir(f"s3://{ENV.s3_staging_bucket}") - await fs._makedirs(f"s3://{ENV.s3_staging_bucket}", exist_ok=True) - - path = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_rmdir/{uuid.uuid4()}" - ) - await fs._pipe_file(path, b"data") - # Only bucket paths can be removed. - with pytest.raises(FileExistsError): - await fs._rmdir(path) - with pytest.raises(FileNotFoundError): - fs.rmdir(f"{path}/nonexistent") - - @pytest.mark.asyncio - async def test_metadata_and_tags(self, fs): - path = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_metadata/{uuid.uuid4()}" - ) - await fs._pipe_file(path, b"data") - - assert fs.metadata(path) == {} - fs.setxattr(path, attr1="value1") - assert fs.metadata(path) == {"attr1": "value1"} - assert fs.getxattr(path, "attr1") == "value1" - assert fs.getxattr(path, "missing") is None - - assert fs.get_tags(path) == {} - fs.put_tags(path, {"tag1": "value1"}) - assert fs.get_tags(path) == {"tag1": "value1"} - def test_pandas_read_csv(self): import pandas @@ -2842,56 +2094,6 @@ def test_pandas_read_csv(self): ) assert [(row["col"],) for _, row in df.iterrows()] == [(123456789,)] - @pytest.mark.parametrize( - "line_count", - [ - 1 * 2**20, - ], - ) - def test_pandas_write_csv(self, line_count): - import pandas - - with tempfile.NamedTemporaryFile("w+t") as tmp: - tmp.write("col1") - tmp.write("\n") - for _ in range(line_count): - tmp.write("a") - tmp.write("\n") - tmp.flush() - - tmp.seek(0) - df = pandas.read_csv(tmp.name) - path = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_pandas_write_csv/{uuid.uuid4()}.csv" - ) - df.to_csv(path, index=False) - - actual = pandas.read_csv(path) - pandas.testing.assert_frame_equal(actual, df) - - def test_sync_wrappers(self, fs): - """Verify that mirror_sync_methods generates working sync wrappers.""" - actual = fs.ls(f"s3://{ENV.s3_staging_bucket}") - assert isinstance(actual, list) - assert len(actual) > 0 - - assert fs.exists(f"s3://{ENV.s3_staging_bucket}") - - def test_invalidate_cache(self, fs): - path = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_invalidate_cache/{uuid.uuid4()}" - ) - fs.pipe_file(path, b"data") - fs.info(path) - - # Cache should be populated - fs.invalidate_cache(path) - # Should not raise after cache invalidation - info = fs.info(path) - assert info.size == 4 - class TestAioS3File: @pytest.mark.parametrize("max_workers", [2, None]) @@ -3016,15 +2218,13 @@ def read(): assert sync_fs.core.get_object.call_count == 4 assert state["peak"] == 2 - @pytest.mark.parametrize("asynchronous", [False, True]) @pytest.mark.asyncio - async def test_open_parts_in_flight(self, asynchronous): + async def test_open_parts_in_flight(self): # GH-1158: a write keeps at most max_workers parts in flight with the - # executor of either kind. The parts waiting for a worker - # used to be queued with their data. - fs = AioS3FileSystem( - connection=mock.MagicMock(), asynchronous=asynchronous, skip_instance_cache=True - ) + # S3AioExecutor of the file; the S3ThreadPoolExecutor of an + # asynchronous filesystem is that of S3File. The parts waiting for a + # worker used to be queued with their data. + fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) block_size = S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE gates: dict[int, threading.Event] = {} @@ -3071,15 +2271,12 @@ def write(): sync_fs.core.complete_multipart_upload.assert_called_once() assert in_flight == [1, 2, 2, 2] - @pytest.mark.parametrize("asynchronous", [False, True]) @pytest.mark.asyncio - async def test_open_failed_part_while_waiting(self, asynchronous): + async def test_open_failed_part_while_waiting(self): # GH-1173: a part found failed when a write waits for a part stops the - # write with its error with the executor of either kind. The write - # used to upload the rest of the data. - fs = AioS3FileSystem( - connection=mock.MagicMock(), asynchronous=asynchronous, skip_instance_cache=True - ) + # write with its error with the S3AioExecutor of the file, as with + # that of S3File. The write used to upload the rest of the data. + fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) block_size = S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE sync_fs = fs._sync_fs sync_fs.core.create_multipart_upload = mock.MagicMock( @@ -3126,26 +2323,9 @@ def test_open_invalid_max_workers(self): with pytest.raises(ValueError, match="max_workers must be greater than 0"): fs.open("s3://bucket/key", "wb", max_workers=0) - def test_open_version_id(self): - fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) - fs._sync_fs.info = mock.MagicMock( - return_value=S3Object( - init={"Key": "key"}, - type=S3ObjectType.S3_OBJECT_TYPE_FILE, - bucket="bucket", - key="key", - ) - ) - fs._sync_fs.info.return_value.size = 4 - - with fs.open("s3://bucket/key", "rb", version_id="v1") as f: - assert isinstance(f, AioS3File) - assert f.version_id == "v1" - assert f.size == 4 - fs._sync_fs.info.assert_called_once_with("bucket/key?versionId=v1", version_id="v1") - def test_open_lookup_parameters(self): - # GH-1004: the lookup of the file sends its lookup parameters. + # GH-1004: the lookup of the file sends its lookup parameters, and + # looks up the version that is opened. fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) fs._sync_fs.info = mock.MagicMock( return_value=S3Object( @@ -3157,38 +2337,14 @@ def test_open_lookup_parameters(self): ) sse_c = {"SSECustomerAlgorithm": "AES256", "SSECustomerKey": "k" * 32} - with fs.open("s3://bucket/key", "rb", ContentType="text/plain", **sse_c) as f: + with fs.open( + "s3://bucket/key", "rb", version_id="v1", ContentType="text/plain", **sse_c + ) as f: assert isinstance(f, AioS3File) - fs._sync_fs.info.assert_called_once_with("bucket/key", version_id=None, **sse_c) - - @pytest.mark.parametrize( - ("start", "end", "max_workers", "worker_block_size", "ranges"), - [ - (42, 1337, 1, 999, [(42, 1337)]), # single worker - (42, 1337, 2, 999, [(42, 42 + 999), (42 + 999, 1337)]), # more workers - ( - 42, - 1337, - 2, - 333, - [ - (42, 42 + 333), - (42 + 333, 42 + 666), - (42 + 666, 42 + 999), - (42 + 999, 1337), - ], - ), - (42, 1337, 2, 1295, [(42, 1337)]), # single block - (42, 1337, 2, 1296, [(42, 1337)]), # single block - (42, 1337, 2, 1294, [(42, 1336), (1336, 1337)]), # single block too small - ], - ) - def test_get_ranges(self, start, end, max_workers, worker_block_size, ranges): - assert ( - S3File._get_ranges( - start, end, max_workers=max_workers, worker_block_size=worker_block_size - ) - == ranges + assert f.version_id == "v1" + assert f.size == 4 + fs._sync_fs.info.assert_called_once_with( + "bucket/key?versionId=v1", version_id="v1", **sse_c ) @@ -3248,7 +2404,9 @@ def upload_part(**kw): written = asyncio.Event() async def release(): + deadline = time.monotonic() + EVENT_TIMEOUT while True: + assert time.monotonic() < deadline with lock: if written.is_set() or len(active) >= 2: # Only the newest part finishes until the data is @@ -3270,7 +2428,8 @@ async def release(): await f.close() await releaser - assert in_flight == [1, 2, 2, 2] + assert len(in_flight) == 4 + assert max(in_flight) <= 2 (_, parts), _ = fs.core.complete_multipart_upload.call_args assert [part.part_number for part in parts] == [1, 2, 3, 4] fs.core.abort_multipart_upload.assert_not_called() @@ -3499,8 +2658,8 @@ async def test_close_cancelled(self, fs): def complete_multipart_upload(*args, **kwargs): started.set() - # The finally block of the test always releases it. - release.wait() + # The finally block of the test releases it. + release.wait(EVENT_TIMEOUT) completed.set() fs.core.complete_multipart_upload = mock.MagicMock(side_effect=complete_multipart_upload) @@ -3703,43 +2862,51 @@ async def test_cancelled_creation_stops_retries(self, fs): assert operation.calls == 1 assert f.closed - @pytest.mark.asyncio - async def test_discard_in_stopped_scope(self, fs): - # GH-1203: a stop scope does not stop the retries of the abort of - # _discard(). - block_size = S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE - f = await fs.open_async("s3://bucket/key", "wb", block_size=block_size) - await f.write(b"x" * block_size) - await f.flush() - # The abort of the core, whose request is stubbed. - del fs.core.abort_multipart_upload - operation = ThrottledOperation(failures=1, operation_name="AbortMultipartUpload") - fs.core.call = retried(operation, RetryConfig(attempt=3, multiplier=0)) - stop = threading.Event() - stop.set() - - def discard(): - with stop_retries_on(stop): - f._discard() - - await asyncio.to_thread(discard) - assert operation.calls == 2 - # Nothing is left to abort. - await f._close_without_commit() - assert operation.calls == 2 - @pytest.mark.asyncio async def test_cancelled_read_frees_ranges(self, fs): # GH-1186: a cancelled parallel read cancels the reads that have not - # finished and frees the data of the others, also while the - # cancellation is held, as a failed read does. + # finished, which stops their retries, and frees the data of the + # others, also while the cancellation is held, as a failed read does. block_size = 2**10 + loop = asyncio.get_running_loop() read = [] + all_read = threading.Barrier(4, timeout=EVENT_TIMEOUT) + held = threading.Event() + # The held range throttles and would retry after EVENT_TIMEOUT. + operation = ThrottledOperation(failures=1) + throttled = retried(operation, RetryConfig(attempt=2, multiplier=EVENT_TIMEOUT)) + + async def others_done(): + def running_reads(): + return [ + t + for t in asyncio.all_tasks() + if t.get_coro().__qualname__ == "S3AioThreadPool.run" and not t.done() + ] + + # At least the held read is found. + assert running_reads() + await wait_until(lambda: len(running_reads()) <= 1) def get_object(path, range_, **kwargs): + if range_[0] == 0: + all_read.wait() + # Retrying when the read is cancelled, once the tasks of the + # other ranges have taken their results. + asyncio.run_coroutine_threadsafe(others_done(), loop).result(EVENT_TIMEOUT) + held.set() + try: + throttled() + except Exception as e: + # The stopped retry keeps its error in a reference cycle + # through the frames of its traceback, which would keep + # this frame and the data, as S3Core.call() avoids. + traceback.clear_frames(e.__traceback__) # Unlike bytes, an array can be weakly referenced. data = array.array("B", bytes(range_[1] - range_[0])) read.append(weakref.ref(data)) + if range_[0] != 0: + all_read.wait() return data fs.core.get_object = mock.MagicMock(side_effect=get_object) @@ -3751,23 +2918,19 @@ def get_object(path, range_, **kwargs): ) info.size = block_size * 4 fs._sync_fs.info = mock.MagicMock(return_value=info) - wait_reads = asyncio.wait - - async def cancelled_wait(reads, **kwargs): - # As asyncio.wait(), whose frame keeps the tasks in a set, - # cancelled once the reads have finished. - reads = set(reads) - await wait_reads(reads) - raise asyncio.CancelledError f = await fs.open_async("s3://bucket/key", "rb", block_size=block_size, max_workers=4) - with ( - mock.patch.object(asyncio, "wait", side_effect=cancelled_wait), - pytest.raises(asyncio.CancelledError) as info_, - ): - await f.read() + task = asyncio.ensure_future(f.read()) + assert await asyncio.to_thread(held.wait, EVENT_TIMEOUT) + started = time.monotonic() + task.cancel() + with pytest.raises(asyncio.CancelledError) as info_: + await asyncio.wait_for(task, EVENT_TIMEOUT) await f.close() - await asyncio.to_thread(f._pool.shutdown, True) + await asyncio.wait_for(asyncio.to_thread(f._pool.shutdown, True), EVENT_TIMEOUT) + # The retries of the held range stopped. + assert time.monotonic() - started < EVENT_TIMEOUT + assert operation.calls == 1 for _ in range(3): await asyncio.sleep(0) assert len(read) == 4 diff --git a/tests/pyathena/filesystem/test_s3_core.py b/tests/pyathena/filesystem/test_s3_core.py index 2a0695b30..7780322a2 100644 --- a/tests/pyathena/filesystem/test_s3_core.py +++ b/tests/pyathena/filesystem/test_s3_core.py @@ -9,12 +9,12 @@ import gc import io import re +import time import weakref from concurrent.futures import Future from datetime import UTC, datetime from itertools import pairwise from typing import Any -from unittest import mock from urllib.parse import parse_qs, urlsplit import boto3 @@ -83,28 +83,11 @@ def test_operation_params(self, method, params, expected): core, _ = _make_core() assert core.operation_params(method, params) == expected - def test_properties(self): + def test_request_kwargs_is_a_copy(self): core, _ = _make_core(request_kwargs={"RequestPayer": "requester"}) - assert core.retry_config.attempt == 1 - assert core.request_kwargs == {"RequestPayer": "requester"} - # A copy, so the core keeps its parameters. + # A copy, so a caller that changes it leaves the core's parameters. core.request_kwargs.clear() assert core.request_kwargs == {"RequestPayer": "requester"} - assert S3Core(core.client).retry_config.attempt == RetryConfig().attempt - - def test_request_kwargs_go_to_operations_that_accept_them(self): - core, stubber = _make_core(request_kwargs={"RequestPayer": "requester"}) - stubber.add_response( - "head_object", - {"ContentLength": 1}, - {"Bucket": "bucket", "Key": "key", "RequestPayer": "requester"}, - ) - # HeadBucket does not accept RequestPayer. - stubber.add_response("head_bucket", {}, {"Bucket": "bucket"}) - with stubber: - core.head_object(S3Path("bucket", "key")) - core.head_bucket("bucket") - stubber.assert_no_pending_responses() def test_call_params_are_sent_as_given(self): # A per-call parameter is not filtered, so botocore rejects a @@ -120,13 +103,18 @@ def test_call_params_are_sent_as_given(self): (404, "NoSuchKey", FileNotFoundError), (403, "403", PermissionError), (500, "InternalError", OSError), + (416, "InvalidRange", OSError), ], ) def test_call_translates_errors(self, status, code, error): core, stubber = _make_core() stubber.add_client_error("head_object", service_error_code=code, http_status_code=status) - with stubber, pytest.raises(error): + with stubber, pytest.raises(error) as e: core.head_object(S3Path("bucket", "key")) + # The ClientError is the cause, whose code callers such as reads of + # an InvalidRange read. + assert isinstance(e.value.__cause__, botocore.exceptions.ClientError) + assert e.value.__cause__.response["Error"]["Code"] == code @pytest.mark.parametrize("translated", [True, False]) def test_call_failure_frees_request(self, translated): @@ -222,11 +210,6 @@ def test_head_object(self): metadata.object_lock_legal_hold_status, ) == ("GOVERNANCE", MODIFIED, "ON") - def test_head_object_requires_a_key(self): - core, _ = _make_core() - with pytest.raises(ValueError, match="no key"): - core.head_object(S3Path("bucket")) - def test_head_bucket(self): core, stubber = _make_core() stubber.add_response("head_bucket", {"BucketRegion": "us-west-2"}, {"Bucket": "bucket"}) @@ -247,15 +230,29 @@ def test_head_bucket(self): ], ) def test_get_object(self, version_id, range_, expected): - core, stubber = _make_core(request_kwargs={"ServerSideEncryption": "AES256"}) - # ServerSideEncryption, which GetObject does not accept, is not sent. + core, stubber = _make_core( + request_kwargs={"ServerSideEncryption": "AES256", "ExpectedBucketOwner": "111122223333"} + ) + # ServerSideEncryption, which GetObject does not accept, is not sent, + # and a parameter of the call takes precedence over request_kwargs. stubber.add_response( "get_object", {"Body": StreamingBody(io.BytesIO(b"data"), 4)}, - {"Bucket": "bucket", "Key": "key", "IfMatch": '"e"', **expected}, + { + "Bucket": "bucket", + "Key": "key", + "IfMatch": '"e"', + "ExpectedBucketOwner": "444455556666", + **expected, + }, ) with stubber: - data = core.get_object(S3Path("bucket", "key", version_id), range_, IfMatch='"e"') + data = core.get_object( + S3Path("bucket", "key", version_id), + range_, + IfMatch='"e"', + ExpectedBucketOwner="444455556666", + ) stubber.assert_no_pending_responses() assert data == b"data" @@ -309,18 +306,6 @@ def test_get_object_closes_body(self): stubber.assert_no_pending_responses() assert [raw.closed for raw in raws] == [True, True] - @pytest.mark.parametrize( - ("status", "code", "error"), - [(404, "NoSuchKey", FileNotFoundError), (416, "InvalidRange", OSError)], - ) - def test_get_object_translates_errors(self, status, code, error): - core, stubber = _make_core() - stubber.add_client_error("get_object", service_error_code=code, http_status_code=status) - with stubber, pytest.raises(error) as e: - core.get_object(S3Path("bucket", "key"), (10, None)) - assert isinstance(e.value.__cause__, botocore.exceptions.ClientError) - assert e.value.__cause__.response["Error"]["Code"] == code - @pytest.mark.parametrize( ("body", "params", "expected"), [ @@ -347,19 +332,6 @@ def test_put_object(self, body, params, expected): stubber.assert_no_pending_responses() assert (result.etag, result.version_id) == ('"e"', "v1") - @pytest.mark.parametrize( - ("path", "match"), - [ - (S3Path("bucket"), "no key"), - (S3Path("bucket", "key", "v1"), "Cannot write to a version"), - (S3Path("bucket", "key", "null"), "Cannot write to a version"), - ], - ) - def test_put_object_rejects_paths(self, path, match): - core, stubber = _make_core() - with stubber, pytest.raises(ValueError, match=match): - core.put_object(path, b"data") - @pytest.mark.parametrize( ("client_region", "kwargs", "expected"), [ @@ -473,13 +445,6 @@ def test_create_bucket_configuration(self, client_region, region_name, configura # The given configuration is not modified. assert given == configuration - @pytest.mark.parametrize("acl", ["invalid", "bucket-owner-full-control"]) - def test_create_bucket_rejects_acls(self, acl): - # An object-only canned ACL is not accepted for a bucket either. - core, stubber = _make_core() - with stubber, pytest.raises(ValueError, match="ACL not in"): - core.create_bucket("bucket", acl=acl) - def test_delete_bucket(self): core, stubber = _make_core() stubber.add_response( @@ -516,20 +481,20 @@ def test_get_bucket_versioning(self, response, expected): stubber.assert_no_pending_responses() assert status == expected - def test_get_bucket_versioning_translates_errors(self): - core, stubber = _make_core() - stubber.add_client_error( - "get_bucket_versioning", service_error_code="NoSuchBucket", http_status_code=404 - ) - with stubber, pytest.raises(FileNotFoundError): - core.get_bucket_versioning("bucket") - def test_list_objects(self): core, stubber = _make_core() stubber.add_response( "list_objects_v2", { - "Contents": [{"Key": "dir/a", "Size": 1, "ETag": '"a"', "LastModified": MODIFIED}], + "Contents": [ + { + "Key": "dir/a", + "Size": 1, + "ETag": '"a"', + "LastModified": MODIFIED, + "StorageClass": "STANDARD_IA", + } + ], "CommonPrefixes": [{"Prefix": "dir/sub/"}], "KeyCount": 2, "IsTruncated": True, @@ -556,7 +521,12 @@ def test_list_objects(self): bucket="bucket", objects=( S3ObjectSummary( - bucket="bucket", key="dir/a", size=1, etag='"a"', last_modified=MODIFIED + bucket="bucket", + key="dir/a", + size=1, + etag='"a"', + last_modified=MODIFIED, + storage_class="STANDARD_IA", ), ), common_prefixes=(S3CommonPrefix("bucket", "dir/sub/"),), @@ -826,11 +796,6 @@ def test_delete_object(self): core.delete_object(S3Path("bucket", "key", "v1"), BypassGovernanceRetention=True) stubber.assert_no_pending_responses() - def test_delete_object_requires_a_key(self): - core, _ = _make_core() - with pytest.raises(ValueError, match="has no key"): - core.delete_object(S3Path("bucket")) - def test_delete_objects(self): core, stubber = _make_core() stubber.add_response( @@ -873,14 +838,6 @@ def test_delete_objects_quiet(self): result = core.delete_objects(S3DeleteBatch("bucket", (S3Path("bucket", "a"),))) assert result == S3DeleteResult("bucket") - def test_delete_objects_translates_errors(self): - core, stubber = _make_core() - stubber.add_client_error( - "delete_objects", service_error_code="AccessDenied", http_status_code=403 - ) - with stubber, pytest.raises(PermissionError): - core.delete_objects(S3DeleteBatch("bucket", (S3Path("bucket", "a"),))) - def test_create_multipart_upload(self): core, stubber = _make_core() stubber.add_response( @@ -895,30 +852,6 @@ def test_create_multipart_upload(self): ) assert upload.upload_id == "u" - def test_upload_part(self): - core, stubber = _make_core() - stubber.add_response( - "upload_part", - {"ETag": '"e1"'}, - { - "Bucket": "bucket", - "Key": "key", - "UploadId": "u", - "PartNumber": 1, - "Body": b"data", - "SSECustomerAlgorithm": "AES256", - }, - ) - with stubber: - part = core.upload_part( - S3MultipartUpload({"Bucket": "bucket", "Key": "key", "UploadId": "u"}), - 1, - b"data", - SSECustomerAlgorithm="AES256", - PartNumber=99, - ) - assert (part.part_number, part.etag) == (1, '"e1"') - @pytest.mark.parametrize( "bucket", ["myap-abc123-s3alias", "arn:aws:s3:us-east-1:123456789012:accesspoint/myap"] ) @@ -1037,6 +970,8 @@ def test_complete_multipart_upload(self): S3MultipartUpload({"Bucket": "bucket", "Key": "key", "UploadId": "u"}), parts, RequestPayer="requester", + # The given parts take precedence. + MultipartUpload={"Parts": []}, ) assert (completed.etag, completed.version_id) == ('"dst"', "v-dst") @@ -1062,37 +997,25 @@ def test_abort_multipart_upload(self): ) @pytest.mark.parametrize( - ("method", "args"), + ("method", "args", "missing", "message"), [ - ("create_multipart_upload", (S3Path("bucket"),)), - ( - "upload_part_copy", - ( - S3MultipartUpload({"Bucket": "bucket", "Key": "dst", "UploadId": "u"}), - 1, - S3Path("bucket"), - ), - ), + ("upload_part", (1, b"data"), "Bucket", "no bucket"), + ("upload_part", (1, b"data"), "Key", "no key"), + ("upload_part", (1, b"data"), "UploadId", "no upload ID"), + # The other operations check the same identity. + ("upload_part_copy", (1, S3Path("source", "key")), "UploadId", "no upload ID"), + ("complete_multipart_upload", ([],), "Key", "no key"), + ("abort_multipart_upload", (), "Bucket", "no bucket"), ], - ) - def test_multipart_upload_requires_keys(self, method, args): - core, _ = _make_core() - with pytest.raises(ValueError, match="has no key"): - getattr(core, method)(*args) - - @pytest.mark.parametrize( - ("method", "args"), - [ - ("upload_part", (1, b"data")), - ("upload_part_copy", (1, S3Path("source", "key"))), - ("complete_multipart_upload", ([],)), - ("abort_multipart_upload", ()), + ids=[ + "upload_part-Bucket", + "upload_part-Key", + "upload_part-UploadId", + "upload_part_copy-UploadId", + "complete_multipart_upload-Key", + "abort_multipart_upload-Bucket", ], ) - @pytest.mark.parametrize( - ("missing", "message"), - [("Bucket", "no bucket"), ("Key", "no key"), ("UploadId", "no upload ID")], - ) def test_multipart_upload_requires_identity(self, method, args, missing, message): core, _ = _make_core() response = {"Bucket": "bucket", "Key": "key", "UploadId": "u"} @@ -1115,9 +1038,14 @@ def test_upload_part_uses_creation_algorithm_and_identity(self): "PartNumber": 1, "Body": b"data", "ChecksumAlgorithm": "SHA256", + "ChecksumSHA256": "c2hh", + "SSECustomerAlgorithm": "AES256", }, ) with stubber: + # The identity and the part number override the per-call fields, + # and the other parameters, such as the checksum of the part, are + # sent. part = core.upload_part( upload, 1, @@ -1125,17 +1053,15 @@ def test_upload_part_uses_creation_algorithm_and_identity(self): Bucket="other", Key="other", UploadId="other", + PartNumber=99, ChecksumAlgorithm="CRC32", + ChecksumSHA256="c2hh", + SSECustomerAlgorithm="AES256", ) stubber.assert_no_pending_responses() + assert (part.part_number, part.etag) == (1, '"part"') assert part.to_api_repr()["ChecksumSHA256"] == "sha" - def test_create_multipart_upload_rejects_versions(self): - # A write replaces the object at the key, not the named version. - core, _ = _make_core() - with pytest.raises(ValueError, match="Cannot write to a version"): - core.create_multipart_upload(S3Path("bucket", "key", "v1")) - @pytest.mark.parametrize( ("size", "block_size", "ranges"), [ @@ -1498,76 +1424,6 @@ def test_copy_object_annotation(self): ) stubber.assert_no_pending_responses() - @pytest.mark.parametrize( - "field", - [ - "ChecksumCRC32", - "ChecksumCRC32C", - "ChecksumCRC64NVME", - "ChecksumSHA1", - "ChecksumSHA256", - "ChecksumSHA512", - "ChecksumMD5", - "ChecksumXXHASH64", - "ChecksumXXHASH3", - "ChecksumXXHASH128", - ], - ) - @pytest.mark.parametrize("copy", [False, True]) - def test_complete_multipart_upload_preserves_part_checksums(self, field, copy): - core, stubber = _make_core() - identity = {"Bucket": "bucket", "Key": "key", "UploadId": "u"} - algorithm = field.removeprefix("Checksum") - stubber.add_response( - "create_multipart_upload", - {**identity, "ChecksumAlgorithm": algorithm}, - {"Bucket": "bucket", "Key": "key", "ChecksumAlgorithm": algorithm}, - ) - expected_parts = [] - for number in (1, 2): - result = {"ETag": f'"e{number}"', field: f"checksum{number}"} - request = {"Bucket": "bucket", "Key": "key", "UploadId": "u", "PartNumber": number} - if copy: - request["CopySource"] = {"Bucket": "bucket", "Key": "source"} - stubber.add_response("upload_part_copy", {"CopyPartResult": result}, request) - else: - request["Body"] = b"data" - request["ChecksumAlgorithm"] = algorithm - request[field] = f"checksum{number}" - stubber.add_response("upload_part", result, request) - expected_parts.append({**result, "PartNumber": number}) - stubber.add_response( - "complete_multipart_upload", - {"ETag": '"done"'}, - { - "Bucket": "bucket", - "Key": "key", - "UploadId": "u", - "MultipartUpload": {"Parts": expected_parts}, - "RequestPayer": "requester", - }, - ) - with stubber: - upload = core.create_multipart_upload( - S3Path("bucket", "key"), ChecksumAlgorithm=algorithm - ) - if copy: - parts = [ - core.upload_part_copy(upload, n, S3Path("bucket", "source")) for n in (1, 2) - ] - else: - parts = [ - core.upload_part(upload, n, b"data", **{field: f"checksum{n}"}) for n in (1, 2) - ] - completed = core.complete_multipart_upload( - upload, - parts, - RequestPayer="requester", - MultipartUpload={"Parts": []}, - ) - stubber.assert_no_pending_responses() - assert completed.etag == '"done"' - @pytest.mark.parametrize( ("algorithm", "checksum_type"), [(None, None), ("SHA256", "COMPOSITE"), ("CRC32", "FULL_OBJECT")], @@ -1685,6 +1541,8 @@ def test_put_bucket_acl(self): ("put_object_acl", (S3Path("bucket", "key"), "invalid")), # An object ACL that a bucket does not accept. ("put_bucket_acl", ("bucket", "bucket-owner-full-control")), + ("create_bucket", ("bucket", "invalid")), + ("create_bucket", ("bucket", "bucket-owner-full-control")), ], ) def test_acl_validation(self, method, args): @@ -1695,6 +1553,31 @@ def test_acl_validation(self, method, args): @pytest.mark.parametrize( ("method", "args", "match"), [ + ("head_object", (S3Path("bucket"),), "has no key"), + ("delete_object", (S3Path("bucket"),), "has no key"), + ("put_object", (S3Path("bucket"), b"data"), "has no key"), + # A write replaces the object at the key, not the named version. + ("put_object", (S3Path("bucket", "key", "v1"), b"data"), "Cannot write to a version"), + ( + "put_object", + (S3Path("bucket", "key", "null"), b"data"), + "Cannot write to a version", + ), + ("create_multipart_upload", (S3Path("bucket"),), "has no key"), + ( + "create_multipart_upload", + (S3Path("bucket", "key", "v1"),), + "Cannot write to a version", + ), + ( + "upload_part_copy", + ( + S3MultipartUpload({"Bucket": "bucket", "Key": "dst", "UploadId": "u"}), + 1, + S3Path("bucket"), + ), + "has no key", + ), ("get_object_tagging", (S3Path("bucket"),), "has no key"), ("put_object_tagging", (S3Path("bucket"), {}), "has no key"), ("put_object_acl", (S3Path("bucket"), "private"), "has no key"), @@ -1776,46 +1659,14 @@ def test_replace_object_metadata(self): core.replace_object_metadata( S3Path("bucket", "key"), head, - {"new": "2"}, + # The user-defined metadata can be any mapping, such as an + # S3Metadata. + S3Metadata({"Metadata": {"new": "2"}}), ContentType="text/plain", SSECustomerAlgorithm="AES256", ) stubber.assert_no_pending_responses() - def test_replace_object_metadata_omits_unset_fields(self): - core, stubber = _make_core() - stubber.add_response( - "copy_object", - {}, - { - "CopySource": {"Bucket": "bucket", "Key": "key"}, - "Bucket": "bucket", - "Key": "key", - "Metadata": {}, - "MetadataDirective": "REPLACE", - # S3Metadata reports STANDARD when HeadObject omits it. - "StorageClass": "STANDARD", - }, - ) - # The user-defined metadata can be any mapping, such as an S3Metadata. - stubber.add_response( - "copy_object", - {}, - { - "CopySource": {"Bucket": "bucket", "Key": "key"}, - "Bucket": "bucket", - "Key": "key", - "Metadata": {"a": "1"}, - "MetadataDirective": "REPLACE", - "StorageClass": "STANDARD", - }, - ) - with stubber: - core.replace_object_metadata(S3Path("bucket", "key"), S3Metadata({}), {}) - head = S3Metadata({"Metadata": {"a": "1"}}) - core.replace_object_metadata(S3Path("bucket", "key"), head, head) - stubber.assert_no_pending_responses() - @pytest.mark.parametrize( "name", ["Metadata", "MetadataDirective", "CopySource", "Bucket", "Key"] ) @@ -1825,53 +1676,69 @@ def test_replace_object_metadata_rejects_fields_of_the_copy(self, name): core.replace_object_metadata(S3Path("bucket", "key"), S3Metadata({}), {}, **{name: "x"}) @pytest.mark.parametrize( - ("path", "client_method", "params", "expected"), + ("path", "client_method", "params", "key", "query"), [ - (S3Path("bucket", "key"), "get_object", {}, {"Bucket": "bucket", "Key": "key"}), + (S3Path("bucket", "key"), "get_object", {}, "key", {}), ( S3Path("bucket", "key", "null"), "get_object", {"ResponseContentType": "text/csv"}, - { - "Bucket": "bucket", - "Key": "key", - "VersionId": "null", - "ResponseContentType": "text/csv", - }, + "key", + {"versionId": ["null"], "response-content-type": ["text/csv"]}, ), # The parameters take precedence over the path. ( S3Path("bucket", "key", "v1"), + "get_object", + {"Key": "other"}, + "other", + {"versionId": ["v1"]}, + ), + # A parameter that only PutObject accepts signs the method given. + ( + S3Path("bucket", "key"), "put_object", - {"Key": "other", "VersionId": None}, - {"Bucket": "bucket", "Key": "other", "VersionId": None}, + {"Key": "other", "ContentType": "text/csv"}, + "other", + {}, ), ], ) - def test_generate_presigned_url(self, path, client_method, params, expected): - core, _ = _make_core(request_kwargs={"RequestPayer": "requester"}) - core.call = mock.MagicMock(return_value="https://signed") - assert ( - core.generate_presigned_url(path, client_method, expires_in=60, **params) - == "https://signed" - ) - # request_kwargs are not signed. - core.call.assert_called_once_with( - core.client.generate_presigned_url, - ClientMethod=client_method, - Params=expected, - ExpiresIn=60, - ) - - def test_generate_presigned_url_signs_locally(self): + def test_generate_presigned_url(self, path, client_method, params, key, query): # No request is sent, so the stubber has no responses to return. core, stubber = _make_core(request_kwargs={"RequestPayer": "requester"}) with stubber: - url = core.generate_presigned_url(S3Path("bucket", "key", "v1")) - query = parse_qs(urlsplit(url).query) - assert query["versionId"] == ["v1"] - # request_kwargs are not signed. - assert not any(k.lower() == "x-amz-request-payer" for k in query) + url = core.generate_presigned_url(path, client_method, expires_in=60, **params) + signed = time.time() + parts = urlsplit(url) + # The addressing style is botocore's. + assert parts.netloc.startswith("bucket.") or parts.path.startswith("/bucket/") + assert parts.path.endswith(f"/{key}") + signed_query = parse_qs(parts.query) + # The expiry, in the fields of the client's signature version. + if "X-Amz-Expires" in signed_query: + assert signed_query["X-Amz-Expires"] == ["60"] + else: + expires = int(signed_query["Expires"][0]) + assert signed + 60 - 30 < expires <= signed + 60 + # The content type is signed: in the query with SigV2, as a header + # with SigV4. + content_type = params.get("ContentType") + if "Signature" in signed_query: + assert signed_query.pop("content-type", None) == ( + [content_type] if content_type else None + ) + else: + signed_headers = signed_query["X-Amz-SignedHeaders"][0].split(";") + assert ("content-type" in signed_headers) is (content_type is not None) + signature_fields = {"AWSAccessKeyId", "Signature", "Expires"} + signed_query = { + k: v + for k, v in signed_query.items() + if k not in signature_fields and not k.startswith("X-Amz-") + } + # request_kwargs, such as RequestPayer, are not signed. + assert signed_query == query class TestS3DeleteBatch: @@ -1890,6 +1757,9 @@ def test_from_paths(self): assert batches[1].objects == (S3Path("b1", f"k{S3DeleteBatch.MAX_KEYS}"),) assert batches[2].objects == (S3Path("b2", "a", "v1"),) assert S3DeleteBatch.from_paths([]) == [] + # A path without a key is rejected, also among other paths. + with pytest.raises(ValueError, match="Not an object of the bucket b1: s3://b1"): + S3DeleteBatch.from_paths([S3Path("b1", "a"), S3Path("b1")]) @pytest.mark.parametrize( ("objects", "match"), @@ -1908,10 +1778,6 @@ def test_invalid(self, objects, match): with pytest.raises(ValueError, match=match): S3DeleteBatch("bucket", objects) - def test_from_paths_rejects_bucket_paths(self): - with pytest.raises(ValueError, match="Not an object of the bucket bucket: s3://bucket"): - S3DeleteBatch.from_paths([S3Path("bucket", "a"), S3Path("bucket")]) - class TestS3DeleteError: def test_from_response(self): @@ -1921,25 +1787,3 @@ def test_from_response(self): assert error == S3DeleteError(S3Path("bucket", "a", "v1"), "InternalError", "Error") assert str(error) == "bucket/a?versionId=v1 (InternalError: Error)" assert str(S3DeleteError.from_response("bucket", {"Key": "a"})) == "bucket/a (None: None)" - - -class TestS3ObjectSummary: - def test_from_response(self): - entry = { - "Key": "dir/a", - "Size": 1, - "ETag": '"a"', - "LastModified": MODIFIED, - "StorageClass": "STANDARD_IA", - } - summary = S3ObjectSummary.from_response("bucket", entry) - assert summary == S3ObjectSummary("bucket", "dir/a", 1, '"a"', MODIFIED, "STANDARD_IA") - assert summary.path == S3Path("bucket", "dir/a") - # Optional fields that the entry does not have are None. - assert S3ObjectSummary.from_response("bucket", {"Key": "k"}) == S3ObjectSummary( - "bucket", "k" - ) - - def test_frozen(self): - with pytest.raises(AttributeError): - S3ObjectSummary("bucket", "k").key = "other" # type: ignore[misc] diff --git a/tests/pyathena/filesystem/test_s3_errors.py b/tests/pyathena/filesystem/test_s3_errors.py index 246f2cb30..68940cd81 100644 --- a/tests/pyathena/filesystem/test_s3_errors.py +++ b/tests/pyathena/filesystem/test_s3_errors.py @@ -23,15 +23,6 @@ def _client_error(code, message="error message", status_code=None, condition=Non class TestS3ClientError: - def test_properties(self): - error = S3ClientError(_client_error("NoSuchKey", status_code=404)) - assert error.code == "NoSuchKey" - assert error.message == "error message" - assert error.http_status_code == 404 - - error = S3ClientError(_client_error("NoSuchKey")) - assert error.http_status_code is None - @pytest.mark.parametrize( ("code", "expected"), [ @@ -51,7 +42,9 @@ def test_properties(self): ], ) def test_os_error_by_error_code(self, code, expected): - actual = S3ClientError(_client_error(code)).os_error + error = S3ClientError(_client_error(code)) + assert (error.code, error.message) == (code, "error message") + actual = error.os_error assert type(actual) is expected assert "error message" in str(actual) @@ -63,10 +56,14 @@ def test_os_error_by_error_code(self, code, expected): ("ServiceUnavailable", errno.EBUSY), ("OperationAborted", errno.EBUSY), ("InternalError", errno.EIO), + # An unknown code without a status code. + ("UnknownCode", errno.EIO), ], ) def test_os_error_by_error_code_errno(self, code, expected_errno): - actual = S3ClientError(_client_error(code)).os_error + error = S3ClientError(_client_error(code)) + assert error.http_status_code is None + actual = error.os_error # OSError specializes itself into a subclass for some errno values # (e.g., EPERM -> PermissionError), so check the errno, not the exact type. assert isinstance(actual, OSError) @@ -87,16 +84,12 @@ def test_os_error_by_error_code_errno(self, code, expected_errno): ], ) def test_os_error_by_http_status_code(self, status_code, expected_errno): - actual = S3ClientError(_client_error("UnknownCode", status_code=status_code)).os_error + error = S3ClientError(_client_error("UnknownCode", status_code=status_code)) + assert error.http_status_code == status_code + actual = error.os_error assert isinstance(actual, OSError) assert actual.errno == expected_errno - def test_os_error_unknown_code(self): - actual = S3ClientError(_client_error("UnknownCode")).os_error - assert type(actual) is OSError - assert actual.errno == errno.EIO - assert "error message" in str(actual) - @pytest.mark.parametrize( ("condition", "expected"), [ diff --git a/tests/pyathena/filesystem/test_s3_executor.py b/tests/pyathena/filesystem/test_s3_executor.py index 601d0d048..e93de804f 100644 --- a/tests/pyathena/filesystem/test_s3_executor.py +++ b/tests/pyathena/filesystem/test_s3_executor.py @@ -9,11 +9,13 @@ import contextvars import gc import os +import signal import threading import time import traceback import weakref from concurrent.futures import Future, ThreadPoolExecutor, wait +from typing import cast from unittest import mock import pytest @@ -25,12 +27,10 @@ class TestS3AioThreadPool: - def test_init(self): - # None takes the default of ThreadPoolExecutor. - cpus = getattr(os, "process_cpu_count", os.cpu_count)() - assert S3AioThreadPool(None).max_workers == min(32, (cpus or 1) + 4) - with pytest.raises(ValueError, match="max_workers must be greater than 0"): - S3AioThreadPool(0) + def test_default_max_workers(self): + # None takes the default of ThreadPoolExecutor, as documented. + with ThreadPoolExecutor() as executor: + assert S3AioThreadPool().max_workers == executor._max_workers def test_run_beyond_default_executor(self): # GH-1159: the calls run in threads of the pool, so that max_workers, @@ -101,11 +101,30 @@ def test_run_cancelled_after_return(self): # not stop the work that the call left running, such as the parts of a # write in a transaction. operation = ThrottledOperation(block=True, failures=1) - returned = threading.Event() left = [] + go = threading.Event() pool = S3AioThreadPool(1) + submit = pool.submit + + def submit_and_cancel_after_return(fn, *args, **kwargs): + future = submit(fn, *args, **kwargs) + task = cast(asyncio.Task[None], asyncio.current_task()) + + def cancel_after_return(): + # Runs once the task waits for the call, and blocks the event + # loop until the call has returned, so that the task is + # cancelled before it resumes. + go.set() + future.result(EVENT_TIMEOUT) + task.cancel() + + asyncio.get_running_loop().call_soon(cancel_after_return) + return future + + pool.submit = submit_and_cancel_after_return # type: ignore[method-assign] def work(): + assert go.wait(EVENT_TIMEOUT) context = contextvars.copy_context() thread = threading.Thread( target=context.run, @@ -113,17 +132,9 @@ def work(): ) thread.start() left.append(thread) - returned.set() async def main(): task = asyncio.ensure_future(pool.run(work)) - await asyncio.sleep(0) - # Blocks the event loop, so that the task is cancelled before it - # resumes after the call has returned, which follows the event at - # once. - assert returned.wait(EVENT_TIMEOUT) - time.sleep(0.1) - task.cancel() with pytest.raises(asyncio.CancelledError): await task @@ -150,14 +161,23 @@ def work(): started.set() release.wait(EVENT_TIMEOUT) + running = pool.submit(work) + submit = pool.submit + async def main(): - running = pool.submit(work) + submitted = asyncio.Event() + + def submit_and_notify(fn, *args, **kwargs): + future = submit(fn, *args, **kwargs) + submitted.set() + return future + + pool.submit = submit_and_notify # type: ignore[method-assign] run = pool.run_shielded if shielded else pool.run pending = asyncio.ensure_future(run(events.append, "pending finished")) - await asyncio.to_thread(started.wait, EVENT_TIMEOUT) - # Lets the task submit the call, which waits for the thread. - while not pending.done() and not pool._get_executor()._work_queue.qsize(): - await asyncio.sleep(0.01) + assert await asyncio.to_thread(started.wait, EVENT_TIMEOUT) + # The call waits for the only thread. + await asyncio.wait_for(submitted.wait(), EVENT_TIMEOUT) pending.cancel() with pytest.raises(asyncio.CancelledError): await pending @@ -169,54 +189,6 @@ async def main(): assert events == (["pending finished"] if shielded else []) - def test_wait_shielded_cancelled(self): - # GH-1186: cancelling the waiting task neither cancels the call nor - # lets asyncio log its failure, which ends after the cancellation; - # the caller can still wait for its future. - started = threading.Event() - release = threading.Event() - pool = S3AioThreadPool(1) - - def work(): - started.set() - # The finally block of the test always releases it. - release.wait() - raise PermissionError("Access Denied") - - async def main(): - logged = [] - asyncio.get_running_loop().set_exception_handler( - lambda _, context: logged.append(context["message"]) - ) - running = pool.submit(work) - try: - waiting = asyncio.ensure_future(pool.wait_shielded(running)) - await asyncio.to_thread(started.wait, EVENT_TIMEOUT) - # Lets the task start waiting before it is cancelled. - deadline = time.monotonic() + EVENT_TIMEOUT - while not running._done_callbacks: - assert time.monotonic() < deadline - await asyncio.sleep(0) - waiting.cancel() - with pytest.raises(asyncio.CancelledError): - await waiting - finally: - release.set() - # The worker has run the callbacks of the call once it is - # shut down, and the event loop runs theirs on the next - # iterations. - await asyncio.to_thread(pool.shutdown, True) - with pytest.raises(PermissionError, match="Access Denied"): - running.result() - # The callbacks of the future of the call keep its waiter. - del waiting, running - gc.collect() - for _ in range(3): - await asyncio.sleep(0) - return logged - - assert asyncio.run(main()) == [] - def test_shutdown(self): # The calls submitted before the shutdown still run. events = [] @@ -225,53 +197,64 @@ def test_shutdown(self): pool.submit(release.wait, EVENT_TIMEOUT) pending = pool.submit(events.append, "pending finished") pool.shutdown(wait=False) - with pytest.raises(RuntimeError, match="after shutdown"): + # The shutdown did not wait for the running call. + assert not pending.done() + with pytest.raises(RuntimeError): pool.submit(sum, [1, 2]) release.set() pending.result(timeout=EVENT_TIMEOUT) assert events == ["pending finished"] + @pytest.mark.skipif(not hasattr(os, "fork"), reason="os.fork() is not available.") + @pytest.mark.filterwarnings("ignore:.*fork\\(\\) may lead to deadlocks:DeprecationWarning") def test_fork(self): # A forked child, whose copy of the pool has no threads, creates new # ones, also if a thread of the parent held the lock of the pool. pool = S3AioThreadPool(1) - parent = pool._get_executor() - pool._lock.acquire() - # As os.register_at_fork() runs it in a forked child. - S3AioThreadPool._reset_after_fork() - # The copied executor is kept, so that its collection does not touch - # the locks that the parent's threads may have held. - assert parent in S3AioThreadPool._forked_executors - child = pool._get_executor() assert pool.submit(sum, [1, 2]).result(timeout=EVENT_TIMEOUT) == 3 - pool.shutdown() - assert child is not parent - S3AioThreadPool._forked_executors.remove(parent) - parent.shutdown() + # No public call holds the lock across a fork; a thread of the parent + # holds it while the main thread forks. + held, release = threading.Event(), threading.Event() + def hold_lock(): + with pool._lock: + held.set() + release.wait(EVENT_TIMEOUT) -class TestS3AioExecutor: - def test_init(self): - # max_workers is optional, as before it was added. - S3AioExecutor(loop=None) - with pytest.raises(ValueError, match="max_workers must be greater than 0"): - S3AioExecutor(loop=None, max_workers=0) - - def test_submit(self): - async def main(): - executor = S3AioExecutor(loop=asyncio.get_running_loop()) - succeeded = executor.submit(sum, [1, 2]) - failed = executor.submit(int, "x") - await asyncio.to_thread(wait, [succeeded, failed]) - return succeeded, failed + holder = threading.Thread(target=hold_lock) + holder.start() + try: + assert held.wait(EVENT_TIMEOUT) + pid = os.fork() + finally: + release.set() + if pid == 0: + # The child exits here, whatever happens, instead of running pytest. + code = 1 + try: + if pool.submit(sum, [3, 4]).result(timeout=EVENT_TIMEOUT) == 7: + code = 0 + finally: + os._exit(code) + status = (0, 0) + try: + # A child whose pool waits for the threads of the parent would + # not exit. + deadline = time.monotonic() + 2 * EVENT_TIMEOUT + while (status := os.waitpid(pid, os.WNOHANG)) == (0, 0): + assert time.monotonic() < deadline, "the child did not exit" + time.sleep(0.05) + finally: + if status == (0, 0): + os.kill(pid, signal.SIGKILL) + os.waitpid(pid, 0) + holder.join(EVENT_TIMEOUT) + pool.shutdown() + assert os.waitstatus_to_exitcode(status[1]) == 0 - succeeded, failed = asyncio.run(main()) - - assert succeeded.result() == 3 - with pytest.raises(ValueError, match="invalid literal"): - failed.result() +class TestS3AioExecutor: def test_submit_context(self): # GH-1203: a function runs with the context of its submitter, as the # task that dispatches it does, so that the stop scope of the @@ -492,54 +475,6 @@ async def main(): assert not running.cancelled() assert pending.cancelled() - def test_loop_shutdown(self): - # A function that has not started when the event loop shuts down is - # never run, and its future is cancelled and settled, so that wait() - # returns, instead of left pending. - events = [] - started = threading.Event() - settled = threading.Event() - - def work(): - started.set() - # Running until the pending future is settled, so that the - # pending function cannot start before the shutdown. - settled.wait(5) - events.append("finished") - - async def main(): - # The pending function waits for the only thread of the executor. - executor = S3AioExecutor(loop=asyncio.get_running_loop(), max_workers=1) - running = executor.submit(work) - pending = executor.submit(events.append, "pending finished") - pending.add_done_callback(lambda _: settled.set()) - while not started.is_set(): - await asyncio.sleep(0.01) - return running, pending - - running, pending = asyncio.run(main()) - - _, not_done = wait([running, pending], timeout=5) - assert not not_done - assert events == ["finished"] - assert not running.cancelled() - assert pending.cancelled() - - def test_executor_shut_down(self): - # An error raised before the function starts is set on the future. - async def main(): - executor = S3AioExecutor(loop=asyncio.get_running_loop()) - # Fails the dispatch of the function to the pool. - executor._pool.shutdown() - future = executor.submit(sum, [1, 2]) - with pytest.raises(RuntimeError, match="after shutdown"): - await asyncio.wrap_future(future) - return future - - future = asyncio.run(main()) - - assert not future.cancelled() - class TestS3Executor: @pytest.mark.parametrize("executor_class", [S3ThreadPoolExecutor, S3AioThreadPool]) @@ -626,32 +561,41 @@ def work(): def test_submit_to_loop_shutdown(self): # A callable that S3AioExecutor drops when its event loop shuts down - # cancels the future of the caller, so that wait() returns. + # is never run, and its future, whether submit() returned it or the + # caller passed it to submit_to(), is cancelled and settled, so that + # wait() returns, instead of left pending. events = [] started = threading.Event() - settled = threading.Event() + settled = threading.Semaphore(0) def work(): started.set() - settled.wait(5) + # Running until both pending futures are settled, so that their + # callables cannot start before the shutdown. + for _ in range(2): + assert settled.acquire(timeout=EVENT_TIMEOUT) events.append("finished") pending: Future[None] = Future() - pending.add_done_callback(lambda _: settled.set()) + pending.add_done_callback(lambda _: settled.release()) async def main(): + # The pending callables wait for the only thread of the executor. executor = S3AioExecutor(loop=asyncio.get_running_loop(), max_workers=1) running = executor.submit(work) + submitted = executor.submit(events.append, "submitted finished") + submitted.add_done_callback(lambda _: settled.release()) executor.submit_to(pending, events.append, "pending finished") - while not started.is_set(): - await asyncio.sleep(0.01) - return running + assert await asyncio.to_thread(started.wait, EVENT_TIMEOUT) + return running, submitted - running = asyncio.run(main()) + running, submitted = asyncio.run(main()) - _, not_done = wait([running, pending], timeout=5) + _, not_done = wait([running, submitted, pending], timeout=EVENT_TIMEOUT) assert not not_done assert events == ["finished"] + assert not running.cancelled() + assert submitted.cancelled() assert pending.cancelled() def test_submit_to_executor_shut_down(self): @@ -665,7 +609,7 @@ async def main(): executor._pool.shutdown() executor.submit_to(future, sum, [1, 2]) with pytest.raises(RuntimeError, match="after shutdown"): - await asyncio.wrap_future(future) + await asyncio.wait_for(asyncio.wrap_future(future), EVENT_TIMEOUT) asyncio.run(main()) @@ -684,7 +628,9 @@ async def main(): executor.submit_to(future, sum, [1, 2]) assert future.cancel() with ThreadPoolExecutor(max_workers=1) as waiter: - return await loop.run_in_executor(waiter, lambda: wait([future], timeout=5)) + return await loop.run_in_executor( + waiter, lambda: wait([future], timeout=EVENT_TIMEOUT) + ) _, not_done = asyncio.run(main()) diff --git a/tests/pyathena/filesystem/test_s3_path.py b/tests/pyathena/filesystem/test_s3_path.py index 6f3a13db0..efa4b906b 100644 --- a/tests/pyathena/filesystem/test_s3_path.py +++ b/tests/pyathena/filesystem/test_s3_path.py @@ -115,9 +115,6 @@ def test_names(self, path, name, string, uri): assert str(path) == string assert path.uri == uri - def test_str_spells_the_version_query_as_version_id(self): - assert str(S3Path.parse("s3a://bucket/key?version_id=v1")) == "bucket/key?versionId=v1" - @pytest.mark.parametrize( ("path", "expected"), [ diff --git a/tests/pyathena/filesystem/test_s3_path_pairing.py b/tests/pyathena/filesystem/test_s3_path_pairing.py index 4b549352b..6685a4dc2 100644 --- a/tests/pyathena/filesystem/test_s3_path_pairing.py +++ b/tests/pyathena/filesystem/test_s3_path_pairing.py @@ -134,32 +134,32 @@ def test_copy_pairs_needs_sources(self): with pytest.raises(ValueError, match="sources"): S3PathPairing("s3://bucket/a/", "s3://bucket/b/").copy_pairs() - def test_move_pairs(self): - # The "null" version of a key moved onto the key stays in place. - pairs = [("s3://bucket/b?versionId=null", "s3://bucket/b"), ("bucket/d/a", "s3://bucket/z")] - - assert _pairing(pairs).conflict_candidates(pairs) == [] - assert _pairing(pairs).move_pairs(pairs) == [("bucket/d/a", "s3://bucket/z")] - @pytest.mark.parametrize( ("enabled_buckets", "expected"), [ + # The "null" version of a key moved onto the key stays in place. ((), []), (("enabled",), [("enabled/key?versionId=null", "enabled/key")]), (("other",), [("other/key?versionId=null", "other/key")]), ], ) def test_move_pairs_null_version_bucket_state(self, enabled_buckets, expected): + moved = ("enabled/d/a", "s3://enabled/z") pairs = [ ("enabled/key?versionId=null", "enabled/key"), ("other/key?versionId=null", "other/key"), ("enabled/unchanged?versionId=null", "s3://enabled/unchanged?versionID=null"), + # An ordinary move. + moved, ] pairing = _pairing(pairs) assert pairing.conflict_candidates(pairs, versioning_enabled_buckets=enabled_buckets) == [] - assert pairing.move_pairs(pairs, versioning_enabled_buckets=enabled_buckets) == expected - assert pairing.move_pairs(pairs) == [] + assert pairing.move_pairs(pairs, versioning_enabled_buckets=enabled_buckets) == [ + *expected, + moved, + ] + assert pairing.move_pairs(pairs) == [moved] def test_move_pairs_distinguishes_null_version_source(self): pairs = [ diff --git a/tests/pyathena/util.py b/tests/pyathena/util.py index 82e3d467c..97b8b0e9c 100644 --- a/tests/pyathena/util.py +++ b/tests/pyathena/util.py @@ -201,28 +201,60 @@ async def wait_until(predicate, timeout=EVENT_TIMEOUT): await asyncio.sleep(0.01) -def overlapping_calls(calls, timeout=5): +async def await_bounded(awaitable, timeout=EVENT_TIMEOUT): + """Await an awaitable, failing if it is not done within the timeout. + + Unlike ``asyncio.wait_for``, it does not wait for the cancellation of an + awaitable that the timeout interrupts, which a regression could also + hold, so that the test fails instead of hanging. + + Args: + awaitable: The coroutine or future to await. + timeout: The longest wait, in seconds. + + Returns: + The result of the awaitable. + """ + task = asyncio.ensure_future(awaitable) + done, _ = await asyncio.wait({task}, timeout=timeout) + assert done, f"not done within {timeout} s" + return task.result() + + +def overlapping_calls(calls, timeout=5, grace=0.5): """Return a callable that records the peak number of its overlapping calls. - Each call waits until ``calls`` calls overlap, or until the timeout. + Until ``calls`` calls have overlapped, each call waits for that, or until + the timeout. The calls of that first overlap then wait up to ``grace`` + seconds for another overlapping call, which only a caller that runs more + than ``calls`` at once makes, so that the peak also shows an exceeded + bound. Later calls do not wait. Args: calls: The number of calls to wait for. - timeout: The longest wait of a call, in seconds. + timeout: The longest wait of a call for the overlap, in seconds. + grace: The longest wait of a call of the first overlap for one more, + in seconds. Returns: The callable, which accepts and ignores any arguments, and a dictionary whose ``peak`` item is the peak number. """ condition = threading.Condition() - state = {"active": 0, "peak": 0} + state = {"active": 0, "peak": 0, "released": False} def call(*args, **kwargs): with condition: state["active"] += 1 state["peak"] = max(state["peak"], state["active"]) condition.notify_all() - condition.wait_for(lambda: state["active"] >= calls, timeout=timeout) + # The peak, unlike the active calls, stays reached when the last + # call returns before the others wake up. + if not state["released"] and condition.wait_for( + lambda: state["peak"] >= calls, timeout=timeout + ): + condition.wait_for(lambda: state["peak"] > calls, timeout=grace) + state["released"] = True with condition: state["active"] -= 1