Skip to content
Open
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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.
Expand All @@ -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
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 @@ -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)
Loading