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

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.
  if self.flush_interval < _MAX_CHUNK_SIZE_BYTES:
  raise exceptions.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
  return self.object_resource
 
+ @property
+ def is_stream_open(self) -> bool:
+ return self._is_stream_open
+
+
  # helper methods.
  async def append_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.
- if self.generation_number is None:
+ # if `generation_number` == 0 new object will be created only if there
+ # isn't any existing object.
+ if self.generation_number is None or self.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 @@
 from google.cloud.storage._experimental.asyncio.async_multi_range_downloader import (
  AsyncMultiRangeDownloader,
 )
+from google.api_core.exceptions import FailedPrecondition
 
 
 pytestmark = pytest.mark.skipif(
@@ -360,3 +361,74 @@ async def test_append_flushes_and_state_lookup(storage_client, blobs_to_delete):
  await mrd.close()
  content = buffer.getvalue()
  assert content == full_data
+
+@pytest.mark.asyncio
+async def test_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.
+ await writer.open()
+ assert writer.is_stream_open
+
+ await writer.close()
+ assert not writer.is_stream_open
+
+
+ with pytest.raises(FailedPrecondition) as exc_info:
+ writer = AsyncAppendableObjectWriter(
+ grpc_client, _ZONAL_BUCKET, object_name, generation=0
+ )
+ await writer.open()
+ assert exc_info.value.code == 400
+
+ # cleanup
+ del writer
+ gc.collect()
+
+ blobs_to_delete.append(storage_client.bucket(_ZONAL_BUCKET).blob(object_name))
+
+@pytest.mark.asyncio
+async def test_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.
+ await writer.open()
+ assert writer.is_stream_open
+ old_gen = writer.generation
+
+
+ await writer.close()
+ assert not writer.is_stream_open
+
+
+
+ new_writer = AsyncAppendableObjectWriter(
+ grpc_client, _ZONAL_BUCKET, object_name, generation=None
+ )
+ await new_writer.open()
+ assert new_writer.generation != old_gen
+
+ # assert exc_info.value.code == 400
+
+ # cleanup
+ del writer
+ del new_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):
  assert writer.object_name == OBJECT
  assert writer.generation is None
  assert writer.write_handle is None
- assert not writer._is_stream_open
+ assert not writer.is_stream_open
  assert writer.offset is None
  assert writer.persisted_size is None
  assert writer.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()
- assert writer._is_stream_open
+ assert writer.is_stream_open
  assert writer.generation == GENERATION
  assert writer.write_handle == WRITE_HANDLE
  assert writer.persisted_size == 0
@@ -255,7 +255,7 @@ async def test_open_appendable_object_writer_existing_object(
 
  # Assert
  mock_stream.open.assert_awaited_once()
- assert writer._is_stream_open
+ assert writer.is_stream_open
  assert writer.generation == GENERATION
  assert writer.write_handle == WRITE_HANDLE
  assert writer.persisted_size == PERSISTED_SIZE
@@ -379,7 +379,7 @@ async def test_close(mock_write_object_stream, mock_client):
  mock_stream.close.assert_awaited_once()
  assert writer.offset is None
  assert persisted_size == 1024
- assert not writer._is_stream_open
+ assert not writer.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()
- assert not writer._is_stream_open
+ assert not writer.is_stream_open
  assert writer.offset is None
  assert writer.object_resource == mock_resource
  assert writer.persisted_size == 2048
@@ -448,7 +448,7 @@ async def test_finalize(mock_write_object_stream, mock_client):
  assert writer.object_resource == mock_resource
  assert writer.persisted_size == 123
  assert gcs_object == mock_resource
- assert writer._is_stream_open is False
+ assert not writer.is_stream_open
  assert writer.offset is None
 
 

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):
  assert stream.persisted_size == 0
 
 
+@pytest.mark.asyncio
+@mock.patch(
+ "google.cloud.storage._experimental.asyncio.async_write_object_stream.AsyncBidiRpc"
+)
+async def test_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
+ await stream.open()
+
+ # Assert
+ mock_async_bidi_rpc.assert_called_once()
+ _, call_kwargs = mock_async_bidi_rpc.call_args
+ initial_request = call_kwargs["initial_request"]
+ assert initial_request.write_object_spec.if_generation_match == 0
+ assert stream._is_stream_open
+ socket_like_rpc.open.assert_called_once()
+ socket_like_rpc.recv.assert_called_once()
+ assert stream.generation_number == GENERATION
+ assert stream.write_handle == WRITE_HANDLE
+ assert stream.persisted_size == 0
+
+
 @pytest.mark.asyncio
 @mock.patch(
  "google.cloud.storage._experimental.asyncio.async_write_object_stream.AsyncBidiRpc"