diff --git a/gcsfs/tests/test_core.py b/gcsfs/tests/test_core.py index d5314cb18..53a3cd93d 100644 --- a/gcsfs/tests/test_core.py +++ b/gcsfs/tests/test_core.py @@ -1745,9 +1745,7 @@ def test_write_if_generation_match(gcs): try: gcs.pipe_file(fn, b"v2", if_generation_match=gen) except NotImplementedError: - pytest.skip( - "zonal CAS needs a google-cloud-storage release with if_generation_match" - ) + pytest.skip("zonal CAS needs a google-cloud-storage release with Preconditions") assert gcs.cat(fn) == b"v2" with pytest.raises(GenerationMismatchError): gcs.pipe_file(fn, b"v3", if_generation_match=gen) diff --git a/gcsfs/tests/test_zb_hns_utils.py b/gcsfs/tests/test_zb_hns_utils.py index 669968fd4..03e5b176b 100644 --- a/gcsfs/tests/test_zb_hns_utils.py +++ b/gcsfs/tests/test_zb_hns_utils.py @@ -1,7 +1,9 @@ import asyncio import collections import concurrent.futures +import dataclasses import logging +from typing import Optional from unittest import mock import pytest @@ -99,18 +101,39 @@ async def test_init_aaow(): assert result is mock_writer_instance +@dataclasses.dataclass +class _FakePreconditions: + if_generation_match: Optional[int] = None + + +# The real class once google-cloud-storage ships it, so CI builds it too. +_Preconditions = zb_hns_utils.Preconditions or _FakePreconditions + + @pytest.mark.asyncio @pytest.mark.parametrize( "if_generation_match, expected_kwargs", - [(0, {"generation": 0}), (7, {"generation": None, "if_generation_match": 7})], + [ + (0, {"generation": 0}), + ( + 7, + { + "generation": None, + "preconditions": _Preconditions(if_generation_match=7), + }, + ), + ], ) async def test_init_aaow_if_generation_match(if_generation_match, expected_kwargs): mock_writer_instance = mock.AsyncMock() - with mock.patch( - "gcsfs.zb_hns_utils.AsyncAppendableObjectWriter", - new_callable=mock.Mock, - return_value=mock_writer_instance, - ) as mock_writer_class: + with ( + mock.patch.object(zb_hns_utils, "Preconditions", _Preconditions), + mock.patch( + "gcsfs.zb_hns_utils.AsyncAppendableObjectWriter", + new_callable=mock.Mock, + return_value=mock_writer_instance, + ) as mock_writer_class, + ): await zb_hns_utils.init_aaow( mock_grpc_client, bucket_name, @@ -127,12 +150,13 @@ async def test_init_aaow_if_generation_match(if_generation_match, expected_kwarg @pytest.mark.asyncio +@pytest.mark.parametrize("generation", [0, 123]) @pytest.mark.parametrize("if_generation_match", [0, 7]) async def test_init_aaow_rejects_generation_with_if_generation_match( - if_generation_match, + generation, if_generation_match ): - # generation > 0 selects an existing object to append to, which carries no - # precondition, so the combination must not silently drop the precondition. + # A set generation selects an existing object (or, for 0, already means + # create-only), so the combination must not silently drop the precondition. with mock.patch( "gcsfs.zb_hns_utils.AsyncAppendableObjectWriter", new_callable=mock.Mock ) as mock_writer_class: @@ -141,7 +165,7 @@ async def test_init_aaow_rejects_generation_with_if_generation_match( mock_grpc_client, bucket_name, object_name, - generation=123, + generation=generation, if_generation_match=if_generation_match, ) mock_writer_class.assert_not_called() @@ -149,26 +173,41 @@ async def test_init_aaow_rejects_generation_with_if_generation_match( @pytest.mark.asyncio async def test_init_aaow_if_generation_match_unsupported_sdk(): - with mock.patch( - "gcsfs.zb_hns_utils.AsyncAppendableObjectWriter", - side_effect=TypeError("unexpected keyword argument 'if_generation_match'"), + with ( + mock.patch.object(zb_hns_utils, "Preconditions", None), + mock.patch( + "gcsfs.zb_hns_utils.AsyncAppendableObjectWriter", + new_callable=mock.Mock, + ) as mock_writer_class, ): with pytest.raises(NotImplementedError): await zb_hns_utils.init_aaow( mock_grpc_client, bucket_name, object_name, if_generation_match=7 ) + mock_writer_class.assert_not_called() @pytest.mark.asyncio -async def test_init_aaow_unrelated_type_error_propagates(): - with mock.patch( - "gcsfs.zb_hns_utils.AsyncAppendableObjectWriter", - side_effect=TypeError("boom"), +async def test_init_aaow_if_generation_match_zero_without_preconditions(): + mock_writer_instance = mock.AsyncMock() + with ( + mock.patch.object(zb_hns_utils, "Preconditions", None), + mock.patch( + "gcsfs.zb_hns_utils.AsyncAppendableObjectWriter", + new_callable=mock.Mock, + return_value=mock_writer_instance, + ) as mock_writer_class, ): - with pytest.raises(TypeError, match="boom"): - await zb_hns_utils.init_aaow( - mock_grpc_client, bucket_name, object_name, if_generation_match=7 - ) + await zb_hns_utils.init_aaow( + mock_grpc_client, bucket_name, object_name, if_generation_match=0 + ) + mock_writer_class.assert_called_once_with( + client=mock_grpc_client, + bucket_name=bucket_name, + object_name=object_name, + generation=0, + writer_options={}, + ) @pytest.mark.asyncio @@ -179,10 +218,13 @@ async def test_init_aaow_unrelated_type_error_propagates(): async def test_init_aaow_failed_precondition(if_generation_match, expected): mock_writer_instance = mock.AsyncMock() mock_writer_instance.open.side_effect = FailedPrecondition("exists") - with mock.patch( - "gcsfs.zb_hns_utils.AsyncAppendableObjectWriter", - new_callable=mock.Mock, - return_value=mock_writer_instance, + with ( + mock.patch.object(zb_hns_utils, "Preconditions", _Preconditions), + mock.patch( + "gcsfs.zb_hns_utils.AsyncAppendableObjectWriter", + new_callable=mock.Mock, + return_value=mock_writer_instance, + ), ): with pytest.raises(expected) as excinfo: await zb_hns_utils.init_aaow( diff --git a/gcsfs/tests/test_zonal.py b/gcsfs/tests/test_zonal.py index a2a90dd3c..62c7f06ee 100644 --- a/gcsfs/tests/test_zonal.py +++ b/gcsfs/tests/test_zonal.py @@ -766,9 +766,22 @@ def test_multiprocess_error_handling_zb(extended_gcsfs): # =========================== Zonal Upload Tests =========================== +# Write kwargs and the if_generation_match the caller must pass to init_aaow. +_generation_match_cases = pytest.mark.parametrize( + "write_kwargs, expected_match", + [({}, None), ({"mode": "create"}, 0), ({"if_generation_match": "7"}, 7)], + ids=["default", "create", "if_generation_match"], +) + + @pytest.mark.asyncio -async def test_simple_upload_zonal(async_gcs, zonal_write_mocks, file_path): +@_generation_match_cases +async def test_simple_upload_zonal( + async_gcs, zonal_write_mocks, file_path, write_kwargs, expected_match +): """Test simple_upload for Zonal buckets calls the correct writer methods.""" + if write_kwargs and not zonal_write_mocks: + pytest.skip("Precondition forwarding is checked with the mocked writer.") data = b"test data for simple_upload" bucket, object_name, _ = async_gcs.split_path(file_path) @@ -778,10 +791,14 @@ async def test_simple_upload_zonal(async_gcs, zonal_write_mocks, file_path): key=object_name, datain=data, finalize_on_close=True, # Finalize on close to make the object immediately readable + **write_kwargs, ) if zonal_write_mocks: zonal_write_mocks["init_aaow"].assert_awaited_once_with( - async_gcs.grpc_client, bucket, object_name, if_generation_match=None + async_gcs.grpc_client, + bucket, + object_name, + if_generation_match=expected_match, ) zonal_write_mocks["aaow"].append.assert_awaited_once_with(data) zonal_write_mocks["aaow"].close.assert_awaited_once_with(finalize_on_close=True) @@ -790,13 +807,23 @@ async def test_simple_upload_zonal(async_gcs, zonal_write_mocks, file_path): @pytest.mark.asyncio -async def test_initiate_upload_zonal(async_gcs, zonal_write_mocks, file_path): +@_generation_match_cases +async def test_initiate_upload_zonal( + async_gcs, zonal_write_mocks, file_path, write_kwargs, expected_match +): """Test initiate_upload for Zonal buckets returns a writer instance.""" + if write_kwargs and not zonal_write_mocks: + pytest.skip("Precondition forwarding is checked with the mocked writer.") bucket, object_name, _ = async_gcs.split_path(file_path) - writer = await initiate_upload(fs=async_gcs, bucket=bucket, key=object_name) + writer = await initiate_upload( + fs=async_gcs, bucket=bucket, key=object_name, **write_kwargs + ) if zonal_write_mocks: zonal_write_mocks["init_aaow"].assert_awaited_once_with( - async_gcs.grpc_client, bucket, object_name, if_generation_match=None + async_gcs.grpc_client, + bucket, + object_name, + if_generation_match=expected_match, ) assert writer is zonal_write_mocks["aaow"] else: @@ -888,8 +915,13 @@ async def test_upload_chunk_zonal_final_chunk(async_gcs, zonal_write_mocks, file @pytest.mark.asyncio -async def test_put_file_zonal(async_gcs, zonal_write_mocks, file_path): +@_generation_match_cases +async def test_put_file_zonal( + async_gcs, zonal_write_mocks, file_path, write_kwargs, expected_match +): """Test _put_file for Zonal buckets.""" + if write_kwargs and not zonal_write_mocks: + pytest.skip("Precondition forwarding is checked with the mocked writer.") data = b"test data for put_file" bucket, object_name, _ = async_gcs.split_path(file_path) with tmpfile() as lpath: @@ -897,12 +929,14 @@ async def test_put_file_zonal(async_gcs, zonal_write_mocks, file_path): f.write(data) # Finalize on close to make the object immediately readable - await async_gcs._put_file(lpath, file_path, finalize_on_close=True) + await async_gcs._put_file( + lpath, file_path, finalize_on_close=True, **write_kwargs + ) if zonal_write_mocks: grpc_client = await async_gcs._get_grpc_client() zonal_write_mocks["init_aaow"].assert_awaited_once_with( - grpc_client, bucket, object_name, if_generation_match=None + grpc_client, bucket, object_name, if_generation_match=expected_match ) zonal_write_mocks["aaow"].append_from_file.assert_awaited_once() args, kwargs = zonal_write_mocks["aaow"].append_from_file.call_args @@ -921,17 +955,22 @@ async def test_put_file_zonal(async_gcs, zonal_write_mocks, file_path): @pytest.mark.asyncio -async def test_pipe_file_zonal(async_gcs, zonal_write_mocks, file_path): +@_generation_match_cases +async def test_pipe_file_zonal( + async_gcs, zonal_write_mocks, file_path, write_kwargs, expected_match +): """Test _pipe_file for Zonal buckets.""" + if write_kwargs and not zonal_write_mocks: + pytest.skip("Precondition forwarding is checked with the mocked writer.") data = b"test data for pipe_file" bucket, object_name, _ = async_gcs.split_path(file_path) # Finalize on close to make the object immediately readable - await async_gcs._pipe_file(file_path, data, finalize_on_close=True) + await async_gcs._pipe_file(file_path, data, finalize_on_close=True, **write_kwargs) if zonal_write_mocks: grpc_client = await async_gcs._get_grpc_client() zonal_write_mocks["init_aaow"].assert_awaited_once_with( - grpc_client, bucket, object_name, if_generation_match=None + grpc_client, bucket, object_name, if_generation_match=expected_match ) zonal_write_mocks["aaow"].append.assert_awaited_once_with(data) zonal_write_mocks["aaow"].close.assert_awaited_once_with(finalize_on_close=True) diff --git a/gcsfs/tests/test_zonal_file.py b/gcsfs/tests/test_zonal_file.py index 965e80b40..5560f539b 100644 --- a/gcsfs/tests/test_zonal_file.py +++ b/gcsfs/tests/test_zonal_file.py @@ -292,13 +292,41 @@ def test_zonal_file_open_write_mode_with_if_generation_match( if_generation_match=7, ) else: - # file_path does not exist, so generation 7 cannot match; SDK releases - # whose writer lacks the parameter raise NotImplementedError instead. - with pytest.raises((NotImplementedError, GenerationMismatchError)): - with extended_gcsfs.open( - file_path, "wb", finalize_on_close=True, if_generation_match=7 - ) as f: - f.write(test_data) + # file_path does not exist, so generation 7 cannot match. + try: + with pytest.raises(GenerationMismatchError): + with extended_gcsfs.open( + file_path, "wb", finalize_on_close=True, if_generation_match=7 + ) as f: + f.write(test_data) + except NotImplementedError: + pytest.skip( + "zonal CAS needs a google-cloud-storage release with Preconditions" + ) + + +@pytest.mark.parametrize( + "mode, kwargs, error", + [ + ("xb", {}, FileExistsError("exists")), + ("wb", {"if_generation_match": 7}, GenerationMismatchError("path", 7)), + ("wb", {"if_generation_match": 7}, NotImplementedError("no Preconditions")), + ], +) +def test_zonal_file_failed_writer_open_is_not_retried_on_close( + extended_gcsfs, zonal_write_mocks, file_path, mode, kwargs, error +): + """A writer that fails to open raises once; close() does not open it again.""" + if not zonal_write_mocks: + pytest.skip("Unit test for writer open failures with mocks.") + zonal_write_mocks["init_aaow"].side_effect = error + + with pytest.raises(type(error)): + with extended_gcsfs.open(file_path, mode, **kwargs) as f: + f.write(test_data) + + zonal_write_mocks["init_aaow"].assert_awaited_once() + assert f.closed @mock.patch("gcsfs.zonal_file.asyn.sync") diff --git a/gcsfs/zb_hns_utils.py b/gcsfs/zb_hns_utils.py index 08483ac70..55ee49426 100644 --- a/gcsfs/zb_hns_utils.py +++ b/gcsfs/zb_hns_utils.py @@ -22,6 +22,13 @@ from gcsfs.retry import GenerationMismatchError +try: + from google.cloud.storage.asyncio.async_appendable_object_writer import ( + Preconditions, + ) +except ImportError: # releases before googleapis/google-cloud-python#18641 + Preconditions = None + MRD_MAX_RANGES = 1000 # MRD supports up to 1000 ranges per request try: DEFAULT_CONCURRENCY = int(os.environ.get("DEFAULT_GCSFS_CONCURRENCY", "4")) @@ -175,7 +182,7 @@ async def init_aaow( ``if_generation_match=0`` creates the object only if it does not exist; a positive value overwrites only that generation (requires a - google-cloud-storage release whose writer accepts ``if_generation_match``). + google-cloud-storage release whose writer accepts ``preconditions``). A failed precondition raises ``FileExistsError`` for ``0`` and ``GenerationMismatchError`` otherwise; without ``if_generation_match`` the SDK's ``FailedPrecondition`` propagates unchanged. @@ -186,9 +193,10 @@ async def init_aaow( if flush_interval_bytes and flush_interval_bytes != _DEFAULT_FLUSH_INTERVAL_BYTES: writer_options["FLUSH_INTERVAL_BYTES"] = flush_interval_bytes kwargs = {} - if generation and if_generation_match is not None: - # generation > 0 appends to an existing object (AppendObjectSpec), which - # carries no precondition, so the SDK would silently ignore it. + if generation is not None and if_generation_match is not None: + # generation > 0 appends to an existing object and generation=0 already + # means create-only, so a generation precondition cannot apply. Check + # here because if_generation_match=0 is mapped to generation=0 below. raise ValueError( "if_generation_match cannot be combined with an explicit generation" ) @@ -196,23 +204,20 @@ async def init_aaow( # Every SDK release maps generation=0 to WriteObjectSpec(if_generation_match=0). generation = 0 elif if_generation_match is not None: - kwargs["if_generation_match"] = if_generation_match - try: - writer = AsyncAppendableObjectWriter( - client=grpc_client, - bucket_name=bucket_name, - object_name=object_name, - generation=generation, - writer_options=writer_options, - **kwargs, - ) - except TypeError as e: - if "if_generation_match" not in str(e): - raise - raise NotImplementedError( - "if_generation_match > 0 on zonal buckets requires a google-cloud-storage " - "release whose AsyncAppendableObjectWriter accepts if_generation_match" - ) from e + if Preconditions is None: + raise NotImplementedError( + "if_generation_match > 0 on zonal buckets requires a google-cloud-storage " + "release whose AsyncAppendableObjectWriter accepts Preconditions" + ) + kwargs["preconditions"] = Preconditions(if_generation_match=if_generation_match) + writer = AsyncAppendableObjectWriter( + client=grpc_client, + bucket_name=bucket_name, + object_name=object_name, + generation=generation, + writer_options=writer_options, + **kwargs, + ) try: await writer.open() except FailedPrecondition as e: diff --git a/gcsfs/zonal_file.py b/gcsfs/zonal_file.py index e4cfc128e..3a3139d07 100644 --- a/gcsfs/zonal_file.py +++ b/gcsfs/zonal_file.py @@ -173,14 +173,20 @@ async def _init_aaow( def _ensure_aaow(self): if self.aaow is None: - self.aaow = asyn.sync( - self.gcsfs.loop, - self._init_aaow, - self.bucket, - self.key, - self.generation, - self.flush_interval_bytes, - ) + try: + self.aaow = asyn.sync( + self.gcsfs.loop, + self._init_aaow, + self.bucket, + self.key, + self.generation, + self.flush_interval_bytes, + ) + except Exception: + # Like fsspec's flush() when _initiate_upload() fails: mark the + # file closed so that close() does not open the writer again. + self.closed = True + raise def _fetch_range( self,