From f730dd1ed19b899e4f48623a2be4412a532a7291 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 4 Oct 2026 13:44:51 +0900 Subject: [PATCH 1/4] Move the multipart upload requests into S3Core (step 3.2 of #1063) Add S3Core.create_multipart_upload(), upload_part(), upload_part_copy(), complete_multipart_upload() and abort_multipart_upload(), one request each, and part_ranges(), the range split of a multipart copy and of an append. The fields that the arguments set still take precedence over inherited parameters of the same name. complete_multipart_upload() takes the typed parts, and abort_multipart_upload() raises; logging and swallowing a failed abort stays in the filesystem's callers. The part limits become S3Core attributes, and the filesystems no longer have the MULTIPART_UPLOAD_* attributes (maintainer decision; breaking change for 4.0.0). S3File, pipe_file(), clear_multipart_uploads() and the sync and aio multipart copies call fs.core.*; the requests sent are unchanged. Co-Authored-By: Claude Opus 5.5 --- docs/filesystem.md | 20 +- pyathena/filesystem/s3.py | 299 +++++-------------- pyathena/filesystem/s3_async.py | 66 ++--- pyathena/filesystem/s3_core.py | 207 ++++++++++++- tests/pyathena/filesystem/test_s3.py | 323 +++++++++------------ tests/pyathena/filesystem/test_s3_async.py | 95 +++--- tests/pyathena/filesystem/test_s3_core.py | 192 ++++++++++++ tests/pyathena/util.py | 5 +- 8 files changed, 702 insertions(+), 505 deletions(-) diff --git a/docs/filesystem.md b/docs/filesystem.md index 9c9e9421..61904928 100644 --- a/docs/filesystem.md +++ b/docs/filesystem.md @@ -280,10 +280,12 @@ directories below the bucket level) and is always a no-op. ## Typed S3 operations `S3FileSystem.core` is an `S3Core`, the typed operations that the filesystem sends -its listing, lookup and delete requests with. It can also be built on a boto3 S3 -client. Each operation sends one request (one per page for the iterators) with the -retry policy, raises `FileNotFoundError` for a missing bucket, or for a missing object -or version that it reads, and caches nothing. +its listing, lookup, delete and multipart upload requests with. It can also be built +on a boto3 S3 client. Each operation sends one request (one per page for the +iterators) with the retry policy, raises `FileNotFoundError` for a missing bucket, or +for a missing object or version that it reads, and caches nothing. Requests sent +through `fs.core` do not invalidate the filesystem's cache: call +`fs.invalidate_cache()` after a change, or make it through the filesystem. ```python import boto3 @@ -306,8 +308,7 @@ for page in core.list_objects("YOUR_S3_BUCKET", prefix="path/to/", delimiter="/" DeleteObjects request accepts. `S3DeleteBatch.from_paths()` groups paths into batches of up to 1,000 objects per bucket. The objects that S3 could not delete are in the `errors` of the returned `S3DeleteResult`, not raised; as in S3, deleting a key that -does not exist is not an error. Requests sent through `fs.core` do not invalidate the -filesystem's cache: call `fs.invalidate_cache()` after them, or delete with `fs.rm()`. +does not exist is not an error. ```python from pyathena.filesystem.s3_core import S3DeleteBatch @@ -322,6 +323,13 @@ for batch in S3DeleteBatch.from_paths(paths): print(error) # path (code: message) ``` +`create_multipart_upload()`, `upload_part()`, `upload_part_copy()`, +`complete_multipart_upload()` and `abort_multipart_upload()` send the requests of a +multipart upload. `part_ranges()` splits an object into the byte ranges of the parts +that copy it, by the part limits `MULTIPART_UPLOAD_MIN_PART_SIZE` (5 MiB), +`MULTIPART_UPLOAD_MAX_PART_SIZE` (5 GiB) and `MULTIPART_UPLOAD_MAX_PARTS` (10,000) of +`S3Core`. + ## Async filesystem `AioS3FileSystem` provides the same functionality on top of fsspec's diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index f632ad00..934dfbf5 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -161,15 +161,6 @@ class S3FileSystem(AbstractFileSystem): stored in S3, but can also be used independently for S3 file operations. """ - # https://docs.aws.amazon.com/AmazonS3/latest/userguide/qfacts.html - # The minimum size of a part in a multipart upload is 5MiB. - MULTIPART_UPLOAD_MIN_PART_SIZE: int = 5 * 2**20 # 5MiB - # 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 per multipart upload is 10,000. - MULTIPART_UPLOAD_MAX_PARTS: int = 10_000 DEFAULT_BLOCK_SIZE: int = 5 * 2**20 # 5MiB # https://docs.aws.amazon.com/AmazonS3/latest/userguide/acl-overview.html#canned-acl OBJECT_ACLS: frozenset[str] = frozenset( @@ -1734,7 +1725,7 @@ def _copy_file(self, path1: str, path2: str, **kwargs) -> bool: return False size1 = info1.get("size", 0) try: - if size1 <= self.MULTIPART_UPLOAD_MAX_PART_SIZE: + if size1 <= self.core.MULTIPART_UPLOAD_MAX_PART_SIZE: self._copy_object( bucket1=source.bucket, key1=source.key, @@ -1829,21 +1820,22 @@ def _copy_object_with_multipart_upload( directive has an invalid value. """ max_workers = max_workers if max_workers else self.max_workers - block_size = block_size if block_size else self.MULTIPART_UPLOAD_MAX_PART_SIZE + block_size = block_size if block_size else self.core.MULTIPART_UPLOAD_MAX_PART_SIZE if ( - block_size < self.MULTIPART_UPLOAD_MIN_PART_SIZE - or block_size > self.MULTIPART_UPLOAD_MAX_PART_SIZE + block_size < self.core.MULTIPART_UPLOAD_MIN_PART_SIZE + or block_size > self.core.MULTIPART_UPLOAD_MAX_PART_SIZE ): raise ValueError( "Block size must be between " - f"5 MiB ({self.MULTIPART_UPLOAD_MIN_PART_SIZE} bytes) and " - f"5 GiB ({self.MULTIPART_UPLOAD_MAX_PART_SIZE} bytes), inclusive: {block_size}." + f"5 MiB ({self.core.MULTIPART_UPLOAD_MIN_PART_SIZE} bytes) and " + f"5 GiB ({self.core.MULTIPART_UPLOAD_MAX_PART_SIZE} bytes), " + f"inclusive: {block_size}." ) create_kwargs, version_id1, head_size = self._get_multipart_copy_kwargs( bucket1, key1, version_id1, kwargs ) - if head_size is not None and head_size <= self.MULTIPART_UPLOAD_MAX_PART_SIZE: + if head_size is not None and head_size <= self.core.MULTIPART_UPLOAD_MAX_PART_SIZE: # The size that the caller found, which may come from a cached # listing, was larger than the copied version, which fits in a # single CopyObject request. @@ -1857,13 +1849,9 @@ def _copy_object_with_multipart_upload( ) return # The size of the copied version, not the one that the caller found. - ranges = self._get_copy_ranges(size1 if head_size is None else head_size, block_size) - copy_source = { - "Bucket": bucket1, - "Key": key1, - } - if version_id1: - copy_source.update({"VersionId": version_id1}) + ranges = self.core.part_ranges(size1 if head_size is None else head_size, block_size) + source = S3Path(bucket1, key1, version_id1) + destination = S3Path(bucket2, key2) # The annotations are listed before anything is written, so that a # missing permission fails first. annotations = ( @@ -1871,21 +1859,16 @@ def _copy_object_with_multipart_upload( if self._copies_annotations(bucket1, kwargs) else [] ) - multipart_upload = self._create_multipart_upload( - bucket=bucket2, - key=key2, - **create_kwargs, - ) + multipart_upload = self.core.create_multipart_upload(destination, **create_kwargs) with self._create_executor(max_workers=max_workers) as executor: futures = [ executor.submit( - self._upload_part_copy, - bucket=bucket2, - key=key2, - copy_source=copy_source, + self.core.upload_part_copy, + path=destination, upload_id=cast(str, multipart_upload.upload_id), part_number=i + 1, - copy_source_ranges=range_, + source=source, + range_=range_, **self.core.operation_params("upload_part_copy", kwargs), ) for i, range_ in enumerate(ranges) @@ -2001,7 +1984,7 @@ def _get_multipart_copy_kwargs( source.update({"VersionId": version_id}) if ( head.content_length is not None - and head.content_length <= self.MULTIPART_UPLOAD_MAX_PART_SIZE + and head.content_length <= self.core.MULTIPART_UPLOAD_MAX_PART_SIZE ): # Copied with CopyObject instead, which applies the directives. return {}, version_id, head.content_length @@ -2156,40 +2139,6 @@ def _copy_object_annotation( **{**self.core.operation_params("put_object_annotation", kwargs), **destination}, ) - def _get_copy_ranges(self, size: int, block_size: int) -> list[tuple[int, int]]: - """Split an object into the source ranges of a multipart copy. - - The object is split into ranges of ``block_size`` bytes, whatever the - number of workers, or of a larger size that splits it into at most - ``MULTIPART_UPLOAD_MAX_PARTS`` ranges. A last range shorter than - ``MULTIPART_UPLOAD_MIN_PART_SIZE`` is merged into the previous one, - which is split in half if the result exceeds - ``MULTIPART_UPLOAD_MAX_PART_SIZE``. Every range is then within the - S3 part size limits, including the last one unless the whole object - is smaller than the minimum part size, so that more parts can follow - the copied ones, as in an append. - - Args: - size: The size of the source object in bytes. - block_size: The size in bytes to split the object by, between - ``MULTIPART_UPLOAD_MIN_PART_SIZE`` and - ``MULTIPART_UPLOAD_MAX_PART_SIZE``. It is raised to - ``size`` divided by ``MULTIPART_UPLOAD_MAX_PARTS``, rounded - up, if smaller. The range that a short last range is merged - into can be longer, up to ``MULTIPART_UPLOAD_MAX_PART_SIZE``. - - Returns: - The ``(start, end)`` byte ranges, with an exclusive end, that - cover the whole object in order. - """ - block_size = max(block_size, math.ceil(size / self.MULTIPART_UPLOAD_MAX_PARTS)) - starts = list(range(0, size, block_size)) - if len(starts) > 1 and size - starts[-1] < self.MULTIPART_UPLOAD_MIN_PART_SIZE: - starts.pop() - if size - starts[-1] > self.MULTIPART_UPLOAD_MAX_PART_SIZE: - starts.append(starts[-1] + (size - starts[-1]) // 2) - return list(zip(starts, [*starts[1:], size], strict=True)) - def _check_multipart_upload_size(self, path: str, size: int, block_size: int) -> None: """Check that data fits in a multipart upload before uploading it. @@ -2200,16 +2149,16 @@ def _check_multipart_upload_size(self, path: str, size: int, block_size: int) -> Raises: ValueError: If the data takes more than - ``MULTIPART_UPLOAD_MAX_PARTS`` blocks. + ``S3Core.MULTIPART_UPLOAD_MAX_PARTS`` blocks. """ - if size > block_size * self.MULTIPART_UPLOAD_MAX_PARTS: + if size > block_size * self.core.MULTIPART_UPLOAD_MAX_PARTS: min_block_size = max( - math.ceil(size / self.MULTIPART_UPLOAD_MAX_PARTS), - self.MULTIPART_UPLOAD_MIN_PART_SIZE, + math.ceil(size / self.core.MULTIPART_UPLOAD_MAX_PARTS), + self.core.MULTIPART_UPLOAD_MIN_PART_SIZE, ) raise ValueError( f"Cannot upload {size} bytes to {path} in " - f"{self.MULTIPART_UPLOAD_MAX_PARTS} parts with a block size of " + f"{self.core.MULTIPART_UPLOAD_MAX_PARTS} parts with a block size of " f"{block_size} bytes. Write the file with a block_size, or a " "default_block_size of the filesystem, of at least " f"{min_block_size} bytes." @@ -2249,7 +2198,7 @@ def pipe_file( committed. ValueError: If the path does not contain a key or specifies a version, if the compression is not supported, or if the data - takes more than ``MULTIPART_UPLOAD_MAX_PARTS`` blocks. + takes more than ``S3Core.MULTIPART_UPLOAD_MAX_PARTS`` blocks. """ # Normalized as open() normalizes it, so that the key written, and # the codec that "infer" takes from its extension, do not depend on @@ -2264,7 +2213,7 @@ def pipe_file( # The size in bytes; the length of a memoryview counts its items. size = memoryview(value).nbytes self._check_multipart_upload_size(path, size, block_size) - if self._intrans or size > min(block_size, self.MULTIPART_UPLOAD_MAX_PART_SIZE): + if self._intrans or size > min(block_size, self.core.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. @@ -2330,13 +2279,11 @@ def _finish_multipart_upload( request_kwargs = request_kwargs or {} 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, + parts = [future.result() for future in futures] + return self.core.complete_multipart_upload( + S3Path(bucket, key), + upload_id, + parts, **self.core.operation_params("complete_multipart_upload", request_kwargs), ) except BaseException: @@ -2365,14 +2312,10 @@ def _abort_multipart_upload( those that it accepts. """ try: - self._call( - self._client.abort_multipart_upload, - **{ - **self.core.operation_params("abort_multipart_upload", request_kwargs), - "Bucket": bucket, - "Key": key, - "UploadId": upload_id, - }, + self.core.abort_multipart_upload( + S3Path(bucket, key), + upload_id, + **self.core.operation_params("abort_multipart_upload", request_kwargs), ) except Exception: _logger.exception( @@ -2495,7 +2438,7 @@ def put_file( exists, or an object is created at it before the upload is committed. ValueError: If the file takes more than - ``MULTIPART_UPLOAD_MAX_PARTS`` blocks. + ``S3Core.MULTIPART_UPLOAD_MAX_PARTS`` blocks. Note: Directories are not supported for upload. If lpath is a directory, @@ -3009,11 +2952,9 @@ def clear_multipart_uploads(self, path: str) -> None: with self._create_executor(max_workers=self.max_workers) as executor: futures = [ executor.submit( - self._call, - self._client.abort_multipart_upload, - Bucket=upload.bucket, - Key=upload.key, - UploadId=upload.upload_id, + self.core.abort_multipart_upload, + S3Path(cast(str, upload.bucket), cast(str, upload.key)), + cast(str, upload.upload_id), ) for upload in uploads ] @@ -3351,97 +3292,6 @@ def _put_object(self, bucket: str, key: str, body: bytes | None, **kwargs) -> S3 ) return S3PutObject(response) - def _create_multipart_upload(self, bucket: str, key: str, **kwargs) -> S3MultipartUpload: - request = { - "Bucket": bucket, - "Key": key, - } - - _logger.debug(f"Create multipart upload to s3://{bucket}/{key}.") - response = self._call( - self._client.create_multipart_upload, - # The fields of the request take precedence over inherited - # parameters of the same name. - **{**kwargs, **request}, - ) - return S3MultipartUpload(response) - - def _upload_part_copy( - self, - bucket: str, - key: str, - copy_source: str | dict[str, Any], - upload_id: str, - part_number: int, - copy_source_ranges: tuple[int, int] | None = None, - **kwargs, - ) -> S3MultipartUploadPart: - request = { - "Bucket": bucket, - "Key": key, - "CopySource": copy_source, - "UploadId": upload_id, - "PartNumber": part_number, - } - if copy_source_ranges: - range_ = S3File._format_ranges(copy_source_ranges) - request.update({"CopySourceRange": range_}) - _logger.debug( - f"Upload part copy from {copy_source} to s3://{bucket}/{key} as part {part_number}." - ) - response = self._call( - self._client.upload_part_copy, - # The fields of the request take precedence over inherited - # parameters of the same name. - **{**kwargs, **request}, - ) - return S3MultipartUploadPart(part_number, response) - - def _upload_part( - self, - bucket: str, - key: str, - upload_id: str, - part_number: int, - body: bytes, - **kwargs, - ) -> S3MultipartUploadPart: - request = { - "Bucket": bucket, - "Key": key, - "UploadId": upload_id, - "PartNumber": part_number, - "Body": body, - } - - _logger.debug(f"Upload part of {upload_id} to s3://{bucket}/{key} as part {part_number}.") - response = self._call( - self._client.upload_part, - # The fields of the request take precedence over inherited - # parameters of the same name. - **{**kwargs, **request}, - ) - return S3MultipartUploadPart(part_number, response) - - def _complete_multipart_upload( - self, bucket: str, key: str, upload_id: str, parts: list[dict[str, Any]], **kwargs - ) -> S3CompleteMultipartUpload: - request = { - "Bucket": bucket, - "Key": key, - "UploadId": upload_id, - "MultipartUpload": {"Parts": parts}, - } - - _logger.debug(f"Complete multipart upload {upload_id} to s3://{bucket}/{key}.") - response = self._call( - self._client.complete_multipart_upload, - # The fields of the request take precedence over inherited - # parameters of the same name. - **{**kwargs, **request}, - ) - return S3CompleteMultipartUpload(response) - def _call(self, method: str | Callable[..., Any], **kwargs) -> dict[str, Any]: """Send a request with the core (see :meth:`S3Core.call`). @@ -3484,7 +3334,7 @@ def __init__( In read mode, the object is looked up with ``info()`` and the reads are made conditional on its ETag (``IfMatch``). In append mode, an - existing object smaller than ``MULTIPART_UPLOAD_MIN_PART_SIZE`` is + existing object smaller than ``S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE`` is read into the write buffer; a larger one is copied with ``UploadPartCopy`` as the first parts of a multipart upload, whatever the block size. In exclusive-create mode, the object must not exist @@ -3504,8 +3354,8 @@ def __init__( executor: The executor for parallel operations. If None, a new ``S3ThreadPoolExecutor`` is created. block_size: The block size for reads and writes. Must be between - ``MULTIPART_UPLOAD_MIN_PART_SIZE`` and - ``MULTIPART_UPLOAD_MAX_PART_SIZE``, inclusive, unless reading. + ``S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE`` and + ``S3Core.MULTIPART_UPLOAD_MAX_PART_SIZE``, inclusive, unless reading. cache_type: The fsspec cache type for reads. autocommit: Whether to commit the written data when the file is closed. If False, :meth:`commit` must be called. @@ -3525,8 +3375,8 @@ def __init__( including when the path is a prefix. ValueError: If the path has no key, the version IDs do not match, a version is given for writing, or the block size is not - between ``MULTIPART_UPLOAD_MIN_PART_SIZE`` and - ``MULTIPART_UPLOAD_MAX_PART_SIZE`` for writing. + between ``S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE`` and + ``S3Core.MULTIPART_UPLOAD_MAX_PART_SIZE`` for writing. """ self.max_workers = max_workers # A new dictionary, so that the caller's is not modified. @@ -3559,14 +3409,16 @@ def __init__( # so that a reopened (e.g., unpickled) file reads the same version. path = f"{path}?versionId={self.version_id}" if "r" not in mode and not ( - fs.MULTIPART_UPLOAD_MIN_PART_SIZE <= block_size <= fs.MULTIPART_UPLOAD_MAX_PART_SIZE + fs.core.MULTIPART_UPLOAD_MIN_PART_SIZE + <= block_size + <= fs.core.MULTIPART_UPLOAD_MAX_PART_SIZE ): # When writing, every full block is uploaded as a part of a # multipart upload. raise ValueError( "Block size for writing must be between " - f"5 MiB ({fs.MULTIPART_UPLOAD_MIN_PART_SIZE} bytes) and " - f"5 GiB ({fs.MULTIPART_UPLOAD_MAX_PART_SIZE} bytes), inclusive: {block_size}." + f"5 MiB ({fs.core.MULTIPART_UPLOAD_MIN_PART_SIZE} bytes) and " + f"5 GiB ({fs.core.MULTIPART_UPLOAD_MAX_PART_SIZE} bytes), inclusive: {block_size}." ) self._details: S3Object | dict[str, Any] = {} @@ -3603,7 +3455,7 @@ def __init__( append_info = fs.info(path, refresh=True, **lookup_kwargs) if ( append_info is not None - and append_info.get("size", 0) < fs.MULTIPART_UPLOAD_MIN_PART_SIZE + and append_info.get("size", 0) < fs.core.MULTIPART_UPLOAD_MIN_PART_SIZE ): # Too small to be a part of a multipart upload: rewritten # from the buffer. Only the lookup parameters are sent, so @@ -3732,45 +3584,41 @@ def _initiate_upload(self) -> None: # a multipart upload, whatever the block size. return - self.multipart_upload = self.fs._create_multipart_upload( - bucket=self.bucket, - key=self.key, - **self._get_request_kwargs("create_multipart_upload"), + self.multipart_upload = self.fs.core.create_multipart_upload( + S3Path(self.bucket, self.key), **self._get_request_kwargs("create_multipart_upload") ) if self.append_block: - if self.tell() > self.fs.MULTIPART_UPLOAD_MAX_PART_SIZE: + if self.tell() > self.fs.core.MULTIPART_UPLOAD_MAX_PART_SIZE: info = self.fs.info( self.path, version_id=self.version_id, **self.fs._get_lookup_kwargs(self.s3_additional_kwargs), ) - ranges = self.fs._get_copy_ranges( + ranges = self.fs.core.part_ranges( # Set copy source file byte size info.get("size", 0), - self.fs.MULTIPART_UPLOAD_MAX_PART_SIZE, + self.fs.core.MULTIPART_UPLOAD_MAX_PART_SIZE, ) for i, range_ in enumerate(ranges): self.multipart_upload_parts.append( self._executor.submit( - self.fs._upload_part_copy, - bucket=self.bucket, - key=self.key, - copy_source=self.path, + self.fs.core.upload_part_copy, + path=S3Path(self.bucket, self.key), upload_id=cast(str, self.multipart_upload.upload_id), part_number=i + 1, - copy_source_ranges=range_, + source=S3Path(self.bucket, self.key), + range_=range_, **self._get_request_kwargs("upload_part_copy"), ) ) else: self.multipart_upload_parts.append( self._executor.submit( - self.fs._upload_part_copy, - bucket=self.bucket, - key=self.key, - copy_source=self.path, + self.fs.core.upload_part_copy, + path=S3Path(self.bucket, self.key), upload_id=cast(str, self.multipart_upload.upload_id), part_number=1, + source=S3Path(self.bucket, self.key), **self._get_request_kwargs("upload_part_copy"), ) ) @@ -3805,10 +3653,10 @@ def _upload_chunk(self, final: bool = False) -> bool: # ahead one block and merge a short last block into this one. next_data = buffer.read(self.blocksize) next_data_size = len(next_data) - if 0 < next_data_size < self.fs.MULTIPART_UPLOAD_MIN_PART_SIZE: + if 0 < next_data_size < self.fs.core.MULTIPART_UPLOAD_MIN_PART_SIZE: upload_data = data + next_data upload_data_size = len(upload_data) - if upload_data_size < self.fs.MULTIPART_UPLOAD_MAX_PART_SIZE: + if upload_data_size < self.fs.core.MULTIPART_UPLOAD_MAX_PART_SIZE: uploads = [upload_data] else: split_size = upload_data_size // 2 @@ -3818,22 +3666,21 @@ def _upload_chunk(self, final: bool = False) -> bool: uploads = [data] for upload in uploads: - if part_number >= self.fs.MULTIPART_UPLOAD_MAX_PARTS: + if part_number >= self.fs.core.MULTIPART_UPLOAD_MAX_PARTS: self._close_without_commit() raise ValueError( - f"Cannot upload more than {self.fs.MULTIPART_UPLOAD_MAX_PARTS} " + f"Cannot upload more than {self.fs.core.MULTIPART_UPLOAD_MAX_PARTS} " f"parts to s3://{self.bucket}/{self.key} with a block size of " f"{self.blocksize} bytes. Write the file with a block_size, or " "a default_block_size of the filesystem, large enough for it to " - f"fit in {self.fs.MULTIPART_UPLOAD_MAX_PARTS} parts, including " + f"fit in {self.fs.core.MULTIPART_UPLOAD_MAX_PARTS} parts, including " "the parts copied from the existing object in an append." ) part_number += 1 self.multipart_upload_parts.append( self._executor.submit( - self.fs._upload_part, - bucket=self.bucket, - key=self.key, + self.fs.core.upload_part, + path=S3Path(self.bucket, self.key), upload_id=cast(str, self.multipart_upload.upload_id), part_number=part_number, body=upload, @@ -3923,14 +3770,10 @@ def discard(self) -> None: # be stored after the abort, so wait for the parts that could not # be cancelled first. wait([f for f in self.multipart_upload_parts if not f.cancel()]) - self.fs._call( - "abort_multipart_upload", - **{ - **self._get_request_kwargs("abort_multipart_upload"), - "Bucket": self.bucket, - "Key": self.key, - "UploadId": self.multipart_upload.upload_id, - }, + self.fs.core.abort_multipart_upload( + S3Path(self.bucket, self.key), + cast(str, self.multipart_upload.upload_id), + **self._get_request_kwargs("abort_multipart_upload"), ) self.multipart_upload = None diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index fe015e9a..88f2d897 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -29,6 +29,7 @@ S3CompleteMultipartUpload, S3Metadata, S3MultipartUpload, + S3MultipartUploadPart, S3Object, S3ObjectType, S3ObjectVersion, @@ -205,7 +206,7 @@ def _pipe_file_in_transaction( FileExistsError: If the mode is "create" and the path already exists. ValueError: If the compression is not supported, or if the data - takes more than ``MULTIPART_UPLOAD_MAX_PARTS`` blocks. + takes more than ``S3Core.MULTIPART_UPLOAD_MAX_PARTS`` blocks. """ # See S3FileSystem.pipe_file. compression = get_compression(self._strip_protocol(path), kwargs.pop("compression", None)) @@ -258,7 +259,7 @@ def _put_file_in_transaction( FileExistsError: If the mode is "create" and the path already exists. ValueError: If the file takes more than - ``MULTIPART_UPLOAD_MAX_PARTS`` blocks. + ``S3Core.MULTIPART_UPLOAD_MAX_PARTS`` blocks. """ if os.path.isdir(lpath): return @@ -530,7 +531,7 @@ async def _copy_file(self, path1: str, path2: str, **kwargs) -> bool: return False size1 = info1.get("size", 0) try: - if size1 <= S3FileSystem.MULTIPART_UPLOAD_MAX_PART_SIZE: + if size1 <= self.core.MULTIPART_UPLOAD_MAX_PART_SIZE: await asyncio.to_thread( self._sync_fs._copy_object, bucket1=source.bucket, @@ -596,22 +597,22 @@ async def _copy_object_with_multipart_upload( directive has an invalid value. """ max_workers = max_workers if max_workers else self._sync_fs.max_workers - block_size = block_size if block_size else S3FileSystem.MULTIPART_UPLOAD_MAX_PART_SIZE + block_size = block_size if block_size else self.core.MULTIPART_UPLOAD_MAX_PART_SIZE if ( - block_size < S3FileSystem.MULTIPART_UPLOAD_MIN_PART_SIZE - or block_size > S3FileSystem.MULTIPART_UPLOAD_MAX_PART_SIZE + block_size < self.core.MULTIPART_UPLOAD_MIN_PART_SIZE + or block_size > self.core.MULTIPART_UPLOAD_MAX_PART_SIZE ): raise ValueError( "Block size must be between " - f"5 MiB ({S3FileSystem.MULTIPART_UPLOAD_MIN_PART_SIZE} bytes) and " - f"5 GiB ({S3FileSystem.MULTIPART_UPLOAD_MAX_PART_SIZE} bytes), " + f"5 MiB ({self.core.MULTIPART_UPLOAD_MIN_PART_SIZE} bytes) and " + f"5 GiB ({self.core.MULTIPART_UPLOAD_MAX_PART_SIZE} bytes), " f"inclusive: {block_size}." ) create_kwargs, version_id1, head_size = await asyncio.to_thread( self._sync_fs._get_multipart_copy_kwargs, bucket1, key1, version_id1, kwargs ) - if head_size is not None and head_size <= S3FileSystem.MULTIPART_UPLOAD_MAX_PART_SIZE: + if head_size is not None and head_size <= self.core.MULTIPART_UPLOAD_MAX_PART_SIZE: # See S3FileSystem._copy_object_with_multipart_upload. await asyncio.to_thread( self._sync_fs._copy_object, @@ -624,15 +625,9 @@ async def _copy_object_with_multipart_upload( ) return # The size of the copied version; see S3FileSystem. - ranges = self._sync_fs._get_copy_ranges( - size1 if head_size is None else head_size, block_size - ) - copy_source: dict[str, Any] = { - "Bucket": bucket1, - "Key": key1, - } - if version_id1: - copy_source["VersionId"] = version_id1 + ranges = self.core.part_ranges(size1 if head_size is None else head_size, block_size) + source = S3Path(bucket1, key1, version_id1) + destination = S3Path(bucket2, key2) # Listed before anything is written; see S3FileSystem. annotations = ( await asyncio.to_thread( @@ -642,42 +637,34 @@ async def _copy_object_with_multipart_upload( else [] ) multipart_upload = await asyncio.to_thread( - self._sync_fs._create_multipart_upload, - bucket=bucket2, - key=key2, - **create_kwargs, + self.core.create_multipart_upload, destination, **create_kwargs ) upload_id = cast(str, multipart_upload.upload_id) semaphore = asyncio.Semaphore(max_workers) - part_kwargs = self._sync_fs.core.operation_params("upload_part_copy", kwargs) + part_kwargs = self.core.operation_params("upload_part_copy", kwargs) failed = False - async def _upload_part(i: int, range_: tuple[int, int]) -> dict[str, Any] | None: + async def _upload_part(i: int, range_: tuple[int, int]) -> S3MultipartUploadPart | None: nonlocal failed async with semaphore: if failed: # The upload is being aborted; do not start more parts. return None try: - result = await asyncio.to_thread( - self._sync_fs._upload_part_copy, - bucket=bucket2, - key=key2, - copy_source=copy_source, + return await asyncio.to_thread( + self.core.upload_part_copy, + path=destination, upload_id=upload_id, part_number=i + 1, - copy_source_ranges=range_, + source=source, + range_=range_, **part_kwargs, ) except Exception: # Set before the semaphore lets a waiting part start. failed = True raise - return { - "ETag": result.etag, - "PartNumber": result.part_number, - } tasks = [asyncio.ensure_future(_upload_part(i, r)) for i, r in enumerate(ranges)] completion: asyncio.Task[S3CompleteMultipartUpload] | None = None @@ -708,12 +695,11 @@ async def _abort() -> None: parts = [task.result() for task in tasks] completion = asyncio.ensure_future( asyncio.to_thread( - self._sync_fs._complete_multipart_upload, - bucket=bucket2, - key=key2, - upload_id=upload_id, - parts=cast(list[dict[str, Any]], parts), - **self._sync_fs.core.operation_params("complete_multipart_upload", kwargs), + self.core.complete_multipart_upload, + destination, + upload_id, + cast(list[S3MultipartUploadPart], parts), + **self.core.operation_params("complete_multipart_upload", kwargs), ) ) # shield keeps a cancellation from cancelling the completion, whose diff --git a/pyathena/filesystem/s3_core.py b/pyathena/filesystem/s3_core.py index b1e55632..bde5cd42 100644 --- a/pyathena/filesystem/s3_core.py +++ b/pyathena/filesystem/s3_core.py @@ -10,7 +10,8 @@ from __future__ import annotations import logging -from collections.abc import Callable, Iterable, Iterator, Mapping +import math +from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence from dataclasses import dataclass from datetime import datetime from typing import Any, ClassVar, cast @@ -19,7 +20,13 @@ from botocore.client import BaseClient from pyathena.filesystem.s3_errors import S3ClientError -from pyathena.filesystem.s3_object import S3Metadata, S3ObjectVersion +from pyathena.filesystem.s3_object import ( + S3CompleteMultipartUpload, + S3Metadata, + S3MultipartUpload, + S3MultipartUploadPart, + S3ObjectVersion, +) from pyathena.filesystem.s3_path import S3Path from pyathena.util import RetryConfig, override, retry_api_call @@ -410,6 +417,14 @@ class S3Core: ... print([o.key for o in page.objects]) """ + # https://docs.aws.amazon.com/AmazonS3/latest/userguide/qfacts.html + # The minimum size of a part in a multipart upload is 5MiB. + MULTIPART_UPLOAD_MIN_PART_SIZE: int = 5 * 2**20 # 5MiB + # The maximum size of a part in a multipart upload is 5GiB. + MULTIPART_UPLOAD_MAX_PART_SIZE: int = 5 * 2**30 # 5GiB + # The maximum number of parts per multipart upload is 10,000. + MULTIPART_UPLOAD_MAX_PARTS: int = 10_000 + def __init__( self, client: BaseClient, @@ -578,6 +593,194 @@ def delete_objects(self, batch: S3DeleteBatch, **params) -> S3DeleteResult: ) return S3DeleteResult.from_response(batch.bucket, response) + def create_multipart_upload(self, path: S3Path, **params) -> S3MultipartUpload: + """Start a multipart upload to an object with CreateMultipartUpload. + + Args: + path: The path of the object to write. + **params: Additional request parameters. The bucket and key of + the path take precedence over parameters of the same name. + + Returns: + The multipart upload. + + Raises: + ValueError: If the path has no key. + """ + if not path.key: + raise ValueError(f"The path has no key: {path.uri}.") + request: dict[str, Any] = {"Bucket": path.bucket, "Key": path.key} + _logger.debug(f"Create multipart upload to {path.uri}.") + response = self.call(self._client.create_multipart_upload, **{**params, **request}) + return S3MultipartUpload(response) + + def upload_part( + self, path: S3Path, upload_id: str, part_number: int, body: bytes, **params + ) -> S3MultipartUploadPart: + """Upload a part of a multipart upload with UploadPart. + + Args: + path: The path of the object that the upload writes. + upload_id: The ID of the multipart upload. + part_number: The number of the part, from 1. + body: The data of the part. + **params: Additional request parameters. The fields that the + other arguments set take precedence over parameters of the + same name. + + Returns: + The uploaded part. + + Raises: + ValueError: If the path has no key. + """ + if not path.key: + raise ValueError(f"The path has no key: {path.uri}.") + request: dict[str, Any] = { + "Bucket": path.bucket, + "Key": path.key, + "UploadId": upload_id, + "PartNumber": part_number, + "Body": body, + } + _logger.debug(f"Upload part of {upload_id} to {path.uri} as part {part_number}.") + response = self.call(self._client.upload_part, **{**params, **request}) + return S3MultipartUploadPart(part_number, response) + + def upload_part_copy( + self, + path: S3Path, + upload_id: str, + part_number: int, + source: S3Path, + range_: tuple[int, int] | None = None, + **params, + ) -> S3MultipartUploadPart: + """Copy a part of a multipart upload from an object with UploadPartCopy. + + Args: + path: The path of the object that the upload writes. + upload_id: The ID of the multipart upload. + part_number: The number of the part, from 1. + source: The path of the object to copy, with the version ID to + copy, if any. + range_: The ``(start, end)`` byte range of the source to copy, + with an exclusive end; None copies the whole source. + **params: Additional request parameters. The fields that the + other arguments set take precedence over parameters of the + same name. + + Returns: + The copied part. + + Raises: + ValueError: If the path or the source has no key. + """ + if not path.key: + raise ValueError(f"The path has no key: {path.uri}.") + if not source.key: + raise ValueError(f"The source has no key: {source.uri}.") + copy_source: dict[str, Any] = {"Bucket": source.bucket, "Key": source.key} + if source.version_id: + copy_source.update({"VersionId": source.version_id}) + request: dict[str, Any] = { + "Bucket": path.bucket, + "Key": path.key, + "CopySource": copy_source, + "UploadId": upload_id, + "PartNumber": part_number, + } + if range_: + request.update({"CopySourceRange": f"bytes={range_[0]}-{range_[1] - 1}"}) + _logger.debug(f"Upload part copy from {source.uri} to {path.uri} as part {part_number}.") + response = self.call(self._client.upload_part_copy, **{**params, **request}) + return S3MultipartUploadPart(part_number, response) + + def complete_multipart_upload( + self, path: S3Path, upload_id: str, parts: Sequence[S3MultipartUploadPart], **params + ) -> S3CompleteMultipartUpload: + """Complete a multipart upload with CompleteMultipartUpload. + + Args: + path: The path of the object that the upload writes. + upload_id: The ID of the multipart upload. + parts: The uploaded parts, in part-number order. + **params: Additional request parameters. The fields that the + other arguments set take precedence over parameters of the + same name. + + Returns: + The completed upload. + + Raises: + ValueError: If the path has no key. + """ + if not path.key: + raise ValueError(f"The path has no key: {path.uri}.") + request: dict[str, Any] = { + "Bucket": path.bucket, + "Key": path.key, + "UploadId": upload_id, + "MultipartUpload": { + "Parts": [{"ETag": p.etag, "PartNumber": p.part_number} for p in parts] + }, + } + _logger.debug(f"Complete multipart upload {upload_id} to {path.uri}.") + response = self.call(self._client.complete_multipart_upload, **{**params, **request}) + return S3CompleteMultipartUpload(response) + + def abort_multipart_upload(self, path: S3Path, upload_id: str, **params) -> None: + """Abort a multipart upload with AbortMultipartUpload. + + Args: + path: The path of the object that the upload writes. + upload_id: The ID of the multipart upload. + **params: Additional request parameters. The fields that the + other arguments set take precedence over parameters of the + same name. + + Raises: + ValueError: If the path has no key. + """ + if not path.key: + raise ValueError(f"The path has no key: {path.uri}.") + request: dict[str, Any] = {"Bucket": path.bucket, "Key": path.key, "UploadId": upload_id} + self.call(self._client.abort_multipart_upload, **{**params, **request}) + + def part_ranges(self, size: int, block_size: int) -> list[tuple[int, int]]: + """Split an object into the source ranges of the parts that copy it. + + The object is split into ranges of ``block_size`` bytes, or of a + larger size that splits it into at most + ``MULTIPART_UPLOAD_MAX_PARTS`` ranges. A last range shorter than + ``MULTIPART_UPLOAD_MIN_PART_SIZE`` is merged into the previous one, + which is split in half if the result exceeds + ``MULTIPART_UPLOAD_MAX_PART_SIZE``. Every range is then within the + S3 part size limits, including the last one unless the whole object + is smaller than the minimum part size, so that more parts can follow + the copied ones, as in an append. + + Args: + size: The size of the source object in bytes. + block_size: The size in bytes to split the object by, between + ``MULTIPART_UPLOAD_MIN_PART_SIZE`` and + ``MULTIPART_UPLOAD_MAX_PART_SIZE``. It is raised to + ``size`` divided by ``MULTIPART_UPLOAD_MAX_PARTS``, rounded + up, if smaller. The range that a short last range is merged + into can be longer, up to ``MULTIPART_UPLOAD_MAX_PART_SIZE``. + + Returns: + The ``(start, end)`` byte ranges, with an exclusive end, that + cover the whole object in order. + """ + block_size = max(block_size, math.ceil(size / self.MULTIPART_UPLOAD_MAX_PARTS)) + starts = list(range(0, size, block_size)) + if len(starts) > 1 and size - starts[-1] < self.MULTIPART_UPLOAD_MIN_PART_SIZE: + starts.pop() + if size - starts[-1] > self.MULTIPART_UPLOAD_MAX_PART_SIZE: + starts.append(starts[-1] + (size - starts[-1]) // 2) + return list(zip(starts, [*starts[1:], size], strict=True)) + def list_objects_page( self, bucket: str, diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index 7fd15cd1..f781bb03 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -17,7 +17,7 @@ import uuid from concurrent.futures import Future, ThreadPoolExecutor, wait from datetime import UTC, datetime -from itertools import chain, pairwise +from itertools import chain from pathlib import Path from types import SimpleNamespace from unittest import mock @@ -1184,15 +1184,15 @@ def test_pipe_file_buffered_s3_parameters(self, transaction): fs = S3FileSystem( key="dummy", secret="dummy", region_name="us-east-1", skip_instance_cache=True ) - fs._create_multipart_upload = mock.MagicMock( + fs.core.create_multipart_upload = mock.MagicMock( return_value=SimpleNamespace(upload_id="uploadid") ) - fs._upload_part = mock.MagicMock( + fs.core.upload_part = mock.MagicMock( side_effect=lambda **kw: SimpleNamespace(etag='"e"', part_number=kw["part_number"]) ) - fs._complete_multipart_upload = mock.MagicMock() + fs.core.complete_multipart_upload = mock.MagicMock() fs._put_object = mock.MagicMock() - data = b"x" * (fs.MULTIPART_UPLOAD_MIN_PART_SIZE + 1) + data = b"x" * (fs.core.MULTIPART_UPLOAD_MIN_PART_SIZE + 1) if transaction: with fs.transaction: @@ -1202,8 +1202,8 @@ def test_pipe_file_buffered_s3_parameters(self, transaction): ) else: fs.pipe_file("s3://bucket/key", data, ContentType="text/csv") - fs._create_multipart_upload.assert_called_once_with( - bucket="bucket", key="key", ContentType="text/csv" + fs.core.create_multipart_upload.assert_called_once_with( + S3Path("bucket", "key"), ContentType="text/csv" ) def test_put_file_open_parameters(self, tmp_path): @@ -1214,7 +1214,7 @@ def test_put_file_open_parameters(self, tmp_path): ) lpath = tmp_path / "data.csv" lpath.write_bytes(b"a") - block_size = fs.MULTIPART_UPLOAD_MIN_PART_SIZE + block_size = fs.core.MULTIPART_UPLOAD_MIN_PART_SIZE with ( mock.patch.object(fs, "open", wraps=fs.open) as open_, @@ -1292,7 +1292,7 @@ def call(method, **request): (0, ["put_object"]), (1, ["put_object"]), ( - S3FileSystem.MULTIPART_UPLOAD_MIN_PART_SIZE + 1, + S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE + 1, ["create_multipart_upload", "upload_part", "complete_multipart_upload"], ), ], @@ -1306,7 +1306,9 @@ def test_open_exclusive_create(self, size, expected): fs.exists = mock.MagicMock(return_value=False) requests = self._record_requests(fs) - with fs.open("s3://bucket/key", "xb", block_size=fs.MULTIPART_UPLOAD_MIN_PART_SIZE) as f: + with fs.open( + "s3://bucket/key", "xb", block_size=fs.core.MULTIPART_UPLOAD_MIN_PART_SIZE + ) as f: f.write(b"a" * size) fs.exists.assert_called_once() @@ -1331,7 +1333,7 @@ def test_open_exclusive_create_existing(self): [ (1, ["put_object"]), ( - S3FileSystem.MULTIPART_UPLOAD_MIN_PART_SIZE + 1, + S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE + 1, [ "create_multipart_upload", "upload_part", @@ -1352,7 +1354,9 @@ def test_open_exclusive_create_created_since(self, size, expected): with ( pytest.raises(FileExistsError), - fs.open("s3://bucket/key", "xb", block_size=fs.MULTIPART_UPLOAD_MIN_PART_SIZE) as f, + fs.open( + "s3://bucket/key", "xb", block_size=fs.core.MULTIPART_UPLOAD_MIN_PART_SIZE + ) as f, ): f.write(b"a" * size) @@ -1393,7 +1397,7 @@ def test_put_file_create_existing(self, tmp_path): ("size", "conditional"), [ (1, "put_object"), - (S3FileSystem.MULTIPART_UPLOAD_MIN_PART_SIZE + 1, "complete_multipart_upload"), + (S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE + 1, "complete_multipart_upload"), ], ) def test_pipe_file_create_created_since(self, size, conditional): @@ -1402,7 +1406,7 @@ def test_pipe_file_create_created_since(self, size, conditional): # created after the existence check is not replaced. fs = self._make_fs() fs.default_cache_type = "bytes" - fs.default_block_size = fs.MULTIPART_UPLOAD_MIN_PART_SIZE + fs.default_block_size = fs.core.MULTIPART_UPLOAD_MIN_PART_SIZE fs.exists = mock.MagicMock(return_value=False) requests = self._record_requests(fs, precondition_failed=True) @@ -1432,7 +1436,7 @@ def call(method, **request): return {"UploadId": "uploadid", "ETag": '"e"'} fs._call.side_effect = call - block_size = fs.MULTIPART_UPLOAD_MIN_PART_SIZE + block_size = fs.core.MULTIPART_UPLOAD_MIN_PART_SIZE with ( pytest.raises(OSError, match="upload failed") if fail else contextlib.nullcontext(), @@ -1460,7 +1464,7 @@ def test_finish_multipart_upload_request_parameters(self): # GH-946: the completion and the abort receive the parameters of the # upload that they accept. fs = self._make_fs() - fs._complete_multipart_upload = mock.MagicMock() + fs.core.complete_multipart_upload = mock.MagicMock() kwargs = { "ContentType": "text/csv", "RequestPayer": "requester", @@ -1483,11 +1487,10 @@ def test_finish_multipart_upload_request_parameters(self): request_kwargs=kwargs, ) - fs._complete_multipart_upload.assert_called_once_with( - bucket="bucket", - key="key", - upload_id="uploadid", - parts=[{"ETag": '"e1"', "PartNumber": 1}], + fs.core.complete_multipart_upload.assert_called_once_with( + S3Path("bucket", "key"), + "uploadid", + [part.result()], RequestPayer="requester", SSECustomerAlgorithm="AES256", ) @@ -1832,12 +1835,12 @@ def test_cp_file_multipart_parameters(self, size): fs.cp_file( "s3://bucket/src", "s3://bucket/dst", - block_size=fs.MULTIPART_UPLOAD_MIN_PART_SIZE, + block_size=fs.core.MULTIPART_UPLOAD_MIN_PART_SIZE, max_workers=2, RequestPayer="requester", ) - if size <= fs.MULTIPART_UPLOAD_MAX_PART_SIZE: + if size <= fs.core.MULTIPART_UPLOAD_MAX_PART_SIZE: fs._copy_object.assert_called_once_with( bucket1="bucket", key1="src", @@ -1855,7 +1858,7 @@ def test_cp_file_multipart_parameters(self, size): bucket2="bucket", key2="dst", max_workers=2, - block_size=fs.MULTIPART_UPLOAD_MIN_PART_SIZE, + block_size=fs.core.MULTIPART_UPLOAD_MIN_PART_SIZE, RequestPayer="requester", ) @@ -1863,10 +1866,10 @@ def test_copy_object_with_multipart_upload_request_parameters(self): # GH-946: the part copies receive the parameters of the copy that # they accept, and the completion and the abort get them all. fs = self._make_fs() - fs._create_multipart_upload = mock.MagicMock( + fs.core.create_multipart_upload = mock.MagicMock( return_value=SimpleNamespace(upload_id="uploadid") ) - fs._upload_part_copy = mock.MagicMock( + fs.core.upload_part_copy = mock.MagicMock( side_effect=lambda **kw: SimpleNamespace(etag='"e"', part_number=kw["part_number"]) ) fs._finish_multipart_upload = mock.MagicMock() @@ -1889,14 +1892,14 @@ def test_copy_object_with_multipart_upload_request_parameters(self): **kwargs, ) - fs._create_multipart_upload.assert_called_once_with( - bucket="bucket", key="dst", ContentType="text/csv", RequestPayer="requester" + fs.core.create_multipart_upload.assert_called_once_with( + S3Path("bucket", "dst"), ContentType="text/csv", RequestPayer="requester" ) # Only the HeadObject of the source, for its version. assert fs._call.call_count == 1 assert all( c.kwargs["RequestPayer"] == "requester" and "ContentType" not in c.kwargs - for c in fs._upload_part_copy.call_args_list + for c in fs.core.upload_part_copy.call_args_list ) assert fs._finish_multipart_upload.call_args.kwargs["request_kwargs"] == kwargs @@ -1971,7 +1974,7 @@ def test_cp_file_failed_multipart_copy_invalidates_cache(self): # written the destination, so its cached entries are removed. fs = self._make_fs() fs.info = mock.MagicMock(return_value=self._file_object("src")) - fs.info.return_value.size = S3FileSystem.MULTIPART_UPLOAD_MAX_PART_SIZE + 1 + fs.info.return_value.size = S3Core.MULTIPART_UPLOAD_MAX_PART_SIZE + 1 fs._copy_object_with_multipart_upload = mock.MagicMock(side_effect=PermissionError) fs.dircache["bucket/dst"] = [self._file_object("dst")] @@ -1986,13 +1989,13 @@ def test_copy_object_with_multipart_upload_head_object_size(self): fs = self._make_fs() block_size = MULTIPART_COPY_BLOCK_SIZE fs._call.return_value = { - "ContentLength": 2 * block_size + S3FileSystem.MULTIPART_UPLOAD_MIN_PART_SIZE, + "ContentLength": 2 * block_size + S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE, "VersionId": "null", } - fs._create_multipart_upload = mock.MagicMock( + fs.core.create_multipart_upload = mock.MagicMock( return_value=SimpleNamespace(upload_id="uploadid") ) - fs._upload_part_copy = mock.MagicMock() + fs.core.upload_part_copy = mock.MagicMock() fs._finish_multipart_upload = mock.MagicMock() self._multipart_copy( @@ -2004,10 +2007,10 @@ def test_copy_object_with_multipart_upload_head_object_size(self): # The parts are copied in parallel, in any order. parts = sorted( - (c.kwargs["copy_source_ranges"], c.kwargs["copy_source"]) - for c in fs._upload_part_copy.call_args_list + (c.kwargs["range_"], c.kwargs["source"]) + for c in fs.core.upload_part_copy.call_args_list ) - source = {"Bucket": "bucket", "Key": "src"} + source = S3Path("bucket", "src") assert parts == [ ((0, block_size), source), ((block_size, 2 * block_size), source), @@ -2022,7 +2025,7 @@ def test_copy_object_with_multipart_upload_small_head_object_size(self, size): fs = self._make_fs() fs._call.return_value = {"ContentLength": size, "VersionId": "v1"} fs._copy_object = mock.MagicMock() - fs._create_multipart_upload = mock.MagicMock() + fs.core.create_multipart_upload = mock.MagicMock() self._multipart_copy(fs, ContentType="text/csv", RequestPayer="requester") @@ -2035,7 +2038,7 @@ def test_copy_object_with_multipart_upload_small_head_object_size(self, size): ContentType="text/csv", RequestPayer="requester", ) - fs._create_multipart_upload.assert_not_called() + fs.core.create_multipart_upload.assert_not_called() # Only HeadObject; the tags are not read for the multipart upload. assert fs._call.call_count == 1 @@ -2176,10 +2179,10 @@ def test_pipe_file_trailing_slash(self, intrans, size): fs.default_cache_type = "bytes" fs._transaction = None fs._put_object = mock.MagicMock() - fs._create_multipart_upload = mock.MagicMock( + fs.core.create_multipart_upload = mock.MagicMock( return_value=SimpleNamespace(upload_id="uploadid") ) - fs._upload_part = mock.MagicMock( + fs.core.upload_part = mock.MagicMock( side_effect=lambda **kw: SimpleNamespace(etag='"e"', part_number=kw["part_number"]) ) fs._finish_multipart_upload = mock.MagicMock() @@ -2187,8 +2190,10 @@ def test_pipe_file_trailing_slash(self, intrans, size): with fs.transaction if intrans else contextlib.nullcontext(): fs.pipe_file("s3://bucket/dir/key/", b"a" * size) - calls = fs._put_object.call_args_list + fs._create_multipart_upload.call_args_list - assert [c.kwargs["key"] for c in calls] == ["dir/key"] + keys = [c.kwargs["key"] for c in fs._put_object.call_args_list] + [ + c.args[0].key for c in fs.core.create_multipart_upload.call_args_list + ] + assert keys == ["dir/key"] def test_pipe_file_memoryview_routed_by_bytes(self): # A memoryview larger than the block size in bytes, but not in items, @@ -2197,10 +2202,10 @@ def test_pipe_file_memoryview_routed_by_bytes(self): fs = self._make_fs() fs.default_cache_type = "bytes" fs._put_object = mock.MagicMock() - fs._create_multipart_upload = mock.MagicMock( + fs.core.create_multipart_upload = mock.MagicMock( return_value=SimpleNamespace(upload_id="uploadid") ) - fs._upload_part = mock.MagicMock( + fs.core.upload_part = mock.MagicMock( side_effect=lambda **kw: SimpleNamespace(etag='"e"', part_number=kw["part_number"]) ) fs._finish_multipart_upload = mock.MagicMock() @@ -2209,7 +2214,7 @@ def test_pipe_file_memoryview_routed_by_bytes(self): fs.pipe_file("s3://bucket/key", memoryview(data).cast("I")) fs._put_object.assert_not_called() - assert b"".join(c.kwargs["body"] for c in fs._upload_part.call_args_list) == data + assert b"".join(c.kwargs["body"] for c in fs.core.upload_part.call_args_list) == data fs._finish_multipart_upload.assert_called_once() def test_pipe_file_small_drops_max_workers(self): @@ -2226,10 +2231,10 @@ def test_pipe_file_buffered_non_contiguous_memoryview(self): # raise BufferError for it. fs = self._make_fs() fs.default_cache_type = "bytes" - fs._create_multipart_upload = mock.MagicMock( + fs.core.create_multipart_upload = mock.MagicMock( return_value=SimpleNamespace(upload_id="uploadid") ) - fs._upload_part = mock.MagicMock( + fs.core.upload_part = mock.MagicMock( side_effect=lambda **kw: SimpleNamespace(etag='"e"', part_number=kw["part_number"]) ) fs._finish_multipart_upload = mock.MagicMock() @@ -2237,7 +2242,7 @@ def test_pipe_file_buffered_non_contiguous_memoryview(self): fs.pipe_file("s3://bucket/key", memoryview(b"ab" * size)[::2]) - assert b"".join(c.kwargs["body"] for c in fs._upload_part.call_args_list) == b"a" * size + assert b"".join(c.kwargs["body"] for c in fs.core.upload_part.call_args_list) == b"a" * size fs._finish_multipart_upload.assert_called_once() fs._call.assert_not_called() @@ -2300,10 +2305,10 @@ def test_pipe_file_compression_multipart(self, intrans): fs = self._make_fs() fs.default_cache_type = "bytes" fs._transaction = None - fs._create_multipart_upload = mock.MagicMock( + fs.core.create_multipart_upload = mock.MagicMock( return_value=SimpleNamespace(upload_id="uploadid") ) - fs._upload_part = mock.MagicMock( + fs.core.upload_part = mock.MagicMock( side_effect=lambda **kw: SimpleNamespace(etag='"e"', part_number=kw["part_number"]) ) fs._finish_multipart_upload = mock.MagicMock() @@ -2313,7 +2318,7 @@ def test_pipe_file_compression_multipart(self, intrans): with fs.transaction if intrans else contextlib.nullcontext(): fs.pipe_file("s3://bucket/key", value, compression="gzip") - body = b"".join(c.kwargs["body"] for c in fs._upload_part.call_args_list) + body = b"".join(c.kwargs["body"] for c in fs.core.upload_part.call_args_list) assert gzip.decompress(body) == value fs._finish_multipart_upload.assert_called_once() @@ -2374,7 +2379,7 @@ def test_pipe_file_failed_write_aborts_multipart_upload(self): # aborts the upload instead of completing it. fs = self._make_fs() fs.default_cache_type = "bytes" - fs._create_multipart_upload = mock.MagicMock( + fs.core.create_multipart_upload = mock.MagicMock( return_value=SimpleNamespace(upload_id="uploadid") ) fs._finish_multipart_upload = mock.MagicMock() @@ -2387,7 +2392,7 @@ def test_pipe_file_failed_write_aborts_multipart_upload(self): fs._finish_multipart_upload.assert_not_called() fs._call.assert_called_once_with( - "abort_multipart_upload", Bucket="bucket", Key="key", UploadId="uploadid" + fs._client.abort_multipart_upload, Bucket="bucket", Key="key", UploadId="uploadid" ) executor.shutdown.assert_called_once() @@ -2429,10 +2434,10 @@ def test_put_file_failed_write_aborts_multipart_upload(self, tmp_path, intrans): fs = self._make_fs() fs.default_cache_type = "bytes" fs._transaction = None - fs._create_multipart_upload = mock.MagicMock( + fs.core.create_multipart_upload = mock.MagicMock( return_value=SimpleNamespace(upload_id="uploadid") ) - fs._upload_part = mock.MagicMock( + fs.core.upload_part = mock.MagicMock( side_effect=lambda **kw: SimpleNamespace(etag='"e"', part_number=kw["part_number"]) ) fs._finish_multipart_upload = mock.MagicMock() @@ -2449,7 +2454,7 @@ def test_put_file_failed_write_aborts_multipart_upload(self, tmp_path, intrans): fs._finish_multipart_upload.assert_not_called() fs._call.assert_called_once_with( - "abort_multipart_upload", Bucket="bucket", Key="key", UploadId="uploadid" + fs._client.abort_multipart_upload, Bucket="bucket", Key="key", UploadId="uploadid" ) @pytest.mark.parametrize( @@ -2465,8 +2470,8 @@ def test_put_file_failed_write_aborts_multipart_upload(self, tmp_path, intrans): ) def test_check_multipart_upload_size(self, size, block_size, min_block_size): fs = self._make_fs() - fs.MULTIPART_UPLOAD_MIN_PART_SIZE = 4 - fs.MULTIPART_UPLOAD_MAX_PARTS = 3 + fs.core.MULTIPART_UPLOAD_MIN_PART_SIZE = 4 + fs.core.MULTIPART_UPLOAD_MAX_PARTS = 3 if min_block_size is None: fs._check_multipart_upload_size("s3://bucket/key", size, block_size) @@ -2479,7 +2484,7 @@ def test_put_file_exceeding_max_parts(self, tmp_path, kwargs): # GH-953: a file that does not fit in the maximum number of parts is # rejected before anything is uploaded. fs = self._make_fs() - fs.MULTIPART_UPLOAD_MAX_PARTS = 3 + fs.core.MULTIPART_UPLOAD_MAX_PARTS = 3 fs.default_block_size = 4 fs.open = mock.MagicMock() lpath = tmp_path / "data" @@ -2542,7 +2547,7 @@ def test_pipe_file_exceeding_max_parts(self, value, kwargs): # GH-953: data that does not fit in the maximum number of parts is # rejected before anything is uploaded. fs = self._make_fs() - fs.MULTIPART_UPLOAD_MAX_PARTS = 3 + fs.core.MULTIPART_UPLOAD_MAX_PARTS = 3 fs.default_block_size = 4 fs.open = mock.MagicMock() fs._put_object = mock.MagicMock() @@ -2628,11 +2633,11 @@ def test_open_version_id_for_writing(self, mode, path, kwargs): # GH-926: the message states the accepted range. ( "s3://bucket/key", - S3FileSystem.MULTIPART_UPLOAD_MIN_PART_SIZE - 1, + S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE - 1, r"between 5 MiB \(5242880 bytes\) and 5 GiB \(5368709120 bytes\), inclusive", ), # GH-952: a part cannot be larger than the maximum part size. - ("s3://bucket/key", S3FileSystem.MULTIPART_UPLOAD_MAX_PART_SIZE + 1, "between"), + ("s3://bucket/key", S3Core.MULTIPART_UPLOAD_MAX_PART_SIZE + 1, "between"), ("s3://bucket", S3FileSystem.DEFAULT_BLOCK_SIZE, "does not contain a key"), ], ) @@ -2904,7 +2909,7 @@ def test_cat_file_range_lookup_parameters(self): @pytest.mark.parametrize( "block_size", - [S3FileSystem.MULTIPART_UPLOAD_MIN_PART_SIZE, S3FileSystem.MULTIPART_UPLOAD_MAX_PART_SIZE], + [S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE, S3Core.MULTIPART_UPLOAD_MAX_PART_SIZE], ) def test_open_block_size_limits_for_writing(self, block_size): fs = self._make_fs() @@ -3244,7 +3249,7 @@ def test_open_directory(self, cache_type): def test_finish_multipart_upload(self): fs = self._make_fs() - fs._complete_multipart_upload = mock.MagicMock() + fs.core.complete_multipart_upload = mock.MagicMock() futures = [] for part_number in (1, 2): future: Future[SimpleNamespace] = Future() @@ -3254,14 +3259,8 @@ def test_finish_multipart_upload(self): 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.core.complete_multipart_upload.assert_called_once_with( + S3Path("bucket", "key"), "uploadid", [f.result() for f in futures] ) fs._call.assert_not_called() @@ -3270,7 +3269,7 @@ def test_finish_multipart_upload(self): @pytest.mark.parametrize("error", [RuntimeError, KeyboardInterrupt]) def test_finish_multipart_upload_aborts_on_failure(self, error): fs = self._make_fs() - fs._complete_multipart_upload = mock.MagicMock() + fs.core.complete_multipart_upload = mock.MagicMock() future: Future[SimpleNamespace] = Future() future.set_exception(error("upload failed")) @@ -3278,7 +3277,7 @@ def test_finish_multipart_upload_aborts_on_failure(self, error): fs._finish_multipart_upload( bucket="bucket", key="key", upload_id="uploadid", futures=[future] ) - fs._complete_multipart_upload.assert_not_called() + fs.core.complete_multipart_upload.assert_not_called() fs._call.assert_called_once_with( fs._client.abort_multipart_upload, Bucket="bucket", @@ -3290,7 +3289,7 @@ def test_finish_multipart_upload_without_abort(self): # A caller that aborts the upload itself, as S3File.commit() does, # gets the original error with the parts and the upload left alone. fs = self._make_fs() - fs._complete_multipart_upload = mock.MagicMock() + fs.core.complete_multipart_upload = mock.MagicMock() failed: Future[SimpleNamespace] = Future() failed.set_exception(RuntimeError("upload failed")) pending: Future[SimpleNamespace] = Future() @@ -3308,7 +3307,7 @@ def test_finish_multipart_upload_without_abort(self): 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.core.complete_multipart_upload = mock.MagicMock() fs._call = mock.MagicMock(side_effect=RuntimeError("abort failed")) future: Future[SimpleNamespace] = Future() future.set_exception(RuntimeError("upload failed")) @@ -3324,7 +3323,7 @@ def test_finish_multipart_upload_waits_for_running_parts(self): # may be stored after the abort, so the abort waits for it. The # parts that have not started are cancelled. fs = self._make_fs() - fs._complete_multipart_upload = mock.MagicMock() + fs.core.complete_multipart_upload = mock.MagicMock() events = [] fs._call.side_effect = lambda *args, **kwargs: events.append("abort") failed: Future[SimpleNamespace] = Future() @@ -3367,7 +3366,7 @@ def test_finish_multipart_upload_does_not_wait_for_cancelled_parts(self): # acknowledge its cancellation, e.g., an event loop blocked by the # caller. fs = self._make_fs() - fs._complete_multipart_upload = mock.MagicMock() + fs.core.complete_multipart_upload = mock.MagicMock() failed: Future[SimpleNamespace] = Future() failed.set_exception(RuntimeError("upload failed")) never_started: Future[SimpleNamespace] = Future() @@ -3393,69 +3392,16 @@ def finish(): assert never_started.cancelled() fs._call.assert_called_once() - @pytest.mark.parametrize( - ("size", "block_size", "ranges"), - [ - # A single range. - (5 * 2**20, 5 * 2**20, [(0, 5 * 2**20)]), - # The size is an exact multiple of the block size. - (10 * 2**30, 5 * 2**30, [(0, 5 * 2**30), (5 * 2**30, 10 * 2**30)]), - # A last range of the minimum part size is kept. - ( - 5 * 2**30 + 5 * 2**20, - 5 * 2**30, - [(0, 5 * 2**30), (5 * 2**30, 5 * 2**30 + 5 * 2**20)], - ), - # GH-951: a last range shorter than the minimum part size is - # merged into the previous one, - (15 * 2**20 - 1, 5 * 2**20, [(0, 5 * 2**20), (5 * 2**20, 15 * 2**20 - 1)]), - # which is split in half if it exceeds the maximum part size. - ( - 5 * 2**30 + 2**20, - 5 * 2**30, - [(0, 5 * 2**29 + 2**19), (5 * 2**29 + 2**19, 5 * 2**30 + 2**20)], - ), - ], - ) - def test_get_copy_ranges(self, size, block_size, ranges): - assert self._make_fs()._get_copy_ranges(size, block_size) == ranges - - @pytest.mark.parametrize( - ("size", "num_ranges"), - [ - # The block size splits the object into the maximum number of parts. - (10_000 * 5 * 2**20, 10_000), - # GH-953: a larger object is split by a larger size instead of - # into more parts than the maximum, - (10_000 * 5 * 2**20 + 1, 9_999), - (50 * 2**30, 10_000), - # including the maximum object size. - (5 * 2**40, 10_000), - ], - ) - def test_get_copy_ranges_max_parts(self, size, num_ranges): - fs = self._make_fs() - ranges = fs._get_copy_ranges(size, 5 * 2**20) - - assert len(ranges) == num_ranges - assert ranges[0][0] == 0 - assert ranges[-1][1] == size - assert all(end == start for (_, end), (start, _) in pairwise(ranges)) - assert all( - fs.MULTIPART_UPLOAD_MIN_PART_SIZE <= end - start <= fs.MULTIPART_UPLOAD_MAX_PART_SIZE - for start, end in ranges - ) - @pytest.mark.parametrize("max_workers", [1, 4]) def test_copy_object_with_multipart_upload_part_sizes(self, max_workers): # GH-951: the parts are within the S3 part size limits whatever the # number of workers; a single worker used to copy the whole object # as one part larger than 5 GiB. fs = self._make_fs() - fs._create_multipart_upload = mock.MagicMock( + fs.core.create_multipart_upload = mock.MagicMock( return_value=SimpleNamespace(upload_id="uploadid") ) - fs._upload_part_copy = mock.MagicMock() + fs.core.upload_part_copy = mock.MagicMock() fs._finish_multipart_upload = mock.MagicMock() fs._call.return_value = {} @@ -3474,8 +3420,8 @@ def test_copy_object_with_multipart_upload_part_sizes(self, max_workers): ) parts = sorted( - (c.kwargs["part_number"], c.kwargs["copy_source_ranges"]) - for c in fs._upload_part_copy.call_args_list + (c.kwargs["part_number"], c.kwargs["range_"]) + for c in fs.core.upload_part_copy.call_args_list ) assert parts == [ (1, (0, 5 * 2**29 + 2**19)), @@ -3485,8 +3431,8 @@ def test_copy_object_with_multipart_upload_part_sizes(self, max_workers): @pytest.mark.parametrize( "block_size", [ - S3FileSystem.MULTIPART_UPLOAD_MIN_PART_SIZE - 1, - S3FileSystem.MULTIPART_UPLOAD_MAX_PART_SIZE + 1, + S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE - 1, + S3Core.MULTIPART_UPLOAD_MAX_PART_SIZE + 1, ], ) def test_copy_object_with_multipart_upload_invalid_block_size(self, block_size): @@ -5141,9 +5087,9 @@ def test_list_and_clear_multipart_uploads(self, fs): ) prefix_path = f"s3://{bucket}/{prefix}" key = f"{prefix}/file" - upload = fs._create_multipart_upload(bucket=bucket, key=key) + upload = fs.core.create_multipart_upload(S3Path(bucket, key)) # A sibling key that starts with the same characters as the prefix. - sibling = fs._create_multipart_upload(bucket=bucket, key=f"{prefix}2/file") + sibling = fs.core.create_multipart_upload(S3Path(bucket, f"{prefix}2/file")) try: uploads = fs.list_multipart_uploads(prefix_path) listed = next((u for u in uploads if u.upload_id == upload.upload_id), None) @@ -5319,6 +5265,16 @@ def _make_mock_fs(): fs = mock.MagicMock(spec=S3FileSystem) fs._client = S3_CLIENT fs.core = S3Core(S3_CLIENT) + # The requests of the core go to the mocked _call, and the + # multipart requests whose results the tests build are mocked. + fs.core.call = fs._call + for name in ( + "create_multipart_upload", + "upload_part", + "upload_part_copy", + "complete_multipart_upload", + ): + setattr(fs.core, name, mock.MagicMock()) fs._get_lookup_kwargs.side_effect = S3FileSystem._get_lookup_kwargs return fs @@ -5333,8 +5289,8 @@ def _make_write_file(data: bytes, autocommit: bool): file.key = "key.txt" file.s3_additional_kwargs = {} file.autocommit = autocommit - file.blocksize = S3FileSystem.MULTIPART_UPLOAD_MIN_PART_SIZE - file.fs.MULTIPART_UPLOAD_MAX_PARTS = S3FileSystem.MULTIPART_UPLOAD_MAX_PARTS + file.blocksize = S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE + file.fs.core.MULTIPART_UPLOAD_MAX_PARTS = S3Core.MULTIPART_UPLOAD_MAX_PARTS file.append_block = False file.multipart_upload = None file.multipart_upload_parts = [] @@ -5350,11 +5306,11 @@ def _make_multipart_write_file(data: bytes, autocommit: bool): # access is needed. file = TestS3File._make_write_file(data, autocommit=autocommit) file.blocksize = 4 - file.fs.MULTIPART_UPLOAD_MIN_PART_SIZE = 4 - file.fs.MULTIPART_UPLOAD_MAX_PART_SIZE = 8 + file.fs.core.MULTIPART_UPLOAD_MIN_PART_SIZE = 4 + file.fs.core.MULTIPART_UPLOAD_MAX_PART_SIZE = 8 file.multipart_upload = SimpleNamespace(upload_id="uploadid") file._executor = ThreadPoolExecutor(max_workers=1) - file.fs._upload_part.side_effect = lambda **kw: SimpleNamespace( + file.fs.core.upload_part.side_effect = lambda **kw: SimpleNamespace( etag=f'"e{kw["part_number"]}"', part_number=kw["part_number"] ) return file @@ -5365,9 +5321,9 @@ def _make_append_fs(existing: bytes): # size of 4 bytes so that the write and append paths can be exercised # with tiny data and no AWS access. fs = TestS3File._make_mock_fs() - fs.MULTIPART_UPLOAD_MIN_PART_SIZE = 4 - fs.MULTIPART_UPLOAD_MAX_PART_SIZE = 64 - fs.MULTIPART_UPLOAD_MAX_PARTS = S3FileSystem.MULTIPART_UPLOAD_MAX_PARTS + fs.core.MULTIPART_UPLOAD_MIN_PART_SIZE = 4 + fs.core.MULTIPART_UPLOAD_MAX_PART_SIZE = 64 + fs.core.MULTIPART_UPLOAD_MAX_PARTS = S3Core.MULTIPART_UPLOAD_MAX_PARTS fs.info.return_value = S3Object( init={"ContentLength": len(existing)}, type=S3ObjectType.S3_OBJECT_TYPE_FILE, @@ -5375,14 +5331,13 @@ def _make_append_fs(existing: bytes): key="key.txt", ) fs.cat_file.return_value = existing - fs._create_multipart_upload.return_value = SimpleNamespace(upload_id="uploadid") + fs.core.create_multipart_upload.return_value = SimpleNamespace(upload_id="uploadid") def part(**kw): return SimpleNamespace(etag=f'"e{kw["part_number"]}"', part_number=kw["part_number"]) - fs._upload_part.side_effect = part - fs._upload_part_copy.side_effect = part - fs._get_copy_ranges.side_effect = functools.partial(S3FileSystem._get_copy_ranges, fs) + fs.core.upload_part.side_effect = part + fs.core.upload_part_copy.side_effect = part return fs @staticmethod @@ -5390,15 +5345,15 @@ def _uploaded_object(fs, existing: bytes) -> bytes: # Rebuild the object S3 would store from the mocked upload calls. # A part copy without a range copies the whole existing object. if fs._put_object.called: - fs._create_multipart_upload.assert_not_called() + fs.core.create_multipart_upload.assert_not_called() return fs._put_object.call_args.kwargs["body"] fs._finish_multipart_upload.assert_called_once() parts = [] - for c in fs._upload_part_copy.call_args_list: - start, end = c.kwargs.get("copy_source_ranges", (0, len(existing))) + for c in fs.core.upload_part_copy.call_args_list: + start, end = c.kwargs.get("range_", (0, len(existing))) parts.append((c.kwargs["part_number"], existing[start:end])) parts += [ - (c.kwargs["part_number"], c.kwargs["body"]) for c in fs._upload_part.call_args_list + (c.kwargs["part_number"], c.kwargs["body"]) for c in fs.core.upload_part.call_args_list ] part_numbers = sorted(n for n, _ in parts) assert part_numbers == list(range(1, len(parts) + 1)) @@ -5426,8 +5381,8 @@ def test_append(self, existing, appended, multipart, part_copy): f.write(appended) assert self._uploaded_object(fs, existing) == existing + appended - assert fs._create_multipart_upload.called is multipart - assert fs._upload_part_copy.called is part_copy + assert fs.core.create_multipart_upload.called is multipart + assert fs.core.upload_part_copy.called is part_copy fs.touch.assert_not_called() @pytest.mark.parametrize("max_workers", [1, 4]) @@ -5446,8 +5401,8 @@ def test_append_part_copy_ranges(self, max_workers): assert self._uploaded_object(fs, existing) == existing + b"b" ranges = sorted( - (c.kwargs["part_number"], c.kwargs["copy_source_ranges"]) - for c in fs._upload_part_copy.call_args_list + (c.kwargs["part_number"], c.kwargs["range_"]) + for c in fs.core.upload_part_copy.call_args_list ) assert ranges == [(1, (0, 64)), (2, (64, 96)), (3, (96, 129))] @@ -5474,11 +5429,12 @@ def test_write_part_sizes(self, writes, block_size): assert self._uploaded_object(fs, b"") == b"".join(writes) parts = sorted( - (c.kwargs["part_number"], len(c.kwargs["body"])) for c in fs._upload_part.call_args_list + (c.kwargs["part_number"], len(c.kwargs["body"])) + for c in fs.core.upload_part.call_args_list ) sizes = [size for _, size in parts] - assert all(size >= fs.MULTIPART_UPLOAD_MIN_PART_SIZE for size in sizes[:-1]) - assert all(size <= fs.MULTIPART_UPLOAD_MAX_PART_SIZE for size in sizes) + assert all(size >= fs.core.MULTIPART_UPLOAD_MIN_PART_SIZE for size in sizes[:-1]) + assert all(size <= fs.core.MULTIPART_UPLOAD_MAX_PART_SIZE for size in sizes) @staticmethod def _write_and_close(f, writes: list[bytes]) -> None: @@ -5500,7 +5456,7 @@ def _write_and_close(f, writes: list[bytes]) -> None: @pytest.mark.parametrize("autocommit", [True, False]) def test_write_max_parts(self, existing, mode, writes, autocommit): fs = self._make_append_fs(existing) - fs.MULTIPART_UPLOAD_MAX_PARTS = 3 + fs.core.MULTIPART_UPLOAD_MAX_PARTS = 3 f = S3File(fs, "s3://bucket/key.txt", mode=mode, block_size=4, autocommit=autocommit) self._write_and_close(f, writes) @@ -5508,7 +5464,7 @@ def test_write_max_parts(self, existing, mode, writes, autocommit): f.commit() assert self._uploaded_object(fs, existing) == existing + b"".join(writes) - assert fs._upload_part_copy.call_count + fs._upload_part.call_count == 3 + assert fs.core.upload_part_copy.call_count + fs.core.upload_part.call_count == 3 @pytest.mark.parametrize( ("existing", "mode", "writes"), @@ -5525,7 +5481,7 @@ def test_write_max_parts(self, existing, mode, writes, autocommit): @pytest.mark.parametrize("autocommit", [True, False]) def test_write_exceeding_max_parts(self, existing, mode, writes, autocommit): fs = self._make_append_fs(existing) - fs.MULTIPART_UPLOAD_MAX_PARTS = 3 + fs.core.MULTIPART_UPLOAD_MAX_PARTS = 3 executor = mock.MagicMock(wraps=S3ThreadPoolExecutor(max_workers=1)) f = S3File( @@ -5548,7 +5504,7 @@ def test_write_exceeding_max_parts(self, existing, mode, writes, autocommit): assert [c.kwargs["part_number"] for c in executor.submit.call_args_list] == [1, 2, 3] executor.shutdown.assert_called() fs._call.assert_called_once_with( - "abort_multipart_upload", Bucket="bucket", Key="key.txt", UploadId="uploadid" + S3_CLIENT.abort_multipart_upload, Bucket="bucket", Key="key.txt", UploadId="uploadid" ) fs._finish_multipart_upload.assert_not_called() fs._put_object.assert_not_called() @@ -5560,7 +5516,7 @@ def test_write_exceeding_max_parts_abort_failure(self, caplog, autocommit): # the upload or completes it. GH-945: the upload is kept, so that # discard() retries the abort. fs = self._make_append_fs(b"") - fs.MULTIPART_UPLOAD_MAX_PARTS = 3 + fs.core.MULTIPART_UPLOAD_MAX_PARTS = 3 fs._call.side_effect = PermissionError("abort failed") executor = mock.MagicMock(wraps=S3ThreadPoolExecutor(max_workers=1)) @@ -5598,7 +5554,7 @@ def test_write_exceeding_max_parts_abort_interrupted(self): # still does not complete the upload; the executor is shut down. # GH-945: the upload is kept, so that discard() retries the abort. fs = self._make_append_fs(b"") - fs.MULTIPART_UPLOAD_MAX_PARTS = 3 + fs.core.MULTIPART_UPLOAD_MAX_PARTS = 3 fs._call.side_effect = KeyboardInterrupt executor = mock.MagicMock(wraps=S3ThreadPoolExecutor(max_workers=1)) @@ -5632,7 +5588,7 @@ def test_write_exceeding_max_parts_without_close(self): # The executor of the closed file is shut down, as fsspec does not # close it again when it is garbage collected. fs = self._make_append_fs(b"") - fs.MULTIPART_UPLOAD_MAX_PARTS = 3 + fs.core.MULTIPART_UPLOAD_MAX_PARTS = 3 executor = mock.MagicMock(wraps=S3ThreadPoolExecutor(max_workers=1)) f = S3File(fs, "s3://bucket/key.txt", mode="wb", block_size=4, executor=executor) @@ -5661,11 +5617,11 @@ def test_multipart_write_request_parameters(self): ) as f: f.write(b"x" * 8) - fs._create_multipart_upload.assert_called_once_with( - bucket="bucket", key="key.txt", **kwargs + fs.core.create_multipart_upload.assert_called_once_with( + S3Path("bucket", "key.txt"), **kwargs ) - assert fs._upload_part.call_count == 2 - for c in fs._upload_part.call_args_list: + assert fs.core.upload_part.call_count == 2 + for c in fs.core.upload_part.call_args_list: assert {k: v for k, v in c.kwargs.items() if k[0].isupper()} == { "RequestPayer": "requester", "SSECustomerAlgorithm": "AES256", @@ -5703,7 +5659,7 @@ def test_append_discard(self): f.discard() fs._call.assert_called_once_with( - "abort_multipart_upload", + S3_CLIENT.abort_multipart_upload, Bucket="bucket", Key="key.txt", UploadId="uploadid", @@ -5809,7 +5765,7 @@ def test_discard(self, multipart): if multipart: file.fs._call.assert_called_once() - assert file.fs._call.call_args.args[0] == "abort_multipart_upload" + assert file.fs._call.call_args.args[0] == S3_CLIENT.abort_multipart_upload else: file.fs._call.assert_not_called() file.fs._put_object.assert_not_called() @@ -5864,7 +5820,7 @@ def test_commit_failure_and_discard(self, caplog, error, abort_fails): file.fs._finish_multipart_upload.side_effect = functools.partial( S3FileSystem._finish_multipart_upload, file.fs ) - file.fs._complete_multipart_upload.side_effect = error("complete failed") + file.fs.core.complete_multipart_upload.side_effect = error("complete failed") if abort_fails: file.fs._call.side_effect = [PermissionError("abort failed"), None] @@ -5879,7 +5835,12 @@ def test_commit_failure_and_discard(self, caplog, error, abort_fails): file.discard() assert file.fs._call.call_args_list == [ - mock.call("abort_multipart_upload", Bucket="bucket", Key="key.txt", UploadId="uploadid") + mock.call( + S3_CLIENT.abort_multipart_upload, + Bucket="bucket", + Key="key.txt", + UploadId="uploadid", + ) ] * (2 if abort_fails else 1) assert file.multipart_upload is None assert file.multipart_upload_parts == [] @@ -5893,7 +5854,7 @@ def test_commit_failure_and_interrupted_abort(self): file.fs._finish_multipart_upload.side_effect = functools.partial( S3FileSystem._finish_multipart_upload, file.fs ) - file.fs._complete_multipart_upload.side_effect = RuntimeError("complete failed") + file.fs.core.complete_multipart_upload.side_effect = RuntimeError("complete failed") file.fs._call.side_effect = [KeyboardInterrupt, None] with pytest.raises(KeyboardInterrupt): @@ -5916,7 +5877,7 @@ def test_discard_on_event_loop_thread(self): async def rollback(): executor = S3AioExecutor(loop=asyncio.get_running_loop()) - parts.extend(executor.submit(file.fs._upload_part) for _ in range(2)) + parts.extend(executor.submit(file.fs.core.upload_part) for _ in range(2)) file.multipart_upload_parts = list(parts) file.discard() @@ -5926,9 +5887,9 @@ async def rollback(): assert not thread.is_alive() assert all(part.cancelled() for part in parts) - file.fs._upload_part.assert_not_called() + file.fs.core.upload_part.assert_not_called() file.fs._call.assert_called_once() - assert file.fs._call.call_args.args[0] == "abort_multipart_upload" + assert file.fs._call.call_args.args[0] == S3_CLIENT.abort_multipart_upload @pytest.mark.parametrize("autocommit", [True, False]) def test_upload_chunk_multipart(self, autocommit): diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index 67925f3d..107d53ac 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -159,13 +159,13 @@ async def test_copy_object_with_multipart_upload_part_sizes(self, max_workers): connection=mock.MagicMock(), max_workers=max_workers, skip_instance_cache=True ) sync_fs = fs._sync_fs - sync_fs._create_multipart_upload = mock.MagicMock( + sync_fs.core.create_multipart_upload = mock.MagicMock( return_value=SimpleNamespace(upload_id="uploadid") ) - sync_fs._upload_part_copy = mock.MagicMock( + sync_fs.core.upload_part_copy = mock.MagicMock( side_effect=lambda **kw: SimpleNamespace(etag='"e"', part_number=kw["part_number"]) ) - sync_fs._complete_multipart_upload = mock.MagicMock() + sync_fs.core.complete_multipart_upload = mock.MagicMock() sync_fs._call = sync_fs._core.call = mock.MagicMock(return_value={}) await fs._copy_object_with_multipart_upload( @@ -182,8 +182,8 @@ async def test_copy_object_with_multipart_upload_part_sizes(self, max_workers): ) parts = sorted( - (c.kwargs["part_number"], c.kwargs["copy_source_ranges"]) - for c in sync_fs._upload_part_copy.call_args_list + (c.kwargs["part_number"], c.kwargs["range_"]) + for c in sync_fs.core.upload_part_copy.call_args_list ) assert parts == [ (1, (0, 5 * 2**29 + 2**19)), @@ -226,7 +226,7 @@ async def test_copy_object_with_multipart_upload_small_head_object_size(self, si return_value={"ContentLength": size, "VersionId": "v1"} ) sync_fs._copy_object = mock.MagicMock() - sync_fs._create_multipart_upload = mock.MagicMock() + sync_fs.core.create_multipart_upload = mock.MagicMock() await fs._copy_object_with_multipart_upload( bucket1="bucket", @@ -247,7 +247,7 @@ async def test_copy_object_with_multipart_upload_small_head_object_size(self, si ContentType="text/csv", RequestPayer="requester", ) - sync_fs._create_multipart_upload.assert_not_called() + sync_fs.core.create_multipart_upload.assert_not_called() # Only HeadObject; the tags are not read for the multipart upload. assert sync_fs._call.call_count == 1 @@ -297,7 +297,7 @@ async def test_cp_file_failed_multipart_copy_invalidates_cache(self): fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) fs._info = mock.AsyncMock( return_value=S3Object( - init={"ContentLength": S3FileSystem.MULTIPART_UPLOAD_MAX_PART_SIZE + 1}, + init={"ContentLength": S3Core.MULTIPART_UPLOAD_MAX_PART_SIZE + 1}, type=S3ObjectType.S3_OBJECT_TYPE_FILE, bucket="bucket", key="src", @@ -316,7 +316,7 @@ async def test_copy_object_with_multipart_upload_waits_for_running_parts(self): # one fails, and no part starts after the failure. fs = AioS3FileSystem(connection=mock.MagicMock(), max_workers=2, skip_instance_cache=True) sync_fs = fs._sync_fs - sync_fs._create_multipart_upload = mock.MagicMock( + sync_fs.core.create_multipart_upload = mock.MagicMock( return_value=SimpleNamespace(upload_id="uploadid") ) events = [] @@ -333,8 +333,8 @@ def upload_part_copy(**kw): failed.set() raise OSError("part failed") - sync_fs._upload_part_copy = mock.MagicMock(side_effect=upload_part_copy) - sync_fs._complete_multipart_upload = mock.MagicMock() + sync_fs.core.upload_part_copy = mock.MagicMock(side_effect=upload_part_copy) + sync_fs.core.complete_multipart_upload = mock.MagicMock() # The HeadObject of the source, for its version. sync_fs._call = sync_fs._core.call = mock.MagicMock(return_value={}) sync_fs._abort_multipart_upload = mock.MagicMock( @@ -345,10 +345,10 @@ def upload_part_copy(**kw): await fs._copy_object_with_multipart_upload( bucket1="bucket", key1="src", - size1=3 * S3FileSystem.MULTIPART_UPLOAD_MIN_PART_SIZE, + size1=3 * S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE, bucket2="bucket", key2="dst", - block_size=S3FileSystem.MULTIPART_UPLOAD_MIN_PART_SIZE, + block_size=S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE, MetadataDirective="REPLACE", TaggingDirective="REPLACE", AnnotationDirective="EXCLUDE", @@ -357,7 +357,7 @@ def upload_part_copy(**kw): # Part 3 waits for a worker and is not started after the failure. assert sorted(events[:2]) == ["start 1", "start 2"] assert events[2:] == ["end 1", "abort"] - sync_fs._complete_multipart_upload.assert_not_called() + sync_fs.core.complete_multipart_upload.assert_not_called() @pytest.mark.parametrize("cancellations", [1, 2]) @pytest.mark.asyncio @@ -368,7 +368,7 @@ async def test_copy_object_with_multipart_upload_cancelled(self, cancellations): # cleanup. fs = AioS3FileSystem(connection=mock.MagicMock(), max_workers=2, skip_instance_cache=True) sync_fs = fs._sync_fs - sync_fs._create_multipart_upload = mock.MagicMock( + sync_fs.core.create_multipart_upload = mock.MagicMock( return_value=SimpleNamespace(upload_id="uploadid") ) events = [] @@ -392,8 +392,8 @@ def abort_multipart_upload(*args): events.append("abort") aborted.set() - sync_fs._upload_part_copy = mock.MagicMock(side_effect=upload_part_copy) - sync_fs._complete_multipart_upload = mock.MagicMock() + sync_fs.core.upload_part_copy = mock.MagicMock(side_effect=upload_part_copy) + sync_fs.core.complete_multipart_upload = mock.MagicMock() # The HeadObject of the source, for its version. sync_fs._call = sync_fs._core.call = mock.MagicMock(return_value={}) sync_fs._abort_multipart_upload = mock.MagicMock(side_effect=abort_multipart_upload) @@ -402,10 +402,10 @@ def abort_multipart_upload(*args): fs._copy_object_with_multipart_upload( bucket1="bucket", key1="src", - size1=3 * S3FileSystem.MULTIPART_UPLOAD_MIN_PART_SIZE, + size1=3 * S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE, bucket2="bucket", key2="dst", - block_size=S3FileSystem.MULTIPART_UPLOAD_MIN_PART_SIZE, + block_size=S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE, MetadataDirective="REPLACE", TaggingDirective="REPLACE", AnnotationDirective="EXCLUDE", @@ -438,7 +438,7 @@ def abort_multipart_upload(*args): assert sorted(events[:2]) == ["start 1", "start 2"] assert sorted(events[2:4]) == ["end 1", "end 2"] assert events[4:] == ["abort"] - sync_fs._complete_multipart_upload.assert_not_called() + sync_fs.core.complete_multipart_upload.assert_not_called() @pytest.mark.parametrize("completion_fails", [False, True]) @pytest.mark.asyncio @@ -447,14 +447,14 @@ async def test_copy_object_with_multipart_upload_cancelled_completion(self, comp # aborts the upload only if it failed, and is re-raised. fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) sync_fs = fs._sync_fs - sync_fs._create_multipart_upload = mock.MagicMock( + sync_fs.core.create_multipart_upload = mock.MagicMock( return_value=SimpleNamespace(upload_id="uploadid") ) events = [] started = threading.Event() release = threading.Event() - def complete_multipart_upload(**kw): + def complete_multipart_upload(*args, **kw): started.set() # The finally blocks of the test always release it. release.wait() @@ -463,10 +463,12 @@ def complete_multipart_upload(**kw): raise OSError("completion failed") return SimpleNamespace() - sync_fs._upload_part_copy = mock.MagicMock( + sync_fs.core.upload_part_copy = mock.MagicMock( side_effect=lambda **kw: SimpleNamespace(etag='"e"', part_number=kw["part_number"]) ) - sync_fs._complete_multipart_upload = mock.MagicMock(side_effect=complete_multipart_upload) + sync_fs.core.complete_multipart_upload = mock.MagicMock( + side_effect=complete_multipart_upload + ) # The HeadObject of the source, for its version. sync_fs._call = sync_fs._core.call = mock.MagicMock(return_value={}) sync_fs._abort_multipart_upload = mock.MagicMock( @@ -477,10 +479,10 @@ def complete_multipart_upload(**kw): fs._copy_object_with_multipart_upload( bucket1="bucket", key1="src", - size1=2 * S3FileSystem.MULTIPART_UPLOAD_MIN_PART_SIZE, + size1=2 * S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE, bucket2="bucket", key2="dst", - block_size=S3FileSystem.MULTIPART_UPLOAD_MIN_PART_SIZE, + block_size=S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE, MetadataDirective="REPLACE", TaggingDirective="REPLACE", AnnotationDirective="EXCLUDE", @@ -507,8 +509,8 @@ def complete_multipart_upload(**kw): @pytest.mark.parametrize( "block_size", [ - S3FileSystem.MULTIPART_UPLOAD_MIN_PART_SIZE - 1, - S3FileSystem.MULTIPART_UPLOAD_MAX_PART_SIZE + 1, + S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE - 1, + S3Core.MULTIPART_UPLOAD_MAX_PART_SIZE + 1, ], ) @pytest.mark.asyncio @@ -681,7 +683,7 @@ def test_transaction_pipe_put_file_exceeding_max_parts(self, tmp_path, kwargs): # reject data that does not fit in the maximum number of parts before # opening the file. fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) - fs._sync_fs.MULTIPART_UPLOAD_MAX_PARTS = 3 + fs._sync_fs.core.MULTIPART_UPLOAD_MAX_PARTS = 3 fs._sync_fs.default_block_size = 4 fs._sync_fs._call = fs._sync_fs._core.call = mock.MagicMock() fs.open = mock.MagicMock() @@ -857,7 +859,7 @@ def test_put_file_in_transaction_open_parameters(self, tmp_path, mode, open_mode "s3://bucket/key", Callback(), mode, - block_size=S3FileSystem.MULTIPART_UPLOAD_MIN_PART_SIZE, + block_size=S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE, max_workers=2, StorageClass="STANDARD_IA", ) @@ -865,7 +867,7 @@ def test_put_file_in_transaction_open_parameters(self, tmp_path, mode, open_mode fs.open.assert_called_once_with( "s3://bucket/key", open_mode, - block_size=S3FileSystem.MULTIPART_UPLOAD_MIN_PART_SIZE, + block_size=S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE, max_workers=2, s3_additional_kwargs={"StorageClass": "STANDARD_IA", "ContentType": "text/csv"}, ) @@ -1008,7 +1010,7 @@ async def test_cp_file_multipart_parameters(self, size): ) sync_fs = fs._sync_fs sync_fs._copy_object = mock.MagicMock() - sync_fs._create_multipart_upload = mock.MagicMock( + sync_fs.core.create_multipart_upload = mock.MagicMock( return_value=SimpleNamespace(upload_id="uploadid") ) running = [] @@ -1021,8 +1023,8 @@ def upload_part_copy(**kw): running.remove(kw["part_number"]) return SimpleNamespace(etag='"e"', part_number=kw["part_number"]) - sync_fs._upload_part_copy = mock.MagicMock(side_effect=upload_part_copy) - sync_fs._complete_multipart_upload = mock.MagicMock() + sync_fs.core.upload_part_copy = mock.MagicMock(side_effect=upload_part_copy) + sync_fs.core.complete_multipart_upload = mock.MagicMock() # The HeadObject of the source, for its version. sync_fs._call = sync_fs._core.call = mock.MagicMock(return_value={}) directives = { @@ -1034,7 +1036,7 @@ def upload_part_copy(**kw): await fs._cp_file( "s3://bucket/src", "s3://bucket/dst", - block_size=S3FileSystem.MULTIPART_UPLOAD_MAX_PART_SIZE // 2, + block_size=S3Core.MULTIPART_UPLOAD_MAX_PART_SIZE // 2, max_workers=1, RequestPayer="requester", ContentType="text/csv", @@ -1043,7 +1045,7 @@ def upload_part_copy(**kw): **directives, ) - if size <= S3FileSystem.MULTIPART_UPLOAD_MAX_PART_SIZE: + if size <= S3Core.MULTIPART_UPLOAD_MAX_PART_SIZE: sync_fs._copy_object.assert_called_once_with( bucket1="bucket", key1="src", @@ -1055,20 +1057,21 @@ def upload_part_copy(**kw): **directives, ) else: - sync_fs._create_multipart_upload.assert_called_once_with( - bucket="bucket", key="dst", RequestPayer="requester", ContentType="text/csv" + sync_fs.core.create_multipart_upload.assert_called_once_with( + S3Path("bucket", "dst"), RequestPayer="requester", ContentType="text/csv" ) # The part copies receive the parameters that they accept, and # max_workers limits how many run at once. # Two parts, the second with the 1-byte tail. - assert sync_fs._upload_part_copy.call_count == 2 + assert sync_fs.core.upload_part_copy.call_count == 2 assert all( c.kwargs["RequestPayer"] == "requester" and "ContentType" not in c.kwargs - for c in sync_fs._upload_part_copy.call_args_list + for c in sync_fs.core.upload_part_copy.call_args_list ) assert max(concurrency) == 1 assert ( - sync_fs._complete_multipart_upload.call_args.kwargs["RequestPayer"] == "requester" + sync_fs.core.complete_multipart_upload.call_args.kwargs["RequestPayer"] + == "requester" ) def test_internal_file_system_not_cached(self): @@ -1886,7 +1889,7 @@ async def test_open_parallel_requests(self, asynchronous): fs = AioS3FileSystem( connection=mock.MagicMock(), asynchronous=asynchronous, skip_instance_cache=True ) - block_size = S3FileSystem.MULTIPART_UPLOAD_MIN_PART_SIZE + block_size = S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE size = block_size * 4 condition = threading.Condition() state = {"active": 0, "peak": 0} @@ -1909,13 +1912,13 @@ def call(**kwargs): return mock.MagicMock(side_effect=call) sync_fs = fs._sync_fs - sync_fs._create_multipart_upload = mock.MagicMock( + sync_fs.core.create_multipart_upload = mock.MagicMock( return_value=SimpleNamespace(upload_id="uploadid") ) - sync_fs._upload_part = track( + sync_fs.core.upload_part = track( lambda **kw: S3MultipartUploadPart(kw["part_number"], {"ETag": '"e"'}) ) - sync_fs._complete_multipart_upload = mock.MagicMock() + sync_fs.core.complete_multipart_upload = mock.MagicMock() sync_fs._get_object = track( lambda **kw: (kw["ranges"][0], b"a" * (kw["ranges"][1] - kw["ranges"][0])) ) @@ -1934,7 +1937,7 @@ def write(): f.write(b"a" * size) await asyncio.to_thread(write) - assert sync_fs._upload_part.call_count == 4 + assert sync_fs.core.upload_part.call_count == 4 assert state["peak"] == 2 def read(): diff --git a/tests/pyathena/filesystem/test_s3_core.py b/tests/pyathena/filesystem/test_s3_core.py index 4b76d461..c74fed8e 100644 --- a/tests/pyathena/filesystem/test_s3_core.py +++ b/tests/pyathena/filesystem/test_s3_core.py @@ -6,6 +6,7 @@ # SPDX-License-Identifier: MIT from datetime import UTC, datetime +from itertools import pairwise import boto3 import botocore.exceptions @@ -24,6 +25,7 @@ S3ListObjectVersionsPage, S3ObjectSummary, ) +from pyathena.filesystem.s3_object import S3MultipartUploadPart from pyathena.filesystem.s3_path import S3Path from pyathena.util import RetryConfig @@ -413,6 +415,196 @@ def test_delete_objects_translates_errors(self): with stubber, pytest.raises(PermissionError): core.delete_objects(S3DeleteBatch("bucket", (S3Path("bucket", "a"),))) + def test_create_multipart_upload(self): + core, stubber = _make_core() + stubber.add_response( + "create_multipart_upload", + {"Bucket": "bucket", "Key": "key", "UploadId": "u"}, + {"Bucket": "bucket", "Key": "key", "ContentType": "text/csv"}, + ) + with stubber: + # The key of the path takes precedence over a parameter. + upload = core.create_multipart_upload( + S3Path("bucket", "key"), ContentType="text/csv", Key="other" + ) + assert upload.upload_id == "u" + + def test_upload_part(self): + core, stubber = _make_core() + stubber.add_response( + "upload_part", + {"ETag": '"e1"'}, + { + "Bucket": "bucket", + "Key": "key", + "UploadId": "u", + "PartNumber": 1, + "Body": b"data", + "SSECustomerAlgorithm": "AES256", + }, + ) + with stubber: + part = core.upload_part( + S3Path("bucket", "key"), + "u", + 1, + b"data", + SSECustomerAlgorithm="AES256", + PartNumber=99, + ) + assert (part.part_number, part.etag) == (1, '"e1"') + + def test_upload_part_copy(self): + core, stubber = _make_core() + stubber.add_response( + "upload_part_copy", + {"CopyPartResult": {"ETag": '"p2"'}}, + { + "Bucket": "bucket", + "Key": "dst", + "CopySource": {"Bucket": "src-bucket", "Key": "src", "VersionId": "v1"}, + "UploadId": "u", + "PartNumber": 2, + "CopySourceRange": "bytes=10-19", + "CopySourceIfMatch": '"src"', + }, + ) + # Without a range, the whole source is copied. + stubber.add_response( + "upload_part_copy", + {"CopyPartResult": {"ETag": '"p1"'}}, + { + "Bucket": "bucket", + "Key": "dst", + "CopySource": {"Bucket": "bucket", "Key": "dst"}, + "UploadId": "u", + "PartNumber": 1, + }, + ) + with stubber: + part = core.upload_part_copy( + S3Path("bucket", "dst"), + "u", + 2, + S3Path("src-bucket", "src", "v1"), + range_=(10, 20), + CopySourceIfMatch='"src"', + ) + whole = core.upload_part_copy(S3Path("bucket", "dst"), "u", 1, S3Path("bucket", "dst")) + stubber.assert_no_pending_responses() + assert (part.part_number, part.etag) == (2, '"p2"') + assert (whole.part_number, whole.etag) == (1, '"p1"') + + def test_complete_multipart_upload(self): + core, stubber = _make_core() + stubber.add_response( + "complete_multipart_upload", + {"ETag": '"dst"', "VersionId": "v-dst"}, + { + "Bucket": "bucket", + "Key": "key", + "UploadId": "u", + "MultipartUpload": { + "Parts": [{"ETag": '"e1"', "PartNumber": 1}, {"ETag": '"e2"', "PartNumber": 2}] + }, + "RequestPayer": "requester", + }, + ) + parts = [S3MultipartUploadPart(n, {"ETag": f'"e{n}"'}) for n in (1, 2)] + with stubber: + completed = core.complete_multipart_upload( + S3Path("bucket", "key"), "u", parts, RequestPayer="requester" + ) + assert (completed.etag, completed.version_id) == ('"dst"', "v-dst") + + def test_abort_multipart_upload(self): + core, stubber = _make_core() + stubber.add_response( + "abort_multipart_upload", + {}, + {"Bucket": "bucket", "Key": "key", "UploadId": "u", "RequestPayer": "requester"}, + ) + # Unlike the callers that log a failed abort, the core raises it. + stubber.add_client_error( + "abort_multipart_upload", service_error_code="NoSuchUpload", http_status_code=404 + ) + with stubber: + core.abort_multipart_upload(S3Path("bucket", "key"), "u", RequestPayer="requester") + with pytest.raises(FileNotFoundError): + core.abort_multipart_upload(S3Path("bucket", "key"), "u") + + @pytest.mark.parametrize( + ("method", "args"), + [ + ("create_multipart_upload", (S3Path("bucket"),)), + ("upload_part", (S3Path("bucket"), "u", 1, b"")), + ("upload_part_copy", (S3Path("bucket"), "u", 1, S3Path("bucket", "src"))), + ("upload_part_copy", (S3Path("bucket", "dst"), "u", 1, S3Path("bucket"))), + ("complete_multipart_upload", (S3Path("bucket"), "u", [])), + ("abort_multipart_upload", (S3Path("bucket"), "u")), + ], + ) + def test_multipart_upload_requires_keys(self, method, args): + core, _ = _make_core() + with pytest.raises(ValueError, match="has no key"): + getattr(core, method)(*args) + + @pytest.mark.parametrize( + ("size", "block_size", "ranges"), + [ + # A single range. + (5 * 2**20, 5 * 2**20, [(0, 5 * 2**20)]), + # The size is an exact multiple of the block size. + (10 * 2**30, 5 * 2**30, [(0, 5 * 2**30), (5 * 2**30, 10 * 2**30)]), + # A last range of the minimum part size is kept. + ( + 5 * 2**30 + 5 * 2**20, + 5 * 2**30, + [(0, 5 * 2**30), (5 * 2**30, 5 * 2**30 + 5 * 2**20)], + ), + # GH-951: a last range shorter than the minimum part size is + # merged into the previous one, + (15 * 2**20 - 1, 5 * 2**20, [(0, 5 * 2**20), (5 * 2**20, 15 * 2**20 - 1)]), + # which is split in half if it exceeds the maximum part size. + ( + 5 * 2**30 + 2**20, + 5 * 2**30, + [(0, 5 * 2**29 + 2**19), (5 * 2**29 + 2**19, 5 * 2**30 + 2**20)], + ), + ], + ) + def test_part_ranges(self, size, block_size, ranges): + core, _ = _make_core() + assert core.part_ranges(size, block_size) == ranges + + @pytest.mark.parametrize( + ("size", "num_ranges"), + [ + # The block size splits the object into the maximum number of parts. + (10_000 * 5 * 2**20, 10_000), + # GH-953: a larger object is split by a larger size instead of + # into more parts than the maximum, + (10_000 * 5 * 2**20 + 1, 9_999), + (50 * 2**30, 10_000), + # including the maximum object size. + (5 * 2**40, 10_000), + ], + ) + def test_part_ranges_max_parts(self, size, num_ranges): + core, _ = _make_core() + ranges = core.part_ranges(size, 5 * 2**20) + + assert len(ranges) == num_ranges + assert ranges[0][0] == 0 + assert ranges[-1][1] == size + assert all(end == start for (_, end), (start, _) in pairwise(ranges)) + assert all( + core.MULTIPART_UPLOAD_MIN_PART_SIZE + <= end - start + <= core.MULTIPART_UPLOAD_MAX_PART_SIZE + for start, end in ranges + ) + class TestS3DeleteBatch: def test_from_paths(self): diff --git a/tests/pyathena/util.py b/tests/pyathena/util.py index cdd77799..3c7fe8d7 100644 --- a/tests/pyathena/util.py +++ b/tests/pyathena/util.py @@ -21,6 +21,7 @@ from pyathena.filesystem.s3 import S3FileSystem from pyathena.filesystem.s3_async import AioS3FileSystem +from pyathena.filesystem.s3_core import S3Core from pyathena.glue import GlueMetadataClient from pyathena.model import AthenaCalculationExecutionStatus, AthenaQueryExecution @@ -249,8 +250,8 @@ def interrupting_wait(futures, timeout=None): # A source object larger than a single CopyObject request allows, copied # from bucket/src to bucket/dst by a multipart copy of two parts with the # maximum block size (GH-973). -MULTIPART_COPY_BLOCK_SIZE = S3FileSystem.MULTIPART_UPLOAD_MAX_PART_SIZE -MULTIPART_COPY_SIZE = MULTIPART_COPY_BLOCK_SIZE + S3FileSystem.MULTIPART_UPLOAD_MIN_PART_SIZE +MULTIPART_COPY_BLOCK_SIZE = S3Core.MULTIPART_UPLOAD_MAX_PART_SIZE +MULTIPART_COPY_SIZE = MULTIPART_COPY_BLOCK_SIZE + S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE MULTIPART_COPY_EXPIRES = datetime(2030, 1, 1, tzinfo=UTC) MULTIPART_COPY_KWARGS = { # Ignored with the default COPY directives, as CopyObject ignores them. From 3825953dfc51b842229f1eb563931176f3797880 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 4 Oct 2026 13:51:16 +0900 Subject: [PATCH 2/4] Reject a versioned destination in create_multipart_upload(), and fix docs A write replaces the object at the key, so the core rejects a path with a version ID instead of ignoring it, as the filesystem already does. Also document that abort_multipart_upload() raises FileNotFoundError for an upload that does not exist (measured), that part_ranges() sends no request, and build the path of an S3File multipart upload once. Co-Authored-By: Claude Opus 5.5 --- docs/filesystem.md | 4 ++-- pyathena/filesystem/s3.py | 12 +++++++----- pyathena/filesystem/s3_core.py | 9 +++++++-- tests/pyathena/filesystem/test_s3_core.py | 6 ++++++ 4 files changed, 22 insertions(+), 9 deletions(-) diff --git a/docs/filesystem.md b/docs/filesystem.md index 61904928..0fd47525 100644 --- a/docs/filesystem.md +++ b/docs/filesystem.md @@ -325,8 +325,8 @@ for batch in S3DeleteBatch.from_paths(paths): `create_multipart_upload()`, `upload_part()`, `upload_part_copy()`, `complete_multipart_upload()` and `abort_multipart_upload()` send the requests of a -multipart upload. `part_ranges()` splits an object into the byte ranges of the parts -that copy it, by the part limits `MULTIPART_UPLOAD_MIN_PART_SIZE` (5 MiB), +multipart upload. `part_ranges()` sends no request: it splits an object into the byte +ranges of the parts that copy it, by the part limits `MULTIPART_UPLOAD_MIN_PART_SIZE` (5 MiB), `MULTIPART_UPLOAD_MAX_PART_SIZE` (5 GiB) and `MULTIPART_UPLOAD_MAX_PARTS` (10,000) of `S3Core`. diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index 934dfbf5..846f56cb 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -3584,8 +3584,9 @@ def _initiate_upload(self) -> None: # a multipart upload, whatever the block size. return + path = S3Path(self.bucket, self.key) self.multipart_upload = self.fs.core.create_multipart_upload( - S3Path(self.bucket, self.key), **self._get_request_kwargs("create_multipart_upload") + path, **self._get_request_kwargs("create_multipart_upload") ) if self.append_block: if self.tell() > self.fs.core.MULTIPART_UPLOAD_MAX_PART_SIZE: @@ -3603,10 +3604,11 @@ def _initiate_upload(self) -> None: self.multipart_upload_parts.append( self._executor.submit( self.fs.core.upload_part_copy, - path=S3Path(self.bucket, self.key), + path=path, upload_id=cast(str, self.multipart_upload.upload_id), part_number=i + 1, - source=S3Path(self.bucket, self.key), + # The existing object is copied into the upload. + source=path, range_=range_, **self._get_request_kwargs("upload_part_copy"), ) @@ -3615,10 +3617,10 @@ def _initiate_upload(self) -> None: self.multipart_upload_parts.append( self._executor.submit( self.fs.core.upload_part_copy, - path=S3Path(self.bucket, self.key), + path=path, upload_id=cast(str, self.multipart_upload.upload_id), part_number=1, - source=S3Path(self.bucket, self.key), + source=path, **self._get_request_kwargs("upload_part_copy"), ) ) diff --git a/pyathena/filesystem/s3_core.py b/pyathena/filesystem/s3_core.py index bde5cd42..4803c9a9 100644 --- a/pyathena/filesystem/s3_core.py +++ b/pyathena/filesystem/s3_core.py @@ -597,7 +597,7 @@ def create_multipart_upload(self, path: S3Path, **params) -> S3MultipartUpload: """Start a multipart upload to an object with CreateMultipartUpload. Args: - path: The path of the object to write. + path: The path of the object to write, without a version ID. **params: Additional request parameters. The bucket and key of the path take precedence over parameters of the same name. @@ -605,10 +605,13 @@ def create_multipart_upload(self, path: S3Path, **params) -> S3MultipartUpload: The multipart upload. Raises: - ValueError: If the path has no key. + ValueError: If the path has no key, or has a version ID, which a + write cannot replace. """ if not path.key: raise ValueError(f"The path has no key: {path.uri}.") + if path.version_id: + raise ValueError(f"Cannot write to a version: {path.uri}.") request: dict[str, Any] = {"Bucket": path.bucket, "Key": path.key} _logger.debug(f"Create multipart upload to {path.uri}.") response = self.call(self._client.create_multipart_upload, **{**params, **request}) @@ -741,6 +744,8 @@ def abort_multipart_upload(self, path: S3Path, upload_id: str, **params) -> None Raises: ValueError: If the path has no key. + FileNotFoundError: If the upload does not exist, for example + because it was completed or aborted. """ if not path.key: raise ValueError(f"The path has no key: {path.uri}.") diff --git a/tests/pyathena/filesystem/test_s3_core.py b/tests/pyathena/filesystem/test_s3_core.py index c74fed8e..b635783a 100644 --- a/tests/pyathena/filesystem/test_s3_core.py +++ b/tests/pyathena/filesystem/test_s3_core.py @@ -549,6 +549,12 @@ def test_multipart_upload_requires_keys(self, method, args): with pytest.raises(ValueError, match="has no key"): getattr(core, method)(*args) + def test_create_multipart_upload_rejects_versions(self): + # A write replaces the object at the key, not the named version. + core, _ = _make_core() + with pytest.raises(ValueError, match="Cannot write to a version"): + core.create_multipart_upload(S3Path("bucket", "key", "v1")) + @pytest.mark.parametrize( ("size", "block_size", "ranges"), [ From 6e2fb2e020f4519f10fb30a9beaffc76f4646654 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 4 Oct 2026 14:07:33 +0900 Subject: [PATCH 3/4] Inject the abort failure through the core in its test, and fix docs The test of a failed abort replaced only fs._call, but the abort now goes through fs.core.call, so the abort succeeded and the test passed without the failure. Set the failure on the shared mock and assert the abort and its log. Also document that a missing multipart upload raises FileNotFoundError, and that range_=None leaves a CopySourceRange of the params in place. Co-Authored-By: Claude Opus 5.5 --- docs/filesystem.md | 5 +++-- pyathena/filesystem/s3_core.py | 8 +++++--- tests/pyathena/filesystem/test_s3.py | 9 +++++++-- 3 files changed, 15 insertions(+), 7 deletions(-) diff --git a/docs/filesystem.md b/docs/filesystem.md index 0fd47525..0202a8a7 100644 --- a/docs/filesystem.md +++ b/docs/filesystem.md @@ -282,8 +282,9 @@ directories below the bucket level) and is always a no-op. `S3FileSystem.core` is an `S3Core`, the typed operations that the filesystem sends its listing, lookup, delete and multipart upload requests with. It can also be built on a boto3 S3 client. Each operation sends one request (one per page for the -iterators) with the retry policy, raises `FileNotFoundError` for a missing bucket, or -for a missing object or version that it reads, and caches nothing. Requests sent +iterators) with the retry policy, raises `FileNotFoundError` for a missing bucket or +multipart upload, or for a missing object or version that it reads, and caches +nothing. Requests sent through `fs.core` do not invalidate the filesystem's cache: call `fs.invalidate_cache()` after a change, or make it through the filesystem. diff --git a/pyathena/filesystem/s3_core.py b/pyathena/filesystem/s3_core.py index 4803c9a9..e47baf5a 100644 --- a/pyathena/filesystem/s3_core.py +++ b/pyathena/filesystem/s3_core.py @@ -406,8 +406,9 @@ class S3Core: Each operation sends one request, or one per page for the iterators, with the retry policy, and translates S3 errors into ``OSError`` subclasses (see :class:`~pyathena.filesystem.s3_errors.S3ClientError`): - a missing bucket, or a missing object or version that an operation reads, - raises ``FileNotFoundError``, and a denied request ``PermissionError``. + a missing bucket or multipart upload, or a missing object or version that + an operation reads, raises ``FileNotFoundError``, and a denied request + ``PermissionError``. As in S3, deleting a missing key is not an error. Nothing is cached. Example: @@ -668,7 +669,8 @@ def upload_part_copy( source: The path of the object to copy, with the version ID to copy, if any. range_: The ``(start, end)`` byte range of the source to copy, - with an exclusive end; None copies the whole source. + with an exclusive end. None sends no range, so the whole + source is copied unless ``params`` has ``CopySourceRange``. **params: Additional request parameters. The fields that the other arguments set take precedence over parameters of the same name. diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index f781bb03..d6c0084d 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -3305,10 +3305,11 @@ def test_finish_multipart_upload_without_abort(self): assert not pending.cancelled() fs._call.assert_not_called() - def test_finish_multipart_upload_abort_failure_does_not_mask_the_original_error(self): + def test_finish_multipart_upload_abort_failure_does_not_mask_the_original_error(self, caplog): fs = self._make_fs() fs.core.complete_multipart_upload = mock.MagicMock() - fs._call = mock.MagicMock(side_effect=RuntimeError("abort failed")) + # The abort is sent through the core, whose call is the same mock. + fs._call.side_effect = RuntimeError("abort failed") future: Future[SimpleNamespace] = Future() future.set_exception(RuntimeError("upload failed")) @@ -3317,6 +3318,10 @@ def test_finish_multipart_upload_abort_failure_does_not_mask_the_original_error( fs._finish_multipart_upload( bucket="bucket", key="key", upload_id="uploadid", futures=[future] ) + fs._call.assert_called_once_with( + fs._client.abort_multipart_upload, Bucket="bucket", Key="key", UploadId="uploadid" + ) + assert "Failed to abort multipart upload uploadid to s3://bucket/key." in caplog.text def test_finish_multipart_upload_waits_for_running_parts(self): # GH-976: a part that is still uploading when the upload is aborted From e8f28532073f37e8c7469b8486fb077b94c19e3f Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 4 Oct 2026 14:12:25 +0900 Subject: [PATCH 4/4] Say that request_kwargs can also give upload_part_copy() its range Co-Authored-By: Claude Opus 5.5 --- pyathena/filesystem/s3_core.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/pyathena/filesystem/s3_core.py b/pyathena/filesystem/s3_core.py index e47baf5a..0a29c195 100644 --- a/pyathena/filesystem/s3_core.py +++ b/pyathena/filesystem/s3_core.py @@ -669,8 +669,9 @@ def upload_part_copy( source: The path of the object to copy, with the version ID to copy, if any. range_: The ``(start, end)`` byte range of the source to copy, - with an exclusive end. None sends no range, so the whole - source is copied unless ``params`` has ``CopySourceRange``. + with an exclusive end. None sends no range of its own, so the + whole source is copied unless ``params`` or ``request_kwargs`` + have ``CopySourceRange``. **params: Additional request parameters. The fields that the other arguments set take precedence over parameters of the same name.