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 1e31b15faf0c..973226972a3f 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 @@ -547,12 +547,30 @@ async def flush(self) -> int: self.bytes_appended_since_last_flush = 0 return self.persisted_size - async def close(self, finalize_on_close=False) -> Union[int, _storage_v2.Object]: + async def close( + self, + finalize_on_close=False, + full_object_checksum: Optional[int] = None, + ) -> Union[int, _storage_v2.Object]: """Closes the underlying bidi-gRPC stream. :type finalize_on_close: bool :param finalize_on_close: Finalizes the Appendable Object. No more data can be appended. + :type full_object_checksum: int + :param full_object_checksum: (Optional) This should be the CRC32C checksum of + the entire contents of the object as a 32-bit integer. + Used only when finalize_on_close is True. + + It can be obtained by running: + + .. code-block:: python + + import google_crc32c + + data = b"Hello, world!" + crc32c_int = google_crc32c.value(data) + print(crc32c_int) rtype: Union[int, _storage_v2.Object] returns: Updated `self.persisted_size` by default after closing the @@ -561,20 +579,32 @@ async def close(self, finalize_on_close=False) -> Union[int, _storage_v2.Object] :raises ValueError: If the stream is not open (i.e., `open()` has not been called). + :raises ValueError: If full_object_checksum is provided but + finalize_on_close is False. + :raises google.api_core.exceptions.InvalidArgument: If the provided + full_object_checksum does not match the checksum computed by the + server. """ if not self._is_stream_open: raise ValueError("Stream is not open. Call open() before close().") + if full_object_checksum is not None and not finalize_on_close: + raise ValueError( + "full_object_checksum can only be provided when finalize_on_close is True." + ) + if finalize_on_close: - return await self.finalize() + return await self.finalize(full_object_checksum=full_object_checksum) await self.write_obj_stream.close() self._is_stream_open = False return self.persisted_size - async def finalize(self) -> _storage_v2.Object: + async def finalize( + self, full_object_checksum: Optional[int] = None + ) -> _storage_v2.Object: """Finalizes the Appendable Object. Note: Once finalized no more data can be appended. @@ -585,18 +615,45 @@ async def finalize(self) -> _storage_v2.Object: However if `.finalize()` is called no more data can be appended to the object. + :type full_object_checksum: int + :param full_object_checksum: (Optional) This should be the CRC32C checksum of + the entire contents of the object as a 32-bit integer. + + It can be obtained by running: + + .. code-block:: python + + import google_crc32c + + data = b"Hello, world!" + crc32c_int = google_crc32c.value(data) + print(crc32c_int) + rtype: google.cloud.storage_v2.types.Object returns: The finalized object resource. :raises ValueError: If the stream is not open (i.e., `open()` has not been called). + :raises google.api_core.exceptions.InvalidArgument: If the provided + full_object_checksum does not match the checksum computed by the + server. """ if not self._is_stream_open: raise ValueError("Stream is not open. Call open() before finalize().") - await self.write_obj_stream.send( - _storage_v2.BidiWriteObjectRequest(finish_write=True) - ) + finalize_req = _storage_v2.BidiWriteObjectRequest(finish_write=True) + + if full_object_checksum is not None: + finalize_req.object_checksums = _storage_v2.ObjectChecksums( + crc32c=full_object_checksum + ) + else: + logger.warning( + "finalize was called without providing full_object_checksum. " + "No checksum validation will be performed on the finalized object." + ) + + await self.write_obj_stream.send(finalize_req) response = await self.write_obj_stream.recv() self.object_resource = response.resource self.persisted_size = self.object_resource.size diff --git a/packages/google-cloud-storage/tests/system/test_zonal.py b/packages/google-cloud-storage/tests/system/test_zonal.py index bcdaeaffb182..7050cab0ea60 100644 --- a/packages/google-cloud-storage/tests/system/test_zonal.py +++ b/packages/google-cloud-storage/tests/system/test_zonal.py @@ -10,6 +10,7 @@ # python additional imports import google_crc32c import pytest +from google.api_core import exceptions from google.api_core.exceptions import FailedPrecondition, NotFound, OutOfRange # current library imports @@ -1046,3 +1047,68 @@ async def _run(): blobs_to_delete.append(storage_client.bucket(_ZONAL_BUCKET).blob(object_name)) event_loop.run_until_complete(_run()) + + +def test_finalize_with_correct_checksum( + storage_client, + blobs_to_delete, + event_loop, + grpc_client_direct, +): + object_name = f"appendabl_cksum_success_{uuid.uuid4()}" + object_data = b"Hello, appendable object with correct checksum!" + object_checksum = google_crc32c.value(object_data) + + async def _run(): + writer = AsyncAppendableObjectWriter( + grpc_client_direct, _ZONAL_BUCKET, object_name + ) + await writer.open() + await writer.append(object_data) + object_metadata = await writer.finalize(full_object_checksum=object_checksum) + + assert object_metadata.size == len(object_data) + assert int(object_metadata.checksums.crc32c) == object_checksum + + # clean up + blobs_to_delete.append(storage_client.bucket(_ZONAL_BUCKET).blob(object_name)) + del writer + gc.collect() + + event_loop.run_until_complete(_run()) + + +def test_finalize_with_incorrect_checksum_fails( + storage_client, + blobs_to_delete, + event_loop, + grpc_client_direct, +): + object_name = f"appendabl_cksum_fail_{uuid.uuid4()}" + object_data = b"Hello, appendable object with incorrect checksum!" + object_checksum = google_crc32c.value(object_data) + incorrect_checksum = object_checksum + 1 + + async def _run(): + writer = AsyncAppendableObjectWriter( + grpc_client_direct, _ZONAL_BUCKET, object_name + ) + await writer.open() + await writer.append(object_data) + + # Finalization should fail with an exception due to mismatch + with pytest.raises(exceptions.InvalidArgument) as excinfo: + await writer.finalize(full_object_checksum=incorrect_checksum) + + # Assert that the error relates to checksum verification/mismatch + error_msg = str(excinfo.value).lower() + assert ( + "crc32c" in error_msg or "checksum" in error_msg or "mismatch" in error_msg + ) + + # clean up + blobs_to_delete.append(storage_client.bucket(_ZONAL_BUCKET).blob(object_name)) + del writer + gc.collect() + + event_loop.run_until_complete(_run()) 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 c0e62eb54242..396fe106e4eb 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 @@ -503,3 +503,52 @@ async def test_methods_require_open_stream_raises(self, mock_appendable_writer): for coro in methods: with pytest.raises(ValueError, match="Stream is not open"): await coro + + @pytest.mark.asyncio + async def test_finalize_with_checksum(self, mock_appendable_writer): + writer = self._make_one(mock_appendable_writer["mock_client"]) + writer._is_stream_open = True + writer.write_obj_stream = mock_appendable_writer["mock_stream"] + + resource = storage_type.Object(size=999) + mock_appendable_writer[ + "mock_stream" + ].recv.return_value = storage_type.BidiWriteObjectResponse(resource=resource) + + checksum = 12345678 + res = await writer.finalize(full_object_checksum=checksum) + + assert res == resource + assert writer.persisted_size == 999 + mock_appendable_writer["mock_stream"].send.assert_awaited_with( + storage_type.BidiWriteObjectRequest( + finish_write=True, + object_checksums=storage_type.ObjectChecksums(crc32c=checksum), + ) + ) + mock_appendable_writer["mock_stream"].close.assert_awaited() + assert not writer._is_stream_open + + @pytest.mark.asyncio + async def test_close_with_checksum_and_finalize(self, mock_appendable_writer): + writer = self._make_one(mock_appendable_writer["mock_client"]) + writer._is_stream_open = True + writer.finalize = AsyncMock() + + checksum = 12345678 + await writer.close(finalize_on_close=True, full_object_checksum=checksum) + writer.finalize.assert_awaited_once_with(full_object_checksum=checksum) + + @pytest.mark.asyncio + async def test_close_with_checksum_without_finalize_raises( + self, mock_appendable_writer + ): + writer = self._make_one(mock_appendable_writer["mock_client"]) + writer._is_stream_open = True + + checksum = 12345678 + with pytest.raises( + ValueError, + match="full_object_checksum can only be provided when finalize_on_close is True", + ): + await writer.close(finalize_on_close=False, full_object_checksum=checksum)