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 @@ -641,6 +650,14 @@ async def finalize(
if not self._is_stream_open:
raise ValueError("Stream is not open. Call open() before finalize().")

if full_object_checksum is not None:
if not isinstance(full_object_checksum, int):
raise TypeError("full_object_checksum must be an integer.")
Comment on lines +654 to +655

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

In Python, bool is a subclass of int, so isinstance(True, int) evaluates to True. If a boolean value is passed as full_object_checksum, it will bypass this type check and be treated as a valid checksum (e.g., 1 or 0).

To prevent this, explicitly reject boolean values.

Suggested change
if not isinstance(full_object_checksum, int):
raise TypeError("full_object_checksum must be an integer.")
if not isinstance(full_object_checksum, int) or isinstance(full_object_checksum, bool):
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."
)

finalize_req = _storage_v2.BidiWriteObjectRequest(finish_write=True)

if full_object_checksum is not None:
Expand All @@ -654,14 +671,15 @@ async def finalize(
)

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
await self.write_obj_stream.close()

self._is_stream_open = False
self.offset = None
return self.object_resource
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
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
Original file line number Diff line number Diff line change
Expand Up @@ -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",
[
Expand Down Expand Up @@ -552,3 +575,50 @@ async def test_close_with_checksum_without_finalize_raises(
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
Loading