From ac2b8842c2496ba46477b5695252f4ff63cadfd4 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 5 Jul 2026 22:33:45 +0900 Subject: [PATCH 01/12] Write pipe_file with direct S3 API calls The inherited pipe_file went through the open() + write() buffered-file path even though the whole payload is already in memory. Override it to issue a single PutObject request for data up to the block size, and a parallel multipart upload (parts sized to fit within the 10,000-part quota, aborted on failure) for larger data. Supports the fsspec "create" mode and passes additional kwargs (e.g., ContentType) to the PutObject / CreateMultipartUpload APIs. AioS3FileSystem._pipe_file already delegates to the sync implementation and inherits the behavior. Part of the filesystem parity checklist in GH-722. Co-Authored-By: Claude Fable 5 --- pyathena/filesystem/s3.py | 86 ++++++++++++++++++++++++ tests/pyathena/filesystem/test_s3.py | 99 +++++++++++++++++++++++++++- 2 files changed, 184 insertions(+), 1 deletion(-) diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index d7da3586..6fa8fc43 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -1,6 +1,7 @@ from __future__ import annotations import logging +import math import mimetypes import os.path import re @@ -101,6 +102,9 @@ class S3FileSystem(AbstractFileSystem): # https://docs.aws.amazon.com/AmazonS3/latest/userguide/qfacts.html # The maximum size of a part in a multipart upload is 5GiB. MULTIPART_UPLOAD_MAX_PART_SIZE: int = 5 * 2**30 # 5GiB + # https://docs.aws.amazon.com/AmazonS3/latest/userguide/qfacts.html + # The maximum number of parts in a multipart upload is 10,000. + MULTIPART_UPLOAD_MAX_PARTS: int = 10000 # https://docs.aws.amazon.com/AmazonS3/latest/API/API_DeleteObjects.html DELETE_OBJECTS_MAX_KEYS: int = 1000 DEFAULT_BLOCK_SIZE: int = 5 * 2**20 # 5MiB @@ -1076,6 +1080,88 @@ def _copy_object_with_multipart_upload( parts=parts, ) + def pipe_file(self, path: str, value: bytes, mode: str = "overwrite", **kwargs) -> None: + """Write bytes into the path with direct S3 API calls. + + Uploads small data with a single PutObject request instead of the + inherited ``open()`` + ``write()`` path, and switches to a parallel + multipart upload when the data exceeds the block size. + + Args: + path: S3 path (s3://bucket/key) to write to. + value: The bytes to write. + mode: "overwrite" (default) or "create". With "create", raise + FileExistsError when the object already exists. + **kwargs: Additional parameters passed to the PutObject or + CreateMultipartUpload API (e.g., ContentType, StorageClass). + + Raises: + FileExistsError: If the mode is "create" and the path already + exists. + ValueError: If the path does not contain a key or specifies a + version. + """ + bucket, key, version_id = self.parse_path(path) + if version_id: + raise ValueError("Cannot write to the file with the version specified.") + if not key: + raise ValueError("Cannot write to a bucket.") + if mode == "create" and self.exists(path): + raise FileExistsError(path) + + size = len(value) + if size <= self.default_block_size: + self._put_object(bucket=bucket, key=key, body=value, **kwargs) + else: + # The part size must be large enough to fit within the maximum + # number of parts of a multipart upload. + part_size = max( + self.default_block_size, math.ceil(size / self.MULTIPART_UPLOAD_MAX_PARTS) + ) + ranges = S3File._get_ranges( + 0, size, max_workers=max(self.max_workers, 2), worker_block_size=part_size + ) + multipart_upload = self._create_multipart_upload(bucket=bucket, key=key, **kwargs) + upload_id = cast(str, multipart_upload.upload_id) + try: + parts: list[dict[str, Any]] = [] + with self._create_executor(max_workers=self.max_workers) as executor: + futures = [ + executor.submit( + self._upload_part, + bucket=bucket, + key=key, + upload_id=upload_id, + part_number=i + 1, + body=value[range_[0] : range_[1]], + ) + for i, range_ in enumerate(ranges) + ] + for future in as_completed(futures): + result = future.result() + parts.append( + { + "ETag": result.etag, + "PartNumber": result.part_number, + } + ) + parts.sort(key=lambda x: x["PartNumber"]) + self._complete_multipart_upload( + bucket=bucket, + key=key, + upload_id=upload_id, + parts=parts, + ) + except Exception: + self._call( + self._client.abort_multipart_upload, + Bucket=bucket, + Key=key, + UploadId=upload_id, + ) + raise + self.invalidate_cache(path) + def cat_file( self, path: str, start: int | None = None, end: int | None = None, **kwargs ) -> bytes: diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index 373214d4..18e76436 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -232,6 +232,89 @@ def test_rmdir_bucket_deletion_disabled(self): fs.rmdir("s3://bucket") fs._call.assert_not_called() + def test_pipe_file_small_uses_put_object(self): + fs = self._make_fs() + fs.default_block_size = S3FileSystem.DEFAULT_BLOCK_SIZE + fs._put_object = mock.MagicMock() + + fs.pipe_file("s3://bucket/key", b"data", ContentType="text/plain") + fs._put_object.assert_called_once_with( + bucket="bucket", key="key", body=b"data", ContentType="text/plain" + ) + + def test_pipe_file_create_mode(self): + fs = self._make_fs() + fs.default_block_size = S3FileSystem.DEFAULT_BLOCK_SIZE + fs._put_object = mock.MagicMock() + + fs.exists = mock.MagicMock(return_value=True) + with pytest.raises(FileExistsError): + fs.pipe_file("s3://bucket/key", b"data", mode="create") + fs._put_object.assert_not_called() + + fs.exists = mock.MagicMock(return_value=False) + fs.pipe_file("s3://bucket/key", b"data", mode="create") + fs._put_object.assert_called_once() + + def test_pipe_file_invalid_path_raises(self): + fs = self._make_fs() + with pytest.raises(ValueError, match="Cannot write to a bucket"): + fs.pipe_file("s3://bucket", b"data") + with pytest.raises(ValueError, match="version"): + fs.pipe_file("s3://bucket/key?versionId=12345abcde", b"data") + + def test_pipe_file_large_uses_multipart_upload(self): + fs = self._make_fs() + fs.default_block_size = 5 + fs._put_object = mock.MagicMock() + fs._create_multipart_upload = mock.MagicMock( + return_value=SimpleNamespace(upload_id="uploadid") + ) + uploaded: dict[int, bytes] = {} + + def upload_part(**kwargs): + uploaded[kwargs["part_number"]] = kwargs["body"] + return SimpleNamespace( + etag=f'"e{kwargs["part_number"]}"', part_number=kwargs["part_number"] + ) + + fs._upload_part = mock.MagicMock(side_effect=upload_part) + fs._complete_multipart_upload = mock.MagicMock() + + fs.pipe_file("s3://bucket/key", b"0123456789ABCDEF", ContentType="text/plain") + + fs._put_object.assert_not_called() + fs._create_multipart_upload.assert_called_once_with( + bucket="bucket", key="key", ContentType="text/plain" + ) + # The data is split into block-size parts and reassembled in order. + assert uploaded == {1: b"01234", 2: b"56789", 3: b"ABCDE", 4: b"F"} + fs._complete_multipart_upload.assert_called_once_with( + bucket="bucket", + key="key", + upload_id="uploadid", + parts=[{"ETag": f'"e{i}"', "PartNumber": i} for i in range(1, 5)], + ) + + def test_pipe_file_aborts_multipart_upload_on_failure(self): + fs = self._make_fs() + fs.default_block_size = 5 + fs._create_multipart_upload = mock.MagicMock( + return_value=SimpleNamespace(upload_id="uploadid") + ) + fs._upload_part = mock.MagicMock(side_effect=RuntimeError("upload failed")) + fs._complete_multipart_upload = mock.MagicMock() + + with pytest.raises(RuntimeError, match="upload failed"): + fs.pipe_file("s3://bucket/key", b"0123456789ABCDEF") + fs._complete_multipart_upload.assert_not_called() + fs._call.assert_called_once_with( + fs._client.abort_multipart_upload, + Bucket="bucket", + Key="key", + UploadId="uploadid", + ) + def test_metadata_with_version_id(self): fs = self._make_fs() fs._call.return_value = {"Metadata": {}} @@ -774,7 +857,8 @@ def test_touch(self, fs): (1, 2**10), # (10, 2**10), # (100, 2**10), - (1, 2**20), + (1, 2**20), # < block size (5 MiB): single PutObject + (6, 2**20), # > block size (5 MiB): parallel multipart upload # (10, 2**20), # (100, 2**20), # (1024, 2**20), @@ -794,6 +878,19 @@ def test_pipe_cat(self, fs, base, exp): fs.pipe(path, data) assert fs.cat(path) == data + def test_pipe_file_create_mode_and_kwargs(self, fs): + path = ( + f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" + f"filesystem/test_pipe_file_create/{uuid.uuid4()}" + ) + data = b"0123456789" + fs.pipe_file(path, data, ContentType="text/plain") + assert fs.cat(path) == data + assert fs.metadata(path).content_type == "text/plain" + + with pytest.raises(FileExistsError): + fs.pipe_file(path, data, mode="create") + def test_cat_ranges(self, fs): data = b"1234567890abcdefghijklmnopqrstuvwxyz" path = ( From aacb1d5652c65e31bae72f3736e5e13fb9f1519c Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 5 Jul 2026 22:56:01 +0900 Subject: [PATCH 02/12] Apply self-review fixes to pipe_file - Preserve fsspec transaction semantics by deferring to the buffered open() path inside a transaction - Apply the filesystem-level s3_additional_kwargs and accept the open()-style block_size / max_worker / s3_additional_kwargs kwargs - Enforce the minimum/maximum part size on the multipart path and cap the single-PutObject branch at the maximum part size - Extract _finish_multipart_upload, shared with _copy_object_with_multipart_upload, which now also cancels remaining parts and aborts the upload on failure without masking the original error (abort failures are logged) - Split parts with a plain offset loop and size the executor by the part count; accept bytes-like values - Fix S3File._get_ranges generating an empty trailing range when the size is an exact multiple of the block size (latent read-path bug) - Add mode="overwrite" to AioS3FileSystem._pipe_file to match the fsspec AsyncFileSystem signature Co-Authored-By: Claude Fable 5 --- pyathena/filesystem/s3.py | 173 +++++++++++++++++---------- pyathena/filesystem/s3_async.py | 6 +- tests/pyathena/filesystem/test_s3.py | 78 +++++++++++- 3 files changed, 194 insertions(+), 63 deletions(-) diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index 6fa8fc43..afba1f98 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -1049,9 +1049,8 @@ def _copy_object_with_multipart_upload( key=key2, **kwargs, ) - parts = [] with self._create_executor(max_workers=max_workers) as executor: - fs = [ + futures = [ executor.submit( self._upload_part_copy, bucket=bucket2, @@ -1063,29 +1062,23 @@ def _copy_object_with_multipart_upload( ) for i, range_ in enumerate(ranges) ] - for f in as_completed(fs): - result = f.result() - parts.append( - { - "ETag": result.etag, - "PartNumber": result.part_number, - } - ) - - parts.sort(key=lambda x: x["PartNumber"]) # type: ignore[arg-type, return-value] - self._complete_multipart_upload( - bucket=bucket2, - key=key2, - upload_id=cast(str, multipart_upload.upload_id), - parts=parts, - ) + self._finish_multipart_upload( + bucket=bucket2, + key=key2, + upload_id=cast(str, multipart_upload.upload_id), + futures=futures, + ) - def pipe_file(self, path: str, value: bytes, mode: str = "overwrite", **kwargs) -> None: + def pipe_file( + self, path: str, value: bytes | bytearray | memoryview, mode: str = "overwrite", **kwargs + ) -> None: """Write bytes into the path with direct S3 API calls. Uploads small data with a single PutObject request instead of the inherited ``open()`` + ``write()`` path, and switches to a parallel - multipart upload when the data exceeds the block size. + multipart upload when the data exceeds the block size. Inside an + fsspec transaction, the write goes through the inherited buffered + path so that the deferred-commit semantics are kept. Args: path: S3 path (s3://bucket/key) to write to. @@ -1094,6 +1087,11 @@ def pipe_file(self, path: str, value: bytes, mode: str = "overwrite", **kwargs) FileExistsError when the object already exists. **kwargs: Additional parameters passed to the PutObject or CreateMultipartUpload API (e.g., ContentType, StorageClass). + The ``block_size``, ``max_worker``, and + ``s3_additional_kwargs`` parameters of the ``open()`` path + are also accepted. Note that parameters that must be + repeated on every UploadPart request (e.g., SSE-C keys) are + not supported on the multipart path. Raises: FileExistsError: If the mode is "create" and the path already @@ -1101,6 +1099,11 @@ def pipe_file(self, path: str, value: bytes, mode: str = "overwrite", **kwargs) ValueError: If the path does not contain a key or specifies a version. """ + if self._intrans: + # Defer to the buffered open() path so that fsspec transactions + # keep their deferred-commit semantics. + super().pipe_file(path, value, mode=mode, **kwargs) + return bucket, key, version_id = self.parse_path(path) if version_id: raise ValueError("Cannot write to the file with the version specified.") @@ -1108,59 +1111,107 @@ def pipe_file(self, path: str, value: bytes, mode: str = "overwrite", **kwargs) raise ValueError("Cannot write to a bucket.") if mode == "create" and self.exists(path): raise FileExistsError(path) + if not isinstance(value, bytes): + # Accept bytes-like values (bytearray, memoryview) as the + # inherited buffered path did. + value = bytes(value) + + block_size = kwargs.pop("block_size", None) or self.default_block_size + max_worker = kwargs.pop("max_worker", None) or self.max_workers + request_kwargs = { + **self.s3_additional_kwargs, + **kwargs.pop("s3_additional_kwargs", {}), + **kwargs, + } size = len(value) - if size <= self.default_block_size: - self._put_object(bucket=bucket, key=key, body=value, **kwargs) + # A single PutObject request accepts up to the maximum part size. + if size <= min(block_size, self.MULTIPART_UPLOAD_MAX_PART_SIZE): + self._put_object(bucket=bucket, key=key, body=value, **request_kwargs) else: - # The part size must be large enough to fit within the maximum - # number of parts of a multipart upload. - part_size = max( - self.default_block_size, math.ceil(size / self.MULTIPART_UPLOAD_MAX_PARTS) + # The part size must satisfy the minimum/maximum part size and + # fit within the maximum number of parts of a multipart upload. + part_size = min( + max( + block_size, + self.MULTIPART_UPLOAD_MIN_PART_SIZE, + math.ceil(size / self.MULTIPART_UPLOAD_MAX_PARTS), + ), + self.MULTIPART_UPLOAD_MAX_PART_SIZE, ) - ranges = S3File._get_ranges( - 0, size, max_workers=max(self.max_workers, 2), worker_block_size=part_size + offsets = range(0, size, part_size) + multipart_upload = self._create_multipart_upload( + bucket=bucket, key=key, **request_kwargs ) - multipart_upload = self._create_multipart_upload(bucket=bucket, key=key, **kwargs) - upload_id = cast(str, multipart_upload.upload_id) - try: - parts: list[dict[str, Any]] = [] - with self._create_executor(max_workers=self.max_workers) as executor: - futures = [ - executor.submit( - self._upload_part, - bucket=bucket, - key=key, - upload_id=upload_id, - part_number=i + 1, - body=value[range_[0] : range_[1]], - ) - for i, range_ in enumerate(ranges) - ] - for future in as_completed(futures): - result = future.result() - parts.append( - { - "ETag": result.etag, - "PartNumber": result.part_number, - } - ) - parts.sort(key=lambda x: x["PartNumber"]) - self._complete_multipart_upload( + with self._create_executor(max_workers=min(len(offsets), max_worker)) as executor: + futures = [ + executor.submit( + self._upload_part, + bucket=bucket, + key=key, + upload_id=cast(str, multipart_upload.upload_id), + part_number=i + 1, + body=value[offset : offset + part_size], + ) + for i, offset in enumerate(offsets) + ] + self._finish_multipart_upload( bucket=bucket, key=key, - upload_id=upload_id, - parts=parts, + upload_id=cast(str, multipart_upload.upload_id), + futures=futures, ) - except Exception: + self.invalidate_cache(path) + + def _finish_multipart_upload( + self, + bucket: str, + key: str, + upload_id: str, + futures: list[Future[S3MultipartUploadPart]], + ) -> S3CompleteMultipartUpload: + """Collect the uploaded parts and complete the multipart upload. + + When any part fails, the remaining parts are cancelled and the + multipart upload is aborted so that no incomplete upload is left + behind, then the original error is re-raised. + + Args: + bucket: S3 bucket name. + key: Object key being uploaded. + upload_id: Unique identifier for the multipart upload. + futures: Futures of the part uploads, in part-number order. + + Returns: + S3CompleteMultipartUpload of the completed upload. + """ + try: + # The futures are in part-number order. + parts = [ + {"ETag": result.etag, "PartNumber": result.part_number} + for result in (future.result() for future in futures) + ] + return self._complete_multipart_upload( + bucket=bucket, + key=key, + upload_id=upload_id, + parts=parts, + ) + except Exception: + for future in futures: + future.cancel() + try: self._call( self._client.abort_multipart_upload, Bucket=bucket, Key=key, UploadId=upload_id, ) - raise - self.invalidate_cache(path) + except Exception: + _logger.exception( + "Failed to abort multipart upload %s to s3://%s/%s.", upload_id, bucket, key + ) + raise def cat_file( self, path: str, start: int | None = None, end: int | None = None, **kwargs @@ -2175,7 +2226,9 @@ def _get_ranges( range_start = start while True: range_end = range_start + worker_block_size - if range_end > end: + if range_end >= end: + # Also when the size is an exact multiple of the block + # size, so that no empty trailing range is generated. ranges.append((range_start, end)) break ranges.append((range_start, range_end)) diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index ceb820fc..4879f3f3 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -111,8 +111,10 @@ async def _exists(self, path: str, **kwargs) -> bool: async def _rm_file(self, path: str, **kwargs) -> None: await asyncio.to_thread(self._sync_fs.rm_file, path, **kwargs) - async def _pipe_file(self, path: str, value: bytes, **kwargs) -> None: - await asyncio.to_thread(self._sync_fs.pipe_file, path, value, **kwargs) + async def _pipe_file( + self, path: str, value: bytes | bytearray | memoryview, mode: str = "overwrite", **kwargs + ) -> None: + await asyncio.to_thread(self._sync_fs.pipe_file, path, value, mode=mode, **kwargs) async def _put_file(self, lpath: str, rpath: str, callback=_DEFAULT_CALLBACK, **kwargs) -> None: await asyncio.to_thread(self._sync_fs.put_file, lpath, rpath, callback=callback, **kwargs) diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index 18e76436..b372eda8 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -142,6 +142,8 @@ def _make_fs(): fs.max_workers = 4 fs.allow_bucket_creation = False fs.allow_bucket_deletion = False + fs.s3_additional_kwargs = {} + fs._intrans = False return fs def test_ls_from_cache_with_cached_object(self): @@ -235,11 +237,18 @@ def test_rmdir_bucket_deletion_disabled(self): def test_pipe_file_small_uses_put_object(self): fs = self._make_fs() fs.default_block_size = S3FileSystem.DEFAULT_BLOCK_SIZE + fs.s3_additional_kwargs = {"ServerSideEncryption": "AES256"} fs._put_object = mock.MagicMock() + # The filesystem-level s3_additional_kwargs are merged with the + # call-level kwargs, as in the open() path. fs.pipe_file("s3://bucket/key", b"data", ContentType="text/plain") fs._put_object.assert_called_once_with( - bucket="bucket", key="key", body=b"data", ContentType="text/plain" + bucket="bucket", + key="key", + body=b"data", + ServerSideEncryption="AES256", + ContentType="text/plain", ) def test_pipe_file_create_mode(self): @@ -266,6 +275,10 @@ def test_pipe_file_invalid_path_raises(self): def test_pipe_file_large_uses_multipart_upload(self): fs = self._make_fs() fs.default_block_size = 5 + # Shrink the part-size limits so the test data is split without + # allocating real 5 MiB payloads. + fs.MULTIPART_UPLOAD_MIN_PART_SIZE = 5 + fs.MULTIPART_UPLOAD_MAX_PART_SIZE = 2**30 fs._put_object = mock.MagicMock() fs._create_multipart_upload = mock.MagicMock( return_value=SimpleNamespace(upload_id="uploadid") @@ -299,6 +312,8 @@ def upload_part(**kwargs): def test_pipe_file_aborts_multipart_upload_on_failure(self): fs = self._make_fs() fs.default_block_size = 5 + fs.MULTIPART_UPLOAD_MIN_PART_SIZE = 5 + fs.MULTIPART_UPLOAD_MAX_PART_SIZE = 2**30 fs._create_multipart_upload = mock.MagicMock( return_value=SimpleNamespace(upload_id="uploadid") ) @@ -315,6 +330,36 @@ def test_pipe_file_aborts_multipart_upload_on_failure(self): UploadId="uploadid", ) + def test_pipe_file_abort_failure_does_not_mask_the_original_error(self): + fs = self._make_fs() + fs.default_block_size = 5 + fs.MULTIPART_UPLOAD_MIN_PART_SIZE = 5 + fs.MULTIPART_UPLOAD_MAX_PART_SIZE = 2**30 + fs._create_multipart_upload = mock.MagicMock( + return_value=SimpleNamespace(upload_id="uploadid") + ) + fs._upload_part = mock.MagicMock(side_effect=RuntimeError("upload failed")) + fs._complete_multipart_upload = mock.MagicMock() + fs._call = mock.MagicMock(side_effect=RuntimeError("abort failed")) + + # The abort failure is logged, and the original error propagates. + with pytest.raises(RuntimeError, match="upload failed"): + fs.pipe_file("s3://bucket/key", b"0123456789ABCDEF") + + def test_pipe_file_in_transaction_uses_buffered_path(self): + fs = self._make_fs() + fs._intrans = True + fs._put_object = mock.MagicMock() + opened = mock.MagicMock() + fs.open = mock.MagicMock(return_value=opened) + + fs.pipe_file("s3://bucket/key", b"data") + # The buffered open() path is used so that fsspec transactions keep + # their deferred-commit semantics; nothing is written directly. + fs.open.assert_called_once() + opened.__enter__.return_value.write.assert_called_once_with(b"data") + fs._put_object.assert_not_called() + def test_metadata_with_version_id(self): fs = self._make_fs() fs._call.return_value = {"Metadata": {}} @@ -891,6 +936,33 @@ def test_pipe_file_create_mode_and_kwargs(self, fs): with pytest.raises(FileExistsError): fs.pipe_file(path, data, mode="create") + def test_pipe_file_transaction(self, fs): + # Inside an fsspec transaction the write is deferred to the commit. + data = b"0123456789" + path = ( + f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" + f"filesystem/test_pipe_file_transaction/{uuid.uuid4()}" + ) + with fs.transaction: + fs.pipe_file(path, data) + assert fs.cat(path) == data + + # Raising inside the transaction must leave no object behind. + path = ( + f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" + f"filesystem/test_pipe_file_transaction/{uuid.uuid4()}" + ) + + def write_then_fail(): + with fs.transaction: + fs.pipe_file(path, data) + raise RuntimeError("rollback") + + with pytest.raises(RuntimeError): + write_then_fail() + fs.invalidate_cache(path) + assert not fs.exists(path) + def test_cat_ranges(self, fs): data = b"1234567890abcdefghijklmnopqrstuvwxyz" path = ( @@ -1375,6 +1447,10 @@ def test_merge_objects(self, objects, target): (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 + # The size is an exact multiple of the block size: + # no empty trailing range is generated. + (0, 10, 2, 5, [(0, 5), (5, 10)]), + (42, 2632, 2, 1295, [(42, 1337), (1337, 2632)]), ], ) def test_get_ranges(self, start, end, max_workers, worker_block_size, ranges): From 9ec4e09f42d44ddc50ac90410c740d774b96052e Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 5 Jul 2026 23:09:43 +0900 Subject: [PATCH 03/12] Respect explicit fsspec registrations when registering the s3 protocol pyathena.pandas / pyathena.polars unconditionally clobbered the fsspec "s3" / "s3a" protocols at import time, silently shadowing s3fs (GH-719). The registration is still required for the pandas/polars result sets, but it now goes through pyathena.filesystem.register_s3_filesystem, which leaves a filesystem class that the user has already registered explicitly (present in fsspec.registry) untouched, logs what it does, and documents how to restore s3fs. Co-Authored-By: Claude Fable 5 --- pyathena/filesystem/__init__.py | 41 ++++++++++++++++++++++ pyathena/pandas/__init__.py | 5 ++- pyathena/polars/__init__.py | 5 ++- tests/pyathena/filesystem/test_registry.py | 39 ++++++++++++++++++++ 4 files changed, 84 insertions(+), 6 deletions(-) create mode 100644 tests/pyathena/filesystem/test_registry.py diff --git a/pyathena/filesystem/__init__.py b/pyathena/filesystem/__init__.py index e69de29b..2c6a9ee4 100644 --- a/pyathena/filesystem/__init__.py +++ b/pyathena/filesystem/__init__.py @@ -0,0 +1,41 @@ +import logging + +import fsspec + +_logger = logging.getLogger(__name__) + +_S3_FILESYSTEM_CLASS = "pyathena.filesystem.s3.S3FileSystem" + + +def register_s3_filesystem(clobber: bool = True) -> None: + """Register PyAthena's S3 filesystem as fsspec's "s3" / "s3a" protocols. + + PyAthena registers its own filesystem so that the pandas/polars result + sets can read query results from S3 without depending on s3fs. The + registration replaces fsspec's default lazy mapping of the "s3" protocol + to s3fs, which means ``fsspec.filesystem("s3")`` returns PyAthena's + implementation and s3fs-specific settings (e.g., the ``S3FS_LOGGING_LEVEL`` + environment variable) have no effect. + + A filesystem class that has already been registered explicitly (present + in ``fsspec.registry``) is left untouched, so registering another + implementation before importing ``pyathena.pandas`` / ``pyathena.polars`` + opts out of the replacement. To restore s3fs afterwards:: + + fsspec.register_implementation("s3", s3fs.S3FileSystem, clobber=True) + + Args: + clobber: Whether to replace an existing lazy registration of the + protocols. + """ + for protocol in ("s3", "s3a"): + if protocol in fsspec.registry: + # A filesystem class has been registered explicitly + # (e.g., the user registered s3fs); leave it alone. + _logger.debug( + "The %r protocol is already registered in the fsspec registry, skipping.", + protocol, + ) + continue + _logger.debug("Registering %s as the fsspec %r protocol.", _S3_FILESYSTEM_CLASS, protocol) + fsspec.register_implementation(protocol, _S3_FILESYSTEM_CLASS, clobber=clobber) diff --git a/pyathena/pandas/__init__.py b/pyathena/pandas/__init__.py index d2a6d61c..5c077197 100644 --- a/pyathena/pandas/__init__.py +++ b/pyathena/pandas/__init__.py @@ -1,4 +1,3 @@ -import fsspec +from pyathena.filesystem import register_s3_filesystem -fsspec.register_implementation("s3", "pyathena.filesystem.s3.S3FileSystem", clobber=True) -fsspec.register_implementation("s3a", "pyathena.filesystem.s3.S3FileSystem", clobber=True) +register_s3_filesystem() diff --git a/pyathena/polars/__init__.py b/pyathena/polars/__init__.py index d2a6d61c..5c077197 100644 --- a/pyathena/polars/__init__.py +++ b/pyathena/polars/__init__.py @@ -1,4 +1,3 @@ -import fsspec +from pyathena.filesystem import register_s3_filesystem -fsspec.register_implementation("s3", "pyathena.filesystem.s3.S3FileSystem", clobber=True) -fsspec.register_implementation("s3a", "pyathena.filesystem.s3.S3FileSystem", clobber=True) +register_s3_filesystem() diff --git a/tests/pyathena/filesystem/test_registry.py b/tests/pyathena/filesystem/test_registry.py new file mode 100644 index 00000000..ad06185e --- /dev/null +++ b/tests/pyathena/filesystem/test_registry.py @@ -0,0 +1,39 @@ +import fsspec +import pytest +from fsspec.registry import _registry, known_implementations + +from pyathena.filesystem import register_s3_filesystem +from pyathena.filesystem.s3 import S3FileSystem + + +@pytest.fixture +def registry_state(): + registry = dict(_registry) + known = {k: dict(v) for k, v in known_implementations.items()} + yield + _registry.clear() + _registry.update(registry) + known_implementations.clear() + known_implementations.update(known) + + +def test_register_s3_filesystem(registry_state): + _registry.pop("s3", None) + _registry.pop("s3a", None) + known_implementations["s3"] = {"class": "s3fs.S3FileSystem", "err": "Install s3fs"} + known_implementations.pop("s3a", None) + + register_s3_filesystem() + assert known_implementations["s3"]["class"] == "pyathena.filesystem.s3.S3FileSystem" + assert known_implementations["s3a"]["class"] == "pyathena.filesystem.s3.S3FileSystem" + + +def test_register_s3_filesystem_respects_explicit_registration(registry_state): + # A filesystem class registered explicitly in the live fsspec registry + # (e.g., the user registered s3fs) is left untouched. + fsspec.register_implementation("s3", S3FileSystem, clobber=True) + known_implementations["s3"] = {"class": "s3fs.S3FileSystem", "err": "Install s3fs"} + + register_s3_filesystem() + assert known_implementations["s3"]["class"] == "s3fs.S3FileSystem" + assert fsspec.registry["s3"] is S3FileSystem From b477d6adf62ae8073fbb9051d3a0940e60822945 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 5 Jul 2026 23:16:54 +0900 Subject: [PATCH 04/12] Add version-aware mode and object version listing - S3FileSystem/AioS3FileSystem accept version_aware (default False): reads pin the object version observed at open/head time, and ls(versions=True) lists all versions of a prefix - object_version_info lists the versions of the objects under a path as typed S3ObjectVersion instances (paginated, optional delete markers) - Add the s3:GetObjectVersion / s3:ListBucketVersions permissions to the GitHub Actions OIDC role for the integration tests Part of the filesystem parity checklist in GH-722. Co-Authored-By: Claude Fable 5 --- pyathena/filesystem/s3.py | 144 +++++++++++++++++- pyathena/filesystem/s3_async.py | 14 +- pyathena/filesystem/s3_object.py | 74 +++++++++ .../cloudformation/github_actions_oidc.yaml | 2 + tests/pyathena/filesystem/test_s3.py | 121 +++++++++++++++ tests/pyathena/filesystem/test_s3_object.py | 49 ++++++ 6 files changed, 402 insertions(+), 2 deletions(-) diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index afba1f98..dbb70d98 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -32,6 +32,7 @@ S3MultipartUploadPart, S3Object, S3ObjectType, + S3ObjectVersion, S3PutObject, S3StorageClass, ) @@ -72,6 +73,10 @@ class S3FileSystem(AbstractFileSystem): Defaults to False. allow_bucket_deletion: Whether rmdir may delete buckets. Defaults to False. + version_aware: Whether reads pin the object version observed at + open time and ls may list all versions. Requires the + s3:GetObjectVersion / s3:ListBucketVersions permissions. + Defaults to False. Example: >>> from pyathena.filesystem.s3 import S3FileSystem @@ -140,6 +145,7 @@ def __init__( s3_additional_kwargs=None, allow_bucket_creation: bool = False, allow_bucket_deletion: bool = False, + version_aware: bool = False, *args, **kwargs, ) -> None: @@ -163,6 +169,7 @@ def __init__( self.s3_additional_kwargs = s3_additional_kwargs if s3_additional_kwargs else {} self.allow_bucket_creation = allow_bucket_creation self.allow_bucket_deletion = allow_bucket_deletion + self.version_aware = version_aware requester_pays = kwargs.pop("requester_pays", False) self.request_kwargs = {"RequestPayer": "requester"} if requester_pays else {} @@ -273,6 +280,10 @@ def _head_object( ) except FileNotFoundError: return None + if self.version_aware and not version_id: + # Pin the version of the object so that subsequent reads see + # the version observed here even if the object is overwritten. + version_id = response.get("VersionId") file = S3Object( init=response, type=S3ObjectType.S3_OBJECT_TYPE_FILE, @@ -388,7 +399,9 @@ def ls( path: S3 path to list (e.g., "s3://bucket" or "s3://bucket/prefix"). detail: If True, return S3Object instances; if False, return paths as strings. refresh: If True, bypass cache and fetch fresh results from S3. - **kwargs: Additional arguments (ignored for S3). + **kwargs: Additional arguments including: + versions: If True, list all versions of the objects. Requires + the filesystem to be constructed with ``version_aware=True``. Returns: List of S3Object instances (if detail=True) or paths as strings (if detail=False). @@ -398,9 +411,16 @@ def ls( >>> fs.ls("s3://my-bucket") # List objects in bucket >>> fs.ls("s3://my-bucket/", detail=True) # Get detailed object info """ + versions = kwargs.pop("versions", False) + if versions and not self.version_aware: + raise ValueError( + "Cannot list the object versions unless the filesystem is version aware." + ) path = self._strip_protocol(path).rstrip("/") if path in ["", "/"]: files = self._ls_buckets(refresh) + elif versions: + files = self._ls_object_versions(path) else: files = self._ls_dirs(path, refresh=refresh) if not files and "/" in path: @@ -409,6 +429,70 @@ def ls( files = [file] return list(files) if detail else [f.name for f in files] + def _ls_object_versions(self, path: str) -> list[S3Object]: + """List a prefix including all versions of the objects. + + The listing is always fetched from S3 and is not cached, because the + dircache stores the current view of a path. + """ + bucket, key, _ = self.parse_path(path) + prefix = f"{key}/" if key else "" + + files: list[S3Object] = [] + next_key_marker: str | None = None + next_version_id_marker: str | None = None + while True: + request: dict[str, Any] = { + "Bucket": bucket, + "Prefix": prefix, + "Delimiter": "/", + } + if next_key_marker: + request.update( + { + "KeyMarker": next_key_marker, + "VersionIdMarker": next_version_id_marker, + } + ) + response = self._call( + self._client.list_object_versions, + **request, + ) + files.extend( + S3Object( + init={ + "ContentLength": 0, + "ContentType": None, + "StorageClass": S3StorageClass.S3_STORAGE_CLASS_DIRECTORY, + "ETag": None, + "LastModified": None, + }, + type=S3ObjectType.S3_OBJECT_TYPE_DIRECTORY, + bucket=bucket, + key=c["Prefix"][:-1].rstrip("/"), + version_id=None, + ) + for c in response.get("CommonPrefixes", []) + ) + files.extend( + S3Object( + init=v, + type=S3ObjectType.S3_OBJECT_TYPE_FILE, + bucket=bucket, + key=v["Key"], + version_id=v.get("VersionId"), + is_latest=v.get("IsLatest", False), + ) + for v in response.get("Versions", []) + ) + if not response.get("IsTruncated"): + break + next_key_marker = response.get("NextKeyMarker") + next_version_id_marker = response.get("NextVersionIdMarker") + if not next_key_marker: + break + return files + def info(self, path: str, **kwargs) -> S3Object: refresh = kwargs.pop("refresh", False) path = self._strip_protocol(path) @@ -1640,6 +1724,60 @@ def list_multipart_uploads(self, path: str) -> list[S3MultipartUpload]: break return uploads + def object_version_info( + self, path: str, delete_markers: bool = False, **kwargs + ) -> list[S3ObjectVersion]: + """List the versions of the objects under the path. + + Args: + path: S3 path (s3://bucket/key or a key prefix) to list the + versions for. + delete_markers: Whether to include delete markers in the result. + **kwargs: Additional parameters passed to the ListObjectVersions + API. + + Returns: + List of S3ObjectVersion instances describing the versions. + """ + bucket, key, _ = self.parse_path(path) + + _logger.debug("List object versions: s3://%s/%s", bucket, key) + versions: list[S3ObjectVersion] = [] + next_key_marker: str | None = None + next_version_id_marker: str | None = None + while True: + request: dict[str, Any] = {"Bucket": bucket} + if key: + request.update({"Prefix": key}) + if next_key_marker: + request.update( + { + "KeyMarker": next_key_marker, + "VersionIdMarker": next_version_id_marker, + } + ) + response = self._call( + self._client.list_object_versions, + **request, + **kwargs, + ) + versions.extend( + S3ObjectVersion(bucket=bucket, is_delete_marker=False, response=v) + for v in response.get("Versions", []) + ) + if delete_markers: + versions.extend( + S3ObjectVersion(bucket=bucket, is_delete_marker=True, response=m) + for m in response.get("DeleteMarkers", []) + ) + if not response.get("IsTruncated"): + break + next_key_marker = response.get("NextKeyMarker") + next_version_id_marker = response.get("NextVersionIdMarker") + if not next_key_marker: + break + return versions + def clear_multipart_uploads(self, path: str) -> None: """Abort any incomplete multipart uploads in the bucket. @@ -1953,6 +2091,10 @@ def __init__( self._details: S3Object | dict[str, Any] if "r" in mode: info = self.fs.info(self.path, version_id=self.version_id) + if self.fs.version_aware and not self.version_id: + # Pin the version observed at open time so that reads are + # consistent even if the object is overwritten. + self.version_id = info.get("version_id") if etag := info.get("etag"): self.s3_additional_kwargs.update({"IfMatch": etag}) self._details = info diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index 4879f3f3..f968adfe 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -10,7 +10,12 @@ from pyathena.filesystem.s3 import S3File, S3FileSystem from pyathena.filesystem.s3_executor import S3AioExecutor -from pyathena.filesystem.s3_object import S3Metadata, S3MultipartUpload, S3Object +from pyathena.filesystem.s3_object import ( + S3Metadata, + S3MultipartUpload, + S3Object, + S3ObjectVersion, +) if TYPE_CHECKING: from datetime import datetime @@ -66,6 +71,7 @@ def __init__( s3_additional_kwargs: dict[str, Any] | None = None, allow_bucket_creation: bool = False, allow_bucket_deletion: bool = False, + version_aware: bool = False, asynchronous: bool = False, loop: Any | None = None, batch_size: int | None = None, @@ -85,6 +91,7 @@ def __init__( s3_additional_kwargs=s3_additional_kwargs, allow_bucket_creation=allow_bucket_creation, allow_bucket_deletion=allow_bucket_deletion, + version_aware=version_aware, **kwargs, ) # Share dircache for cache coherence between async and sync instances @@ -355,6 +362,11 @@ def put_tags(self, path: str, tags: dict[str, str], mode: str = "o") -> None: def chmod(self, path: str, acl: str, recursive: bool = False, **kwargs) -> None: self._sync_fs.chmod(path, acl, recursive=recursive, **kwargs) + def object_version_info( + self, path: str, delete_markers: bool = False, **kwargs + ) -> list[S3ObjectVersion]: + return self._sync_fs.object_version_info(path, delete_markers=delete_markers, **kwargs) + def list_multipart_uploads(self, path: str) -> list[S3MultipartUpload]: return self._sync_fs.list_multipart_uploads(path) diff --git a/pyathena/filesystem/s3_object.py b/pyathena/filesystem/s3_object.py index 0c55d960..192b7373 100644 --- a/pyathena/filesystem/s3_object.py +++ b/pyathena/filesystem/s3_object.py @@ -318,6 +318,80 @@ def user_metadata(self) -> dict[str, str]: return dict(self._user_metadata) +class S3ObjectVersion: + """Represents a version of an S3 object as returned by ListObjectVersions. + + Attributes: + bucket: S3 bucket name. + key: Object key. + version_id: The version ID of the object. + is_latest: Whether the version is the latest version of the object. + is_delete_marker: Whether the version is a delete marker. + last_modified: Date and time when the version was last modified. + etag: Entity tag of the version. None for delete markers. + size: Size in bytes of the version. None for delete markers. + storage_class: Storage class of the version. None for delete markers. + owner: Owner of the version. + """ + + def __init__(self, bucket: str, is_delete_marker: bool, response: dict[str, Any]) -> None: + self._bucket = bucket + self._is_delete_marker = is_delete_marker + self._key: str = response["Key"] + self._version_id: str | None = response.get("VersionId") + self._is_latest: bool = response.get("IsLatest", False) + self._last_modified: datetime | None = response.get("LastModified") + self._etag: str | None = response.get("ETag") + self._size: int | None = response.get("Size") + self._storage_class: str | None = response.get("StorageClass") + owner = response.get("Owner") + self._owner: S3Owner | None = S3Owner(owner) if owner else None + + @property + def bucket(self) -> str: + return self._bucket + + @property + def key(self) -> str: + return self._key + + @property + def name(self) -> str: + return f"{self._bucket}/{self._key}" + + @property + def version_id(self) -> str | None: + return self._version_id + + @property + def is_latest(self) -> bool: + return self._is_latest + + @property + def is_delete_marker(self) -> bool: + return self._is_delete_marker + + @property + def last_modified(self) -> datetime | None: + return self._last_modified + + @property + def etag(self) -> str | None: + return self._etag + + @property + def size(self) -> int | None: + return self._size + + @property + def storage_class(self) -> str | None: + return self._storage_class + + @property + def owner(self) -> S3Owner | None: + return self._owner + + class S3PutObject: """Represents the response from an S3 PUT object operation. diff --git a/scripts/cloudformation/github_actions_oidc.yaml b/scripts/cloudformation/github_actions_oidc.yaml index e6ac5a0b..0212439f 100644 --- a/scripts/cloudformation/github_actions_oidc.yaml +++ b/scripts/cloudformation/github_actions_oidc.yaml @@ -105,7 +105,9 @@ Resources: "s3:GetBucketLocation", "s3:GetObject", "s3:GetObjectTagging", + "s3:GetObjectVersion", "s3:ListBucket", + "s3:ListBucketVersions", "s3:ListBucketMultipartUploads", "s3:ListMultipartUploadParts", "s3:AbortMultipartUpload", diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index b372eda8..8d666287 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -144,6 +144,7 @@ def _make_fs(): fs.allow_bucket_deletion = False fs.s3_additional_kwargs = {} fs._intrans = False + fs.version_aware = False return fs def test_ls_from_cache_with_cached_object(self): @@ -360,6 +361,98 @@ def test_pipe_file_in_transaction_uses_buffered_path(self): opened.__enter__.return_value.write.assert_called_once_with(b"data") fs._put_object.assert_not_called() + def test_head_object_version_aware(self): + fs = self._make_fs() + fs._call.return_value = {"ContentLength": 4, "ETag": '"etag"', "VersionId": "v1"} + + # The observed version is pinned only in version-aware mode. + assert fs._head_object("bucket/key").version_id is None + + fs = self._make_fs() + fs.version_aware = True + fs._call.return_value = {"ContentLength": 4, "ETag": '"etag"', "VersionId": "v1"} + assert fs._head_object("bucket/key").version_id == "v1" + + def test_object_version_info_paginates(self): + fs = self._make_fs() + fs._call.side_effect = [ + { + "Versions": [ + {"Key": "key", "VersionId": "v2", "IsLatest": True, "Size": 4}, + ], + "DeleteMarkers": [ + {"Key": "key", "VersionId": "m1", "IsLatest": False}, + ], + "IsTruncated": True, + "NextKeyMarker": "key", + "NextVersionIdMarker": "v2", + }, + { + "Versions": [ + {"Key": "key", "VersionId": "v1", "IsLatest": False, "Size": 2}, + ], + "IsTruncated": False, + }, + ] + + actual = fs.object_version_info("s3://bucket/key") + assert [(v.key, v.version_id, v.is_latest, v.size) for v in actual] == [ + ("key", "v2", True, 4), + ("key", "v1", False, 2), + ] + assert all(not v.is_delete_marker for v in actual) + fs._call.assert_any_call(fs._client.list_object_versions, Bucket="bucket", Prefix="key") + fs._call.assert_any_call( + fs._client.list_object_versions, + Bucket="bucket", + Prefix="key", + KeyMarker="key", + VersionIdMarker="v2", + ) + + def test_object_version_info_with_delete_markers(self): + fs = self._make_fs() + fs._call.return_value = { + "Versions": [{"Key": "key", "VersionId": "v1", "IsLatest": False}], + "DeleteMarkers": [{"Key": "key", "VersionId": "m1", "IsLatest": True}], + "IsTruncated": False, + } + + actual = fs.object_version_info("s3://bucket/key", delete_markers=True) + assert [(v.version_id, v.is_delete_marker) for v in actual] == [ + ("v1", False), + ("m1", True), + ] + + def test_ls_versions_requires_version_aware(self): + fs = self._make_fs() + with pytest.raises(ValueError, match="version aware"): + fs.ls("s3://bucket/path", versions=True) + + def test_ls_versions(self): + fs = self._make_fs() + fs.version_aware = True + fs._call.return_value = { + "CommonPrefixes": [{"Prefix": "path/dir/"}], + "Versions": [ + {"Key": "path/key", "VersionId": "v2", "IsLatest": True, "Size": 4}, + {"Key": "path/key", "VersionId": "v1", "IsLatest": False, "Size": 2}, + ], + "IsTruncated": False, + } + + actual = fs.ls("s3://bucket/path", detail=True, versions=True) + fs._call.assert_called_once_with( + fs._client.list_object_versions, Bucket="bucket", Prefix="path/", Delimiter="/" + ) + assert [(f.name, f.version_id, f.is_latest) for f in actual] == [ + ("bucket/path/dir", None, None), + ("bucket/path/key", "v2", True), + ("bucket/path/key", "v1", False), + ] + assert actual[1].size == 4 + assert actual[2].size == 2 + def test_metadata_with_version_id(self): fs = self._make_fs() fs._call.return_value = {"Metadata": {}} @@ -1310,6 +1403,34 @@ def test_list_and_clear_multipart_uploads(self, fs): uploads = fs.list_multipart_uploads(prefix_path) assert not any(u.upload_id == upload.upload_id for u in uploads) + def test_object_version_info(self, fs): + path = ( + f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" + f"filesystem/test_object_version_info/{uuid.uuid4()}" + ) + fs.pipe(path, b"data") + + versions = fs.object_version_info(path) + assert len(versions) == 1 + version = versions[0] + bucket, key, _ = fs.parse_path(path) + assert version.bucket == bucket + assert version.key == key + assert version.name == f"{bucket}/{key}" + assert version.is_latest + assert not version.is_delete_marker + assert version.size == 4 + # An unversioned bucket reports the "null" version. + assert version.version_id + + @pytest.mark.parametrize("fs", [{"version_aware": True}], indirect=True) + def test_version_aware_read(self, fs): + # On an unversioned bucket, the version-aware mode is a no-op for + # reads: HeadObject returns no version to pin. + path = f"s3://{ENV.s3_staging_bucket}/{ENV.s3_filesystem_test_file_key}" + with fs.open(path, "rb") as f: + assert f.read() == b"0123456789" + def test_file_url_metadata_getxattr_setxattr(self, fs): path = ( f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" diff --git a/tests/pyathena/filesystem/test_s3_object.py b/tests/pyathena/filesystem/test_s3_object.py index 9521ac73..bb901fca 100644 --- a/tests/pyathena/filesystem/test_s3_object.py +++ b/tests/pyathena/filesystem/test_s3_object.py @@ -7,6 +7,7 @@ S3MultipartUploadPart, S3Object, S3ObjectType, + S3ObjectVersion, S3PutObject, S3StorageClass, ) @@ -166,6 +167,54 @@ def test_mapping_interface(self): assert "content_type" not in actual +class TestS3ObjectVersion: + def test_init(self): + actual = S3ObjectVersion( + bucket="test_bucket", + is_delete_marker=False, + response={ + "Key": "test_key", + "VersionId": "test_version_id", + "IsLatest": True, + "LastModified": datetime(2015, 1, 1, 0, 0, 0), + "ETag": '"test_etag"', + "Size": 100, + "StorageClass": "STANDARD", + "Owner": {"DisplayName": "test_owner", "ID": "test_owner_id"}, + }, + ) + assert actual.bucket == "test_bucket" + assert actual.key == "test_key" + assert actual.name == "test_bucket/test_key" + assert actual.version_id == "test_version_id" + assert actual.is_latest is True + assert actual.is_delete_marker is False + assert actual.last_modified == datetime(2015, 1, 1, 0, 0, 0) + assert actual.etag == '"test_etag"' + assert actual.size == 100 + assert actual.storage_class == "STANDARD" + assert actual.owner + assert actual.owner.display_name == "test_owner" + assert actual.owner.id == "test_owner_id" + + def test_init_delete_marker(self): + actual = S3ObjectVersion( + bucket="test_bucket", + is_delete_marker=True, + response={ + "Key": "test_key", + "VersionId": "test_version_id", + "IsLatest": True, + "LastModified": datetime(2015, 1, 1, 0, 0, 0), + }, + ) + assert actual.is_delete_marker is True + assert actual.etag is None + assert actual.size is None + assert actual.storage_class is None + assert actual.owner is None + + class TestS3PutObject: def test_init(self): actual = S3PutObject( From 15b8e6cc79f72d6702aa34b20b3d2ab284dede68 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 5 Jul 2026 23:36:01 +0900 Subject: [PATCH 05/12] Apply review feedback: version-aware fixes, f-string logging, registry overwrite - Version-aware opens bypass the dircache so the pinned version is the one actually observed at open time, not a stale or version-less cached entry - ls(versions=True) on an object path falls back to listing the versions of the key itself instead of returning an empty list - Guard the ListObjectVersions pagination against responses without a NextVersionIdMarker - register_s3_filesystem now always registers PyAthena's filesystem class (live fsspec registry), logging a warning when it overwrites an explicitly registered implementation - Use f-strings for log messages in the filesystem modules and ignore ruff's G004 accordingly Co-Authored-By: Claude Fable 5 --- pyathena/filesystem/__init__.py | 30 ++++--- pyathena/filesystem/s3.py | 93 ++++++++++++---------- pyproject.toml | 1 + tests/pyathena/filesystem/test_registry.py | 31 +++++--- tests/pyathena/filesystem/test_s3.py | 26 ++++++ 5 files changed, 110 insertions(+), 71 deletions(-) diff --git a/pyathena/filesystem/__init__.py b/pyathena/filesystem/__init__.py index 2c6a9ee4..b393e7c2 100644 --- a/pyathena/filesystem/__init__.py +++ b/pyathena/filesystem/__init__.py @@ -2,9 +2,9 @@ import fsspec -_logger = logging.getLogger(__name__) +from pyathena.filesystem.s3 import S3FileSystem -_S3_FILESYSTEM_CLASS = "pyathena.filesystem.s3.S3FileSystem" +_logger = logging.getLogger(__name__) def register_s3_filesystem(clobber: bool = True) -> None: @@ -17,25 +17,23 @@ def register_s3_filesystem(clobber: bool = True) -> None: implementation and s3fs-specific settings (e.g., the ``S3FS_LOGGING_LEVEL`` environment variable) have no effect. - A filesystem class that has already been registered explicitly (present - in ``fsspec.registry``) is left untouched, so registering another - implementation before importing ``pyathena.pandas`` / ``pyathena.polars`` - opts out of the replacement. To restore s3fs afterwards:: + A filesystem class that has already been registered explicitly is also + overwritten, with a warning log. To restore another implementation, + re-register it after importing ``pyathena.pandas`` / ``pyathena.polars``:: fsspec.register_implementation("s3", s3fs.S3FileSystem, clobber=True) Args: - clobber: Whether to replace an existing lazy registration of the + clobber: Whether to replace an existing registration of the protocols. """ for protocol in ("s3", "s3a"): - if protocol in fsspec.registry: - # A filesystem class has been registered explicitly - # (e.g., the user registered s3fs); leave it alone. - _logger.debug( - "The %r protocol is already registered in the fsspec registry, skipping.", - protocol, + registered = fsspec.registry.get(protocol) + if registered is not None and registered is not S3FileSystem: + _logger.warning( + f"The fsspec {protocol!r} protocol is already registered as " + f"{registered.__module__}.{registered.__qualname__} and will be overwritten by " + f"{S3FileSystem.__module__}.{S3FileSystem.__qualname__}." ) - continue - _logger.debug("Registering %s as the fsspec %r protocol.", _S3_FILESYSTEM_CLASS, protocol) - fsspec.register_implementation(protocol, _S3_FILESYSTEM_CLASS, clobber=clobber) + _logger.debug(f"Registering {S3FileSystem} as the fsspec {protocol!r} protocol.") + fsspec.register_implementation(protocol, S3FileSystem, clobber=clobber) diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index dbb70d98..e633f88b 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -488,9 +488,30 @@ def _ls_object_versions(self, path: str) -> list[S3Object]: if not response.get("IsTruncated"): break next_key_marker = response.get("NextKeyMarker") - next_version_id_marker = response.get("NextVersionIdMarker") + next_version_id_marker = response.get("NextVersionIdMarker", "") if not next_key_marker: break + + if not files and key: + # The path may point at an object rather than a key prefix. + files = [ + S3Object( + init={ + "ContentLength": version.size or 0, + "ContentType": None, + "StorageClass": version.storage_class, + "ETag": version.etag, + "LastModified": version.last_modified, + }, + type=S3ObjectType.S3_OBJECT_TYPE_FILE, + bucket=bucket, + key=version.key, + version_id=version.version_id, + is_latest=version.is_latest, + ) + for version in self.object_version_info(path) + if version.key == key + ] return files def info(self, path: str, **kwargs) -> S3Object: @@ -794,7 +815,7 @@ def _delete_object( if version_id: request.update({"VersionId": version_id}) - _logger.debug("Delete object: s3://%s/%s?versionId=%s", bucket, key, version_id) + _logger.debug(f"Delete object: s3://{bucket}/{key}?versionId={version_id}") self._call( self._client.delete_object, **request, @@ -906,7 +927,7 @@ def mkdir(self, path: str, create_parents: bool = True, **kwargs) -> None: # us-east-1 does not accept a location constraint. request.update({"CreateBucketConfiguration": {"LocationConstraint": region_name}}) - _logger.debug("Create bucket: s3://%s", bucket) + _logger.debug(f"Create bucket: s3://{bucket}") try: self._call( self._client.create_bucket, @@ -981,7 +1002,7 @@ def rmdir(self, path: str) -> None: "Set allow_bucket_deletion=True on the filesystem to enable it." ) - _logger.debug("Delete bucket: s3://%s", bucket) + _logger.debug(f"Delete bucket: s3://{bucket}") self._call( self._client.delete_bucket, Bucket=bucket, @@ -1086,12 +1107,8 @@ def _copy_object( } _logger.debug( - "Copy object from s3://%s/%s?versionId=%s to s3://%s/%s.", - bucket1, - key1, - version_id1, - bucket2, - key2, + f"Copy object from s3://{bucket1}/{key1}?versionId={version_id1} " + f"to s3://{bucket2}/{key2}." ) self._call(self._client.copy_object, **request, **kwargs) @@ -1293,7 +1310,7 @@ def _finish_multipart_upload( ) except Exception: _logger.exception( - "Failed to abort multipart upload %s to s3://%s/%s.", upload_id, bucket, key + f"Failed to abort multipart upload {upload_id} to s3://{bucket}/{key}." ) raise @@ -1466,7 +1483,7 @@ def sign(self, path: str, expiration: int = 3600, **kwargs): "ExpiresIn": expiration, } - _logger.debug("Generate signed url: s3://%s/%s?versionId=%s", bucket, key, version_id) + _logger.debug(f"Generate signed url: s3://{bucket}/{key}?versionId={version_id}") return self._call( self._client.generate_presigned_url, **request, @@ -1492,7 +1509,7 @@ def metadata(self, path: str, **kwargs) -> S3Metadata: if version_id: request.update({"VersionId": version_id}) - _logger.debug("Head object metadata: s3://%s/%s?versionId=%s", bucket, key, version_id) + _logger.debug(f"Head object metadata: s3://{bucket}/{key}?versionId={version_id}") response = self._call( self._client.head_object, **request, @@ -1551,7 +1568,7 @@ def setxattr(self, path: str, copy_kwargs: dict[str, Any] | None = None, **kw_ar if version_id: copy_source.update({"VersionId": version_id}) - _logger.debug("Set object metadata: s3://%s/%s?versionId=%s", bucket, key, version_id) + _logger.debug(f"Set object metadata: s3://{bucket}/{key}?versionId={version_id}") self._call( self._client.copy_object, CopySource=copy_source, @@ -1579,7 +1596,7 @@ def get_tags(self, path: str) -> dict[str, str]: if version_id: request.update({"VersionId": version_id}) - _logger.debug("Get object tagging: s3://%s/%s?versionId=%s", bucket, key, version_id) + _logger.debug(f"Get object tagging: s3://{bucket}/{key}?versionId={version_id}") response = self._call( self._client.get_object_tagging, **request, @@ -1621,7 +1638,7 @@ def put_tags(self, path: str, tags: dict[str, str], mode: str = "o") -> None: if version_id: request.update({"VersionId": version_id}) - _logger.debug("Put object tagging: s3://%s/%s?versionId=%s", bucket, key, version_id) + _logger.debug(f"Put object tagging: s3://{bucket}/{key}?versionId={version_id}") self._call( self._client.put_object_tagging, **request, @@ -1664,14 +1681,14 @@ def chmod(self, path: str, acl: str, recursive: bool = False, **kwargs) -> None: if version_id: request.update({"VersionId": version_id}) - _logger.debug("Put object acl: s3://%s/%s?versionId=%s", bucket, key, version_id) + _logger.debug(f"Put object acl: s3://{bucket}/{key}?versionId={version_id}") self._call( self._client.put_object_acl, **request, **kwargs, ) else: - _logger.debug("Put bucket acl: s3://%s", bucket) + _logger.debug(f"Put bucket acl: s3://{bucket}") self._call( self._client.put_bucket_acl, Bucket=bucket, @@ -1697,7 +1714,7 @@ def list_multipart_uploads(self, path: str) -> list[S3MultipartUpload]: """ bucket, key, _ = self.parse_path(path) - _logger.debug("List multipart uploads: s3://%s/%s", bucket, key) + _logger.debug(f"List multipart uploads: s3://{bucket}/{key}") uploads: list[S3MultipartUpload] = [] next_key_marker: str | None = None next_upload_id_marker: str | None = None @@ -1741,7 +1758,7 @@ def object_version_info( """ bucket, key, _ = self.parse_path(path) - _logger.debug("List object versions: s3://%s/%s", bucket, key) + _logger.debug(f"List object versions: s3://{bucket}/{key}") versions: list[S3ObjectVersion] = [] next_key_marker: str | None = None next_version_id_marker: str | None = None @@ -1773,7 +1790,7 @@ def object_version_info( if not response.get("IsTruncated"): break next_key_marker = response.get("NextKeyMarker") - next_version_id_marker = response.get("NextVersionIdMarker") + next_version_id_marker = response.get("NextVersionIdMarker", "") if not next_key_marker: break return versions @@ -1898,13 +1915,7 @@ def _get_object( if version_id: request.update({"VersionId": version_id}) - _logger.debug( - "Get object: s3://%s/%s?versionId=%s&range=%s", - bucket, - key, - version_id, - range_, - ) + _logger.debug(f"Get object: s3://{bucket}/{key}?versionId={version_id}&range={range_}") response = self._call( self._client.get_object, **request, @@ -1917,7 +1928,7 @@ def _put_object(self, bucket: str, key: str, body: bytes | None, **kwargs) -> S3 if body: request.update({"Body": body}) - _logger.debug("Put object: s3://%s/%s", bucket, key) + _logger.debug(f"Put object: s3://{bucket}/{key}") response = self._call( self._client.put_object, **request, @@ -1931,7 +1942,7 @@ def _create_multipart_upload(self, bucket: str, key: str, **kwargs) -> S3Multipa "Key": key, } - _logger.debug("Create multipart upload to s3://%s/%s.", bucket, key) + _logger.debug(f"Create multipart upload to s3://{bucket}/{key}.") response = self._call( self._client.create_multipart_upload, **request, @@ -1960,11 +1971,7 @@ def _upload_part_copy( range_ = S3File._format_ranges(copy_source_ranges) request.update({"CopySourceRange": range_}) _logger.debug( - "Upload part copy from %s to s3://%s/%s as part %s.", - copy_source, - bucket, - key, - part_number, + f"Upload part copy from {copy_source} to s3://{bucket}/{key} as part {part_number}." ) response = self._call( self._client.upload_part_copy, @@ -1990,13 +1997,7 @@ def _upload_part( "Body": body, } - _logger.debug( - "Upload part of %s to s3://%s/%s as part %s.", - upload_id, - bucket, - key, - part_number, - ) + _logger.debug(f"Upload part of {upload_id} to s3://{bucket}/{key} as part {part_number}.") response = self._call( self._client.upload_part, **request, @@ -2014,7 +2015,7 @@ def _complete_multipart_upload( "MultipartUpload": {"Parts": parts}, } - _logger.debug("Complete multipart upload %s to s3://%s/%s.", upload_id, bucket, key) + _logger.debug(f"Complete multipart upload {upload_id} to s3://{bucket}/{key}.") response = self._call( self._client.complete_multipart_upload, **request, @@ -2090,7 +2091,11 @@ def __init__( self.append_block = False self._details: S3Object | dict[str, Any] if "r" in mode: - info = self.fs.info(self.path, version_id=self.version_id) + # In version-aware mode, bypass the dircache (which may have been + # populated by a listing without version information) so that the + # version observed here is the one that is pinned. + refresh = self.fs.version_aware and not self.version_id + info = self.fs.info(self.path, version_id=self.version_id, refresh=refresh) if self.fs.version_aware and not self.version_id: # Pin the version observed at open time so that reads are # consistent even if the object is overwritten. diff --git a/pyproject.toml b/pyproject.toml index a2ad98af..c9d198bf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -165,6 +165,7 @@ select = [ ] ignore = [ "RUF059", # unused-unpacked-variable (too noisy for interface-heavy code) + "G004", # logging-f-string (f-strings are preferred for log messages) ] [tool.ruff.lint.per-file-ignores] diff --git a/tests/pyathena/filesystem/test_registry.py b/tests/pyathena/filesystem/test_registry.py index ad06185e..a5595503 100644 --- a/tests/pyathena/filesystem/test_registry.py +++ b/tests/pyathena/filesystem/test_registry.py @@ -1,3 +1,5 @@ +import logging + import fsspec import pytest from fsspec.registry import _registry, known_implementations @@ -6,6 +8,10 @@ from pyathena.filesystem.s3 import S3FileSystem +class _DummyFileSystem: + pass + + @pytest.fixture def registry_state(): registry = dict(_registry) @@ -20,20 +26,23 @@ def registry_state(): def test_register_s3_filesystem(registry_state): _registry.pop("s3", None) _registry.pop("s3a", None) - known_implementations["s3"] = {"class": "s3fs.S3FileSystem", "err": "Install s3fs"} - known_implementations.pop("s3a", None) register_s3_filesystem() - assert known_implementations["s3"]["class"] == "pyathena.filesystem.s3.S3FileSystem" - assert known_implementations["s3a"]["class"] == "pyathena.filesystem.s3.S3FileSystem" + assert fsspec.registry["s3"] is S3FileSystem + assert fsspec.registry["s3a"] is S3FileSystem -def test_register_s3_filesystem_respects_explicit_registration(registry_state): - # A filesystem class registered explicitly in the live fsspec registry - # (e.g., the user registered s3fs) is left untouched. - fsspec.register_implementation("s3", S3FileSystem, clobber=True) - known_implementations["s3"] = {"class": "s3fs.S3FileSystem", "err": "Install s3fs"} +def test_register_s3_filesystem_overwrites_existing_registration(registry_state, caplog): + fsspec.register_implementation("s3", _DummyFileSystem, clobber=True) - register_s3_filesystem() - assert known_implementations["s3"]["class"] == "s3fs.S3FileSystem" + with caplog.at_level(logging.WARNING, logger="pyathena.filesystem"): + register_s3_filesystem() + # An explicitly registered filesystem class is overwritten with a warning. assert fsspec.registry["s3"] is S3FileSystem + assert "will be overwritten" in caplog.text + + # Re-registering PyAthena's own filesystem does not warn. + caplog.clear() + with caplog.at_level(logging.WARNING, logger="pyathena.filesystem"): + register_s3_filesystem() + assert not caplog.text diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index 8d666287..17ed70ba 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -453,6 +453,32 @@ def test_ls_versions(self): assert actual[1].size == 4 assert actual[2].size == 2 + def test_ls_versions_object_path_falls_back_to_the_key(self): + fs = self._make_fs() + fs.version_aware = True + fs._call.side_effect = [ + # The prefix listing returns nothing: the path is an object. + {"IsTruncated": False}, + { + "Versions": [ + {"Key": "path/key", "VersionId": "v2", "IsLatest": True, "Size": 4}, + {"Key": "path/key", "VersionId": "v1", "IsLatest": False, "Size": 2}, + # A sibling key sharing the prefix is excluded. + {"Key": "path/key2", "VersionId": "x1", "IsLatest": True, "Size": 1}, + ], + "IsTruncated": False, + }, + ] + + actual = fs.ls("s3://bucket/path/key", detail=True, versions=True) + fs._call.assert_any_call( + fs._client.list_object_versions, Bucket="bucket", Prefix="path/key" + ) + assert [(f.name, f.version_id, f.size) for f in actual] == [ + ("bucket/path/key", "v2", 4), + ("bucket/path/key", "v1", 2), + ] + def test_metadata_with_version_id(self): fs = self._make_fs() fs._call.return_value = {"Metadata": {}} From 3cb42a2d0966f18852806c915f10cfb4b187c0e0 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Mon, 6 Jul 2026 00:06:39 +0900 Subject: [PATCH 06/12] Apply /simplify findings and add filesystem documentation - Extract _list_object_versions_pages, shared by ls(versions=True) and object_version_info, and _directory_object/_versioned_file_object factories replacing the repeated S3Object literals - Move the version-aware staleness decision into info(): a version-less cached entry is re-headed there, removing the file-layer refresh workaround and the doubled HEAD on version-aware opens - Route S3File.commit's multipart branch through _finish_multipart_upload so the complete-or-abort protocol has a single owner (commit now also aborts on failure) - Bound the ls(versions=True) object-path fallback listing with a delimiter; drop the S3ObjectVersion round-trip in the fallback - Drop the dead clobber parameter of register_s3_filesystem and reuse the helper in the test fixture; remove mocked tests duplicating integration coverage - Add the docs/filesystem.md user documentation page and fix the RST list formatting in the S3FileSystem class docstring Co-Authored-By: Claude Fable 5 --- docs/filesystem.md | 174 ++++++++++++++++++ docs/index.md | 1 + pyathena/filesystem/__init__.py | 8 +- pyathena/filesystem/s3.py | 260 ++++++++++----------------- tests/pyathena/filesystem/test_s3.py | 39 +--- 5 files changed, 279 insertions(+), 203 deletions(-) create mode 100644 docs/filesystem.md diff --git a/docs/filesystem.md b/docs/filesystem.md new file mode 100644 index 00000000..8cabdd08 --- /dev/null +++ b/docs/filesystem.md @@ -0,0 +1,174 @@ +(filesystem)= + +# S3 filesystem + +PyAthena ships its own [fsspec](https://filesystem-spec.readthedocs.io/en/latest/)-compatible +filesystem implementation for Amazon S3 (`S3FileSystem`), instead of depending on +[s3fs](https://github.com/fsspec/s3fs). This avoids the dependency and version-conflict +problems between s3fs/aiobotocore and boto3/botocore, while providing an API surface +compatible with s3fs for users migrating from it. + +The filesystem is used internally by the pandas/polars result sets to read query results +from S3, and can also be used independently for S3 file operations. + +## fsspec registration + +Importing `pyathena.pandas` or `pyathena.polars` registers `S3FileSystem` as the fsspec +`s3` / `s3a` protocols via `pyathena.filesystem.register_s3_filesystem`. This replaces +fsspec's default lazy mapping of the `s3` protocol to s3fs, which means +`fsspec.filesystem("s3")` returns PyAthena's implementation and s3fs-specific settings +(such as the `S3FS_LOGGING_LEVEL` environment variable) have no effect. + +A filesystem class that has already been registered explicitly is also overwritten, +with a warning log, so that the replacement is diagnosable. To restore another +implementation, re-register it afterwards: + +```python +import fsspec +import s3fs + +import pyathena.pandas # Registers PyAthena's S3FileSystem. + +fsspec.register_implementation("s3", s3fs.S3FileSystem, clobber=True) +``` + +## Basic usage + +The filesystem can be constructed from a PyAthena connection, or directly with +s3fs-compatible credential arguments: + +```python +from pyathena import connect +from pyathena.filesystem.s3 import S3FileSystem + +fs = S3FileSystem(connect(region_name="us-west-2")) + +# Or with direct credentials (s3fs-compatible arguments). +fs = S3FileSystem(key="YOUR_ACCESS_KEY", secret="YOUR_SECRET_KEY") + +# Or anonymously for public buckets. +fs = S3FileSystem(anon=True) +``` + +Standard fsspec operations work as expected: + +```python +fs.ls("s3://YOUR_S3_BUCKET/path/to/") +fs.find("s3://YOUR_S3_BUCKET/path/to/") +fs.exists("s3://YOUR_S3_BUCKET/path/to/object") +fs.info("s3://YOUR_S3_BUCKET/path/to/object") + +with fs.open("s3://YOUR_S3_BUCKET/path/to/object", "rb") as f: + data = f.read() + +fs.pipe("s3://YOUR_S3_BUCKET/path/to/object", b"data") +fs.cat("s3://YOUR_S3_BUCKET/path/to/object") +fs.cp("s3://YOUR_S3_BUCKET/src", "s3://YOUR_S3_BUCKET/dst") +fs.rm("s3://YOUR_S3_BUCKET/path/to/", recursive=True) +``` + +Writes with `pipe`/`pipe_file` issue a single PutObject request for data up to the +block size (5 MiB by default) and switch to a parallel multipart upload for larger +data. Inside an [fsspec transaction](https://filesystem-spec.readthedocs.io/en/latest/features.html#transactions), +writes are deferred until the transaction commits and are discarded on rollback. + +## Error translation + +S3 error responses are translated into standard Python exceptions, so filesystem +operations raise natural errors instead of botocore's `ClientError`: + +| S3 error | Python exception | +| --- | --- | +| `404` / `NoSuchKey` / `NoSuchBucket` | `FileNotFoundError` | +| `403` / `AccessDenied` | `PermissionError` | +| `BucketAlreadyExists` / `BucketAlreadyOwnedByYou` | `FileExistsError` | +| `RequestTimeout` | `TimeoutError` | +| Others | `OSError` with the matching `errno` | + +## Object metadata, tags, and ACLs + +```python +# User-defined metadata (x-amz-meta-*). +fs.setxattr("s3://YOUR_S3_BUCKET/path/to/object", attr1="value1") +metadata = fs.metadata("s3://YOUR_S3_BUCKET/path/to/object") +metadata["attr1"] # User-defined metadata via the mapping interface. +metadata.content_type # System-defined metadata as typed properties. +fs.getxattr("s3://YOUR_S3_BUCKET/path/to/object", "attr1") + +# Object tagging. +fs.put_tags("s3://YOUR_S3_BUCKET/path/to/object", {"tag1": "value1"}) +fs.put_tags("s3://YOUR_S3_BUCKET/path/to/object", {"tag2": "value2"}, mode="m") # Merge. +fs.get_tags("s3://YOUR_S3_BUCKET/path/to/object") + +# Canned ACLs. +fs.chmod("s3://YOUR_S3_BUCKET/path/to/object", "bucket-owner-full-control") +fs.chmod("s3://YOUR_S3_BUCKET/path/to/", "private", recursive=True) +``` + +Note that `setxattr` rewrites the object by copying it onto itself (S3 does not allow +updating the metadata of an existing object in place), which updates its last-modified +time. + +## Multipart upload management + +Incomplete multipart uploads continue to accrue storage costs until they are completed +or aborted. The filesystem can discover and abort them: + +```python +uploads = fs.list_multipart_uploads("s3://YOUR_S3_BUCKET") +for upload in uploads: + print(upload.key, upload.upload_id, upload.initiated) + +# Abort all incomplete uploads under a bucket or key prefix. +fs.clear_multipart_uploads("s3://YOUR_S3_BUCKET/path/to/") +``` + +## Versioning + +With `version_aware=True`, reads pin the object version observed at open time, so a +file handle keeps returning consistent data even if the object is overwritten while +reading. Explicit versions can always be read with the `?versionId=` suffix or the +`version_id` argument. + +```python +fs = S3FileSystem(connect(region_name="us-west-2"), version_aware=True) + +with fs.open("s3://YOUR_S3_BUCKET/path/to/object", "rb") as f: + data = f.read() # Pinned to the version observed at open time. + +# List all versions of the objects under a prefix. +fs.ls("s3://YOUR_S3_BUCKET/path/to/", versions=True, detail=True) + +# Typed version information, including delete markers if requested. +versions = fs.object_version_info("s3://YOUR_S3_BUCKET/path/to/object") +for version in versions: + print(version.version_id, version.is_latest, version.last_modified) +``` + +Version-aware operations require the `s3:GetObjectVersion` and +`s3:ListBucketVersions` permissions. + +## Bucket lifecycle + +Bucket creation and deletion are infrastructure-level changes and are disabled by +default: `mkdir`/`makedirs` and `rmdir` raise `PermissionError` when they would +create or delete a bucket. Pass the opt-in flags to enable them: + +```python +fs = S3FileSystem( + connect(region_name="us-west-2"), + allow_bucket_creation=True, + allow_bucket_deletion=True, +) +fs.mkdir("s3://YOUR_NEW_BUCKET") +fs.rmdir("s3://YOUR_NEW_BUCKET") # The bucket must be empty. +``` + +Creating a key prefix under an existing bucket requires no operation (S3 has no real +directories below the bucket level) and is always a no-op. + +## Async filesystem + +`AioS3FileSystem` provides the same functionality on top of fsspec's +`AsyncFileSystem`, dispatching parallel operations through the asyncio event loop. +See {ref}`aio-s3-filesystem` for details. diff --git a/docs/index.md b/docs/index.md index 42adab15..3198a6dc 100644 --- a/docs/index.md +++ b/docs/index.md @@ -74,6 +74,7 @@ sqlalchemy :maxdepth: 2 :caption: Advanced Topics +filesystem null_handling testing ``` diff --git a/pyathena/filesystem/__init__.py b/pyathena/filesystem/__init__.py index b393e7c2..8ab27c25 100644 --- a/pyathena/filesystem/__init__.py +++ b/pyathena/filesystem/__init__.py @@ -7,7 +7,7 @@ _logger = logging.getLogger(__name__) -def register_s3_filesystem(clobber: bool = True) -> None: +def register_s3_filesystem() -> None: """Register PyAthena's S3 filesystem as fsspec's "s3" / "s3a" protocols. PyAthena registers its own filesystem so that the pandas/polars result @@ -22,10 +22,6 @@ def register_s3_filesystem(clobber: bool = True) -> None: re-register it after importing ``pyathena.pandas`` / ``pyathena.polars``:: fsspec.register_implementation("s3", s3fs.S3FileSystem, clobber=True) - - Args: - clobber: Whether to replace an existing registration of the - protocols. """ for protocol in ("s3", "s3a"): registered = fsspec.registry.get(protocol) @@ -36,4 +32,4 @@ def register_s3_filesystem(clobber: bool = True) -> None: f"{S3FileSystem.__module__}.{S3FileSystem.__qualname__}." ) _logger.debug(f"Registering {S3FileSystem} as the fsspec {protocol!r} protocol.") - fsspec.register_implementation(protocol, S3FileSystem, clobber=clobber) + fsspec.register_implementation(protocol, S3FileSystem, clobber=True) diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index e633f88b..3427b9ae 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -5,7 +5,7 @@ import mimetypes import os.path import re -from collections.abc import Callable +from collections.abc import Callable, Iterator from concurrent.futures import Future, as_completed from copy import deepcopy from datetime import datetime @@ -52,12 +52,14 @@ class S3FileSystem(AbstractFileSystem): designed to be compatible with s3fs while offering PyAthena-specific optimizations. The filesystem supports standard S3 operations including: + - Listing objects and directories - Reading and writing files - Copying and moving objects - Reading and writing object metadata, tags, and canned ACLs - Multipart uploads for large files, including management of incomplete uploads + - Version-aware reads and object version listing (see ``version_aware``) - Creating and removing buckets (disabled by default; see ``allow_bucket_creation`` / ``allow_bucket_deletion``) - Various S3 storage classes and encryption options @@ -234,6 +236,35 @@ def parse_path(path: str) -> tuple[str, str | None, str | None]: return match.group("bucket"), match.group("key"), match.group("version_id") raise ValueError(f"Invalid S3 path format {path}.") + @staticmethod + def _directory_object(bucket: str, key: str | None, version_id: str | None = None) -> S3Object: + """Build an S3Object representing a directory entry.""" + return S3Object( + init={ + "ContentLength": 0, + "ContentType": None, + "StorageClass": S3StorageClass.S3_STORAGE_CLASS_DIRECTORY, + "ETag": None, + "LastModified": None, + }, + type=S3ObjectType.S3_OBJECT_TYPE_DIRECTORY, + bucket=bucket, + key=key, + version_id=version_id, + ) + + @staticmethod + def _versioned_file_object(bucket: str, version: dict[str, Any]) -> S3Object: + """Build an S3Object from a ListObjectVersions Versions entry.""" + return S3Object( + init=version, + type=S3ObjectType.S3_OBJECT_TYPE_FILE, + bucket=bucket, + key=version["Key"], + version_id=version.get("VersionId"), + is_latest=version.get("IsLatest", False), + ) + def _head_bucket(self, bucket, refresh: bool = False) -> S3Object | None: if bucket not in self.dircache or refresh: try: @@ -356,19 +387,7 @@ def _ls_dirs( **request, ) files.extend( - S3Object( - init={ - "ContentLength": 0, - "ContentType": None, - "StorageClass": S3StorageClass.S3_STORAGE_CLASS_DIRECTORY, - "ETag": None, - "LastModified": None, - }, - type=S3ObjectType.S3_OBJECT_TYPE_DIRECTORY, - bucket=bucket, - key=c["Prefix"][:-1].rstrip("/"), - version_id=version_id, - ) + self._directory_object(bucket, c["Prefix"][:-1].rstrip("/"), version_id) for c in response.get("CommonPrefixes", []) ) files.extend( @@ -439,14 +458,35 @@ def _ls_object_versions(self, path: str) -> list[S3Object]: prefix = f"{key}/" if key else "" files: list[S3Object] = [] + for response in self._list_object_versions_pages(bucket, prefix=prefix, delimiter="/"): + files.extend( + self._directory_object(bucket, c["Prefix"][:-1].rstrip("/")) + for c in response.get("CommonPrefixes", []) + ) + files.extend( + self._versioned_file_object(bucket, v) for v in response.get("Versions", []) + ) + + if not files and key: + # The path may point at an object rather than a key prefix. + files = [ + self._versioned_file_object(bucket, v) + for response in self._list_object_versions_pages(bucket, prefix=key, delimiter="/") + for v in response.get("Versions", []) + if v["Key"] == key + ] + return files + + def _list_object_versions_pages( + self, bucket: str, prefix: str, delimiter: str | None = None, **kwargs + ) -> Iterator[dict[str, Any]]: + """Iterate over the pages of a ListObjectVersions request.""" next_key_marker: str | None = None next_version_id_marker: str | None = None while True: - request: dict[str, Any] = { - "Bucket": bucket, - "Prefix": prefix, - "Delimiter": "/", - } + request: dict[str, Any] = {"Bucket": bucket, "Prefix": prefix} + if delimiter: + request.update({"Delimiter": delimiter}) if next_key_marker: request.update( { @@ -457,34 +497,9 @@ def _ls_object_versions(self, path: str) -> list[S3Object]: response = self._call( self._client.list_object_versions, **request, + **kwargs, ) - files.extend( - S3Object( - init={ - "ContentLength": 0, - "ContentType": None, - "StorageClass": S3StorageClass.S3_STORAGE_CLASS_DIRECTORY, - "ETag": None, - "LastModified": None, - }, - type=S3ObjectType.S3_OBJECT_TYPE_DIRECTORY, - bucket=bucket, - key=c["Prefix"][:-1].rstrip("/"), - version_id=None, - ) - for c in response.get("CommonPrefixes", []) - ) - files.extend( - S3Object( - init=v, - type=S3ObjectType.S3_OBJECT_TYPE_FILE, - bucket=bucket, - key=v["Key"], - version_id=v.get("VersionId"), - is_latest=v.get("IsLatest", False), - ) - for v in response.get("Versions", []) - ) + yield response if not response.get("IsTruncated"): break next_key_marker = response.get("NextKeyMarker") @@ -492,28 +507,6 @@ def _ls_object_versions(self, path: str) -> list[S3Object]: if not next_key_marker: break - if not files and key: - # The path may point at an object rather than a key prefix. - files = [ - S3Object( - init={ - "ContentLength": version.size or 0, - "ContentType": None, - "StorageClass": version.storage_class, - "ETag": version.etag, - "LastModified": version.last_modified, - }, - type=S3ObjectType.S3_OBJECT_TYPE_FILE, - bucket=bucket, - key=version.key, - version_id=version.version_id, - is_latest=version.is_latest, - ) - for version in self.object_version_info(path) - if version.key == key - ] - return files - def info(self, path: str, **kwargs) -> S3Object: refresh = kwargs.pop("refresh", False) path = self._strip_protocol(path) @@ -544,20 +537,22 @@ def info(self, path: str, **kwargs) -> S3Object: cache = None if cache: - return cache - return S3Object( - init={ - "ContentLength": 0, - "ContentType": None, - "StorageClass": S3StorageClass.S3_STORAGE_CLASS_DIRECTORY, - "ETag": None, - "LastModified": None, - }, - type=S3ObjectType.S3_OBJECT_TYPE_DIRECTORY, - bucket=bucket, - key=key.rstrip("/") if key else None, - version_id=version_id, - ) + if ( + self.version_aware + and not version_id + and cache.get("type") == S3ObjectType.S3_OBJECT_TYPE_FILE + and not cache.get("version_id") + ): + # A version-aware lookup needs the version to pin; + # treat a version-less cached entry (e.g., populated + # by a listing) as stale and head the object again. + refresh = True + else: + return cache + else: + return self._directory_object( + bucket, key.rstrip("/") if key else None, version_id + ) if key: object_info = self._head_object(path, refresh=refresh, version_id=version_id) if object_info: @@ -580,19 +575,7 @@ def info(self, path: str, **kwargs) -> S3Object: or response.get("Contents", []) or response.get("CommonPrefixes", []) ): - return S3Object( - init={ - "ContentLength": 0, - "ContentType": None, - "StorageClass": S3StorageClass.S3_STORAGE_CLASS_DIRECTORY, - "ETag": None, - "LastModified": None, - }, - type=S3ObjectType.S3_OBJECT_TYPE_DIRECTORY, - bucket=bucket, - key=key.rstrip("/") if key else None, - version_id=version_id, - ) + return self._directory_object(bucket, key.rstrip("/") if key else None, version_id) raise FileNotFoundError(path) def _extract_parent_directories( @@ -634,25 +617,7 @@ def _extract_parent_directories( dir_path = "/".join(parts[:i]) dirs.add(dir_path) - # Create S3Object instances for directories - directory_objects = [] - for dir_path in dirs: - dir_obj = S3Object( - init={ - "ContentLength": 0, - "ContentType": None, - "StorageClass": S3StorageClass.S3_STORAGE_CLASS_DIRECTORY, - "ETag": None, - "LastModified": None, - }, - type=S3ObjectType.S3_OBJECT_TYPE_DIRECTORY, - bucket=bucket, - key=dir_path, - version_id=None, - ) - directory_objects.append(dir_obj) - - return directory_objects + return [self._directory_object(bucket, dir_path) for dir_path in dirs] def _find( self, @@ -1288,10 +1253,8 @@ def _finish_multipart_upload( """ try: # The futures are in part-number order. - parts = [ - {"ETag": result.etag, "PartNumber": result.part_number} - for result in (future.result() for future in futures) - ] + results = [future.result() for future in futures] + parts = [{"ETag": r.etag, "PartNumber": r.part_number} for r in results] return self._complete_multipart_upload( bucket=bucket, key=key, @@ -1760,24 +1723,7 @@ def object_version_info( _logger.debug(f"List object versions: s3://{bucket}/{key}") versions: list[S3ObjectVersion] = [] - next_key_marker: str | None = None - next_version_id_marker: str | None = None - while True: - request: dict[str, Any] = {"Bucket": bucket} - if key: - request.update({"Prefix": key}) - if next_key_marker: - request.update( - { - "KeyMarker": next_key_marker, - "VersionIdMarker": next_version_id_marker, - } - ) - response = self._call( - self._client.list_object_versions, - **request, - **kwargs, - ) + for response in self._list_object_versions_pages(bucket, prefix=key or "", **kwargs): versions.extend( S3ObjectVersion(bucket=bucket, is_delete_marker=False, response=v) for v in response.get("Versions", []) @@ -1787,12 +1733,6 @@ def object_version_info( S3ObjectVersion(bucket=bucket, is_delete_marker=True, response=m) for m in response.get("DeleteMarkers", []) ) - if not response.get("IsTruncated"): - break - next_key_marker = response.get("NextKeyMarker") - next_version_id_marker = response.get("NextVersionIdMarker", "") - if not next_key_marker: - break return versions def clear_multipart_uploads(self, path: str) -> None: @@ -2091,14 +2031,11 @@ def __init__( self.append_block = False self._details: S3Object | dict[str, Any] if "r" in mode: - # In version-aware mode, bypass the dircache (which may have been - # populated by a listing without version information) so that the - # version observed here is the one that is pinned. - refresh = self.fs.version_aware and not self.version_id - info = self.fs.info(self.path, version_id=self.version_id, refresh=refresh) + info = self.fs.info(self.path, version_id=self.version_id) if self.fs.version_aware and not self.version_id: # Pin the version observed at open time so that reads are - # consistent even if the object is overwritten. + # consistent even if the object is overwritten. info() heads + # the object when the cached entry carries no version. self.version_id = info.get("version_id") if etag := info.get("etag"): self.s3_additional_kwargs.update({"IfMatch": etag}) @@ -2243,22 +2180,19 @@ def commit(self) -> None: if not self.multipart_upload: raise RuntimeError("Multipart upload is not initialized.") - parts: list[dict[str, Any]] = [] - for f in as_completed(self.multipart_upload_parts): - result = f.result() - parts.append( - { - "ETag": result.etag, - "PartNumber": result.part_number, - } + try: + self.fs._finish_multipart_upload( + bucket=self.bucket, + key=self.key, + upload_id=cast(str, self.multipart_upload.upload_id), + futures=self.multipart_upload_parts, ) - parts.sort(key=lambda x: x["PartNumber"]) - self.fs._complete_multipart_upload( - bucket=self.bucket, - key=self.key, - upload_id=cast(str, self.multipart_upload.upload_id), - parts=parts, - ) + except Exception: + # The multipart upload has been aborted by the helper; + # prevent discard() from aborting it again. + self.multipart_upload = None + self.multipart_upload_parts = [] + raise self.fs.invalidate_cache(self.path) diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index 17ed70ba..dd3da8ae 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -12,10 +12,10 @@ from types import SimpleNamespace from unittest import mock -import fsspec import pytest from fsspec import Callback +from pyathena.filesystem import register_s3_filesystem from pyathena.filesystem.s3 import S3File, S3FileSystem from pyathena.filesystem.s3_object import S3Object, S3ObjectType, S3StorageClass from pyathena.util import RetryConfig @@ -25,8 +25,7 @@ @pytest.fixture(scope="class") def register_filesystem(): - fsspec.register_implementation("s3", "pyathena.filesystem.s3.S3FileSystem", clobber=True) - fsspec.register_implementation("s3a", "pyathena.filesystem.s3.S3FileSystem", clobber=True) + register_s3_filesystem() @pytest.mark.usefixtures("register_filesystem") @@ -252,20 +251,6 @@ def test_pipe_file_small_uses_put_object(self): ContentType="text/plain", ) - def test_pipe_file_create_mode(self): - fs = self._make_fs() - fs.default_block_size = S3FileSystem.DEFAULT_BLOCK_SIZE - fs._put_object = mock.MagicMock() - - fs.exists = mock.MagicMock(return_value=True) - with pytest.raises(FileExistsError): - fs.pipe_file("s3://bucket/key", b"data", mode="create") - fs._put_object.assert_not_called() - - fs.exists = mock.MagicMock(return_value=False) - fs.pipe_file("s3://bucket/key", b"data", mode="create") - fs._put_object.assert_called_once() - def test_pipe_file_invalid_path_raises(self): fs = self._make_fs() with pytest.raises(ValueError, match="Cannot write to a bucket"): @@ -347,20 +332,6 @@ def test_pipe_file_abort_failure_does_not_mask_the_original_error(self): with pytest.raises(RuntimeError, match="upload failed"): fs.pipe_file("s3://bucket/key", b"0123456789ABCDEF") - def test_pipe_file_in_transaction_uses_buffered_path(self): - fs = self._make_fs() - fs._intrans = True - fs._put_object = mock.MagicMock() - opened = mock.MagicMock() - fs.open = mock.MagicMock(return_value=opened) - - fs.pipe_file("s3://bucket/key", b"data") - # The buffered open() path is used so that fsspec transactions keep - # their deferred-commit semantics; nothing is written directly. - fs.open.assert_called_once() - opened.__enter__.return_value.write.assert_called_once_with(b"data") - fs._put_object.assert_not_called() - def test_head_object_version_aware(self): fs = self._make_fs() fs._call.return_value = {"ContentLength": 4, "ETag": '"etag"', "VersionId": "v1"} @@ -472,7 +443,7 @@ def test_ls_versions_object_path_falls_back_to_the_key(self): actual = fs.ls("s3://bucket/path/key", detail=True, versions=True) fs._call.assert_any_call( - fs._client.list_object_versions, Bucket="bucket", Prefix="path/key" + fs._client.list_object_versions, Bucket="bucket", Prefix="path/key", Delimiter="/" ) assert [(f.name, f.version_id, f.size) for f in actual] == [ ("bucket/path/key", "v2", 4), @@ -1678,7 +1649,7 @@ def test_upload_chunk_multipart(self, autocommit): assert file._upload_chunk(final=True) is False # final -> buffer kept if not autocommit: # Deferred: the multipart upload is completed by commit(). - file.fs._complete_multipart_upload.assert_not_called() + file.fs._finish_multipart_upload.assert_not_called() file.commit() - file.fs._complete_multipart_upload.assert_called_once() + file.fs._finish_multipart_upload.assert_called_once() file.fs._put_object.assert_not_called() From da13a76a81319c586be4e78a17ab252e9ca433d9 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Mon, 6 Jul 2026 00:13:55 +0900 Subject: [PATCH 07/12] Move the registry tests into test_s3.py Keep the test file layout mirroring the source layout instead of adding a separate test_registry.py. Co-Authored-By: Claude Fable 5 --- tests/pyathena/filesystem/test_registry.py | 48 ---------------------- tests/pyathena/filesystem/test_s3.py | 43 +++++++++++++++++++ 2 files changed, 43 insertions(+), 48 deletions(-) delete mode 100644 tests/pyathena/filesystem/test_registry.py diff --git a/tests/pyathena/filesystem/test_registry.py b/tests/pyathena/filesystem/test_registry.py deleted file mode 100644 index a5595503..00000000 --- a/tests/pyathena/filesystem/test_registry.py +++ /dev/null @@ -1,48 +0,0 @@ -import logging - -import fsspec -import pytest -from fsspec.registry import _registry, known_implementations - -from pyathena.filesystem import register_s3_filesystem -from pyathena.filesystem.s3 import S3FileSystem - - -class _DummyFileSystem: - pass - - -@pytest.fixture -def registry_state(): - registry = dict(_registry) - known = {k: dict(v) for k, v in known_implementations.items()} - yield - _registry.clear() - _registry.update(registry) - known_implementations.clear() - known_implementations.update(known) - - -def test_register_s3_filesystem(registry_state): - _registry.pop("s3", None) - _registry.pop("s3a", None) - - register_s3_filesystem() - assert fsspec.registry["s3"] is S3FileSystem - assert fsspec.registry["s3a"] is S3FileSystem - - -def test_register_s3_filesystem_overwrites_existing_registration(registry_state, caplog): - fsspec.register_implementation("s3", _DummyFileSystem, clobber=True) - - with caplog.at_level(logging.WARNING, logger="pyathena.filesystem"): - register_s3_filesystem() - # An explicitly registered filesystem class is overwritten with a warning. - assert fsspec.registry["s3"] is S3FileSystem - assert "will be overwritten" in caplog.text - - # Re-registering PyAthena's own filesystem does not warn. - caplog.clear() - with caplog.at_level(logging.WARNING, logger="pyathena.filesystem"): - register_s3_filesystem() - assert not caplog.text diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index dd3da8ae..78431684 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -1,4 +1,5 @@ import io +import logging import os import tempfile import time @@ -12,8 +13,10 @@ from types import SimpleNamespace from unittest import mock +import fsspec import pytest from fsspec import Callback +from fsspec.registry import _registry, known_implementations from pyathena.filesystem import register_s3_filesystem from pyathena.filesystem.s3 import S3File, S3FileSystem @@ -1653,3 +1656,43 @@ def test_upload_chunk_multipart(self, autocommit): file.commit() file.fs._finish_multipart_upload.assert_called_once() file.fs._put_object.assert_not_called() + + +class _DummyFileSystem: + pass + + +@pytest.fixture +def registry_state(): + registry = dict(_registry) + known = {k: dict(v) for k, v in known_implementations.items()} + yield + _registry.clear() + _registry.update(registry) + known_implementations.clear() + known_implementations.update(known) + + +def test_register_s3_filesystem(registry_state): + _registry.pop("s3", None) + _registry.pop("s3a", None) + + register_s3_filesystem() + assert fsspec.registry["s3"] is S3FileSystem + assert fsspec.registry["s3a"] is S3FileSystem + + +def test_register_s3_filesystem_overwrites_existing_registration(registry_state, caplog): + fsspec.register_implementation("s3", _DummyFileSystem, clobber=True) + + with caplog.at_level(logging.WARNING, logger="pyathena.filesystem"): + register_s3_filesystem() + # An explicitly registered filesystem class is overwritten with a warning. + assert fsspec.registry["s3"] is S3FileSystem + assert "will be overwritten" in caplog.text + + # Re-registering PyAthena's own filesystem does not warn. + caplog.clear() + with caplog.at_level(logging.WARNING, logger="pyathena.filesystem"): + register_s3_filesystem() + assert not caplog.text From 0cf6d5cd7caa733115297d06e504405255eccca8 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Mon, 6 Jul 2026 09:36:28 +0900 Subject: [PATCH 08/12] Mirror the __init__ module tests as test_init.py Co-Authored-By: Claude Fable 5 --- tests/pyathena/filesystem/test_init.py | 48 ++++++++++++++++++++++++++ tests/pyathena/filesystem/test_s3.py | 43 ----------------------- 2 files changed, 48 insertions(+), 43 deletions(-) create mode 100644 tests/pyathena/filesystem/test_init.py diff --git a/tests/pyathena/filesystem/test_init.py b/tests/pyathena/filesystem/test_init.py new file mode 100644 index 00000000..a5595503 --- /dev/null +++ b/tests/pyathena/filesystem/test_init.py @@ -0,0 +1,48 @@ +import logging + +import fsspec +import pytest +from fsspec.registry import _registry, known_implementations + +from pyathena.filesystem import register_s3_filesystem +from pyathena.filesystem.s3 import S3FileSystem + + +class _DummyFileSystem: + pass + + +@pytest.fixture +def registry_state(): + registry = dict(_registry) + known = {k: dict(v) for k, v in known_implementations.items()} + yield + _registry.clear() + _registry.update(registry) + known_implementations.clear() + known_implementations.update(known) + + +def test_register_s3_filesystem(registry_state): + _registry.pop("s3", None) + _registry.pop("s3a", None) + + register_s3_filesystem() + assert fsspec.registry["s3"] is S3FileSystem + assert fsspec.registry["s3a"] is S3FileSystem + + +def test_register_s3_filesystem_overwrites_existing_registration(registry_state, caplog): + fsspec.register_implementation("s3", _DummyFileSystem, clobber=True) + + with caplog.at_level(logging.WARNING, logger="pyathena.filesystem"): + register_s3_filesystem() + # An explicitly registered filesystem class is overwritten with a warning. + assert fsspec.registry["s3"] is S3FileSystem + assert "will be overwritten" in caplog.text + + # Re-registering PyAthena's own filesystem does not warn. + caplog.clear() + with caplog.at_level(logging.WARNING, logger="pyathena.filesystem"): + register_s3_filesystem() + assert not caplog.text diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index 78431684..dd3da8ae 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -1,5 +1,4 @@ import io -import logging import os import tempfile import time @@ -13,10 +12,8 @@ from types import SimpleNamespace from unittest import mock -import fsspec import pytest from fsspec import Callback -from fsspec.registry import _registry, known_implementations from pyathena.filesystem import register_s3_filesystem from pyathena.filesystem.s3 import S3File, S3FileSystem @@ -1656,43 +1653,3 @@ def test_upload_chunk_multipart(self, autocommit): file.commit() file.fs._finish_multipart_upload.assert_called_once() file.fs._put_object.assert_not_called() - - -class _DummyFileSystem: - pass - - -@pytest.fixture -def registry_state(): - registry = dict(_registry) - known = {k: dict(v) for k, v in known_implementations.items()} - yield - _registry.clear() - _registry.update(registry) - known_implementations.clear() - known_implementations.update(known) - - -def test_register_s3_filesystem(registry_state): - _registry.pop("s3", None) - _registry.pop("s3a", None) - - register_s3_filesystem() - assert fsspec.registry["s3"] is S3FileSystem - assert fsspec.registry["s3a"] is S3FileSystem - - -def test_register_s3_filesystem_overwrites_existing_registration(registry_state, caplog): - fsspec.register_implementation("s3", _DummyFileSystem, clobber=True) - - with caplog.at_level(logging.WARNING, logger="pyathena.filesystem"): - register_s3_filesystem() - # An explicitly registered filesystem class is overwritten with a warning. - assert fsspec.registry["s3"] is S3FileSystem - assert "will be overwritten" in caplog.text - - # Re-registering PyAthena's own filesystem does not warn. - caplog.clear() - with caplog.at_level(logging.WARNING, logger="pyathena.filesystem"): - register_s3_filesystem() - assert not caplog.text From 8810a6a94ae52ee4160bd084f66f45c337f4dd74 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Mon, 6 Jul 2026 09:40:16 +0900 Subject: [PATCH 09/12] Remove the dependency-rationale sentence from the filesystem docs Co-Authored-By: Claude Fable 5 --- docs/filesystem.md | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/docs/filesystem.md b/docs/filesystem.md index 8cabdd08..55a0c40f 100644 --- a/docs/filesystem.md +++ b/docs/filesystem.md @@ -3,10 +3,8 @@ # S3 filesystem PyAthena ships its own [fsspec](https://filesystem-spec.readthedocs.io/en/latest/)-compatible -filesystem implementation for Amazon S3 (`S3FileSystem`), instead of depending on -[s3fs](https://github.com/fsspec/s3fs). This avoids the dependency and version-conflict -problems between s3fs/aiobotocore and boto3/botocore, while providing an API surface -compatible with s3fs for users migrating from it. +filesystem implementation for Amazon S3 (`S3FileSystem`), built on boto3, with an API +surface compatible with [s3fs](https://github.com/fsspec/s3fs) for users migrating from it. The filesystem is used internally by the pandas/polars result sets to read query results from S3, and can also be used independently for S3 file operations. From 17326533ef9716e9c46fbf8b16350d0d241f5c5c Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Mon, 6 Jul 2026 09:58:01 +0900 Subject: [PATCH 10/12] Keep multipart uploads owned by the buffered file path pipe_file now writes only the single-PutObject case directly (the checklist item in GH-722) and defers larger data to the buffered open() path, so the multipart upload orchestration has a single owner in S3File instead of a second copy in pipe_file. The pipe-specific multipart tests are replaced by unit tests of the shared _finish_multipart_upload helper, which remains in use by the copy path and S3File.commit. Co-Authored-By: Claude Fable 5 --- docs/filesystem.md | 5 +- pyathena/filesystem/s3.py | 79 +++++++--------------------- tests/pyathena/filesystem/test_s3.py | 75 ++++++++++---------------- 3 files changed, 50 insertions(+), 109 deletions(-) diff --git a/docs/filesystem.md b/docs/filesystem.md index 55a0c40f..38ea8224 100644 --- a/docs/filesystem.md +++ b/docs/filesystem.md @@ -66,8 +66,9 @@ fs.rm("s3://YOUR_S3_BUCKET/path/to/", recursive=True) ``` Writes with `pipe`/`pipe_file` issue a single PutObject request for data up to the -block size (5 MiB by default) and switch to a parallel multipart upload for larger -data. Inside an [fsspec transaction](https://filesystem-spec.readthedocs.io/en/latest/features.html#transactions), +block size (5 MiB by default); larger data is uploaded as a parallel multipart upload +through the buffered file path. Inside an +[fsspec transaction](https://filesystem-spec.readthedocs.io/en/latest/features.html#transactions), writes are deferred until the transaction commits and are discarded on rollback. ## Error translation diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index 3427b9ae..d8bcf2f3 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -1,7 +1,6 @@ from __future__ import annotations import logging -import math import mimetypes import os.path import re @@ -109,9 +108,6 @@ class S3FileSystem(AbstractFileSystem): # https://docs.aws.amazon.com/AmazonS3/latest/userguide/qfacts.html # The maximum size of a part in a multipart upload is 5GiB. MULTIPART_UPLOAD_MAX_PART_SIZE: int = 5 * 2**30 # 5GiB - # https://docs.aws.amazon.com/AmazonS3/latest/userguide/qfacts.html - # The maximum number of parts in a multipart upload is 10,000. - MULTIPART_UPLOAD_MAX_PARTS: int = 10000 # https://docs.aws.amazon.com/AmazonS3/latest/API/API_DeleteObjects.html DELETE_OBJECTS_MAX_KEYS: int = 1000 DEFAULT_BLOCK_SIZE: int = 5 * 2**20 # 5MiB @@ -1138,26 +1134,24 @@ def _copy_object_with_multipart_upload( def pipe_file( self, path: str, value: bytes | bytearray | memoryview, mode: str = "overwrite", **kwargs ) -> None: - """Write bytes into the path with direct S3 API calls. + """Write bytes into the path. - Uploads small data with a single PutObject request instead of the - inherited ``open()`` + ``write()`` path, and switches to a parallel - multipart upload when the data exceeds the block size. Inside an - fsspec transaction, the write goes through the inherited buffered - path so that the deferred-commit semantics are kept. + Writes data up to the block size with a single PutObject request + instead of the inherited ``open()`` + ``write()`` path. Larger data + and writes inside an fsspec transaction go through the buffered + path, which uploads the data as a parallel multipart upload and + keeps the deferred-commit semantics of transactions. Args: path: S3 path (s3://bucket/key) to write to. value: The bytes to write. mode: "overwrite" (default) or "create". With "create", raise FileExistsError when the object already exists. - **kwargs: Additional parameters passed to the PutObject or - CreateMultipartUpload API (e.g., ContentType, StorageClass). - The ``block_size``, ``max_worker``, and + **kwargs: Additional parameters passed to the PutObject API + (e.g., ContentType, StorageClass) on the single-request + path. The ``block_size``, ``max_worker``, and ``s3_additional_kwargs`` parameters of the ``open()`` path - are also accepted. Note that parameters that must be - repeated on every UploadPart request (e.g., SSE-C keys) are - not supported on the multipart path. + are also accepted. Raises: FileExistsError: If the mode is "create" and the path already @@ -1165,9 +1159,11 @@ def pipe_file( ValueError: If the path does not contain a key or specifies a version. """ - if self._intrans: - # Defer to the buffered open() path so that fsspec transactions - # keep their deferred-commit semantics. + block_size = kwargs.get("block_size") or self.default_block_size + if self._intrans or len(value) > min(block_size, self.MULTIPART_UPLOAD_MAX_PART_SIZE): + # Defer to the buffered open() path, which keeps the + # deferred-commit semantics of fsspec transactions and uploads + # large data as a parallel multipart upload. super().pipe_file(path, value, mode=mode, **kwargs) return bucket, key, version_id = self.parse_path(path) @@ -1179,54 +1175,17 @@ def pipe_file( raise FileExistsError(path) if not isinstance(value, bytes): # Accept bytes-like values (bytearray, memoryview) as the - # inherited buffered path did. + # buffered path does. value = bytes(value) - block_size = kwargs.pop("block_size", None) or self.default_block_size - max_worker = kwargs.pop("max_worker", None) or self.max_workers + kwargs.pop("block_size", None) + kwargs.pop("max_worker", None) request_kwargs = { **self.s3_additional_kwargs, **kwargs.pop("s3_additional_kwargs", {}), **kwargs, } - - size = len(value) - # A single PutObject request accepts up to the maximum part size. - if size <= min(block_size, self.MULTIPART_UPLOAD_MAX_PART_SIZE): - self._put_object(bucket=bucket, key=key, body=value, **request_kwargs) - else: - # The part size must satisfy the minimum/maximum part size and - # fit within the maximum number of parts of a multipart upload. - part_size = min( - max( - block_size, - self.MULTIPART_UPLOAD_MIN_PART_SIZE, - math.ceil(size / self.MULTIPART_UPLOAD_MAX_PARTS), - ), - self.MULTIPART_UPLOAD_MAX_PART_SIZE, - ) - offsets = range(0, size, part_size) - multipart_upload = self._create_multipart_upload( - bucket=bucket, key=key, **request_kwargs - ) - with self._create_executor(max_workers=min(len(offsets), max_worker)) as executor: - futures = [ - executor.submit( - self._upload_part, - bucket=bucket, - key=key, - upload_id=cast(str, multipart_upload.upload_id), - part_number=i + 1, - body=value[offset : offset + part_size], - ) - for i, offset in enumerate(offsets) - ] - self._finish_multipart_upload( - bucket=bucket, - key=key, - upload_id=cast(str, multipart_upload.upload_id), - futures=futures, - ) + self._put_object(bucket=bucket, key=key, body=value, **request_kwargs) self.invalidate_cache(path) def _finish_multipart_upload( diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index dd3da8ae..a17d18a2 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -5,7 +5,7 @@ import urllib.parse import urllib.request import uuid -from concurrent.futures import ThreadPoolExecutor +from concurrent.futures import Future, ThreadPoolExecutor from datetime import datetime, timezone from itertools import chain from pathlib import Path @@ -139,6 +139,7 @@ def _make_fs(): fs._retry_config = RetryConfig() fs.request_kwargs = {} fs.max_workers = 4 + fs.default_block_size = S3FileSystem.DEFAULT_BLOCK_SIZE fs.allow_bucket_creation = False fs.allow_bucket_deletion = False fs.s3_additional_kwargs = {} @@ -258,56 +259,39 @@ def test_pipe_file_invalid_path_raises(self): with pytest.raises(ValueError, match="version"): fs.pipe_file("s3://bucket/key?versionId=12345abcde", b"data") - def test_pipe_file_large_uses_multipart_upload(self): + def test_finish_multipart_upload(self): fs = self._make_fs() - fs.default_block_size = 5 - # Shrink the part-size limits so the test data is split without - # allocating real 5 MiB payloads. - fs.MULTIPART_UPLOAD_MIN_PART_SIZE = 5 - fs.MULTIPART_UPLOAD_MAX_PART_SIZE = 2**30 - fs._put_object = mock.MagicMock() - fs._create_multipart_upload = mock.MagicMock( - return_value=SimpleNamespace(upload_id="uploadid") - ) - uploaded: dict[int, bytes] = {} - - def upload_part(**kwargs): - uploaded[kwargs["part_number"]] = kwargs["body"] - return SimpleNamespace( - etag=f'"e{kwargs["part_number"]}"', part_number=kwargs["part_number"] - ) - - fs._upload_part = mock.MagicMock(side_effect=upload_part) fs._complete_multipart_upload = mock.MagicMock() - - fs.pipe_file("s3://bucket/key", b"0123456789ABCDEF", ContentType="text/plain") - - fs._put_object.assert_not_called() - fs._create_multipart_upload.assert_called_once_with( - bucket="bucket", key="key", ContentType="text/plain" + futures = [] + for part_number in (1, 2): + future: Future[SimpleNamespace] = Future() + future.set_result(SimpleNamespace(etag=f'"e{part_number}"', part_number=part_number)) + futures.append(future) + + fs._finish_multipart_upload( + bucket="bucket", key="key", upload_id="uploadid", futures=futures ) - # The data is split into block-size parts and reassembled in order. - assert uploaded == {1: b"01234", 2: b"56789", 3: b"ABCDE", 4: b"F"} fs._complete_multipart_upload.assert_called_once_with( bucket="bucket", key="key", upload_id="uploadid", - parts=[{"ETag": f'"e{i}"', "PartNumber": i} for i in range(1, 5)], + parts=[ + {"ETag": '"e1"', "PartNumber": 1}, + {"ETag": '"e2"', "PartNumber": 2}, + ], ) + fs._call.assert_not_called() - def test_pipe_file_aborts_multipart_upload_on_failure(self): + def test_finish_multipart_upload_aborts_on_failure(self): fs = self._make_fs() - fs.default_block_size = 5 - fs.MULTIPART_UPLOAD_MIN_PART_SIZE = 5 - fs.MULTIPART_UPLOAD_MAX_PART_SIZE = 2**30 - fs._create_multipart_upload = mock.MagicMock( - return_value=SimpleNamespace(upload_id="uploadid") - ) - fs._upload_part = mock.MagicMock(side_effect=RuntimeError("upload failed")) fs._complete_multipart_upload = mock.MagicMock() + future: Future[SimpleNamespace] = Future() + future.set_exception(RuntimeError("upload failed")) with pytest.raises(RuntimeError, match="upload failed"): - fs.pipe_file("s3://bucket/key", b"0123456789ABCDEF") + fs._finish_multipart_upload( + bucket="bucket", key="key", upload_id="uploadid", futures=[future] + ) fs._complete_multipart_upload.assert_not_called() fs._call.assert_called_once_with( fs._client.abort_multipart_upload, @@ -316,21 +300,18 @@ def test_pipe_file_aborts_multipart_upload_on_failure(self): UploadId="uploadid", ) - def test_pipe_file_abort_failure_does_not_mask_the_original_error(self): + def test_finish_multipart_upload_abort_failure_does_not_mask_the_original_error(self): fs = self._make_fs() - fs.default_block_size = 5 - fs.MULTIPART_UPLOAD_MIN_PART_SIZE = 5 - fs.MULTIPART_UPLOAD_MAX_PART_SIZE = 2**30 - fs._create_multipart_upload = mock.MagicMock( - return_value=SimpleNamespace(upload_id="uploadid") - ) - fs._upload_part = mock.MagicMock(side_effect=RuntimeError("upload failed")) fs._complete_multipart_upload = mock.MagicMock() fs._call = mock.MagicMock(side_effect=RuntimeError("abort failed")) + future: Future[SimpleNamespace] = Future() + future.set_exception(RuntimeError("upload failed")) # The abort failure is logged, and the original error propagates. with pytest.raises(RuntimeError, match="upload failed"): - fs.pipe_file("s3://bucket/key", b"0123456789ABCDEF") + fs._finish_multipart_upload( + bucket="bucket", key="key", upload_id="uploadid", futures=[future] + ) def test_head_object_version_aware(self): fs = self._make_fs() From eee98afb1f934a60c9a2fedf3c49edcd956f9a7a Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Mon, 6 Jul 2026 10:05:28 +0900 Subject: [PATCH 11/12] Document the onerror workaround with its removal condition fsspec's AbstractFileSystem.mv() passed the typo'd "onerror" keyword (instead of "on_error", which copy() consumes) until it was fixed in fsspec 2026.6.0, so it leaks through copy(**kwargs) into cp_file on older releases and must not reach the S3 API. Replace the stale TODO with the actual mechanism and the fsspec version at which the pop can be removed. Co-Authored-By: Claude Fable 5 --- pyathena/filesystem/s3.py | 8 +++++--- pyathena/filesystem/s3_async.py | 2 ++ 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index d8bcf2f3..29dd5dc7 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -1011,10 +1011,12 @@ def cp_file( to optimize performance for large files. The copy operation is performed entirely on the S3 service without data transfer. """ - # TODO: Delete the value that seems to be a typo, onerror=false. + # fsspec < 2026.6.0: AbstractFileSystem.mv() passed the typo'd + # "onerror" keyword (instead of "on_error", which copy() consumes), + # so it leaked through copy(**kwargs) into cp_file and must not + # reach the S3 API. Remove this once the fsspec requirement is + # >= 2026.6.0, where mv() passes on_error correctly. # https://github.com/fsspec/filesystem_spec/commit/346a589fef9308550ffa3d0d510f2db67281bb05 - # https://github.com/fsspec/filesystem_spec/blob/2024.10.0/fsspec/spec.py#L1185 - # https://github.com/fsspec/filesystem_spec/blob/2024.10.0/fsspec/spec.py#L1077 kwargs.pop("onerror", None) bucket1, key1, version_id1 = self.parse_path(path1) bucket2, key2, version_id2 = self.parse_path(path2) diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index f968adfe..834ffd92 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -191,6 +191,8 @@ async def _delete_chunk(chunk: list[dict[str, Any]]) -> None: async def _cp_file(self, path1: str, path2: str, **kwargs) -> None: """Copy an S3 object, using async parallel multipart upload for large files.""" + # fsspec < 2026.6.0 leaks the typo'd "onerror" keyword from mv(); + # see S3FileSystem.cp_file. kwargs.pop("onerror", None) bucket1, key1, version_id1 = self.parse_path(path1) bucket2, key2, version_id2 = self.parse_path(path2) From bef67e67c96e6fc6e9bb7c16c61e81411984878a Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Mon, 6 Jul 2026 10:14:41 +0900 Subject: [PATCH 12/12] Clean up the s3fs-compatible client construction - Import Connection at the top of the module; the runtime import guarded against an import cycle that no longer exists - Replace the bare s3fs source link with a docstring describing the accepted arguments - Honor use_ssl=False, which the previous truthiness check silently dropped, and stop injecting None credential values into the session and client arguments Co-Authored-By: Claude Fable 5 --- pyathena/filesystem/s3.py | 70 ++++++++++++++-------------- tests/pyathena/filesystem/test_s3.py | 23 +++++++++ 2 files changed, 59 insertions(+), 34 deletions(-) diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index 29dd5dc7..ab941678 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -10,7 +10,7 @@ from datetime import datetime from multiprocessing import cpu_count from re import Pattern -from typing import TYPE_CHECKING, Any, cast +from typing import Any, cast import botocore.exceptions from boto3 import Session @@ -22,6 +22,7 @@ from fsspec.utils import tokenize import pyathena +from pyathena.connection import Connection from pyathena.filesystem.s3_errors import S3ClientError from pyathena.filesystem.s3_executor import S3Executor, S3ThreadPoolExecutor from pyathena.filesystem.s3_object import ( @@ -37,9 +38,6 @@ ) from pyathena.util import RetryConfig, retry_api_call -if TYPE_CHECKING: - from pyathena.connection import Connection - _logger = logging.getLogger(__name__) @@ -173,55 +171,59 @@ def __init__( self.request_kwargs = {"RequestPayer": "requester"} if requester_pays else {} def _get_client_compatible_with_s3fs(self, **kwargs) -> BaseClient: - """ - https://github.com/fsspec/s3fs/blob/2023.4.0/s3fs/core.py#L457-L535 - """ - from pyathena.connection import Connection + """Build a boto3 S3 client from s3fs-compatible constructor arguments. + + Accepts the constructor arguments that s3fs users pass through fsspec + storage options — ``key``/``username``, ``secret``/``password``, + ``token``, ``anon``, ``use_ssl``, ``endpoint_url``, + ``connect_timeout``/``read_timeout``, and the ``client_kwargs`` / + ``config_kwargs`` dictionaries — in addition to boto3 session + arguments such as ``region_name`` and ``profile_name``. + Args: + **kwargs: The filesystem constructor arguments. + + Returns: + A boto3 S3 client configured from the arguments. + """ config_kwargs = deepcopy(kwargs.pop("config_kwargs", {})) + client_kwargs = deepcopy(kwargs.pop("client_kwargs", {})) + user_agent_extra = config_kwargs.pop("user_agent_extra", None) - if user_agent_extra: - if pyathena.user_agent_extra not in user_agent_extra: - config_kwargs.update( - {"user_agent_extra": f"{pyathena.user_agent_extra} {user_agent_extra}"} - ) - else: - config_kwargs.update({"user_agent_extra": user_agent_extra}) - else: - config_kwargs.update({"user_agent_extra": pyathena.user_agent_extra}) - connect_timeout = kwargs.pop("connect_timeout", None) - if connect_timeout: + if user_agent_extra and pyathena.user_agent_extra not in user_agent_extra: + user_agent_extra = f"{pyathena.user_agent_extra} {user_agent_extra}" + config_kwargs.update({"user_agent_extra": user_agent_extra or pyathena.user_agent_extra}) + if connect_timeout := kwargs.pop("connect_timeout", None): config_kwargs.update({"connect_timeout": connect_timeout}) - read_timeout = kwargs.pop("read_timeout", None) - if read_timeout: + if read_timeout := kwargs.pop("read_timeout", None): config_kwargs.update({"read_timeout": read_timeout}) - client_kwargs = deepcopy(kwargs.pop("client_kwargs", {})) use_ssl = kwargs.pop("use_ssl", None) - if use_ssl: + if use_ssl is not None: client_kwargs.update({"use_ssl": use_ssl}) - endpoint_url = kwargs.pop("endpoint_url", None) - if endpoint_url: + if endpoint_url := kwargs.pop("endpoint_url", None): client_kwargs.update({"endpoint_url": endpoint_url}) - anon = kwargs.pop("anon", False) - if anon: + if kwargs.pop("anon", False): config_kwargs.update({"signature_version": UNSIGNED}) else: creds = { - "aws_access_key_id": kwargs.pop("key", kwargs.pop("username", None)), - "aws_secret_access_key": kwargs.pop("secret", kwargs.pop("password", None)), - "aws_session_token": kwargs.pop("token", None), + key: value + for key, value in { + "aws_access_key_id": kwargs.pop("key", kwargs.pop("username", None)), + "aws_secret_access_key": kwargs.pop("secret", kwargs.pop("password", None)), + "aws_session_token": kwargs.pop("token", None), + }.items() + if value is not None } - kwargs.update(**creds) - client_kwargs.update(**creds) + kwargs.update(creds) + client_kwargs.update(creds) - config = Config(**config_kwargs) session = Session( **{k: v for k, v in kwargs.items() if k in Connection._SESSION_PASSING_ARGS} ) return session.client( "s3", - config=config, + config=Config(**config_kwargs), **{k: v for k, v in client_kwargs.items() if k in Connection._CLIENT_PASSING_ARGS}, ) diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index a17d18a2..14ae956b 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -15,6 +15,7 @@ import pytest from fsspec import Callback +import pyathena from pyathena.filesystem import register_s3_filesystem from pyathena.filesystem.s3 import S3File, S3FileSystem from pyathena.filesystem.s3_object import S3Object, S3ObjectType, S3StorageClass @@ -147,6 +148,28 @@ def _make_fs(): fs.version_aware = False return fs + def test_get_client_compatible_with_s3fs(self): + # Only constructs a boto3 client; no AWS access. + fs = S3FileSystem( + key="test_access_key", + secret="test_secret_key", + region_name="us-east-1", + use_ssl=False, + skip_instance_cache=True, + ) + # use_ssl=False is honored (previously dropped by a truthiness check). + assert fs._client.meta.endpoint_url.startswith("http://") + assert pyathena.user_agent_extra in fs._client.meta.config.user_agent_extra + + fs = S3FileSystem( + key="test_access_key", + secret="test_secret_key", + region_name="us-east-1", + endpoint_url="http://localhost:9000", + skip_instance_cache=True, + ) + assert fs._client.meta.endpoint_url == "http://localhost:9000" + def test_ls_from_cache_with_cached_object(self): fs = self._make_fs() obj = S3Object(