From c50294b8bece71c1eb68ce35384b2cf0132a142b Mon Sep 17 00:00:00 2001 From: yuxin00j Date: Mon, 12 Oct 2026 09:57:34 +0000 Subject: [PATCH 1/6] feat(storage): add if_generation_match to AsyncAppendableObjectWriter `AsyncAppendableObjectWriter` and `_AsyncWriteObjectStream` overloaded a single `generation` argument: `None`/`0` created a new object via `WriteObjectSpec(if_generation_match=generation)`, while `N > 0` reopened an existing object via `AppendObjectSpec(generation=N)`. There was no way to request `WriteObjectSpec(appendable=True, if_generation_match=N)`, i.e. overwrite an object only if its current generation is still `N`, although the gRPC API and the other GCS client libraries support it. This adds an explicit `if_generation_match` parameter. It is set on `WriteObjectSpec` when creating a new object (`generation` is `None` or `0`; `generation=0` keeps meaning `if_generation_match=0` when the new parameter is not given) and is ignored when `generation > 0` selects an existing object to append to, so the redirect-retry path that reopens with the redirected generation is unaffected. Test Plan: ``` cd packages/google-cloud-storage PYTHONPATH=$PWD python -m pytest tests/unit/asyncio/test_async_write_object_stream.py tests/unit/asyncio/test_async_appendable_object_writer.py # 73 passed ruff check --select I --target-version=py310 --line-length=88 ruff format --check --target-version=py310 --line-length=88 flake8 ``` Authored with an AI assistant. --- .../asyncio/async_appendable_object_writer.py | 12 ++++ .../asyncio/async_write_object_stream.py | 16 +++++- .../test_async_appendable_object_writer.py | 15 +++++ .../asyncio/test_async_write_object_stream.py | 55 +++++++++++++++++++ 4 files changed, 96 insertions(+), 2 deletions(-) diff --git a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_appendable_object_writer.py b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_appendable_object_writer.py index 1af54d6ffe43..b4a7fcfed5ca 100644 --- a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_appendable_object_writer.py +++ b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_appendable_object_writer.py @@ -112,6 +112,7 @@ def __init__( write_handle: Optional[_storage_v2.BidiWriteHandle] = None, writer_options: Optional[dict] = None, storage_class: Optional[str] = None, + if_generation_match: Optional[int] = None, ): """ Class for appending data to a GCS Appendable Object. @@ -184,6 +185,15 @@ def __init__( :param storage_class: (Optional) Storage class of the object bytes. If specified, it overrides the bucket's `storage_class`. If not, object storage class will be the same as bucket's storage_class. + + :type if_generation_match: Optional[int] + :param if_generation_match: (Optional) Only create the object if its + current generation matches this value; ``0`` means the object must + not exist. Use this for compare-and-swap overwrites of an existing + object. Ignored when ``generation`` selects an existing object to + append to. A mismatch raises + :class:`google.api_core.exceptions.FailedPrecondition` from + ``open()``. """ _utils.raise_if_no_fast_crc32c() self.client = client @@ -192,6 +202,7 @@ def __init__( self.write_handle = write_handle self.generation = generation self.storage_class = storage_class + self.if_generation_match = if_generation_match self.write_obj_stream: Optional[_AsyncWriteObjectStream] = None self._is_stream_open: bool = False @@ -369,6 +380,7 @@ async def _do_open(): write_handle=self.write_handle, routing_token=self._routing_token, storage_class=self.storage_class, + if_generation_match=self.if_generation_match, ) if self._routing_token: diff --git a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_write_object_stream.py b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_write_object_stream.py index 2d78a8c20fd1..25cda251fa7b 100644 --- a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_write_object_stream.py +++ b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_write_object_stream.py @@ -61,6 +61,12 @@ class _AsyncWriteObjectStream(_AsyncAbstractObjectStream): :type storage_class: Optional[str] :param storage_class: (Optional) The storage class of the object. + + :type if_generation_match: Optional[int] + :param if_generation_match: (Optional) Precondition for creating a new + object: the write succeeds only if the object's current generation + matches this value (``0`` means the object must not exist). Ignored + when ``generation_number`` selects an existing object to append to. """ def __init__( @@ -73,6 +79,7 @@ def __init__( routing_token: Optional[str] = None, blob: Optional[Blob] = None, storage_class: Optional[str] = None, + if_generation_match: Optional[int] = None, ) -> None: if client is None: raise ValueError("client must be provided") @@ -91,6 +98,7 @@ def __init__( self.routing_token: Optional[str] = routing_token self.blob: Optional[Blob] = blob self.storage_class: Optional[str] = storage_class + self.if_generation_match: Optional[int] = if_generation_match self._full_bucket_name = f"projects/_/buckets/{self.bucket_name}" self.rpc = self.client._client._transport._wrapped_methods[ @@ -114,7 +122,8 @@ async def open(self, metadata: Optional[List[Tuple[str, str]]] = None) -> None: :rtype: None :raises ValueError: If the stream is already open. :raises google.api_core.exceptions.FailedPrecondition: - if `generation_number` is 0 and object already exists. + if `if_generation_match` (or `generation_number` is 0) does not + match the object's current generation. """ if self._is_stream_open: raise ValueError("Stream is already open") @@ -133,11 +142,14 @@ async def open(self, metadata: Optional[List[Tuple[str, str]]] = None) -> None: bucket=self._full_bucket_name, storage_class=self.storage_class, ) + if_generation_match = self.if_generation_match + if if_generation_match is None: + if_generation_match = self.generation_number self.first_bidi_write_req = _storage_v2.BidiWriteObjectRequest( write_object_spec=_storage_v2.WriteObjectSpec( resource=resource, appendable=True, - if_generation_match=self.generation_number, + if_generation_match=if_generation_match, ), ) else: diff --git a/packages/google-cloud-storage/tests/unit/asyncio/test_async_appendable_object_writer.py b/packages/google-cloud-storage/tests/unit/asyncio/test_async_appendable_object_writer.py index e6fcd90f0b84..ec6175b8d52c 100644 --- a/packages/google-cloud-storage/tests/unit/asyncio/test_async_appendable_object_writer.py +++ b/packages/google-cloud-storage/tests/unit/asyncio/test_async_appendable_object_writer.py @@ -173,6 +173,7 @@ async def test_init_with_storage_class(self, mock_appendable_writer, storage_cla write_handle=None, routing_token=None, storage_class=storage_class, + if_generation_match=None, ) assert writer.write_obj_stream.storage_class == storage_class @@ -272,6 +273,7 @@ async def test_from_blob_open_passes_storage_class( write_handle=None, routing_token=None, storage_class=storage_class, + if_generation_match=None, ) assert writer.write_obj_stream.storage_class == storage_class @@ -312,6 +314,19 @@ async def test_open_success(self, mock_appendable_writer): assert mock_stream_cls.call_args.kwargs["storage_class"] is None assert writer.write_obj_stream.storage_class is None + @pytest.mark.asyncio + async def test_open_passes_if_generation_match(self, mock_appendable_writer): + writer = self._make_one( + mock_appendable_writer["mock_client"], if_generation_match=123 + ) + assert writer.if_generation_match == 123 + + await writer.open() + + mock_stream_cls = mock_appendable_writer["mock_stream_cls"] + assert mock_stream_cls.call_args.kwargs["if_generation_match"] == 123 + assert mock_stream_cls.call_args.kwargs["generation_number"] is None + def test_on_open_error_redirection(self, mock_appendable_writer): """Verify redirect info is extracted from helper.""" writer = self._make_one(mock_appendable_writer["mock_client"]) diff --git a/packages/google-cloud-storage/tests/unit/asyncio/test_async_write_object_stream.py b/packages/google-cloud-storage/tests/unit/asyncio/test_async_write_object_stream.py index 049c9f6608a9..827590e024f8 100644 --- a/packages/google-cloud-storage/tests/unit/asyncio/test_async_write_object_stream.py +++ b/packages/google-cloud-storage/tests/unit/asyncio/test_async_write_object_stream.py @@ -186,6 +186,61 @@ async def test_open_existing_object_with_token(self, mock_rpc_cls, mock_client): assert initial_request.append_object_spec.routing_token == "token-123" assert stream.persisted_size == 1024 + @mock.patch("google.cloud.storage.asyncio.async_write_object_stream.AsyncBidiRpc") + @pytest.mark.asyncio + @pytest.mark.parametrize( + "generation_number, if_generation_match, expected", + [(None, None, None), (0, None, 0), (None, 0, 0), (None, 123, 123)], + ) + async def test_open_new_object_if_generation_match( + self, + mock_rpc_cls, + mock_client, + generation_number, + if_generation_match, + expected, + ): + mock_rpc = mock_rpc_cls.return_value + mock_rpc.open = AsyncMock() + mock_rpc.recv = AsyncMock(return_value=MagicMock(resource=None)) + + stream = _AsyncWriteObjectStream( + mock_client, + BUCKET, + OBJECT, + generation_number=generation_number, + if_generation_match=if_generation_match, + ) + await stream.open() + + spec = mock_rpc_cls.call_args.kwargs["initial_request"].write_object_spec + assert spec.appendable + if expected is None: + assert "if_generation_match" not in spec + else: + assert spec.if_generation_match == expected + + @mock.patch("google.cloud.storage.asyncio.async_write_object_stream.AsyncBidiRpc") + @pytest.mark.asyncio + async def test_open_existing_object_ignores_if_generation_match( + self, mock_rpc_cls, mock_client + ): + mock_rpc = mock_rpc_cls.return_value + mock_rpc.open = AsyncMock() + mock_rpc.recv = AsyncMock(return_value=MagicMock(resource=None)) + + stream = _AsyncWriteObjectStream( + mock_client, + BUCKET, + OBJECT, + generation_number=GENERATION, + if_generation_match=123, + ) + await stream.open() + + initial_request = mock_rpc_cls.call_args.kwargs["initial_request"] + assert initial_request.append_object_spec.generation == GENERATION + @mock.patch("google.cloud.storage.asyncio.async_write_object_stream.AsyncBidiRpc") @pytest.mark.asyncio async def test_open_metadata_merging(self, mock_rpc_cls, mock_client): From 092eb437bf545e653e3717659775e1b066edc620 Mon Sep 17 00:00:00 2001 From: yuxin00j Date: Thu, 29 Oct 2026 07:44:31 +0000 Subject: [PATCH 2/6] feat(storage): replace if_generation_match with Preconditions dataclass on AsyncAppendableObjectWriter --- .../asyncio/async_appendable_object_writer.py | 75 +++++- .../asyncio/async_write_object_stream.py | 62 +++-- .../tests/system/test_rapid.py | 66 +++++ .../test_async_appendable_object_writer.py | 237 +++++++++++++++++- .../asyncio/test_async_write_object_stream.py | 69 ++++- 5 files changed, 469 insertions(+), 40 deletions(-) diff --git a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_appendable_object_writer.py b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_appendable_object_writer.py index b4a7fcfed5ca..cd797333fe64 100644 --- a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_appendable_object_writer.py +++ b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_appendable_object_writer.py @@ -14,6 +14,7 @@ import io import logging +from dataclasses import dataclass from io import BufferedReader from typing import Dict, List, Optional, Tuple, Union @@ -100,6 +101,44 @@ def _is_write_retryable(exc): return False +@dataclass +class Preconditions: + """Preconditions for :class:`AsyncAppendableObjectWriter` operations. + + Generation preconditions (``if_generation_match``, + ``if_generation_not_match``) apply only when creating or overwriting an + object (``WriteObjectSpec``). Metageneration preconditions + (``if_metageneration_match``, ``if_metageneration_not_match``) apply to both + creating/overwriting (``WriteObjectSpec``) and appending to an existing + appendable object (``AppendObjectSpec``). + + :type if_generation_match: Optional[int] + :param if_generation_match: (Optional) Makes the operation conditional on + whether the object's current generation matches the given value. Setting + to ``0`` makes the operation succeed only if there are no live versions + of the object. + + :type if_generation_not_match: Optional[int] + :param if_generation_not_match: (Optional) Makes the operation conditional + on whether the object's current generation does not match the given + value. If no live object exists, the precondition fails. + + :type if_metageneration_match: Optional[int] + :param if_metageneration_match: (Optional) Makes the operation conditional + on whether the object's current metageneration matches the given value. + + :type if_metageneration_not_match: Optional[int] + :param if_metageneration_not_match: (Optional) Makes the operation + conditional on whether the object's current metageneration does not + match the given value. + """ + + if_generation_match: Optional[int] = None + if_generation_not_match: Optional[int] = None + if_metageneration_match: Optional[int] = None + if_metageneration_not_match: Optional[int] = None + + class AsyncAppendableObjectWriter: """Class for appending data to a GCS Appendable Object asynchronously.""" @@ -112,7 +151,7 @@ def __init__( write_handle: Optional[_storage_v2.BidiWriteHandle] = None, writer_options: Optional[dict] = None, storage_class: Optional[str] = None, - if_generation_match: Optional[int] = None, + preconditions: Optional[Preconditions] = None, ): """ Class for appending data to a GCS Appendable Object. @@ -166,7 +205,8 @@ def __init__( object doesn't exist in the bucket (useful for not accidentally overwriting existing objects). - Warning: If `None`, a new object is created. If an object with the + Warning: If `None` and no generation precondition is set in + `preconditions`, a new object is created. If an object with the same name already exists, it will be overwritten the moment `writer.open()` is called. @@ -186,23 +226,31 @@ def __init__( If specified, it overrides the bucket's `storage_class`. If not, object storage class will be the same as bucket's storage_class. - :type if_generation_match: Optional[int] - :param if_generation_match: (Optional) Only create the object if its - current generation matches this value; ``0`` means the object must - not exist. Use this for compare-and-swap overwrites of an existing - object. Ignored when ``generation`` selects an existing object to - append to. A mismatch raises + :type preconditions: Optional[:class:`Preconditions`] + :param preconditions: (Optional) Preconditions for the write operation. + Generation preconditions (``if_generation_match``, + ``if_generation_not_match``) apply only when creating or overwriting + an object (``generation`` is ``None`` or ``0``) and are ignored when + appending to an existing object (``generation > 0``). Metageneration + preconditions (``if_metageneration_match``, + ``if_metageneration_not_match``) apply to both creating/overwriting + and appending. A mismatch raises :class:`google.api_core.exceptions.FailedPrecondition` from ``open()``. + + :raises TypeError: If ``preconditions`` is not a :class:`Preconditions` + instance. """ _utils.raise_if_no_fast_crc32c() + if preconditions is not None and not isinstance(preconditions, Preconditions): + raise TypeError("preconditions must be an instance of Preconditions") self.client = client self.bucket_name = bucket_name self.object_name = object_name self.write_handle = write_handle self.generation = generation self.storage_class = storage_class - self.if_generation_match = if_generation_match + self.preconditions = preconditions self.write_obj_stream: Optional[_AsyncWriteObjectStream] = None self._is_stream_open: bool = False @@ -238,6 +286,7 @@ def from_blob( blob: Blob, write_handle: Optional[_storage_v2.BidiWriteHandle] = None, writer_options: Optional[dict] = None, + preconditions: Optional[Preconditions] = None, ) -> "AsyncAppendableObjectWriter": """Creates an AsyncAppendableObjectWriter from an existing Blob object. @@ -270,6 +319,9 @@ def from_blob( :param writer_options: (Optional) Configuration settings for the underlying appendable writer. + :type preconditions: Optional[:class:`Preconditions`] + :param preconditions: (Optional) Preconditions for the write operation. + :rtype: :class:`AsyncAppendableObjectWriter` :returns: An initialized writer instance. """ @@ -281,6 +333,7 @@ def from_blob( write_handle=write_handle, writer_options=writer_options, storage_class=blob.storage_class, + preconditions=preconditions, ) instance.blob = blob return instance @@ -349,6 +402,8 @@ async def open( """Opens the underlying bidi-gRPC stream. :raises ValueError: If the stream is already open. + :raises google.api_core.exceptions.FailedPrecondition: If a precondition + in ``preconditions`` (or ``generation=0``) is not met. """ if self._is_stream_open: @@ -380,7 +435,7 @@ async def _do_open(): write_handle=self.write_handle, routing_token=self._routing_token, storage_class=self.storage_class, - if_generation_match=self.if_generation_match, + preconditions=self.preconditions, ) if self._routing_token: diff --git a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_write_object_stream.py b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_write_object_stream.py index 25cda251fa7b..ab1c57c44976 100644 --- a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_write_object_stream.py +++ b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_write_object_stream.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import List, Optional, Tuple +from typing import TYPE_CHECKING, List, Optional, Tuple import grpc from google.api_core.bidi_async import AsyncBidiRpc @@ -25,6 +25,11 @@ ) from google.cloud.storage.asyncio.async_grpc_client import AsyncGrpcClient +if TYPE_CHECKING: + from google.cloud.storage.asyncio.async_appendable_object_writer import ( + Preconditions, + ) + class _AsyncWriteObjectStream(_AsyncAbstractObjectStream): """Class representing a gRPC bidi-stream for writing data from a GCS @@ -51,7 +56,8 @@ class _AsyncWriteObjectStream(_AsyncAbstractObjectStream): object doesn't exist in the bucket (useful for not accidentally overwriting existing objects). - Warning: If `None`, a new object is created. If an object with the + Warning: If `None` and no generation precondition is set in + `preconditions`, a new object is created. If an object with the same name already exists, it will be overwritten the moment `writer.open()` is called. @@ -62,11 +68,11 @@ class _AsyncWriteObjectStream(_AsyncAbstractObjectStream): :type storage_class: Optional[str] :param storage_class: (Optional) The storage class of the object. - :type if_generation_match: Optional[int] - :param if_generation_match: (Optional) Precondition for creating a new - object: the write succeeds only if the object's current generation - matches this value (``0`` means the object must not exist). Ignored - when ``generation_number`` selects an existing object to append to. + :type preconditions: Optional[:class:`~google.cloud.storage.asyncio.async_appendable_object_writer.Preconditions`] + :param preconditions: (Optional) Preconditions for the write operation. + Generation preconditions apply only when creating or overwriting an + object (``generation_number`` is ``None`` or ``0``); metageneration + preconditions apply to both creating/overwriting and appending. """ def __init__( @@ -79,7 +85,7 @@ def __init__( routing_token: Optional[str] = None, blob: Optional[Blob] = None, storage_class: Optional[str] = None, - if_generation_match: Optional[int] = None, + preconditions: Optional["Preconditions"] = None, ) -> None: if client is None: raise ValueError("client must be provided") @@ -98,7 +104,7 @@ def __init__( self.routing_token: Optional[str] = routing_token self.blob: Optional[Blob] = blob self.storage_class: Optional[str] = storage_class - self.if_generation_match: Optional[int] = if_generation_match + self.preconditions: Optional["Preconditions"] = preconditions self._full_bucket_name = f"projects/_/buckets/{self.bucket_name}" self.rpc = self.client._client._transport._wrapped_methods[ @@ -122,8 +128,8 @@ async def open(self, metadata: Optional[List[Tuple[str, str]]] = None) -> None: :rtype: None :raises ValueError: If the stream is already open. :raises google.api_core.exceptions.FailedPrecondition: - if `if_generation_match` (or `generation_number` is 0) does not - match the object's current generation. + If a precondition in `preconditions` (or `generation_number=0`) is + not met. """ if self._is_stream_open: raise ValueError("Stream is already open") @@ -142,14 +148,32 @@ async def open(self, metadata: Optional[List[Tuple[str, str]]] = None) -> None: bucket=self._full_bucket_name, storage_class=self.storage_class, ) - if_generation_match = self.if_generation_match - if if_generation_match is None: - if_generation_match = self.generation_number + if_generation_match = ( + self.preconditions.if_generation_match + if self.preconditions + and self.preconditions.if_generation_match is not None + else self.generation_number + ) self.first_bidi_write_req = _storage_v2.BidiWriteObjectRequest( write_object_spec=_storage_v2.WriteObjectSpec( resource=resource, appendable=True, if_generation_match=if_generation_match, + if_generation_not_match=( + self.preconditions.if_generation_not_match + if self.preconditions + else None + ), + if_metageneration_match=( + self.preconditions.if_metageneration_match + if self.preconditions + else None + ), + if_metageneration_not_match=( + self.preconditions.if_metageneration_not_match + if self.preconditions + else None + ), ), ) else: @@ -160,6 +184,16 @@ async def open(self, metadata: Optional[List[Tuple[str, str]]] = None) -> None: generation=self.generation_number, write_handle=self.write_handle if self.write_handle else None, routing_token=self.routing_token if self.routing_token else None, + if_metageneration_match=( + self.preconditions.if_metageneration_match + if self.preconditions + else None + ), + if_metageneration_not_match=( + self.preconditions.if_metageneration_not_match + if self.preconditions + else None + ), ), ) diff --git a/packages/google-cloud-storage/tests/system/test_rapid.py b/packages/google-cloud-storage/tests/system/test_rapid.py index 52a3c1fba464..5a09e2f55ac3 100644 --- a/packages/google-cloud-storage/tests/system/test_rapid.py +++ b/packages/google-cloud-storage/tests/system/test_rapid.py @@ -30,6 +30,7 @@ from google.cloud.storage.asyncio.async_appendable_object_writer import ( _DEFAULT_FLUSH_INTERVAL_BYTES, AsyncAppendableObjectWriter, + Preconditions, ) from google.cloud.storage.asyncio.async_grpc_client import AsyncGrpcClient from google.cloud.storage.asyncio.async_multi_range_downloader import ( @@ -825,6 +826,71 @@ async def _run(): event_loop.run_until_complete(_run()) +def test_open_with_preconditions( + storage_client, blobs_to_delete, event_loop, grpc_client +): + """Tests create-only and compare-and-swap writes using `Preconditions`.""" + object_name = f"test_append_with_preconditions-{uuid.uuid4()}" + + async def _run(): + writer = AsyncAppendableObjectWriter( + grpc_client, + _BUCKET_UNDER_TEST, + object_name, + preconditions=Preconditions(if_generation_match=0), + storage_class="RAPID" if RCU_SYSTEM_TESTS else None, + ) + await writer.open() + assert writer.is_stream_open + gen_1 = writer.generation + await writer.close() + assert not writer.is_stream_open + + with pytest.raises(FailedPrecondition) as exc_info: + writer_create_fail = AsyncAppendableObjectWriter( + grpc_client, + _BUCKET_UNDER_TEST, + object_name, + preconditions=Preconditions(if_generation_match=0), + storage_class="RAPID" if RCU_SYSTEM_TESTS else None, + ) + await writer_create_fail.open() + assert exc_info.value.code == 400 + + writer_cas = AsyncAppendableObjectWriter( + grpc_client, + _BUCKET_UNDER_TEST, + object_name, + preconditions=Preconditions(if_generation_match=gen_1), + storage_class="RAPID" if RCU_SYSTEM_TESTS else None, + ) + await writer_cas.open() + assert writer_cas.generation != gen_1 + await writer_cas.close() + + with pytest.raises(FailedPrecondition) as exc_info: + writer_stale_cas = AsyncAppendableObjectWriter( + grpc_client, + _BUCKET_UNDER_TEST, + object_name, + preconditions=Preconditions(if_generation_match=gen_1), + storage_class="RAPID" if RCU_SYSTEM_TESTS else None, + ) + await writer_stale_cas.open() + assert exc_info.value.code == 400 + + # cleanup + del writer + del writer_cas + gc.collect() + + blobs_to_delete.append( + storage_client.bucket(_BUCKET_UNDER_TEST).blob(object_name) + ) + + event_loop.run_until_complete(_run()) + + def test_open_existing_object_with_gen_None_overrides_existing( storage_client, blobs_to_delete, event_loop, grpc_client ): diff --git a/packages/google-cloud-storage/tests/unit/asyncio/test_async_appendable_object_writer.py b/packages/google-cloud-storage/tests/unit/asyncio/test_async_appendable_object_writer.py index ec6175b8d52c..88a548276bd6 100644 --- a/packages/google-cloud-storage/tests/unit/asyncio/test_async_appendable_object_writer.py +++ b/packages/google-cloud-storage/tests/unit/asyncio/test_async_appendable_object_writer.py @@ -27,6 +27,7 @@ _DEFAULT_FLUSH_INTERVAL_BYTES, _MAX_CHUNK_SIZE_BYTES, AsyncAppendableObjectWriter, + Preconditions, _is_write_retryable, ) @@ -173,7 +174,7 @@ async def test_init_with_storage_class(self, mock_appendable_writer, storage_cla write_handle=None, routing_token=None, storage_class=storage_class, - if_generation_match=None, + preconditions=None, ) assert writer.write_obj_stream.storage_class == storage_class @@ -273,10 +274,61 @@ async def test_from_blob_open_passes_storage_class( write_handle=None, routing_token=None, storage_class=storage_class, - if_generation_match=None, + preconditions=None, ) assert writer.write_obj_stream.storage_class == storage_class + def test_init_invalid_preconditions_type_raises(self, mock_appendable_writer): + with pytest.raises( + TypeError, match="preconditions must be an instance of Preconditions" + ): + self._make_one( + mock_appendable_writer["mock_client"], + preconditions={"if_generation_match": 0}, + ) + + def test_init_generation_with_preconditions(self, mock_appendable_writer): + preconditions = Preconditions( + if_generation_match=123, + if_metageneration_match=1, + if_metageneration_not_match=2, + ) + writer = self._make_one( + mock_appendable_writer["mock_client"], + generation=GENERATION, + preconditions=preconditions, + ) + assert writer.generation == GENERATION + assert writer.preconditions == preconditions + + @pytest.mark.asyncio + @pytest.mark.parametrize("blob_generation", [None, GENERATION]) + async def test_from_blob_with_preconditions( + self, mock_appendable_writer, blob_generation + ): + mock_blob = mock.Mock(spec=Blob) + mock_blob.name = OBJECT + mock_blob.bucket.name = BUCKET + mock_blob.generation = blob_generation + mock_blob.storage_class = "RAPID" + + preconditions = Preconditions( + if_generation_match=123, if_metageneration_match=1 + ) + writer = AsyncAppendableObjectWriter.from_blob( + mock_appendable_writer["mock_client"], + mock_blob, + preconditions=preconditions, + ) + assert writer.preconditions == preconditions + + await writer.open() + + mock_stream_cls = mock_appendable_writer["mock_stream_cls"] + assert mock_stream_cls.call_args.kwargs["preconditions"] == preconditions + assert mock_stream_cls.call_args.kwargs["generation_number"] == blob_generation + assert mock_stream_cls.call_args.kwargs["blob"] == mock_blob + # ------------------------------------------------------------------------- # Stream Lifecycle Tests # ------------------------------------------------------------------------- @@ -315,18 +367,191 @@ async def test_open_success(self, mock_appendable_writer): assert writer.write_obj_stream.storage_class is None @pytest.mark.asyncio - async def test_open_passes_if_generation_match(self, mock_appendable_writer): + async def test_open_passes_preconditions(self, mock_appendable_writer): + preconditions = Preconditions(if_generation_match=123) writer = self._make_one( - mock_appendable_writer["mock_client"], if_generation_match=123 + mock_appendable_writer["mock_client"], + preconditions=preconditions, ) - assert writer.if_generation_match == 123 + assert writer.preconditions == preconditions await writer.open() mock_stream_cls = mock_appendable_writer["mock_stream_cls"] - assert mock_stream_cls.call_args.kwargs["if_generation_match"] == 123 + assert mock_stream_cls.call_args.kwargs["preconditions"] == preconditions assert mock_stream_cls.call_args.kwargs["generation_number"] is None + @pytest.mark.asyncio + async def test_open_failed_precondition_not_retried(self, mock_appendable_writer): + preconditions = Preconditions(if_generation_match=123) + writer = self._make_one( + mock_appendable_writer["mock_client"], + preconditions=preconditions, + ) + mock_appendable_writer[ + "mock_stream" + ].open.side_effect = exceptions.FailedPrecondition("precondition failed") + + with pytest.raises(exceptions.FailedPrecondition): + await writer.open() + + mock_appendable_writer["mock_stream"].open.assert_awaited_once() + assert not writer.is_stream_open + + @pytest.mark.asyncio + async def test_open_redirect_without_generation_preserves_preconditions(self): + mock_client = mock.MagicMock() + mock_transport = mock.MagicMock() + mock_transport.bidi_write_object = mock.sentinel.bidi_write_object + mock_transport._wrapped_methods = { + mock.sentinel.bidi_write_object: mock.sentinel.wrapped_bidi_write_object + } + mock_client.grpc_client._client._transport = mock_transport + + preconditions = Preconditions( + if_generation_match=123, if_metageneration_match=2 + ) + writer = self._make_one(mock_client, preconditions=preconditions) + + redirect = BidiWriteObjectRedirectedError(routing_token="rt-1") + redirect_exc = exceptions.Aborted("redirect", errors=[redirect]) + open_resp = storage_type.BidiWriteObjectResponse( + resource=storage_type.Object(generation=888, size=0), + write_handle=storage_type.BidiWriteHandle(handle=b"wh-1"), + ) + + with mock.patch( + "google.cloud.storage.asyncio.async_write_object_stream.AsyncBidiRpc" + ) as mock_rpc_cls: + rpc_1 = mock.MagicMock() + rpc_1.open = AsyncMock(side_effect=redirect_exc) + rpc_2 = mock.MagicMock() + rpc_2.open = AsyncMock() + rpc_2.recv = AsyncMock(return_value=open_resp) + mock_rpc_cls.side_effect = [rpc_1, rpc_2] + + await writer.open() + + assert mock_rpc_cls.call_count == 2 + for call in mock_rpc_cls.call_args_list: + spec = call.kwargs["initial_request"].write_object_spec + assert spec.appendable + assert "if_generation_match" in spec + assert spec.if_generation_match == 123 + assert "if_metageneration_match" in spec + assert spec.if_metageneration_match == 2 + + second_meta = dict(mock_rpc_cls.call_args_list[1].kwargs["metadata"]) + assert "routing_token=rt-1" in second_meta["x-goog-request-params"] + assert writer.generation == 888 + + @pytest.mark.asyncio + async def test_open_redirect_with_generation_switches_to_append_spec(self): + mock_client = mock.MagicMock() + mock_transport = mock.MagicMock() + mock_transport.bidi_write_object = mock.sentinel.bidi_write_object + mock_transport._wrapped_methods = { + mock.sentinel.bidi_write_object: mock.sentinel.wrapped_bidi_write_object + } + mock_client.grpc_client._client._transport = mock_transport + + preconditions = Preconditions( + if_generation_match=123, if_metageneration_match=2 + ) + writer = self._make_one(mock_client, preconditions=preconditions) + + redirect = BidiWriteObjectRedirectedError( + routing_token="rt-2", + write_handle=storage_type.BidiWriteHandle(handle=b"wh-2"), + generation=777, + ) + redirect_exc = exceptions.Aborted("redirect", errors=[redirect]) + open_resp = storage_type.BidiWriteObjectResponse( + persisted_size=0, + write_handle=storage_type.BidiWriteHandle(handle=b"wh-2"), + ) + + with mock.patch( + "google.cloud.storage.asyncio.async_write_object_stream.AsyncBidiRpc" + ) as mock_rpc_cls: + rpc_1 = mock.MagicMock() + rpc_1.open = AsyncMock(side_effect=redirect_exc) + rpc_2 = mock.MagicMock() + rpc_2.open = AsyncMock() + rpc_2.recv = AsyncMock(return_value=open_resp) + mock_rpc_cls.side_effect = [rpc_1, rpc_2] + + await writer.open() + + assert mock_rpc_cls.call_count == 2 + first_req = mock_rpc_cls.call_args_list[0].kwargs["initial_request"] + assert "write_object_spec" in first_req + assert "if_generation_match" in first_req.write_object_spec + assert first_req.write_object_spec.if_generation_match == 123 + + second_req = mock_rpc_cls.call_args_list[1].kwargs["initial_request"] + assert "append_object_spec" in second_req + assert second_req.append_object_spec.generation == 777 + assert second_req.append_object_spec.routing_token == "rt-2" + assert "if_metageneration_match" in second_req.append_object_spec + assert second_req.append_object_spec.if_metageneration_match == 2 + assert writer.generation == 777 + + @pytest.mark.asyncio + async def test_reopen_after_create_sends_append_spec(self): + mock_client = mock.MagicMock() + mock_transport = mock.MagicMock() + mock_transport.bidi_write_object = mock.sentinel.bidi_write_object + mock_transport._wrapped_methods = { + mock.sentinel.bidi_write_object: mock.sentinel.wrapped_bidi_write_object + } + mock_client.grpc_client._client._transport = mock_transport + + preconditions = Preconditions(if_generation_match=0, if_metageneration_match=1) + writer = self._make_one(mock_client, preconditions=preconditions) + + first_open_resp = storage_type.BidiWriteObjectResponse( + resource=storage_type.Object(generation=999, size=0), + write_handle=storage_type.BidiWriteHandle(handle=b"wh-init"), + ) + reopen_resp = storage_type.BidiWriteObjectResponse( + persisted_size=0, + write_handle=storage_type.BidiWriteHandle(handle=b"wh-reopen"), + ) + + with mock.patch( + "google.cloud.storage.asyncio.async_write_object_stream.AsyncBidiRpc" + ) as mock_rpc_cls: + rpc_1 = mock.MagicMock() + rpc_1.open = AsyncMock() + rpc_1.recv = AsyncMock(return_value=first_open_resp) + rpc_1.close = AsyncMock() + rpc_1.send = AsyncMock() + rpc_2 = mock.MagicMock() + rpc_2.open = AsyncMock() + rpc_2.recv = AsyncMock(return_value=reopen_resp) + mock_rpc_cls.side_effect = [rpc_1, rpc_2] + + await writer.open() + assert writer.generation == 999 + + # Simulate reconnect re-opening the writer + writer.write_obj_stream._is_stream_open = False + writer._is_stream_open = False + await writer.open() + + assert mock_rpc_cls.call_count == 2 + first_req = mock_rpc_cls.call_args_list[0].kwargs["initial_request"] + assert "write_object_spec" in first_req + assert "if_generation_match" in first_req.write_object_spec + assert first_req.write_object_spec.if_generation_match == 0 + + second_req = mock_rpc_cls.call_args_list[1].kwargs["initial_request"] + assert "append_object_spec" in second_req + assert second_req.append_object_spec.generation == 999 + assert "if_metageneration_match" in second_req.append_object_spec + assert second_req.append_object_spec.if_metageneration_match == 1 + def test_on_open_error_redirection(self, mock_appendable_writer): """Verify redirect info is extracted from helper.""" writer = self._make_one(mock_appendable_writer["mock_client"]) diff --git a/packages/google-cloud-storage/tests/unit/asyncio/test_async_write_object_stream.py b/packages/google-cloud-storage/tests/unit/asyncio/test_async_write_object_stream.py index 827590e024f8..4645e3d48aae 100644 --- a/packages/google-cloud-storage/tests/unit/asyncio/test_async_write_object_stream.py +++ b/packages/google-cloud-storage/tests/unit/asyncio/test_async_write_object_stream.py @@ -21,6 +21,9 @@ from google.cloud import _storage_v2 from google.cloud.storage import Blob, Bucket +from google.cloud.storage.asyncio.async_appendable_object_writer import ( + Preconditions, +) from google.cloud.storage.asyncio.async_write_object_stream import ( _AsyncWriteObjectStream, ) @@ -189,15 +192,20 @@ async def test_open_existing_object_with_token(self, mock_rpc_cls, mock_client): @mock.patch("google.cloud.storage.asyncio.async_write_object_stream.AsyncBidiRpc") @pytest.mark.asyncio @pytest.mark.parametrize( - "generation_number, if_generation_match, expected", - [(None, None, None), (0, None, 0), (None, 0, 0), (None, 123, 123)], + "generation_number, preconditions, expected", + [ + (None, None, None), + (0, None, 0), + (None, Preconditions(if_generation_match=0), 0), + (None, Preconditions(if_generation_match=123), 123), + ], ) async def test_open_new_object_if_generation_match( self, mock_rpc_cls, mock_client, generation_number, - if_generation_match, + preconditions, expected, ): mock_rpc = mock_rpc_cls.return_value @@ -209,7 +217,7 @@ async def test_open_new_object_if_generation_match( BUCKET, OBJECT, generation_number=generation_number, - if_generation_match=if_generation_match, + preconditions=preconditions, ) await stream.open() @@ -218,28 +226,69 @@ async def test_open_new_object_if_generation_match( if expected is None: assert "if_generation_match" not in spec else: + assert "if_generation_match" in spec assert spec.if_generation_match == expected @mock.patch("google.cloud.storage.asyncio.async_write_object_stream.AsyncBidiRpc") @pytest.mark.asyncio - async def test_open_existing_object_ignores_if_generation_match( - self, mock_rpc_cls, mock_client - ): + async def test_open_new_object_all_preconditions(self, mock_rpc_cls, mock_client): mock_rpc = mock_rpc_cls.return_value mock_rpc.open = AsyncMock() mock_rpc.recv = AsyncMock(return_value=MagicMock(resource=None)) + preconditions = Preconditions( + if_generation_match=10, + if_generation_not_match=20, + if_metageneration_match=3, + if_metageneration_not_match=4, + ) stream = _AsyncWriteObjectStream( mock_client, BUCKET, OBJECT, - generation_number=GENERATION, + preconditions=preconditions, + ) + await stream.open() + + spec = mock_rpc_cls.call_args.kwargs["initial_request"].write_object_spec + assert spec.appendable + assert "if_generation_match" in spec + assert spec.if_generation_match == 10 + assert "if_generation_not_match" in spec + assert spec.if_generation_not_match == 20 + assert "if_metageneration_match" in spec + assert spec.if_metageneration_match == 3 + assert "if_metageneration_not_match" in spec + assert spec.if_metageneration_not_match == 4 + + @mock.patch("google.cloud.storage.asyncio.async_write_object_stream.AsyncBidiRpc") + @pytest.mark.asyncio + async def test_open_existing_object_preconditions(self, mock_rpc_cls, mock_client): + mock_rpc = mock_rpc_cls.return_value + mock_rpc.open = AsyncMock() + mock_rpc.recv = AsyncMock(return_value=MagicMock(resource=None)) + + preconditions = Preconditions( if_generation_match=123, + if_generation_not_match=456, + if_metageneration_match=2, + if_metageneration_not_match=3, + ) + stream = _AsyncWriteObjectStream( + mock_client, + BUCKET, + OBJECT, + generation_number=GENERATION, + preconditions=preconditions, ) await stream.open() - initial_request = mock_rpc_cls.call_args.kwargs["initial_request"] - assert initial_request.append_object_spec.generation == GENERATION + spec = mock_rpc_cls.call_args.kwargs["initial_request"].append_object_spec + assert spec.generation == GENERATION + assert "if_metageneration_match" in spec + assert spec.if_metageneration_match == 2 + assert "if_metageneration_not_match" in spec + assert spec.if_metageneration_not_match == 3 @mock.patch("google.cloud.storage.asyncio.async_write_object_stream.AsyncBidiRpc") @pytest.mark.asyncio From b90cbfbb17b4f603926eb4d282f0f61aa497e9ca Mon Sep 17 00:00:00 2001 From: yuxin00j Date: Fri, 30 Oct 2026 03:29:44 +0000 Subject: [PATCH 3/6] fix(storage): reject generation preconditions with generation and stop resending preconditions after create --- .../asyncio/async_appendable_object_writer.py | 43 +++++-- .../test_async_appendable_object_writer.py | 113 ++++++++++++++---- 2 files changed, 122 insertions(+), 34 deletions(-) diff --git a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_appendable_object_writer.py b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_appendable_object_writer.py index cd797333fe64..7f4cb5a95e1c 100644 --- a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_appendable_object_writer.py +++ b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_appendable_object_writer.py @@ -229,21 +229,36 @@ def __init__( :type preconditions: Optional[:class:`Preconditions`] :param preconditions: (Optional) Preconditions for the write operation. Generation preconditions (``if_generation_match``, - ``if_generation_not_match``) apply only when creating or overwriting - an object (``generation`` is ``None`` or ``0``) and are ignored when - appending to an existing object (``generation > 0``). Metageneration - preconditions (``if_metageneration_match``, - ``if_metageneration_not_match``) apply to both creating/overwriting - and appending. A mismatch raises + ``if_generation_not_match``) require ``generation`` to be ``None``. + Metageneration preconditions (``if_metageneration_match``, + ``if_metageneration_not_match``) apply to the existing object: the + one being overwritten, or the one being appended to when + ``generation > 0``. A mismatch raises :class:`google.api_core.exceptions.FailedPrecondition` from - ``open()``. + ``open()``. Preconditions are not re-sent when the writer reconnects + to a generation it created. :raises TypeError: If ``preconditions`` is not a :class:`Preconditions` instance. + :raises ValueError: If ``generation`` is set together with a generation + precondition in ``preconditions``. """ _utils.raise_if_no_fast_crc32c() if preconditions is not None and not isinstance(preconditions, Preconditions): raise TypeError("preconditions must be an instance of Preconditions") + if ( + generation is not None + and preconditions is not None + and ( + preconditions.if_generation_match is not None + or preconditions.if_generation_not_match is not None + ) + ): + raise ValueError( + "if_generation_match/if_generation_not_match require generation=None; " + "generation=0 already means if_generation_match=0 and generation > 0 " + "(including from_blob's blob.generation) appends to an existing object." + ) self.client = client self.bucket_name = bucket_name self.object_name = object_name @@ -251,6 +266,9 @@ def __init__( self.generation = generation self.storage_class = storage_class self.preconditions = preconditions + # Set once this writer creates a new generation; `preconditions` then no + # longer apply to the object being reconnected to. + self._created_generation = False self.write_obj_stream: Optional[_AsyncWriteObjectStream] = None self._is_stream_open: bool = False @@ -321,9 +339,12 @@ def from_blob( :type preconditions: Optional[:class:`Preconditions`] :param preconditions: (Optional) Preconditions for the write operation. + Generation preconditions require ``blob.generation`` to be ``None``. :rtype: :class:`AsyncAppendableObjectWriter` :returns: An initialized writer instance. + :raises ValueError: If ``blob.generation`` is set together with a + generation precondition in ``preconditions``. """ instance = cls( client=client, @@ -368,6 +389,9 @@ def _on_open_error(self, exc): if redirect_proto.write_handle: self.write_handle = redirect_proto.write_handle if redirect_proto.generation: + if not self.generation: + # The redirected create already produced this generation. + self._created_generation = True self.generation = redirect_proto.generation def _merge_retry_policy( @@ -413,6 +437,7 @@ async def open( async def _do_open(): current_metadata = list(metadata) if metadata else [] + creating = not self.generation # Cleanup stream from previous failed attempt, if any. if self.write_obj_stream: @@ -435,7 +460,7 @@ async def _do_open(): write_handle=self.write_handle, routing_token=self._routing_token, storage_class=self.storage_class, - preconditions=self.preconditions, + preconditions=None if self._created_generation else self.preconditions, ) if self._routing_token: @@ -448,6 +473,8 @@ async def _do_open(): ) if self.write_obj_stream.generation_number: + if creating: + self._created_generation = True self.generation = self.write_obj_stream.generation_number if self.write_obj_stream.write_handle: self.write_handle = self.write_obj_stream.write_handle diff --git a/packages/google-cloud-storage/tests/unit/asyncio/test_async_appendable_object_writer.py b/packages/google-cloud-storage/tests/unit/asyncio/test_async_appendable_object_writer.py index 88a548276bd6..501424161ea7 100644 --- a/packages/google-cloud-storage/tests/unit/asyncio/test_async_appendable_object_writer.py +++ b/packages/google-cloud-storage/tests/unit/asyncio/test_async_appendable_object_writer.py @@ -287,34 +287,34 @@ def test_init_invalid_preconditions_type_raises(self, mock_appendable_writer): preconditions={"if_generation_match": 0}, ) - def test_init_generation_with_preconditions(self, mock_appendable_writer): - preconditions = Preconditions( - if_generation_match=123, - if_metageneration_match=1, - if_metageneration_not_match=2, - ) - writer = self._make_one( - mock_appendable_writer["mock_client"], - generation=GENERATION, - preconditions=preconditions, - ) - assert writer.generation == GENERATION - assert writer.preconditions == preconditions + @pytest.mark.parametrize("generation", [0, GENERATION]) + @pytest.mark.parametrize( + "preconditions", + [ + Preconditions(if_generation_match=0), + Preconditions(if_generation_match=456), + Preconditions(if_generation_not_match=456), + ], + ) + def test_init_generation_and_generation_precondition_raises( + self, mock_appendable_writer, generation, preconditions + ): + with pytest.raises(ValueError, match="require generation=None"): + self._make_one( + mock_appendable_writer["mock_client"], + generation=generation, + preconditions=preconditions, + ) @pytest.mark.asyncio - @pytest.mark.parametrize("blob_generation", [None, GENERATION]) - async def test_from_blob_with_preconditions( - self, mock_appendable_writer, blob_generation - ): + async def test_from_blob_with_preconditions(self, mock_appendable_writer): mock_blob = mock.Mock(spec=Blob) mock_blob.name = OBJECT mock_blob.bucket.name = BUCKET - mock_blob.generation = blob_generation + mock_blob.generation = None mock_blob.storage_class = "RAPID" - preconditions = Preconditions( - if_generation_match=123, if_metageneration_match=1 - ) + preconditions = Preconditions(if_generation_match=456) writer = AsyncAppendableObjectWriter.from_blob( mock_appendable_writer["mock_client"], mock_blob, @@ -326,9 +326,25 @@ async def test_from_blob_with_preconditions( mock_stream_cls = mock_appendable_writer["mock_stream_cls"] assert mock_stream_cls.call_args.kwargs["preconditions"] == preconditions - assert mock_stream_cls.call_args.kwargs["generation_number"] == blob_generation + assert mock_stream_cls.call_args.kwargs["generation_number"] is None assert mock_stream_cls.call_args.kwargs["blob"] == mock_blob + def test_from_blob_with_generation_and_generation_precondition_raises( + self, mock_appendable_writer + ): + mock_blob = mock.Mock(spec=Blob) + mock_blob.name = OBJECT + mock_blob.bucket.name = BUCKET + mock_blob.generation = GENERATION + mock_blob.storage_class = "RAPID" + + with pytest.raises(ValueError, match="require generation=None"): + AsyncAppendableObjectWriter.from_blob( + mock_appendable_writer["mock_client"], + mock_blob, + preconditions=Preconditions(if_generation_match=456), + ) + # ------------------------------------------------------------------------- # Stream Lifecycle Tests # ------------------------------------------------------------------------- @@ -493,8 +509,9 @@ async def test_open_redirect_with_generation_switches_to_append_spec(self): assert "append_object_spec" in second_req assert second_req.append_object_spec.generation == 777 assert second_req.append_object_spec.routing_token == "rt-2" - assert "if_metageneration_match" in second_req.append_object_spec - assert second_req.append_object_spec.if_metageneration_match == 2 + # The metageneration precondition described the replaced object, not + # the generation created by this writer. + assert "if_metageneration_match" not in second_req.append_object_spec assert writer.generation == 777 @pytest.mark.asyncio @@ -549,8 +566,52 @@ async def test_reopen_after_create_sends_append_spec(self): second_req = mock_rpc_cls.call_args_list[1].kwargs["initial_request"] assert "append_object_spec" in second_req assert second_req.append_object_spec.generation == 999 - assert "if_metageneration_match" in second_req.append_object_spec - assert second_req.append_object_spec.if_metageneration_match == 1 + # The new generation starts at metageneration 1, so the precondition + # for the replaced object must not be re-sent. + assert "if_metageneration_match" not in second_req.append_object_spec + + @pytest.mark.asyncio + async def test_reopen_append_mode_keeps_metageneration_preconditions(self): + mock_client = mock.MagicMock() + mock_transport = mock.MagicMock() + mock_transport.bidi_write_object = mock.sentinel.bidi_write_object + mock_transport._wrapped_methods = { + mock.sentinel.bidi_write_object: mock.sentinel.wrapped_bidi_write_object + } + mock_client.grpc_client._client._transport = mock_transport + + writer = self._make_one( + mock_client, + generation=GENERATION, + preconditions=Preconditions(if_metageneration_match=3), + ) + open_resp = storage_type.BidiWriteObjectResponse( + resource=storage_type.Object(generation=GENERATION, size=0), + write_handle=storage_type.BidiWriteHandle(handle=b"wh"), + ) + + with mock.patch( + "google.cloud.storage.asyncio.async_write_object_stream.AsyncBidiRpc" + ) as mock_rpc_cls: + rpc_1 = mock.MagicMock() + rpc_1.open = AsyncMock() + rpc_1.recv = AsyncMock(return_value=open_resp) + rpc_2 = mock.MagicMock() + rpc_2.open = AsyncMock() + rpc_2.recv = AsyncMock(return_value=open_resp) + mock_rpc_cls.side_effect = [rpc_1, rpc_2] + + await writer.open() + # Simulate reconnect re-opening the writer + writer.write_obj_stream._is_stream_open = False + writer._is_stream_open = False + await writer.open() + + assert mock_rpc_cls.call_count == 2 + for call in mock_rpc_cls.call_args_list: + spec = call.kwargs["initial_request"].append_object_spec + assert spec.generation == GENERATION + assert spec.if_metageneration_match == 3 def test_on_open_error_redirection(self, mock_appendable_writer): """Verify redirect info is extracted from helper.""" From 8ce6bf81ebbf64fc0aa83943012856bf4224e272 Mon Sep 17 00:00:00 2001 From: yuxin00j Date: Fri, 30 Oct 2026 06:28:23 +0000 Subject: [PATCH 4/6] fix(storage): make Preconditions immutable --- .../cloud/storage/asyncio/async_appendable_object_writer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_appendable_object_writer.py b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_appendable_object_writer.py index 7f4cb5a95e1c..6955792cdcf1 100644 --- a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_appendable_object_writer.py +++ b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_appendable_object_writer.py @@ -101,7 +101,7 @@ def _is_write_retryable(exc): return False -@dataclass +@dataclass(frozen=True) class Preconditions: """Preconditions for :class:`AsyncAppendableObjectWriter` operations. From 15d221e5d4f2b64d2e1bff68aeb6c78417b8efbb Mon Sep 17 00:00:00 2001 From: yuxin00j Date: Fri, 30 Oct 2026 08:37:22 +0000 Subject: [PATCH 5/6] fix(storage): add from_blob generation override and reject write_handle with metageneration preconditions Passing generation=None to from_blob overwrites a fetched Blob (for example with if_generation_match=blob.generation) while keeping its metadata. Metageneration preconditions are rejected together with write_handle because GCS only checks them when the handle is empty. Also consolidates the precondition unit tests. --- .../asyncio/async_appendable_object_writer.py | 54 ++- .../tests/system/test_rapid.py | 100 ++++++ .../test_async_appendable_object_writer.py | 318 +++++++++--------- .../asyncio/test_async_write_object_stream.py | 69 ++-- 4 files changed, 321 insertions(+), 220 deletions(-) diff --git a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_appendable_object_writer.py b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_appendable_object_writer.py index 6955792cdcf1..da399659a343 100644 --- a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_appendable_object_writer.py +++ b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_appendable_object_writer.py @@ -50,6 +50,7 @@ _BIDI_WRITE_REDIRECTED_TYPE_URL = ( "type.googleapis.com/google.storage.v2.BidiWriteObjectRedirectedError" ) +_marker = object() logger = logging.getLogger(__name__) @@ -110,7 +111,9 @@ class Preconditions: object (``WriteObjectSpec``). Metageneration preconditions (``if_metageneration_match``, ``if_metageneration_not_match``) apply to both creating/overwriting (``WriteObjectSpec``) and appending to an existing - appendable object (``AppendObjectSpec``). + appendable object (``AppendObjectSpec``). GCS only checks metageneration + preconditions when no ``write_handle`` is sent, so they cannot be combined + with ``write_handle``. :type if_generation_match: Optional[int] :param if_generation_match: (Optional) Makes the operation conditional on @@ -233,15 +236,17 @@ def __init__( Metageneration preconditions (``if_metageneration_match``, ``if_metageneration_not_match``) apply to the existing object: the one being overwritten, or the one being appended to when - ``generation > 0``. A mismatch raises - :class:`google.api_core.exceptions.FailedPrecondition` from - ``open()``. Preconditions are not re-sent when the writer reconnects - to a generation it created. + ``generation > 0``. They cannot be combined with ``write_handle``, + because GCS only checks them when ``write_handle`` is empty. A + mismatch raises :class:`google.api_core.exceptions.FailedPrecondition` + from ``open()``. Preconditions are not re-sent when the writer + reconnects to a generation it created. :raises TypeError: If ``preconditions`` is not a :class:`Preconditions` instance. :raises ValueError: If ``generation`` is set together with a generation - precondition in ``preconditions``. + precondition, or ``write_handle`` together with a metageneration + precondition, in ``preconditions``. """ _utils.raise_if_no_fast_crc32c() if preconditions is not None and not isinstance(preconditions, Preconditions): @@ -257,7 +262,21 @@ def __init__( raise ValueError( "if_generation_match/if_generation_not_match require generation=None; " "generation=0 already means if_generation_match=0 and generation > 0 " - "(including from_blob's blob.generation) appends to an existing object." + "appends to an existing object. To overwrite using a fetched Blob, " + "pass generation=None to from_blob." + ) + if ( + write_handle + and preconditions is not None + and ( + preconditions.if_metageneration_match is not None + or preconditions.if_metageneration_not_match is not None + ) + ): + raise ValueError( + "if_metageneration_match/if_metageneration_not_match cannot be " + "combined with write_handle: GCS only checks metageneration " + "preconditions when write_handle is empty." ) self.client = client self.bucket_name = bucket_name @@ -305,6 +324,7 @@ def from_blob( write_handle: Optional[_storage_v2.BidiWriteHandle] = None, writer_options: Optional[dict] = None, preconditions: Optional[Preconditions] = None, + generation=_marker, ) -> "AsyncAppendableObjectWriter": """Creates an AsyncAppendableObjectWriter from an existing Blob object. @@ -339,18 +359,30 @@ def from_blob( :type preconditions: Optional[:class:`Preconditions`] :param preconditions: (Optional) Preconditions for the write operation. - Generation preconditions require ``blob.generation`` to be ``None``. + Generation preconditions require ``generation`` to be ``None``. + Warning: a Blob fetched from GCS has ``blob.generation`` set, so + unless ``generation=None`` is passed the writer appends to that + generation, including when only metageneration preconditions are + given. + + :type generation: Optional[int] + :param generation: (Optional) Overrides ``blob.generation``, which is + used by default. Pass ``None`` to overwrite the object using a + fetched Blob, for example with + ``Preconditions(if_generation_match=blob.generation)``. :rtype: :class:`AsyncAppendableObjectWriter` :returns: An initialized writer instance. - :raises ValueError: If ``blob.generation`` is set together with a - generation precondition in ``preconditions``. + :raises ValueError: If ``generation`` (``blob.generation`` by default) + is set together with a generation precondition, or + ``write_handle`` together with a metageneration precondition, in + ``preconditions``. """ instance = cls( client=client, bucket_name=blob.bucket.name, object_name=blob.name, - generation=blob.generation, + generation=blob.generation if generation is _marker else generation, write_handle=write_handle, writer_options=writer_options, storage_class=blob.storage_class, diff --git a/packages/google-cloud-storage/tests/system/test_rapid.py b/packages/google-cloud-storage/tests/system/test_rapid.py index 5a09e2f55ac3..b90011816419 100644 --- a/packages/google-cloud-storage/tests/system/test_rapid.py +++ b/packages/google-cloud-storage/tests/system/test_rapid.py @@ -891,6 +891,106 @@ async def _run(): event_loop.run_until_complete(_run()) +@pytest.mark.skipif( + RCU_SYSTEM_TESTS, reason="Write from blob for RCU not supported in SDK yet" +) +def test_from_blob_overwrite_fetched_blob_with_if_generation_match( + storage_client, blobs_to_delete, event_loop, grpc_client +): + """A fetched Blob is overwritten with CAS when `generation=None` is passed.""" + object_name = f"test_from_blob_cas-{uuid.uuid4()}" + + async def _run(): + blob = storage_client.bucket(_BUCKET_UNDER_TEST).blob(object_name) + blob.content_type = "text/plain" + blob.metadata = {"environment": "system-test"} + writer = AsyncAppendableObjectWriter.from_blob(grpc_client, blob) + await writer.open() + await writer.append(b"v1") + await writer.close(finalize_on_close=True) + blobs_to_delete.append(blob) + + fetched = storage_client.bucket(_BUCKET_UNDER_TEST).get_blob(object_name) + writer_cas = AsyncAppendableObjectWriter.from_blob( + grpc_client, + fetched, + generation=None, + preconditions=Preconditions(if_generation_match=fetched.generation), + ) + await writer_cas.open() + await writer_cas.append(b"v2") + await writer_cas.close(finalize_on_close=True) + + obj = await grpc_client.get_object( + bucket_name=_BUCKET_UNDER_TEST, object_name=object_name + ) + assert obj.generation != fetched.generation + assert obj.content_type == "text/plain" + assert obj.metadata["environment"] == "system-test" + + # cleanup + del writer + del writer_cas + gc.collect() + + event_loop.run_until_complete(_run()) + + +def test_append_with_metageneration_precondition( + storage_client, blobs_to_delete, event_loop, grpc_client +): + """Appending without a write_handle enforces metageneration preconditions.""" + object_name = f"test_append_metageneration-{uuid.uuid4()}" + + async def _run(): + writer = AsyncAppendableObjectWriter( + grpc_client, + _BUCKET_UNDER_TEST, + object_name, + storage_class="RAPID" if RCU_SYSTEM_TESTS else None, + ) + await writer.open() + await writer.append(b"v1") + await writer.close() + blobs_to_delete.append( + storage_client.bucket(_BUCKET_UNDER_TEST).blob(object_name) + ) + obj = await grpc_client.get_object( + bucket_name=_BUCKET_UNDER_TEST, object_name=object_name + ) + + with pytest.raises(FailedPrecondition) as exc_info: + writer_stale = AsyncAppendableObjectWriter( + grpc_client, + _BUCKET_UNDER_TEST, + object_name, + generation=writer.generation, + preconditions=Preconditions( + if_metageneration_match=obj.metageneration + 1 + ), + ) + await writer_stale.open() + assert exc_info.value.code == 400 + + writer_ok = AsyncAppendableObjectWriter( + grpc_client, + _BUCKET_UNDER_TEST, + object_name, + generation=writer.generation, + preconditions=Preconditions(if_metageneration_match=obj.metageneration), + ) + await writer_ok.open() + await writer_ok.append(b"v2") + await writer_ok.close() + + # cleanup + del writer + del writer_ok + gc.collect() + + event_loop.run_until_complete(_run()) + + def test_open_existing_object_with_gen_None_overrides_existing( storage_client, blobs_to_delete, event_loop, grpc_client ): diff --git a/packages/google-cloud-storage/tests/unit/asyncio/test_async_appendable_object_writer.py b/packages/google-cloud-storage/tests/unit/asyncio/test_async_appendable_object_writer.py index 501424161ea7..5c38ec8cb058 100644 --- a/packages/google-cloud-storage/tests/unit/asyncio/test_async_appendable_object_writer.py +++ b/packages/google-cloud-storage/tests/unit/asyncio/test_async_appendable_object_writer.py @@ -22,7 +22,7 @@ from google.cloud._storage_v2.types import storage as storage_type from google.cloud._storage_v2.types.storage import BidiWriteObjectRedirectedError -from google.cloud.storage import Blob +from google.cloud.storage import Blob, Bucket from google.cloud.storage.asyncio.async_appendable_object_writer import ( _DEFAULT_FLUSH_INTERVAL_BYTES, _MAX_CHUNK_SIZE_BYTES, @@ -131,6 +131,20 @@ def _create_stream(*args, **kwargs): stream_patcher.stop() +def _make_grpc_client(): + """Returns a client whose transport resolves `bidi_write_object`. + + Use it with a patched `AsyncBidiRpc` to inspect the requests sent on open. + """ + mock_client = mock.MagicMock() + transport = mock_client.grpc_client._client._transport + transport.bidi_write_object = mock.sentinel.bidi_write_object + transport._wrapped_methods = { + mock.sentinel.bidi_write_object: mock.sentinel.wrapped_bidi_write_object + } + return mock_client + + class TestAsyncAppendableObjectWriter: def _make_one(self, mock_client, **kwargs): return AsyncAppendableObjectWriter(mock_client, BUCKET, OBJECT, **kwargs) @@ -278,73 +292,133 @@ async def test_from_blob_open_passes_storage_class( ) assert writer.write_obj_stream.storage_class == storage_class - def test_init_invalid_preconditions_type_raises(self, mock_appendable_writer): - with pytest.raises( - TypeError, match="preconditions must be an instance of Preconditions" - ): - self._make_one( - mock_appendable_writer["mock_client"], - preconditions={"if_generation_match": 0}, - ) - - @pytest.mark.parametrize("generation", [0, GENERATION]) @pytest.mark.parametrize( - "preconditions", + "kwargs, error, match", [ - Preconditions(if_generation_match=0), - Preconditions(if_generation_match=456), - Preconditions(if_generation_not_match=456), + ( + {"preconditions": {"if_generation_match": 0}}, + TypeError, + "must be an instance of Preconditions", + ), + ( + { + "generation": 0, + "preconditions": Preconditions(if_generation_match=456), + }, + ValueError, + "require generation=None", + ), + ( + { + "generation": GENERATION, + "preconditions": Preconditions(if_generation_not_match=456), + }, + ValueError, + "require generation=None", + ), + ( + { + "write_handle": WRITE_HANDLE, + "preconditions": Preconditions(if_metageneration_match=1), + }, + ValueError, + "write_handle", + ), + ( + { + "write_handle": WRITE_HANDLE, + "preconditions": Preconditions(if_metageneration_not_match=1), + }, + ValueError, + "write_handle", + ), ], ) - def test_init_generation_and_generation_precondition_raises( - self, mock_appendable_writer, generation, preconditions + def test_init_rejects_invalid_preconditions( + self, mock_appendable_writer, kwargs, error, match ): - with pytest.raises(ValueError, match="require generation=None"): - self._make_one( - mock_appendable_writer["mock_client"], - generation=generation, - preconditions=preconditions, - ) - - @pytest.mark.asyncio - async def test_from_blob_with_preconditions(self, mock_appendable_writer): - mock_blob = mock.Mock(spec=Blob) - mock_blob.name = OBJECT - mock_blob.bucket.name = BUCKET - mock_blob.generation = None - mock_blob.storage_class = "RAPID" - - preconditions = Preconditions(if_generation_match=456) - writer = AsyncAppendableObjectWriter.from_blob( - mock_appendable_writer["mock_client"], - mock_blob, - preconditions=preconditions, - ) - assert writer.preconditions == preconditions - - await writer.open() + with pytest.raises(error, match=match): + self._make_one(mock_appendable_writer["mock_client"], **kwargs) - mock_stream_cls = mock_appendable_writer["mock_stream_cls"] - assert mock_stream_cls.call_args.kwargs["preconditions"] == preconditions - assert mock_stream_cls.call_args.kwargs["generation_number"] is None - assert mock_stream_cls.call_args.kwargs["blob"] == mock_blob - - def test_from_blob_with_generation_and_generation_precondition_raises( - self, mock_appendable_writer + @pytest.mark.parametrize( + "kwargs, match", + [ + # blob.generation is used unless `generation` overrides it. + ( + {"preconditions": Preconditions(if_generation_match=456)}, + "require generation=None", + ), + ( + { + "write_handle": WRITE_HANDLE, + "preconditions": Preconditions(if_metageneration_match=1), + }, + "write_handle", + ), + ], + ) + def test_from_blob_rejects_invalid_preconditions( + self, mock_appendable_writer, kwargs, match ): mock_blob = mock.Mock(spec=Blob) mock_blob.name = OBJECT mock_blob.bucket.name = BUCKET mock_blob.generation = GENERATION - mock_blob.storage_class = "RAPID" + mock_blob.storage_class = None - with pytest.raises(ValueError, match="require generation=None"): + with pytest.raises(ValueError, match=match): AsyncAppendableObjectWriter.from_blob( - mock_appendable_writer["mock_client"], - mock_blob, - preconditions=Preconditions(if_generation_match=456), + mock_appendable_writer["mock_client"], mock_blob, **kwargs ) + @pytest.mark.asyncio + @pytest.mark.parametrize( + "preconditions, field, value", + [ + ( + Preconditions(if_generation_match=GENERATION), + "if_generation_match", + GENERATION, + ), + (Preconditions(if_metageneration_match=4), "if_metageneration_match", 4), + ], + ) + async def test_from_blob_generation_none_overwrites_fetched_blob( + self, preconditions, field, value + ): + # A Blob fetched from GCS carries its generation and metadata. + blob = Blob(OBJECT, bucket=Bucket(client=None, name=BUCKET)) + blob._set_properties( + { + "name": OBJECT, + "generation": str(GENERATION), + "metageneration": "4", + "contentType": "text/plain", + "metadata": {"k": "v"}, + } + ) + writer = AsyncAppendableObjectWriter.from_blob( + _make_grpc_client(), blob, generation=None, preconditions=preconditions + ) + open_resp = storage_type.BidiWriteObjectResponse( + resource=storage_type.Object(generation=456, size=0), + write_handle=storage_type.BidiWriteHandle(handle=b"wh"), + ) + + with mock.patch( + "google.cloud.storage.asyncio.async_write_object_stream.AsyncBidiRpc" + ) as mock_rpc_cls: + mock_rpc_cls.return_value.open = AsyncMock() + mock_rpc_cls.return_value.recv = AsyncMock(return_value=open_resp) + await writer.open() + + request = mock_rpc_cls.call_args.kwargs["initial_request"] + assert "append_object_spec" not in request + spec = request.write_object_spec + assert getattr(spec, field) == value + assert spec.resource.content_type == "text/plain" + assert dict(spec.resource.metadata) == {"k": "v"} + # ------------------------------------------------------------------------- # Stream Lifecycle Tests # ------------------------------------------------------------------------- @@ -382,21 +456,6 @@ async def test_open_success(self, mock_appendable_writer): assert mock_stream_cls.call_args.kwargs["storage_class"] is None assert writer.write_obj_stream.storage_class is None - @pytest.mark.asyncio - async def test_open_passes_preconditions(self, mock_appendable_writer): - preconditions = Preconditions(if_generation_match=123) - writer = self._make_one( - mock_appendable_writer["mock_client"], - preconditions=preconditions, - ) - assert writer.preconditions == preconditions - - await writer.open() - - mock_stream_cls = mock_appendable_writer["mock_stream_cls"] - assert mock_stream_cls.call_args.kwargs["preconditions"] == preconditions - assert mock_stream_cls.call_args.kwargs["generation_number"] is None - @pytest.mark.asyncio async def test_open_failed_precondition_not_retried(self, mock_appendable_writer): preconditions = Preconditions(if_generation_match=123) @@ -416,18 +475,10 @@ async def test_open_failed_precondition_not_retried(self, mock_appendable_writer @pytest.mark.asyncio async def test_open_redirect_without_generation_preserves_preconditions(self): - mock_client = mock.MagicMock() - mock_transport = mock.MagicMock() - mock_transport.bidi_write_object = mock.sentinel.bidi_write_object - mock_transport._wrapped_methods = { - mock.sentinel.bidi_write_object: mock.sentinel.wrapped_bidi_write_object - } - mock_client.grpc_client._client._transport = mock_transport - preconditions = Preconditions( if_generation_match=123, if_metageneration_match=2 ) - writer = self._make_one(mock_client, preconditions=preconditions) + writer = self._make_one(_make_grpc_client(), preconditions=preconditions) redirect = BidiWriteObjectRedirectedError(routing_token="rt-1") redirect_exc = exceptions.Aborted("redirect", errors=[redirect]) @@ -463,18 +514,10 @@ async def test_open_redirect_without_generation_preserves_preconditions(self): @pytest.mark.asyncio async def test_open_redirect_with_generation_switches_to_append_spec(self): - mock_client = mock.MagicMock() - mock_transport = mock.MagicMock() - mock_transport.bidi_write_object = mock.sentinel.bidi_write_object - mock_transport._wrapped_methods = { - mock.sentinel.bidi_write_object: mock.sentinel.wrapped_bidi_write_object - } - mock_client.grpc_client._client._transport = mock_transport - preconditions = Preconditions( if_generation_match=123, if_metageneration_match=2 ) - writer = self._make_one(mock_client, preconditions=preconditions) + writer = self._make_one(_make_grpc_client(), preconditions=preconditions) redirect = BidiWriteObjectRedirectedError( routing_token="rt-2", @@ -515,75 +558,25 @@ async def test_open_redirect_with_generation_switches_to_append_spec(self): assert writer.generation == 777 @pytest.mark.asyncio - async def test_reopen_after_create_sends_append_spec(self): - mock_client = mock.MagicMock() - mock_transport = mock.MagicMock() - mock_transport.bidi_write_object = mock.sentinel.bidi_write_object - mock_transport._wrapped_methods = { - mock.sentinel.bidi_write_object: mock.sentinel.wrapped_bidi_write_object - } - mock_client.grpc_client._client._transport = mock_transport - - preconditions = Preconditions(if_generation_match=0, if_metageneration_match=1) - writer = self._make_one(mock_client, preconditions=preconditions) - - first_open_resp = storage_type.BidiWriteObjectResponse( - resource=storage_type.Object(generation=999, size=0), - write_handle=storage_type.BidiWriteHandle(handle=b"wh-init"), - ) - reopen_resp = storage_type.BidiWriteObjectResponse( - persisted_size=0, - write_handle=storage_type.BidiWriteHandle(handle=b"wh-reopen"), - ) - - with mock.patch( - "google.cloud.storage.asyncio.async_write_object_stream.AsyncBidiRpc" - ) as mock_rpc_cls: - rpc_1 = mock.MagicMock() - rpc_1.open = AsyncMock() - rpc_1.recv = AsyncMock(return_value=first_open_resp) - rpc_1.close = AsyncMock() - rpc_1.send = AsyncMock() - rpc_2 = mock.MagicMock() - rpc_2.open = AsyncMock() - rpc_2.recv = AsyncMock(return_value=reopen_resp) - mock_rpc_cls.side_effect = [rpc_1, rpc_2] - - await writer.open() - assert writer.generation == 999 - - # Simulate reconnect re-opening the writer - writer.write_obj_stream._is_stream_open = False - writer._is_stream_open = False - await writer.open() - - assert mock_rpc_cls.call_count == 2 - first_req = mock_rpc_cls.call_args_list[0].kwargs["initial_request"] - assert "write_object_spec" in first_req - assert "if_generation_match" in first_req.write_object_spec - assert first_req.write_object_spec.if_generation_match == 0 - - second_req = mock_rpc_cls.call_args_list[1].kwargs["initial_request"] - assert "append_object_spec" in second_req - assert second_req.append_object_spec.generation == 999 - # The new generation starts at metageneration 1, so the precondition - # for the replaced object must not be re-sent. - assert "if_metageneration_match" not in second_req.append_object_spec - - @pytest.mark.asyncio - async def test_reopen_append_mode_keeps_metageneration_preconditions(self): - mock_client = mock.MagicMock() - mock_transport = mock.MagicMock() - mock_transport.bidi_write_object = mock.sentinel.bidi_write_object - mock_transport._wrapped_methods = { - mock.sentinel.bidi_write_object: mock.sentinel.wrapped_bidi_write_object - } - mock_client.grpc_client._client._transport = mock_transport - + @pytest.mark.parametrize( + "generation, preconditions, resent", + [ + # The writer created the generation, so the precondition described + # the replaced object (the new generation starts at metageneration 1). + ( + None, + Preconditions(if_generation_match=0, if_metageneration_match=1), + False, + ), + # The caller opened an existing generation, which it still describes. + (GENERATION, Preconditions(if_metageneration_match=1), True), + ], + ) + async def test_reopen_resends_metageneration_only_in_append_mode( + self, generation, preconditions, resent + ): writer = self._make_one( - mock_client, - generation=GENERATION, - preconditions=Preconditions(if_metageneration_match=3), + _make_grpc_client(), generation=generation, preconditions=preconditions ) open_resp = storage_type.BidiWriteObjectResponse( resource=storage_type.Object(generation=GENERATION, size=0), @@ -593,14 +586,8 @@ async def test_reopen_append_mode_keeps_metageneration_preconditions(self): with mock.patch( "google.cloud.storage.asyncio.async_write_object_stream.AsyncBidiRpc" ) as mock_rpc_cls: - rpc_1 = mock.MagicMock() - rpc_1.open = AsyncMock() - rpc_1.recv = AsyncMock(return_value=open_resp) - rpc_2 = mock.MagicMock() - rpc_2.open = AsyncMock() - rpc_2.recv = AsyncMock(return_value=open_resp) - mock_rpc_cls.side_effect = [rpc_1, rpc_2] - + mock_rpc_cls.return_value.open = AsyncMock() + mock_rpc_cls.return_value.recv = AsyncMock(return_value=open_resp) await writer.open() # Simulate reconnect re-opening the writer writer.write_obj_stream._is_stream_open = False @@ -608,10 +595,9 @@ async def test_reopen_append_mode_keeps_metageneration_preconditions(self): await writer.open() assert mock_rpc_cls.call_count == 2 - for call in mock_rpc_cls.call_args_list: - spec = call.kwargs["initial_request"].append_object_spec - assert spec.generation == GENERATION - assert spec.if_metageneration_match == 3 + reopen_req = mock_rpc_cls.call_args_list[1].kwargs["initial_request"] + assert reopen_req.append_object_spec.generation == GENERATION + assert ("if_metageneration_match" in reopen_req.append_object_spec) == resent def test_on_open_error_redirection(self, mock_appendable_writer): """Verify redirect info is extracted from helper.""" diff --git a/packages/google-cloud-storage/tests/unit/asyncio/test_async_write_object_stream.py b/packages/google-cloud-storage/tests/unit/asyncio/test_async_write_object_stream.py index 4645e3d48aae..297f6e6ee597 100644 --- a/packages/google-cloud-storage/tests/unit/asyncio/test_async_write_object_stream.py +++ b/packages/google-cloud-storage/tests/unit/asyncio/test_async_write_object_stream.py @@ -194,13 +194,27 @@ async def test_open_existing_object_with_token(self, mock_rpc_cls, mock_client): @pytest.mark.parametrize( "generation_number, preconditions, expected", [ - (None, None, None), - (0, None, 0), - (None, Preconditions(if_generation_match=0), 0), - (None, Preconditions(if_generation_match=123), 123), + (None, None, {}), + (0, None, {"if_generation_match": 0}), + (None, Preconditions(if_generation_match=0), {"if_generation_match": 0}), + ( + None, + Preconditions( + if_generation_match=10, + if_generation_not_match=20, + if_metageneration_match=3, + if_metageneration_not_match=4, + ), + { + "if_generation_match": 10, + "if_generation_not_match": 20, + "if_metageneration_match": 3, + "if_metageneration_not_match": 4, + }, + ), ], ) - async def test_open_new_object_if_generation_match( + async def test_open_new_object_preconditions( self, mock_rpc_cls, mock_client, @@ -223,43 +237,14 @@ async def test_open_new_object_if_generation_match( spec = mock_rpc_cls.call_args.kwargs["initial_request"].write_object_spec assert spec.appendable - if expected is None: - assert "if_generation_match" not in spec - else: - assert "if_generation_match" in spec - assert spec.if_generation_match == expected - - @mock.patch("google.cloud.storage.asyncio.async_write_object_stream.AsyncBidiRpc") - @pytest.mark.asyncio - async def test_open_new_object_all_preconditions(self, mock_rpc_cls, mock_client): - mock_rpc = mock_rpc_cls.return_value - mock_rpc.open = AsyncMock() - mock_rpc.recv = AsyncMock(return_value=MagicMock(resource=None)) - - preconditions = Preconditions( - if_generation_match=10, - if_generation_not_match=20, - if_metageneration_match=3, - if_metageneration_not_match=4, - ) - stream = _AsyncWriteObjectStream( - mock_client, - BUCKET, - OBJECT, - preconditions=preconditions, + fields = ( + "if_generation_match", + "if_generation_not_match", + "if_metageneration_match", + "if_metageneration_not_match", ) - await stream.open() - - spec = mock_rpc_cls.call_args.kwargs["initial_request"].write_object_spec - assert spec.appendable - assert "if_generation_match" in spec - assert spec.if_generation_match == 10 - assert "if_generation_not_match" in spec - assert spec.if_generation_not_match == 20 - assert "if_metageneration_match" in spec - assert spec.if_metageneration_match == 3 - assert "if_metageneration_not_match" in spec - assert spec.if_metageneration_not_match == 4 + # `in` checks field presence, so an unset field is not mistaken for 0. + assert {f: getattr(spec, f) for f in fields if f in spec} == expected @mock.patch("google.cloud.storage.asyncio.async_write_object_stream.AsyncBidiRpc") @pytest.mark.asyncio @@ -269,8 +254,6 @@ async def test_open_existing_object_preconditions(self, mock_rpc_cls, mock_clien mock_rpc.recv = AsyncMock(return_value=MagicMock(resource=None)) preconditions = Preconditions( - if_generation_match=123, - if_generation_not_match=456, if_metageneration_match=2, if_metageneration_not_match=3, ) From e931fc3d88b19d27ff1406c9f58f78ef93d40168 Mon Sep 17 00:00:00 2001 From: yuxin00j Date: Mon, 2 Nov 2026 03:06:59 +0000 Subject: [PATCH 6/6] fix(storage): strip the KMS key version when writing a fetched Blob A Blob fetched from a CMEK object reports kms_key_name with a /cryptoKeyVersions/N suffix, which GCS rejects when creating an object. Strip it in blob_to_proto so from_blob(..., generation=None) overwrites keep the object under the same key. --- .../google/cloud/storage/_grpc_conversions.py | 9 +++++++++ .../google-cloud-storage/tests/system/test_rapid.py | 10 +++++++--- .../asyncio/test_async_appendable_object_writer.py | 7 ++++++- 3 files changed, 22 insertions(+), 4 deletions(-) diff --git a/packages/google-cloud-storage/google/cloud/storage/_grpc_conversions.py b/packages/google-cloud-storage/google/cloud/storage/_grpc_conversions.py index 84ae8e7679d8..12cf4c9d99c5 100644 --- a/packages/google-cloud-storage/google/cloud/storage/_grpc_conversions.py +++ b/packages/google-cloud-storage/google/cloud/storage/_grpc_conversions.py @@ -12,6 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +import re + from google.protobuf import timestamp_pb2 from google.cloud import _storage_v2 @@ -30,6 +32,10 @@ "storage_class": "storage_class", } +# A Blob fetched from a CMEK object reports the key version it was encrypted +# with, which GCS rejects when creating an object. +_KMS_KEY_VERSION_SUFFIX = re.compile(r"/cryptoKeyVersions/[^/]+$") + def blob_to_proto(blob): """Converts a Blob instance to a GCS V2 Object proto message.""" @@ -44,6 +50,9 @@ def blob_to_proto(blob): for attr_name, proto_field in _BLOB_ATTR_TO_PROTO_FIELD.items(): value = getattr(blob, attr_name, None) if value is not None: + if proto_field == "kms_key": + # Write the new generation under the same key. + value = _KMS_KEY_VERSION_SUFFIX.sub("", value) resource_params[proto_field] = value custom_time = getattr(blob, "custom_time", None) diff --git a/packages/google-cloud-storage/tests/system/test_rapid.py b/packages/google-cloud-storage/tests/system/test_rapid.py index b90011816419..038203cb5103 100644 --- a/packages/google-cloud-storage/tests/system/test_rapid.py +++ b/packages/google-cloud-storage/tests/system/test_rapid.py @@ -895,13 +895,15 @@ async def _run(): RCU_SYSTEM_TESTS, reason="Write from blob for RCU not supported in SDK yet" ) def test_from_blob_overwrite_fetched_blob_with_if_generation_match( - storage_client, blobs_to_delete, event_loop, grpc_client + storage_client, blobs_to_delete, event_loop, grpc_client, rapid_kms_key ): - """A fetched Blob is overwritten with CAS when `generation=None` is passed.""" + """A fetched CMEK Blob is overwritten with CAS when `generation=None` is passed.""" object_name = f"test_from_blob_cas-{uuid.uuid4()}" async def _run(): - blob = storage_client.bucket(_BUCKET_UNDER_TEST).blob(object_name) + blob = storage_client.bucket(_BUCKET_UNDER_TEST).blob( + object_name, kms_key_name=rapid_kms_key + ) blob.content_type = "text/plain" blob.metadata = {"environment": "system-test"} writer = AsyncAppendableObjectWriter.from_blob(grpc_client, blob) @@ -910,6 +912,7 @@ async def _run(): await writer.close(finalize_on_close=True) blobs_to_delete.append(blob) + # The fetched kms_key_name includes the key version. fetched = storage_client.bucket(_BUCKET_UNDER_TEST).get_blob(object_name) writer_cas = AsyncAppendableObjectWriter.from_blob( grpc_client, @@ -927,6 +930,7 @@ async def _run(): assert obj.generation != fetched.generation assert obj.content_type == "text/plain" assert obj.metadata["environment"] == "system-test" + assert obj.kms_key.startswith(rapid_kms_key) # cleanup del writer diff --git a/packages/google-cloud-storage/tests/unit/asyncio/test_async_appendable_object_writer.py b/packages/google-cloud-storage/tests/unit/asyncio/test_async_appendable_object_writer.py index 5c38ec8cb058..93b78f07db1f 100644 --- a/packages/google-cloud-storage/tests/unit/asyncio/test_async_appendable_object_writer.py +++ b/packages/google-cloud-storage/tests/unit/asyncio/test_async_appendable_object_writer.py @@ -386,7 +386,9 @@ def test_from_blob_rejects_invalid_preconditions( async def test_from_blob_generation_none_overwrites_fetched_blob( self, preconditions, field, value ): - # A Blob fetched from GCS carries its generation and metadata. + kms_key = "projects/p/locations/l/keyRings/r/cryptoKeys/k" + # A Blob fetched from GCS carries its generation, metadata and, for CMEK + # objects, the key version used to encrypt it. blob = Blob(OBJECT, bucket=Bucket(client=None, name=BUCKET)) blob._set_properties( { @@ -395,6 +397,7 @@ async def test_from_blob_generation_none_overwrites_fetched_blob( "metageneration": "4", "contentType": "text/plain", "metadata": {"k": "v"}, + "kmsKeyName": f"{kms_key}/cryptoKeyVersions/1", } ) writer = AsyncAppendableObjectWriter.from_blob( @@ -418,6 +421,8 @@ async def test_from_blob_generation_none_overwrites_fetched_blob( assert getattr(spec, field) == value assert spec.resource.content_type == "text/plain" assert dict(spec.resource.metadata) == {"k": "v"} + # GCS rejects a key version when creating an object. + assert spec.resource.kms_key == kms_key # ------------------------------------------------------------------------- # Stream Lifecycle Tests