feat(ZonalBuckets): add support for generation=0 to avoid overwriting existing objects and add is_stream_open support (#1709)

feat(ZonalBuckets): add support for `generation=0` to prevent
overwriting existing objects
feat(ZonalBuckets): add `is_stream_open` property to
AsyncAppendableObjectWriter for stream status check
---------
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Author: chandra-siri

google/cloud/storage/_experimental/asyncio/async_appendable_object_writer.py

@@ -97,15 +97,29 @@ def __init__(
  :param object_name: The name of the GCS Appendable Object to be written.
 
  :type generation: int
- :param generation: (Optional) If present, selects a specific revision of
- that object.
- If None, a new object is created.
- If None and Object already exists then it'll will be
- overwritten.
+ :param generation: (Optional) If present, creates writer for that
+ specific revision of that object. Use this to append data to an
+ existing Appendable Object.
+
+ Setting to ``0`` makes the `writer.open()` succeed only if
+ object doesn't exist in the bucket (useful for not accidentally
+ overwriting existing objects).
+
+ Warning: If `None`, a new object is created. If an object with the
+ same name already exists, it will be overwritten the moment 
+ `writer.open()` is called.
 
  :type write_handle: bytes
- :param write_handle: (Optional) An existing handle for writing the object.
- If provided, opening the bidi-gRPC connection will be faster.
+ :param write_handle: (Optional) An handle for writing the object.
+ If provided, opening the bidi-gRPC connection will be faster.
+
+ :type writer_options: dict
+ :param writer_options: (Optional) A dictionary of writer options.
+ Supported options:
+ - "FLUSH_INTERVAL_BYTES": int
+ The number of bytes to append before "persisting" data in GCS
+ servers. Default is `_DEFAULT_FLUSH_INTERVAL_BYTES`.
+ Must be a multiple of `_MAX_CHUNK_SIZE_BYTES`.
  """
 raise_if_no_fast_crc32c()
 self.client=client
@@ -133,7 +147,6 @@ def __init__(
 self.flush_interval=writer_options.get(
 "FLUSH_INTERVAL_BYTES", _DEFAULT_FLUSH_INTERVAL_BYTES
  )
-# TODO: add test case for this.
 ifself.flush_interval<_MAX_CHUNK_SIZE_BYTES:
 raiseexceptions.OutOfRange(
 f"flush_interval must be >= {_MAX_CHUNK_SIZE_BYTES} , but provided {self.flush_interval}"
@@ -346,6 +359,11 @@ async def finalize(self) -> _storage_v2.Object:
 self.offset=None
 returnself.object_resource
 
+@property
+defis_stream_open(self) ->bool:
+returnself._is_stream_open
+
+
 # helper methods.
 asyncdefappend_from_string(self, data: str):
 """

google/cloud/storage/_experimental/asyncio/async_write_object_stream.py

@@ -47,8 +47,17 @@ class _AsyncWriteObjectStream(_AsyncAbstractObjectStream):
  :param object_name: The name of the GCS ``Appendable Object`` to be write.
 
  :type generation_number: int
- :param generation_number: (Optional) If present, selects a specific revision of
- this object. If None, a new object is created.
+ :param generation_number: (Optional) If present, creates writer for that
+ specific revision of that object. Use this to append data to an
+ existing Appendable Object.
+
+ Setting to ``0`` makes the `writer.open()` succeed only if
+ object doesn't exist in the bucket (useful for not accidentally
+ overwriting existing objects).
+
+ Warning: If `None`, a new object is created. If an object with the
+ same name already exists, it will be overwritten the moment
+ `writer.open()` is called.
 
  :type write_handle: bytes
  :param write_handle: (Optional) An existing handle for writing the object.
@@ -101,13 +110,16 @@ async def open(self) -> None:
 # Create a new object or overwrite existing one if generation_number
 # is None. This makes it consistent with GCS JSON API behavior.
 # Created object type would be Appendable Object.
