diff --git a/docs/filesystem.md b/docs/filesystem.md new file mode 100644 index 00000000..38ea8224 --- /dev/null +++ b/docs/filesystem.md @@ -0,0 +1,173 @@ +(filesystem)= + +# S3 filesystem + +PyAthena ships its own [fsspec](https://filesystem-spec.readthedocs.io/en/latest/)-compatible +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. + +## 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); 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 + +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 e69de29b..8ab27c25 100644 --- a/pyathena/filesystem/__init__.py +++ b/pyathena/filesystem/__init__.py @@ -0,0 +1,35 @@ +import logging + +import fsspec + +from pyathena.filesystem.s3 import S3FileSystem + +_logger = logging.getLogger(__name__) + + +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 + 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 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) + """ + for protocol in ("s3", "s3a"): + 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__}." + ) + _logger.debug(f"Registering {S3FileSystem} as the fsspec {protocol!r} protocol.") + fsspec.register_implementation(protocol, S3FileSystem, clobber=True) diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index d7da3586..ab941678 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -4,13 +4,13 @@ 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 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 ( @@ -31,14 +32,12 @@ S3MultipartUploadPart, S3Object, S3ObjectType, + S3ObjectVersion, S3PutObject, S3StorageClass, ) from pyathena.util import RetryConfig, retry_api_call -if TYPE_CHECKING: - from pyathena.connection import Connection - _logger = logging.getLogger(__name__) @@ -50,12 +49,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 @@ -71,6 +72,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 @@ -136,6 +141,7 @@ def __init__( s3_additional_kwargs=None, allow_bucket_creation: bool = False, allow_bucket_deletion: bool = False, + version_aware: bool = False, *args, **kwargs, ) -> None: @@ -159,60 +165,65 @@ 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 {} 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}, ) @@ -223,6 +234,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: @@ -269,6 +309,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, @@ -341,19 +385,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( @@ -384,7 +416,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). @@ -394,9 +428,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: @@ -405,6 +446,65 @@ 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] = [] + 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} + if delimiter: + request.update({"Delimiter": 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, + **kwargs, + ) + yield response + 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 + def info(self, path: str, **kwargs) -> S3Object: refresh = kwargs.pop("refresh", False) path = self._strip_protocol(path) @@ -435,20 +535,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: @@ -471,19 +573,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( @@ -525,25 +615,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, @@ -706,7 +778,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, @@ -818,7 +890,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, @@ -893,7 +965,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, @@ -941,10 +1013,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) @@ -998,12 +1072,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) @@ -1045,9 +1115,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, @@ -1059,22 +1128,117 @@ 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, - } - ) + self._finish_multipart_upload( + bucket=bucket2, + key=key2, + upload_id=cast(str, multipart_upload.upload_id), + futures=futures, + ) - 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, - ) + def pipe_file( + self, path: str, value: bytes | bytearray | memoryview, mode: str = "overwrite", **kwargs + ) -> None: + """Write bytes into the path. + + 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 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. + + 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. + """ + 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) + 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) + if not isinstance(value, bytes): + # Accept bytes-like values (bytearray, memoryview) as the + # buffered path does. + value = bytes(value) + + kwargs.pop("block_size", None) + kwargs.pop("max_worker", None) + request_kwargs = { + **self.s3_additional_kwargs, + **kwargs.pop("s3_additional_kwargs", {}), + **kwargs, + } + self._put_object(bucket=bucket, key=key, body=value, **request_kwargs) + 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. + 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, + 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, + ) + except Exception: + _logger.exception( + f"Failed to abort multipart upload {upload_id} to s3://{bucket}/{key}." + ) + raise def cat_file( self, path: str, start: int | None = None, end: int | None = None, **kwargs @@ -1245,7 +1409,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, @@ -1271,7 +1435,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, @@ -1330,7 +1494,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, @@ -1358,7 +1522,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, @@ -1400,7 +1564,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, @@ -1443,14 +1607,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, @@ -1476,7 +1640,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 @@ -1503,6 +1667,37 @@ 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(f"List object versions: s3://{bucket}/{key}") + versions: list[S3ObjectVersion] = [] + 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", []) + ) + if delete_markers: + versions.extend( + S3ObjectVersion(bucket=bucket, is_delete_marker=True, response=m) + for m in response.get("DeleteMarkers", []) + ) + return versions + def clear_multipart_uploads(self, path: str) -> None: """Abort any incomplete multipart uploads in the bucket. @@ -1623,13 +1818,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, @@ -1642,7 +1831,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, @@ -1656,7 +1845,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, @@ -1685,11 +1874,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, @@ -1715,13 +1900,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, @@ -1739,7 +1918,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, @@ -1816,6 +1995,11 @@ 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. 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}) self._details = info @@ -1959,22 +2143,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) @@ -2089,7 +2270,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..834ffd92 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 @@ -111,8 +118,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) @@ -182,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) @@ -353,6 +364,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/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/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/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_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 373214d4..14ae956b 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -5,17 +5,18 @@ 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 from types import SimpleNamespace from unittest import mock -import fsspec 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 from pyathena.util import RetryConfig @@ -25,8 +26,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") @@ -140,10 +140,36 @@ 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 = {} + fs._intrans = False + 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( @@ -232,6 +258,202 @@ 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.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", + ServerSideEncryption="AES256", + ContentType="text/plain", + ) + + 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_finish_multipart_upload(self): + fs = self._make_fs() + fs._complete_multipart_upload = mock.MagicMock() + 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 + ) + fs._complete_multipart_upload.assert_called_once_with( + bucket="bucket", + key="key", + upload_id="uploadid", + parts=[ + {"ETag": '"e1"', "PartNumber": 1}, + {"ETag": '"e2"', "PartNumber": 2}, + ], + ) + fs._call.assert_not_called() + + def test_finish_multipart_upload_aborts_on_failure(self): + fs = self._make_fs() + fs._complete_multipart_upload = mock.MagicMock() + future: Future[SimpleNamespace] = Future() + future.set_exception(RuntimeError("upload failed")) + + with pytest.raises(RuntimeError, match="upload failed"): + 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, + Bucket="bucket", + Key="key", + UploadId="uploadid", + ) + + def test_finish_multipart_upload_abort_failure_does_not_mask_the_original_error(self): + fs = self._make_fs() + 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._finish_multipart_upload( + bucket="bucket", key="key", upload_id="uploadid", futures=[future] + ) + + 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_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", Delimiter="/" + ) + 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": {}} @@ -774,7 +996,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 +1017,46 @@ 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_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 = ( @@ -1141,6 +1404,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}/" @@ -1278,6 +1569,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): @@ -1358,7 +1653,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() 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(