diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index 4104f9354..64812f37b 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -2111,8 +2111,9 @@ 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 - read into the write buffer; a larger one is copied as the first parts - once a multipart upload starts. + read into the write buffer; a larger one is copied with + ``UploadPartCopy`` as the first parts of a multipart upload, whatever + the block size. Args: fs: The filesystem that the file belongs to. @@ -2188,11 +2189,15 @@ def __init__( self.s3_additional_kwargs.update({"IfMatch": etag}) self._details = info elif "a" in mode and self.fs.exists(path): - self.append_block = True info = self.fs.info(self.path, version_id=self.version_id) loc = info.get("size", 0) if loc < self.fs.MULTIPART_UPLOAD_MIN_PART_SIZE: + # Too small to be a part of a multipart upload: rewrite it + # from the buffer. self.write(self.fs.cat(self.path)) + else: + # Copied with UploadPartCopy as the leading part(s). + self.append_block = True self.loc = loc self.s3_additional_kwargs.update(info.to_api_repr()) self._details = info @@ -2208,8 +2213,10 @@ def close(self) -> None: self._executor.shutdown() def _initiate_upload(self) -> None: - if self.tell() < self.blocksize: + if not self.append_block and self.tell() < self.blocksize: # Files smaller than block size in size cannot be multipart uploaded. + # An append to an object copied with UploadPartCopy always uses + # a multipart upload, whatever the block size. return self.multipart_upload = self.fs._create_multipart_upload( @@ -2259,7 +2266,7 @@ def _upload_chunk(self, final: bool = False) -> bool: # can still read the bytes; resetting it there would upload an empty # object for small files. Mid-stream chunks (final=False) return True so # fsspec clears the already-uploaded buffer between parts. - if self.tell() < self.blocksize: + if not self.append_block and self.tell() < self.blocksize: # Files smaller than block size in size cannot be multipart uploaded. if self.autocommit and final: self.commit() @@ -2360,12 +2367,19 @@ def discard(self) -> None: if self.multipart_upload: for f in self.multipart_upload_parts: f.cancel() + # s3_additional_kwargs also holds object parameters (e.g., the + # existing object's metadata in append mode) that + # AbortMultipartUpload rejects. self.fs._call( "abort_multipart_upload", Bucket=self.bucket, Key=self.key, UploadId=self.multipart_upload.upload_id, - **self.s3_additional_kwargs, + **{ + k: v + for k, v in self.s3_additional_kwargs.items() + if k in ("RequestPayer", "ExpectedBucketOwner") + }, ) self.multipart_upload = None diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index a2ead0af2..9a9fe578e 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -676,6 +676,65 @@ def test_append(self, fs, base, exp): assert len(actual) == len(data + extra) assert actual == data + extra + @pytest.mark.parametrize( + ("size", "extra_size", "block_size"), + [ + # GH-921: an existing object of at least 5 MiB, appended within a + # larger block size, is copied with UploadPartCopy. + (6 * 2**20, 5, 16 * 2**20), + # An existing object smaller than 5 MiB is rewritten from the + # buffer, not copied as well, when the append crosses the block size. + (2**10, 5 * 2**20, None), + ], + ) + def test_append_with_block_size(self, fs, size, extra_size, block_size): + data = b"a" * size + extra = b"b" * extra_size + path = ( + f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" + f"filesystem/test_append_with_block_size/{uuid.uuid4()}" + ) + fs.pipe_file(path, data) + with fs.open(path, "ab", block_size=block_size) as f: + f.write(extra) + # Check the size and the bytes at the ends and around the boundary + # instead of reading the whole object back, to keep the transfer small. + assert fs.info(path, refresh=True).size == size + extra_size + assert fs.cat_file(path, start=0, end=1) == b"a" + assert fs.cat_file(path, start=size - 1, end=size + 1) == b"ab" + assert fs.cat_file(path, start=-1) == b"b" + + @pytest.mark.parametrize("block_size", [None, 16 * 2**20]) + def test_append_transaction_rollback(self, fs, block_size): + # Raising inside the transaction aborts the multipart upload that + # copies the existing object and leaves the object unchanged. + data = b"a" * (6 * 2**20) + path = ( + f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" + f"filesystem/test_append_transaction_rollback/{uuid.uuid4()}" + ) + fs.pipe_file(path, data) + before = fs.info(path, refresh=True) + + def append_then_fail(): + with fs.transaction: + f = fs.open(path, "ab", block_size=block_size) + f.write(b"b" * 5) + f.close() + raise RuntimeError("rollback") + + with pytest.raises(RuntimeError): + append_then_fail() + # A committed append (a multipart upload, or the appended bytes alone) + # would change the ETag and the size, so the object is not read back. + after = fs.info(path, refresh=True) + assert (after.etag, after.last_modified, after.size) == ( + before.etag, + before.last_modified, + before.size, + ) + assert not fs.list_multipart_uploads(path) + def test_ls_buckets(self, fs): fs.invalidate_cache() actual = fs.ls("s3://") @@ -1511,6 +1570,7 @@ def _make_write_file(data: bytes, autocommit: bool): file.s3_additional_kwargs = {} file.autocommit = autocommit file.blocksize = S3FileSystem.MULTIPART_UPLOAD_MIN_PART_SIZE + file.append_block = False file.multipart_upload = None file.multipart_upload_parts = [] file.buffer = io.BytesIO(data) @@ -1534,6 +1594,102 @@ def _make_multipart_write_file(data: bytes, autocommit: bool): ) return file + @staticmethod + def _make_append_fs(existing: bytes): + # A mocked filesystem holding an existing object, with a minimum part + # size of 4 bytes so that the append paths can be exercised with tiny + # data and no AWS access. + fs = mock.MagicMock(spec=S3FileSystem) + fs.MULTIPART_UPLOAD_MIN_PART_SIZE = 4 + fs.MULTIPART_UPLOAD_MAX_PART_SIZE = 64 + fs.exists.return_value = True + fs.info.return_value = S3Object( + init={"ContentLength": len(existing)}, + type=S3ObjectType.S3_OBJECT_TYPE_FILE, + bucket="bucket", + key="key.txt", + ) + fs.cat.return_value = existing + fs._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 + return fs + + @staticmethod + 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() + return fs._put_object.call_args.kwargs["body"] + fs._finish_multipart_upload.assert_called_once() + parts = [(c.kwargs["part_number"], existing) for c in fs._upload_part_copy.call_args_list] + parts += [ + (c.kwargs["part_number"], c.kwargs["body"]) for c in fs._upload_part.call_args_list + ] + part_numbers = sorted(n for n, _ in parts) + assert part_numbers == list(range(1, len(parts) + 1)) + return b"".join(body for _, body in sorted(parts)) + + @pytest.mark.parametrize( + ("existing", "appended", "multipart", "part_copy"), + [ + # Smaller than the minimum part size: read into the buffer. + (b"aa", b"bb", False, False), + # GH-921: an existing object of at least the minimum part size is + # copied with UploadPartCopy even when the block size is larger + # than the whole object. + (b"a" * 6, b"bb", True, True), + (b"a" * 6, b"", True, True), + # An existing object read into the buffer is not copied again + # when the append crosses the block size. + (b"aa", b"b" * 16, True, False), + ], + ) + def test_append(self, existing, appended, multipart, part_copy): + fs = self._make_append_fs(existing) + + with S3File(fs, "s3://bucket/key.txt", mode="ab", block_size=16) as f: + 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 + fs.touch.assert_not_called() + + def test_append_discard(self): + # Rolling back an append aborts its multipart upload without the + # existing object's metadata, which AbortMultipartUpload rejects, + # but with the request parameters it accepts. + fs = self._make_append_fs(b"a" * 6) + f = S3File( + fs, + "s3://bucket/key.txt", + mode="ab", + block_size=16, + autocommit=False, + s3_additional_kwargs={"RequestPayer": "requester", "ExpectedBucketOwner": "123"}, + ) + f.write(b"bb") + f.close() + + f.discard() + + fs._call.assert_called_once_with( + "abort_multipart_upload", + Bucket="bucket", + Key="key.txt", + UploadId="uploadid", + RequestPayer="requester", + ExpectedBucketOwner="123", + ) + fs._finish_multipart_upload.assert_not_called() + fs._put_object.assert_not_called() + @pytest.mark.parametrize( ("objects", "target"), [