-ifself.generation_numberisNone:
+# if `generation_number` == 0 new object will be created only if there
+# isn't any existing object.
+ifself.generation_numberisNoneorself.generation_number==0:
 self.first_bidi_write_req=_storage_v2.BidiWriteObjectRequest(
 write_object_spec=_storage_v2.WriteObjectSpec(
 resource=_storage_v2.Object(
 name=self.object_name, bucket=self._full_bucket_name
  ),
 appendable=True,
+if_generation_match=self.generation_number,
  ),
  )
 else:
@@ -118,7 +130,6 @@ async def open(self) -> None:
 generation=self.generation_number,
  ),
  )
-
 self.socket_like_rpc=AsyncBidiRpc(
 self.rpc, initial_request=self.first_bidi_write_req, metadata=self.metadata
  )

tests/system/test_zonal.py

@@ -18,6 +18,7 @@
 fromgoogle.cloud.storage._experimental.asyncio.async_multi_range_downloaderimport (
 AsyncMultiRangeDownloader,
 )
+fromgoogle.api_core.exceptionsimportFailedPrecondition
 
 
 pytestmark=pytest.mark.skipif(
@@ -360,3 +361,74 @@ async def test_append_flushes_and_state_lookup(storage_client, blobs_to_delete):
 awaitmrd.close()
 content=buffer.getvalue()
 assertcontent==full_data
+
+@pytest.mark.asyncio
+asyncdeftest_open_with_generation_zero(storage_client, blobs_to_delete):
+"""Tests that using `generation=0` fails if the object already exists.
+
+ This test verifies that:
+ 1. An object can be created using `AsyncAppendableObjectWriter` with `generation=0`.
+ 2. Attempting to create the same object again with `generation=0` raises a
+ `FailedPrecondition` error with a 400 status code, because the
+ precondition (object must not exist) is not met.
+ """
+object_name=f"test_append_with_generation-{uuid.uuid4()}"
+grpc_client=AsyncGrpcClient().grpc_client
+writer=AsyncAppendableObjectWriter(grpc_client, _ZONAL_BUCKET, object_name, generation=0)
+
+# Empty object is created.
+awaitwriter.open()
+assertwriter.is_stream_open
+
+awaitwriter.close()
+assertnotwriter.is_stream_open
+
+
+withpytest.raises(FailedPrecondition) asexc_info:
+writer=AsyncAppendableObjectWriter(
+grpc_client, _ZONAL_BUCKET, object_name, generation=0
+ )
+awaitwriter.open()
+assertexc_info.value.code==400
+
+# cleanup
+delwriter
+gc.collect()
+
+blobs_to_delete.append(storage_client.bucket(_ZONAL_BUCKET).blob(object_name))
+
+@pytest.mark.asyncio
+asyncdeftest_open_existing_object_with_gen_None_overrides_existing(storage_client, blobs_to_delete):
+"""
+ Test that a new writer when specifies `None` overrides the existing object.
+ """
+object_name=f"test_append_with_generation-{uuid.uuid4()}"
+
+grpc_client=AsyncGrpcClient().grpc_client
+writer=AsyncAppendableObjectWriter(grpc_client, _ZONAL_BUCKET, object_name, generation=0)
+
+# Empty object is created.
+awaitwriter.open()
+assertwriter.is_stream_open
+old_gen=writer.generation
+
+
+awaitwriter.close()
+assertnotwriter.is_stream_open
+
+
+
+new_writer=AsyncAppendableObjectWriter(
+grpc_client, _ZONAL_BUCKET, object_name, generation=None
+ )
+awaitnew_writer.open()
+assertnew_writer.generation!=old_gen
+
+# assert exc_info.value.code == 400
+
+# cleanup
+delwriter
+delnew_writer
+gc.collect()
+
+blobs_to_delete.append(storage_client.bucket(_ZONAL_BUCKET).blob(object_name))

tests/unit/asyncio/test_async_appendable_object_writer.py

@@ -55,7 +55,7 @@ def test_init(mock_write_object_stream, mock_client):
 assertwriter.object_name==OBJECT
 assertwriter.generationisNone
 assertwriter.write_handleisNone
-assertnotwriter._is_stream_open
+assertnotwriter.is_stream_open
 assertwriter.offsetisNone
 assertwriter.persisted_sizeisNone
 assertwriter.bytes_appended_since_last_flush==0
@@ -225,7 +225,7 @@ async def test_open_appendable_object_writer(mock_write_object_stream, mock_clie
 
 # Assert
 mock_stream.open.assert_awaited_once()
-assertwriter._is_stream_open
+assertwriter.is_stream_open
 assertwriter.generation==GENERATION
 assertwriter.write_handle==WRITE_HANDLE
 assertwriter.persisted_size==0
@@ -255,7 +255,7 @@ async def test_open_appendable_object_writer_existing_object(
 
 # Assert
 mock_stream.open.assert_awaited_once()
-assertwriter._is_stream_open
+assertwriter.is_stream_open
 assertwriter.generation==GENERATION
 assertwriter.write_handle==WRITE_HANDLE
 assertwriter.persisted_size==PERSISTED_SIZE
@@ -379,7 +379,7 @@ async def test_close(mock_write_object_stream, mock_client):
 mock_stream.close.assert_awaited_once()
 assertwriter.offsetisNone
 assertpersisted_size==1024
-assertnotwriter._is_stream_open
+assertnotwriter.is_stream_open
 
 
 @pytest.mark.asyncio
@@ -415,7 +415,7 @@ async def test_finalize_on_close(mock_write_object_stream, mock_client):
 
 # Assert
 mock_stream.close.assert_awaited_once()
-assertnotwriter._is_stream_open
+assertnotwriter.is_stream_open
 assertwriter.offsetisNone
 assertwriter.object_resource==mock_resource
 assertwriter.persisted_size==2048
@@ -448,7 +448,7 @@ async def test_finalize(mock_write_object_stream, mock_client):
 assertwriter.object_resource==mock_resource
 assertwriter.persisted_size==123
 assertgcs_object==mock_resource
-assertwriter._is_stream_openisFalse
+assertnotwriter.is_stream_open
 assertwriter.offsetisNone
 
 

tests/unit/asyncio/test_async_write_object_stream.py

@@ -148,6 +148,42 @@ async def test_open_for_new_object(mock_async_bidi_rpc, mock_client):
 assertstream.persisted_size==0
 
 
+@pytest.mark.asyncio
+@mock.patch(
+"google.cloud.storage._experimental.asyncio.async_write_object_stream.AsyncBidiRpc"
+)
+asyncdeftest_open_for_new_object_with_generation_zero(mock_async_bidi_rpc, mock_client):
+"""Test opening a stream for a new object."""
+# Arrange
+socket_like_rpc=mock.AsyncMock()
+mock_async_bidi_rpc.return_value=socket_like_rpc
+socket_like_rpc.open=mock.AsyncMock()
+
+mock_response=mock.MagicMock(spec=_storage_v2.BidiWriteObjectResponse)
+mock_response.resource=mock.MagicMock(spec=_storage_v2.Object)
+mock_response.resource.generation=GENERATION
+mock_response.resource.size=0
+mock_response.write_handle=WRITE_HANDLE
+socket_like_rpc.recv=mock.AsyncMock(return_value=mock_response)
+
+stream=_AsyncWriteObjectStream(mock_client, BUCKET, OBJECT, generation_number=0)
+
+# Act
+awaitstream.open()
+
+# Assert
+mock_async_bidi_rpc.assert_called_once()
+_, call_kwargs=mock_async_bidi_rpc.call_args
+initial_request=call_kwargs["initial_request"]
+assertinitial_request.write_object_spec.if_generation_match==0
+assertstream._is_stream_open
+socket_like_rpc.open.assert_called_once()
+socket_like_rpc.recv.assert_called_once()
+assertstream.generation_number==GENERATION
+assertstream.write_handle==WRITE_HANDLE
+assertstream.persisted_size==0
+
+
 @pytest.mark.asyncio
 @mock.patch(
 "google.cloud.storage._experimental.asyncio.async_write_object_stream.AsyncBidiRpc"