diff --git a/docs/filesystem.md b/docs/filesystem.md index 9c9e9421..0202a8a7 100644 --- a/docs/filesystem.md +++ b/docs/filesystem.md @@ -280,10 +280,13 @@ 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 +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. ```python import boto3 @@ -306,8 +309,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 +324,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()` 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`. + ## 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..846f56cb 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,43 @@ 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"), + path = S3Path(self.bucket, self.key) + self.multipart_upload = self.fs.core.create_multipart_upload( + path, **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=path, upload_id=cast(str, self.multipart_upload.upload_id), part_number=i + 1, - copy_source_ranges=range_, + # The existing object is copied into the upload. + source=path, + 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=path, upload_id=cast(str, self.multipart_upload.upload_id), part_number=1, + source=path, **self._get_request_kwargs("upload_part_copy"), ) ) @@ -3805,10 +3655,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 +3668,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 +3772,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..0a29c195 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 @@ -399,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: @@ -410,6 +418,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 +594,201 @@ 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, without a version ID. + **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, 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}) + 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 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. + + 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. + 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}.") + 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..d6c0084d 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() @@ -3306,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._complete_multipart_upload = mock.MagicMock() - fs._call = mock.MagicMock(side_effect=RuntimeError("abort failed")) + fs.core.complete_multipart_upload = mock.MagicMock() + # 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")) @@ -3318,13 +3318,17 @@ 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 # 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 +3371,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 +3397,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 +3425,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 +3436,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 +5092,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 +5270,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 +5294,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 +5311,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 +5326,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 +5336,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 +5350,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 +5386,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 +5406,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 +5434,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 +5461,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 +5469,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 +5486,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 +5509,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 +5521,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 +5559,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 +5593,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 +5622,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 +5664,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 +5770,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 +5825,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 +5840,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 +5859,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 +5882,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 +5892,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..b635783a 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,202 @@ 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) + + 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"), + [ + # 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.