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..6551cc880605 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 @@ -386,6 +386,7 @@ async def append( data: bytes, retry_policy: Optional[AsyncRetry] = None, metadata: Optional[List[Tuple[str, str]]] = None, + enable_checksum: bool = True, ) -> None: """Appends data to the Appendable object with automatic retries. @@ -406,6 +407,9 @@ async def append( :type metadata: List[Tuple[str, str]] :param metadata: (Optional) The metadata to be sent with the request. + :type enable_checksum: bool + :param enable_checksum: (Optional) If True, calculates and checks checksums for each chunk. Defaults to True. + :raises ValueError: If the stream is not open. """ if not self._is_stream_open: @@ -487,7 +491,12 @@ async def generator(): return generator() # State initialization - write_state = _WriteState(_MAX_CHUNK_SIZE_BYTES, buffer, self.flush_interval) + write_state = _WriteState( + _MAX_CHUNK_SIZE_BYTES, + buffer, + self.flush_interval, + enable_checksum=enable_checksum, + ) write_state.write_handle = self.write_handle write_state.persisted_size = self.persisted_size # offset is set during `open()` call. @@ -547,12 +556,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 +588,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,26 +624,62 @@ 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) - ) - response = await self.write_obj_stream.recv() - self.object_resource = response.resource - self.persisted_size = self.object_resource.size - await self.write_obj_stream.close() + if full_object_checksum is not None: + if not isinstance(full_object_checksum, int): + raise TypeError("full_object_checksum must be an integer.") + if not (0 <= full_object_checksum <= 0xFFFFFFFF): + raise ValueError( + "full_object_checksum must be a 32-bit unsigned integer." + ) - self._is_stream_open = False - self.offset = None - return self.object_resource + 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) + try: + response = await self.write_obj_stream.recv() + self.object_resource = response.resource + self.persisted_size = self.object_resource.size + return self.object_resource + finally: + await self.write_obj_stream.close() + self._is_stream_open = False + self.offset = None @property def is_stream_open(self) -> bool: diff --git a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_multi_range_downloader.py b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_multi_range_downloader.py index 6d3f5e2fab4b..4c0dc87fd099 100644 --- a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_multi_range_downloader.py +++ b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_multi_range_downloader.py @@ -407,6 +407,12 @@ async def download_ranges( :type retry_policy: :class:`~google.api_core.retry_async.AsyncRetry` :param retry_policy: (Optional) The retry policy to use for the operation. + :type metadata: List[Tuple[str, str]] + :param metadata: (Optional) The metadata to be sent with the request. + + :type enable_checksum: bool + :param enable_checksum: (Optional) If True, checksums are verified for downloaded data. Defaults to True. + :raises ValueError: if the underlying bidi-GRPC stream is not open. :raises ValueError: if the length of read_ranges is more than 1000. :raises DataCorruption: if a checksum mismatch is detected while reading data. diff --git a/packages/google-cloud-storage/google/cloud/storage/asyncio/retry/writes_resumption_strategy.py b/packages/google-cloud-storage/google/cloud/storage/asyncio/retry/writes_resumption_strategy.py index df243a0a6eb2..bec0b31c5407 100644 --- a/packages/google-cloud-storage/google/cloud/storage/asyncio/retry/writes_resumption_strategy.py +++ b/packages/google-cloud-storage/google/cloud/storage/asyncio/retry/writes_resumption_strategy.py @@ -44,6 +44,7 @@ def __init__( chunk_size: int, user_buffer: IO[bytes], flush_interval: int, + enable_checksum: bool = True, ): self.chunk_size = chunk_size self.user_buffer = user_buffer @@ -59,6 +60,7 @@ def __init__( self.write_handle: Union[bytes, storage_type.BidiWriteHandle, None] = None self.routing_token: Optional[str] = None self.is_finalized: bool = False + self.enable_checksum: bool = enable_checksum class _WriteResumptionStrategy(_BaseResumptionStrategy): @@ -84,7 +86,8 @@ def generate_requests( break checksummed_data = storage_type.ChecksummedData(content=chunk) - checksummed_data.crc32c = google_crc32c.value(chunk) + if write_state.enable_checksum: + checksummed_data.crc32c = google_crc32c.value(chunk) request = storage_type.BidiWriteObjectRequest( write_offset=write_state.bytes_sent, 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/retry/test_writes_resumption_strategy.py b/packages/google-cloud-storage/tests/unit/asyncio/retry/test_writes_resumption_strategy.py index 4a2e3d192f87..2af036f80cb2 100644 --- a/packages/google-cloud-storage/tests/unit/asyncio/retry/test_writes_resumption_strategy.py +++ b/packages/google-cloud-storage/tests/unit/asyncio/retry/test_writes_resumption_strategy.py @@ -128,6 +128,22 @@ def test_generate_requests_checksum_verification(self, strategy): expected_int = google_crc32c.value(chunk_data) assert requests[0].checksummed_data.crc32c == expected_int + def test_generate_requests_checksum_disabled(self, strategy): + """Verify CRC32C is not calculated if enable_checksum is False.""" + chunk_data = b"test_data" + mock_buffer = io.BytesIO(chunk_data) + write_state = _WriteState( + chunk_size=10, + user_buffer=mock_buffer, + flush_interval=10, + enable_checksum=False, + ) + state = {"write_state": write_state} + + requests = strategy.generate_requests(state) + + assert not requests[0].checksummed_data.crc32c + def test_generate_requests_flush_logic_exact_interval(self, strategy): """Verify the flush bit is set exactly when the interval is reached.""" mock_buffer = io.BytesIO(b"A" * 12) 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..e831d433c6fe 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 @@ -299,6 +299,29 @@ async def test_append_data_less_than_flush_interval(self, mock_appendable_writer assert writer.offset == data_len assert writer.bytes_appended_since_last_flush == data_len + @pytest.mark.asyncio + async def test_append_with_checksum_disabled(self, mock_appendable_writer): + """Verify append propagates enable_checksum=False to the _WriteState.""" + writer = self._make_one(mock_appendable_writer["mock_client"]) + writer._is_stream_open = True + writer.persisted_size = 0 + writer.write_obj_stream = mock_appendable_writer["mock_stream"] + writer.write_obj_stream.send = AsyncMock() + + with mock.patch( + "google.cloud.storage.asyncio.async_appendable_object_writer._BidiStreamRetryManager" + ) as MockManager: + mock_execute = AsyncMock() + MockManager.return_value.execute = mock_execute + + await writer.append(DATA_LESS_THAN_FLUSH_INTERVAL, enable_checksum=False) + + # Check that execute was called with a state dictionary containing write_state having enable_checksum=False + mock_execute.assert_called_once() + state_arg = mock_execute.call_args[0][0] + assert "write_state" in state_arg + assert state_arg["write_state"].enable_checksum is False + @pytest.mark.parametrize( "data_len", [ @@ -503,3 +526,99 @@ 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) + + @pytest.mark.asyncio + async def test_finalize_invalid_checksum_type(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"] + + with pytest.raises(TypeError, match="full_object_checksum must be an integer"): + await writer.finalize(full_object_checksum="not-an-int") + + @pytest.mark.asyncio + async def test_finalize_invalid_checksum_range(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"] + + # negative + with pytest.raises( + ValueError, match="full_object_checksum must be a 32-bit unsigned integer" + ): + await writer.finalize(full_object_checksum=-1) + + # overflow + with pytest.raises( + ValueError, match="full_object_checksum must be a 32-bit unsigned integer" + ): + await writer.finalize(full_object_checksum=0x100000000) + + @pytest.mark.asyncio + async def test_finalize_mismatch_closes_stream(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"] + + # Mock recv to raise an exception (like server rejecting checksum mismatch) + from google.api_core.exceptions import InvalidArgument + + mock_appendable_writer["mock_stream"].recv.side_effect = InvalidArgument( + "checksum mismatch" + ) + + with pytest.raises(InvalidArgument): + await writer.finalize(full_object_checksum=12345) + + # Assert stream was closed and local state reset despite exception + mock_appendable_writer["mock_stream"].close.assert_awaited() + assert not writer._is_stream_open