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/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..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 @@ -14,6 +14,7 @@ import io import logging +from dataclasses import dataclass from io import BufferedReader from typing import Dict, List, Optional, Tuple, Union @@ -49,6 +50,7 @@ _BIDI_WRITE_REDIRECTED_TYPE_URL = ( "type.googleapis.com/google.storage.v2.BidiWriteObjectRedirectedError" ) +_marker = object() logger = logging.getLogger(__name__) @@ -100,6 +102,46 @@ def _is_write_retryable(exc): return False +@dataclass(frozen=True) +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``). 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 + 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,6 +154,7 @@ def __init__( write_handle: Optional[_storage_v2.BidiWriteHandle] = None, writer_options: Optional[dict] = None, storage_class: Optional[str] = None, + preconditions: Optional[Preconditions] = None, ): """ Class for appending data to a GCS Appendable Object. @@ -165,7 +208,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. @@ -184,14 +228,66 @@ 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 preconditions: Optional[:class:`Preconditions`] + :param preconditions: (Optional) Preconditions for the write operation. + Generation preconditions (``if_generation_match``, + ``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``. 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, 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): + 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 " + "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 self.object_name = object_name self.write_handle = write_handle 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 @@ -227,6 +323,8 @@ def from_blob( blob: 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. @@ -259,17 +357,36 @@ 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. + 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 ``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, + preconditions=preconditions, ) instance.blob = blob return instance @@ -304,6 +421,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( @@ -338,6 +458,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: @@ -347,6 +469,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: @@ -369,6 +492,7 @@ async def _do_open(): write_handle=self.write_handle, routing_token=self._routing_token, storage_class=self.storage_class, + preconditions=None if self._created_generation else self.preconditions, ) if self._routing_token: @@ -381,6 +505,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/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..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. @@ -61,6 +67,12 @@ class _AsyncWriteObjectStream(_AsyncAbstractObjectStream): :type storage_class: Optional[str] :param storage_class: (Optional) The storage class of the object. + + :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__( @@ -73,6 +85,7 @@ def __init__( routing_token: Optional[str] = None, blob: Optional[Blob] = None, storage_class: Optional[str] = None, + preconditions: Optional["Preconditions"] = None, ) -> None: if client is None: raise ValueError("client must be provided") @@ -91,6 +104,7 @@ def __init__( self.routing_token: Optional[str] = routing_token self.blob: Optional[Blob] = blob self.storage_class: Optional[str] = storage_class + self.preconditions: Optional["Preconditions"] = preconditions self._full_bucket_name = f"projects/_/buckets/{self.bucket_name}" self.rpc = self.client._client._transport._wrapped_methods[ @@ -114,7 +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 `generation_number` is 0 and object already exists. + If a precondition in `preconditions` (or `generation_number=0`) is + not met. """ if self._is_stream_open: raise ValueError("Stream is already open") @@ -133,11 +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.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=self.generation_number, + 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: @@ -148,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..038203cb5103 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,175 @@ 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()) + + +@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, rapid_kms_key +): + """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, kms_key_name=rapid_kms_key + ) + 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) + + # 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, + 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" + assert obj.kms_key.startswith(rapid_kms_key) + + # 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 e6fcd90f0b84..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 @@ -22,11 +22,12 @@ 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, AsyncAppendableObjectWriter, + Preconditions, _is_write_retryable, ) @@ -130,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) @@ -173,6 +188,7 @@ async def test_init_with_storage_class(self, mock_appendable_writer, storage_cla write_handle=None, routing_token=None, storage_class=storage_class, + preconditions=None, ) assert writer.write_obj_stream.storage_class == storage_class @@ -272,9 +288,142 @@ async def test_from_blob_open_passes_storage_class( write_handle=None, routing_token=None, storage_class=storage_class, + preconditions=None, ) assert writer.write_obj_stream.storage_class == storage_class + @pytest.mark.parametrize( + "kwargs, error, match", + [ + ( + {"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_rejects_invalid_preconditions( + self, mock_appendable_writer, kwargs, error, match + ): + with pytest.raises(error, match=match): + self._make_one(mock_appendable_writer["mock_client"], **kwargs) + + @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 = None + + with pytest.raises(ValueError, match=match): + AsyncAppendableObjectWriter.from_blob( + 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 + ): + 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( + { + "name": OBJECT, + "generation": str(GENERATION), + "metageneration": "4", + "contentType": "text/plain", + "metadata": {"k": "v"}, + "kmsKeyName": f"{kms_key}/cryptoKeyVersions/1", + } + ) + 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"} + # GCS rejects a key version when creating an object. + assert spec.resource.kms_key == kms_key + # ------------------------------------------------------------------------- # Stream Lifecycle Tests # ------------------------------------------------------------------------- @@ -312,6 +461,149 @@ 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_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): + preconditions = Preconditions( + if_generation_match=123, if_metageneration_match=2 + ) + writer = self._make_one(_make_grpc_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): + preconditions = Preconditions( + if_generation_match=123, if_metageneration_match=2 + ) + writer = self._make_one(_make_grpc_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" + # 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 + @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( + _make_grpc_client(), generation=generation, preconditions=preconditions + ) + 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: + 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 + writer._is_stream_open = False + await writer.open() + + assert mock_rpc_cls.call_count == 2 + 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.""" 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..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 @@ -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, ) @@ -186,6 +189,90 @@ 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, preconditions, expected", + [ + (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_preconditions( + self, + mock_rpc_cls, + mock_client, + generation_number, + preconditions, + 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, + preconditions=preconditions, + ) + await stream.open() + + spec = mock_rpc_cls.call_args.kwargs["initial_request"].write_object_spec + assert spec.appendable + fields = ( + "if_generation_match", + "if_generation_not_match", + "if_metageneration_match", + "if_metageneration_not_match", + ) + # `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 + 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_metageneration_match=2, + if_metageneration_not_match=3, + ) + stream = _AsyncWriteObjectStream( + mock_client, + BUCKET, + OBJECT, + generation_number=GENERATION, + preconditions=preconditions, + ) + await stream.open() + + 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 async def test_open_metadata_merging(self, mock_rpc_cls, mock_client):