diff --git a/docs/filesystem.md b/docs/filesystem.md index 9159fd6ef..f3ce1af80 100644 --- a/docs/filesystem.md +++ b/docs/filesystem.md @@ -135,6 +135,7 @@ operations raise natural errors instead of botocore's `ClientError`: | `404` / `NoSuchKey` / `NoSuchBucket` | `FileNotFoundError` | | `403` / `AccessDenied` | `PermissionError` | | `BucketAlreadyExists` / `BucketAlreadyOwnedByYou` | `FileExistsError` | +| `PreconditionFailed` of an `If-None-Match` condition | `FileExistsError` | | `RequestTimeout` | `TimeoutError` | | Others | `OSError` with the matching `errno` | diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index a7a608771..ffddaa8cc 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -1568,7 +1568,8 @@ def pipe_file( path: S3 path (s3://bucket/key) to write to. value: The bytes to write. mode: "overwrite" (default) or "create". With "create", raise - FileExistsError when the object already exists. + FileExistsError when the object already exists, including + one created during the write, which is not replaced. **kwargs: Additional parameters passed to the PutObject API (e.g., ContentType, StorageClass) on the single-request path. The ``block_size``, ``max_workers``, and @@ -1577,7 +1578,8 @@ def pipe_file( Raises: FileExistsError: If the mode is "create" and the path already - exists. + exists, or an object is created at it before the write is + committed. ValueError: If the path does not contain a key or specifies a version, or if the data takes more than ``MULTIPART_UPLOAD_MAX_PARTS`` blocks. @@ -1589,15 +1591,20 @@ def pipe_file( # Defer to the buffered open() path, which keeps the # deferred-commit semantics of fsspec transactions and uploads # large data as a parallel multipart upload. - super().pipe_file(path, value, mode=mode, **kwargs) + with self.open(path, "xb" if mode == "create" else "wb", **kwargs) as f: + f.write(value) return bucket, key, version_id = self.parse_path(path) if version_id: raise ValueError("Cannot write to the file with the version specified.") if not key: raise ValueError("Cannot write to a bucket.") - if mode == "create" and self.exists(path): - raise FileExistsError(path) + if mode == "create": + # Checked up front, as open() does in "xb" mode, and with + # IfNoneMatch for an object created since. + if self.exists(path): + raise FileExistsError(path) + kwargs["IfNoneMatch"] = "*" if not isinstance(value, bytes): # Accept bytes-like values (bytearray, memoryview) as the # buffered path does. @@ -1753,7 +1760,14 @@ def cat_file( return b"" raise - def put_file(self, lpath: str, rpath: str, callback=_DEFAULT_CALLBACK, **kwargs): + def put_file( + self, + lpath: str, + rpath: str, + callback=_DEFAULT_CALLBACK, + mode: str = "overwrite", + **kwargs, + ): """Upload a local file to S3. Uploads a file from the local filesystem to an S3 location. Supports @@ -1764,11 +1778,18 @@ def put_file(self, lpath: str, rpath: str, callback=_DEFAULT_CALLBACK, **kwargs) lpath: Local file path to upload. rpath: S3 destination path (s3://bucket/key). callback: Progress callback for tracking upload progress. + mode: "overwrite" (default) or "create". With "create", the file + is written as with ``open()`` in ``xb`` mode: raise + FileExistsError when the object already exists, including + one created during the upload, which is not replaced. **kwargs: Additional S3 parameters (e.g., ContentType, StorageClass). The ``block_size``, ``max_workers``, and ``s3_additional_kwargs`` parameters of ``open()`` are also accepted. Raises: + FileExistsError: If the mode is "create" and the path already + 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. @@ -1801,7 +1822,7 @@ def put_file(self, lpath: str, rpath: str, callback=_DEFAULT_CALLBACK, **kwargs) with ( self.open( rpath, - "wb", + "xb" if mode == "create" else "wb", block_size=block_size, max_workers=max_workers, s3_additional_kwargs=s3_additional_kwargs, @@ -2589,12 +2610,15 @@ def __init__( existing object smaller than ``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. + the block size. In exclusive-create mode, the object must not exist + when the file is opened, and the upload is committed with + ``IfNoneMatch="*"`` so that it does not replace an object created in + the meantime. Args: fs: The filesystem that the file belongs to. path: S3 path (s3://bucket/key) of the file. - mode: The file mode, such as ``rb``, ``wb`` or ``ab``. + mode: The file mode: ``rb``, ``wb``, ``ab``, or ``xb``. version_id: The version ID to read. Must match the version ID in the path if both are given. A version cannot be given, in either form, for writing or appending. @@ -2618,6 +2642,8 @@ def __init__( which take precedence over ``s3_additional_kwargs``. Raises: + FileExistsError: If an object exists at the path in + exclusive-create mode. FileNotFoundError: If no object exists at the path when reading, including when the path is a prefix. ValueError: If the path has no key, the version IDs do not match, @@ -2691,6 +2717,12 @@ def __init__( # Too small to be a part of a multipart upload: rewritten # from the buffer. append_data = fs.cat(path) + elif "x" in mode: + # Checked up front so that no data is uploaded for an existing + # object, and on commit with IfNoneMatch for one created since. + if fs.exists(path): + raise FileExistsError(path) + self.s3_additional_kwargs.update({"IfNoneMatch": "*"}) self._executor: S3Executor = executor or S3ThreadPoolExecutor(max_workers=max_workers) super().__init__( @@ -2880,6 +2912,8 @@ def commit(self) -> None: completion fails. Invalidates the cache of the path afterwards. Raises: + FileExistsError: If an object was created at the path after the + file was opened in exclusive-create mode. RuntimeError: If parts were submitted but no multipart upload is initialized. """ diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index 1993d2c3d..8c0947f83 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -180,8 +180,10 @@ def _pipe_file_in_transaction( Args: path: S3 path (s3://bucket/key) to write to. value: The bytes to write. - mode: "overwrite" or "create". With "create", raise - FileExistsError when the object already exists. + mode: "overwrite" or "create". With "create", the file is + opened in ``xb`` mode: raise FileExistsError when the object + already exists, including one created before the + transaction is committed, which is not replaced. **kwargs: Additional parameters passed to ``open()``. Raises: @@ -193,19 +195,30 @@ def _pipe_file_in_transaction( block_size = kwargs.get("block_size") or self._sync_fs.default_block_size # The size in bytes; the length of a memoryview counts its items. self._sync_fs._check_multipart_upload_size(path, memoryview(value).nbytes, block_size) - if mode == "create" and self._sync_fs.exists(path): - raise FileExistsError(path) - with self.open(path, "wb", **kwargs) as f: + with self.open(path, "xb" if mode == "create" else "wb", **kwargs) as f: f.write(value) - async def _put_file(self, lpath: str, rpath: str, callback=_DEFAULT_CALLBACK, **kwargs) -> None: + async def _put_file( + self, + lpath: str, + rpath: str, + callback=_DEFAULT_CALLBACK, + mode: str = "overwrite", + **kwargs, + ) -> None: if self._intrans: # See _pipe_file. - await asyncio.to_thread(self._put_file_in_transaction, lpath, rpath, callback, **kwargs) + await asyncio.to_thread( + self._put_file_in_transaction, lpath, rpath, callback, mode, **kwargs + ) return - await asyncio.to_thread(self._sync_fs.put_file, lpath, rpath, callback=callback, **kwargs) + await asyncio.to_thread( + self._sync_fs.put_file, lpath, rpath, callback=callback, mode=mode, **kwargs + ) - def _put_file_in_transaction(self, lpath: str, rpath: str, callback, **kwargs) -> None: + def _put_file_in_transaction( + self, lpath: str, rpath: str, callback, mode: str, **kwargs + ) -> None: """Upload a local file as a file of this filesystem's transaction. Mirrors :meth:`S3FileSystem.put_file`, but writes through ``open()`` @@ -215,11 +228,17 @@ def _put_file_in_transaction(self, lpath: str, rpath: str, callback, **kwargs) - lpath: Local file path to upload. rpath: S3 destination path (s3://bucket/key). callback: Progress callback for tracking upload progress. + mode: "overwrite" or "create". With "create", the file is + opened in ``xb`` mode: raise FileExistsError when the object + already exists, including one created before the + transaction is committed, which is not replaced. **kwargs: Additional S3 parameters (e.g., ContentType, StorageClass). The ``block_size``, ``max_workers``, and ``s3_additional_kwargs`` parameters of ``open()`` are also accepted. Raises: + FileExistsError: If the mode is "create" and the path already + exists. ValueError: If the file takes more than ``MULTIPART_UPLOAD_MAX_PARTS`` blocks. """ @@ -244,7 +263,7 @@ def _put_file_in_transaction(self, lpath: str, rpath: str, callback, **kwargs) - with ( self.open( rpath, - "wb", + "xb" if mode == "create" else "wb", block_size=block_size, max_workers=max_workers, s3_additional_kwargs=s3_additional_kwargs, diff --git a/pyathena/filesystem/s3_errors.py b/pyathena/filesystem/s3_errors.py index 0929ffc4c..f77623501 100644 --- a/pyathena/filesystem/s3_errors.py +++ b/pyathena/filesystem/s3_errors.py @@ -33,7 +33,9 @@ class S3ClientError: as properties, along with :attr:`os_error`, the equivalent standard Python exception. The error is mapped by its S3 error code first, then by its HTTP status code; if neither is recognized, a generic ``OSError`` - with the original error message is used. + with the original error message is used. A ``PreconditionFailed`` error + of an ``If-None-Match`` condition, from a write that must not replace an + existing object, is mapped to ``FileExistsError``. Example: >>> try: @@ -105,11 +107,15 @@ def __init__(self, error: botocore.exceptions.ClientError) -> None: error_info = error.response.get("Error", {}) self._code: str = str(error_info.get("Code", "")) self._message: str = str(error_info.get("Message", error)) + self._condition: str = str(error_info.get("Condition", "")) status_code = error.response.get("ResponseMetadata", {}).get("HTTPStatusCode") self._http_status_code: int | None = int(status_code) if status_code is not None else None self._os_error: OSError = self._translate() def _translate(self) -> OSError: + if self._code == "PreconditionFailed" and self._condition == "If-None-Match": + # A conditional write (IfNoneMatch="*") found an existing object. + return FileExistsError(self._message) exception = self._ERROR_CODE_TO_EXCEPTION.get(self._code) if exception: return exception(self._message) diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index ad1859301..668d01315 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -30,6 +30,7 @@ import pyathena from pyathena.filesystem import register_s3_filesystem from pyathena.filesystem.s3 import S3File, S3FileSystem +from pyathena.filesystem.s3_errors import S3ClientError from pyathena.filesystem.s3_executor import S3AioExecutor, S3ThreadPoolExecutor from pyathena.filesystem.s3_object import S3Object, S3ObjectType, S3StorageClass from pyathena.util import RetryConfig @@ -887,6 +888,165 @@ def test_put_file_open_parameters(self, tmp_path): s3_additional_kwargs={"StorageClass": "STANDARD_IA", "ContentType": "text/csv"}, ) + @staticmethod + def _record_requests(fs, precondition_failed=False): + # Record the S3 requests of the filesystem by operation name. With + # precondition_failed, the conditional writes fail as S3 fails them + # when an object exists. + requests = [] + + def call(method, **request): + name = method if isinstance(method, str) else method._extract_mock_name() + name = name.split(".")[-1] + requests.append((name, request)) + if precondition_failed and name in {"put_object", "complete_multipart_upload"}: + error = botocore.exceptions.ClientError( + { + "Error": { + "Code": "PreconditionFailed", + "Message": "At least one of the pre-conditions you specified " + "did not hold", + "Condition": "If-None-Match", + }, + "ResponseMetadata": {"HTTPStatusCode": 412}, + }, + name, + ) + raise S3ClientError(error).os_error from error + return {"UploadId": "uploadid", "ETag": '"e"'} + + fs._call.side_effect = call + return requests + + @pytest.mark.parametrize( + ("size", "expected"), + [ + # An empty file is created by touch(). + (0, ["put_object"]), + (1, ["put_object"]), + ( + S3FileSystem.MULTIPART_UPLOAD_MIN_PART_SIZE + 1, + ["create_multipart_upload", "upload_part", "complete_multipart_upload"], + ), + ], + ) + def test_open_exclusive_create(self, size, expected): + # GH-972: "xb" used to replace an existing object. The upload is + # committed only if no object exists, with IfNoneMatch on the + # requests that accept it. + fs = self._make_fs() + fs.default_cache_type = "bytes" + 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: + f.write(b"a" * size) + + fs.exists.assert_called_once() + assert [name for name, _ in requests] == expected + for name, request in requests: + conditional = name in {"put_object", "complete_multipart_upload"} + assert request.get("IfNoneMatch") == ("*" if conditional else None) + + def test_open_exclusive_create_existing(self): + # GH-972: an existing object is found when the file is opened, before + # any data is uploaded. + fs = self._make_fs() + fs.default_cache_type = "bytes" + fs.exists = mock.MagicMock(return_value=True) + + with pytest.raises(FileExistsError): + fs.open("s3://bucket/key", "xb") + fs._call.assert_not_called() + + @pytest.mark.parametrize( + ("size", "expected"), + [ + (1, ["put_object"]), + ( + S3FileSystem.MULTIPART_UPLOAD_MIN_PART_SIZE + 1, + [ + "create_multipart_upload", + "upload_part", + "complete_multipart_upload", + "abort_multipart_upload", + ], + ), + ], + ) + def test_open_exclusive_create_created_since(self, size, expected): + # GH-972: an object created after the file was opened is not + # replaced. S3 rejects the conditional write, which raises + # FileExistsError, and the multipart upload is aborted. + fs = self._make_fs() + fs.default_cache_type = "bytes" + fs.exists = mock.MagicMock(return_value=False) + requests = self._record_requests(fs, precondition_failed=True) + + with ( + pytest.raises(FileExistsError), + fs.open("s3://bucket/key", "xb", block_size=fs.MULTIPART_UPLOAD_MIN_PART_SIZE) as f, + ): + f.write(b"a" * size) + + assert [name for name, _ in requests] == expected + + @pytest.mark.parametrize(("mode", "open_mode"), [("overwrite", "wb"), ("create", "xb")]) + def test_put_file_mode(self, tmp_path, mode, open_mode): + # GH-972: fsspec's mode argument used to be sent to PutObject. It + # selects the mode of the remote file instead. + fs = self._make_fs() + fs.default_cache_type = "bytes" + fs.exists = mock.MagicMock(return_value=False) + requests = self._record_requests(fs) + lpath = tmp_path / "data" + lpath.write_bytes(b"a") + + with mock.patch.object(fs, "open", wraps=fs.open) as open_: + fs.put_file(str(lpath), "s3://bucket/key", mode=mode) + + assert open_.call_args.args == ("s3://bucket/key", open_mode) + ((name, request),) = requests + assert name == "put_object" + assert "mode" not in request + assert request.get("IfNoneMatch") == ("*" if mode == "create" else None) + + def test_put_file_create_existing(self, tmp_path): + fs = self._make_fs() + fs.default_cache_type = "bytes" + fs.exists = mock.MagicMock(return_value=True) + lpath = tmp_path / "data" + lpath.write_bytes(b"a") + + with pytest.raises(FileExistsError): + fs.put_file(str(lpath), "s3://bucket/key", mode="create") + fs._call.assert_not_called() + + @pytest.mark.parametrize( + ("size", "conditional"), + [ + (1, "put_object"), + (S3FileSystem.MULTIPART_UPLOAD_MIN_PART_SIZE + 1, "complete_multipart_upload"), + ], + ) + def test_pipe_file_create_created_since(self, size, conditional): + # GH-972: pipe_file(mode="create") also writes conditionally, on the + # single-request path as on the buffered one, so that an object + # 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.exists = mock.MagicMock(return_value=False) + requests = self._record_requests(fs, precondition_failed=True) + + with pytest.raises(FileExistsError): + fs.pipe_file("s3://bucket/key", b"a" * size, mode="create") + + fs.exists.assert_called_once() + assert [name for name, request in requests if request.get("IfNoneMatch") == "*"] == [ + conditional + ] + @pytest.mark.parametrize("fail", [False, True]) def test_open_parameters_named_as_request_fields(self, fail): # Parameters of a file named like the fields that a request sets @@ -2602,6 +2762,34 @@ def test_pipe_cat(self, fs, base, exp): fs.pipe(path, data) assert fs.cat(path) == data + def test_exclusive_create(self, fs, tmp_path): + # GH-972: "xb" and put_file(mode="create") used to replace an + # existing object. + prefix = ( + f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" + f"filesystem/test_exclusive_create/{uuid.uuid4()}" + ) + path = f"{prefix}/existing" + with fs.open(path, "xb") as f: + f.write(b"old") + lpath = tmp_path / "data" + lpath.write_bytes(b"new") + with pytest.raises(FileExistsError): + fs.open(path, "xb") + with pytest.raises(FileExistsError): + fs.put_file(str(lpath), path, mode="create") + assert fs.cat(path) == b"old" + + # S3 rejects the conditional write of an object created after the + # file was opened. + path = f"{prefix}/created_since" + f = fs.open(path, "xb") + f.write(b"new") + fs.pipe_file(path, b"old") + with pytest.raises(FileExistsError): + f.close() + assert fs.cat(path) == b"old" + def test_pipe_file_create_mode_and_kwargs(self, fs): path = ( f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index 59604bb97..50b9e1851 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -12,6 +12,7 @@ from types import SimpleNamespace from unittest import mock +import boto3 import fsspec import pytest from fsspec import Callback @@ -232,12 +233,40 @@ def write(): write() put_object.assert_not_called() - def test_transaction_pipe_file_create_existing(self): + def test_transaction_pipe_put_file_create_existing(self, tmp_path): + # GH-972: in a transaction, put_file(mode="create") also raises, when + # the file is opened, for an existing object. fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) fs._sync_fs.exists = mock.MagicMock(return_value=True) - with fs.transaction, pytest.raises(FileExistsError): - fs.pipe_file("s3://bucket/key", b"data", mode="create") - fs._sync_fs.exists.assert_called_once_with("s3://bucket/key") + fs._sync_fs._call = mock.MagicMock() + local = tmp_path / "local" + local.write_bytes(b"a") + with fs.transaction: + with pytest.raises(FileExistsError): + fs.pipe_file("s3://bucket/k1", b"data", mode="create") + with pytest.raises(FileExistsError): + fs.put_file(str(local), "s3://bucket/k2", mode="create") + assert fs._sync_fs.exists.call_count == 2 + fs._sync_fs._call.assert_not_called() + + @pytest.mark.parametrize("mode", ["overwrite", "create"]) + def test_put_file_mode(self, tmp_path, mode): + # GH-972: fsspec's mode argument used to be sent to PutObject. + # A real client selects the request parameters of each operation. + fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) + fs._sync_fs._client = boto3.client( + "s3", region_name="us-east-1", aws_access_key_id="dummy", aws_secret_access_key="dummy" + ) + fs._sync_fs.exists = mock.MagicMock(return_value=False) + fs._sync_fs._call = mock.MagicMock(return_value={"ETag": '"e"'}) + local = tmp_path / "local" + local.write_bytes(b"a") + + fs.put_file(str(local), "s3://bucket/key", mode=mode) + + (call,) = fs._sync_fs._call.call_args_list + assert "mode" not in call.kwargs + assert call.kwargs.get("IfNoneMatch") == ("*" if mode == "create" else None) @pytest.mark.parametrize("kwargs", [{"block_size": 4}, {}]) def test_transaction_pipe_put_file_exceeding_max_parts(self, tmp_path, kwargs): @@ -407,9 +436,11 @@ async def test_rm_maxdepth(self): (call,) = sync_fs._call.call_args_list assert call.kwargs["Delete"]["Objects"] == [{"Key": "dir"}, {"Key": "dir/a"}] - def test_put_file_in_transaction_open_parameters(self, tmp_path): + @pytest.mark.parametrize(("mode", "open_mode"), [("overwrite", "wb"), ("create", "xb")]) + def test_put_file_in_transaction_open_parameters(self, tmp_path, mode, open_mode): # GH-969: the open() parameters of put_file() go to open(), and the # other parameters, also in s3_additional_kwargs, to S3. + # GH-972: fsspec's mode argument selects the mode of the file. fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) fs.open = mock.MagicMock() fs.open.return_value.__enter__.return_value.blocksize = 4 @@ -420,6 +451,7 @@ def test_put_file_in_transaction_open_parameters(self, tmp_path): str(lpath), "s3://bucket/key", Callback(), + mode, block_size=S3FileSystem.MULTIPART_UPLOAD_MIN_PART_SIZE, max_workers=2, StorageClass="STANDARD_IA", @@ -427,7 +459,7 @@ def test_put_file_in_transaction_open_parameters(self, tmp_path): fs.open.assert_called_once_with( "s3://bucket/key", - "wb", + open_mode, block_size=S3FileSystem.MULTIPART_UPLOAD_MIN_PART_SIZE, max_workers=2, s3_additional_kwargs={"StorageClass": "STANDARD_IA", "ContentType": "text/csv"}, diff --git a/tests/pyathena/filesystem/test_s3_errors.py b/tests/pyathena/filesystem/test_s3_errors.py index 152692b60..246f2cb30 100644 --- a/tests/pyathena/filesystem/test_s3_errors.py +++ b/tests/pyathena/filesystem/test_s3_errors.py @@ -13,8 +13,10 @@ from pyathena.filesystem.s3_errors import S3ClientError -def _client_error(code, message="error message", status_code=None): +def _client_error(code, message="error message", status_code=None, condition=None): error_response = {"Error": {"Code": code, "Message": message}} + if condition is not None: + error_response["Error"]["Condition"] = condition if status_code is not None: error_response["ResponseMetadata"] = {"HTTPStatusCode": status_code} return botocore.exceptions.ClientError(error_response, "TestOperation") @@ -94,3 +96,23 @@ def test_os_error_unknown_code(self): assert type(actual) is OSError assert actual.errno == errno.EIO assert "error message" in str(actual) + + @pytest.mark.parametrize( + ("condition", "expected"), + [ + # GH-972: a conditional write (IfNoneMatch="*") found an existing + # object. + ("If-None-Match", FileExistsError), + # A read conditional on the ETag (IfMatch) found another object. + ("If-Match", OSError), + (None, OSError), + ], + ) + def test_os_error_precondition_failed(self, condition, expected): + actual = S3ClientError( + _client_error("PreconditionFailed", status_code=412, condition=condition) + ).os_error + assert type(actual) is expected + if expected is OSError: + assert actual.errno == errno.EINVAL + assert "error message" in str(actual)