Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand All @@ -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:
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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
Expand All @@ -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.
Expand All @@ -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:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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):
Expand All @@ -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,
Expand Down
66 changes: 66 additions & 0 deletions packages/google-cloud-storage/tests/system/test_zonal.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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())
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
Loading
Loading