feat(zb-experimental): add async write object stream (#1612)

feat(zb-experimental): add async write object stream

---------

Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>

Author: chandra-siri

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

@@ -0,0 +1,134 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+NOTE:
+This is _experimental module for upcoming support for Rapid Storage.
+(https://cloud.google.com/blog/products/storage-data-transfer/high-performance-storage-innovations-for-ai-hpc#:~:text=your%20AI%20workloads%3A-,Rapid%20Storage,-%3A%20A%20new)
+
+APIs may not work as intended and are not stable yet. Feature is not
+GA(Generally Available) yet, please contact your TAM(Technical Account Manager)
+if you want to use these Rapid Storage APIs.
+
+"""
+from typing import Optional
+from google.cloud import _storage_v2
+from google.cloud.storage._experimental.asyncio.async_grpc_client import AsyncGrpcClient
+from google.cloud.storage._experimental.asyncio.async_abstract_object_stream import (
+    _AsyncAbstractObjectStream,
+)
+from google.api_core.bidi_async import AsyncBidiRpc
+
+
+class _AsyncWriteObjectStream(_AsyncAbstractObjectStream):
+    """Class representing a gRPC bidi-stream for writing data from a GCS
+      ``Appendable Object``.
+
+    This class provides a unix socket-like interface to a GCS ``Object``, with
+    methods like ``open``, ``close``, ``send``, and ``recv``.
+
+    :type client: :class:`~google.cloud.storage._experimental.asyncio.async_grpc_client.AsyncGrpcClient.grpc_client`
+    :param client: async grpc client to use for making API requests.
+
+    :type bucket_name: str
+    :param bucket_name: The name of the GCS ``bucket`` containing the object.
+
+    :type object_name: str
+    :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.
+
+    :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.
+    """
+
+    def __init__(
+        self,
+        client: AsyncGrpcClient.grpc_client,
+        bucket_name: str,
+        object_name: str,
+        generation_number: Optional[int] = None,  # None means new object
+        write_handle: Optional[bytes] = None,
+    ) -> None:
+        if client is None:
+            raise ValueError("client must be provided")
+        if bucket_name is None:
+            raise ValueError("bucket_name must be provided")
+        if object_name is None:
+            raise ValueError("object_name must be provided")
+
+        super().__init__(
+            bucket_name=bucket_name,
+            object_name=object_name,
+            generation_number=generation_number,
+        )
+        self.client: AsyncGrpcClient.grpc_client = client
+        self.write_handle: Optional[bytes] = write_handle
+
+        self._full_bucket_name = f"projects/_/buckets/{self.bucket_name}"
+
+        self.rpc = self.client._client._transport._wrapped_methods[
+            self.client._client._transport.bidi_write_object
+        ]
+
+        self.metadata = (("x-goog-request-params", f"bucket={self._full_bucket_name}"),)
+        self.socket_like_rpc: Optional[AsyncBidiRpc] = None
+        self._is_stream_open: bool = False
+        self.first_bidi_write_req = None
+        self.persisted_size = 0
+        self.object_resource: Optional[_storage_v2.Object] = None
+
+    async def open(self) -> None:
+        """Opening an object for write , should do it's state lookup
+        to know what's the persisted size is.
+        """
+        raise NotImplementedError(
+            "open() is not implemented yet in _AsyncWriteObjectStream"
+        )
+
+    async def close(self) -> None:
+        """Closes the bidi-gRPC connection."""
+        raise NotImplementedError(
+            "close() is not implemented yet in _AsyncWriteObjectStream"
+        )
+
+    async def send(
+        self, bidi_write_object_request: _storage_v2.BidiWriteObjectRequest
+    ) -> None:
+        """Sends a request message on the stream.
+
+        Args:
+            bidi_write_object_request (:class:`~google.cloud._storage_v2.types.BidiReadObjectRequest`):
+                The request message to send. This is typically used to specify
+                the read offset and limit.
+        """
+        raise NotImplementedError(
+            "send() is not implemented yet in _AsyncWriteObjectStream"
+        )
+
+    async def recv(self) -> _storage_v2.BidiWriteObjectResponse:
+        """Receives a response from the stream.
+
+        This method waits for the next message from the server, which could
+        contain object data or metadata.
+
+        Returns:
+            :class:`~google.cloud._storage_v2.types.BidiWriteObjectResponse`:
+                The response message from the server.
+        """
+        raise NotImplementedError(
+            "recv() is not implemented yet in _AsyncWriteObjectStream"
+        )

google/cloud/storage/client.py

@@ -73,6 +73,7 @@ def _buckets_page_start(iterator, page, response):
         )
     page.unreachable = unreachable
 
+
 class Client(ClientWithProject):
     """Client to bundle configuration needed for API requests.
 

tests/unit/asyncio/test_async_write_object_stream.py

@@ -0,0 +1,108 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+from unittest import mock
+
+from google.cloud.storage._experimental.asyncio.async_write_object_stream import (
+    _AsyncWriteObjectStream,
+)
+from google.cloud import _storage_v2
+
+BUCKET = "my-bucket"
+OBJECT = "my-object"
+
+
+@pytest.fixture
+def mock_client():
+    """Mock the async gRPC client."""
+    mock_transport = mock.AsyncMock()
+    mock_transport.bidi_write_object = mock.sentinel.bidi_write_object
+    mock_transport._wrapped_methods = {
+        mock.sentinel.bidi_write_object: mock.sentinel.wrapped_bidi_write_object
+    }
+
+    mock_gapic_client = mock.AsyncMock()
+    mock_gapic_client._transport = mock_transport
+
+    client = mock.AsyncMock()
+    client._client = mock_gapic_client
+    return client
+
+
+def test_async_write_object_stream_init(mock_client):
+    """Test the constructor of _AsyncWriteObjectStream."""
+    stream = _AsyncWriteObjectStream(mock_client, BUCKET, OBJECT)
+
+    assert stream.client == mock_client
+    assert stream.bucket_name == BUCKET
+    assert stream.object_name == OBJECT
+    assert stream.generation_number is None
+    assert stream.write_handle is None
+    assert stream._full_bucket_name == f"projects/_/buckets/{BUCKET}"
+    assert stream.rpc == mock.sentinel.wrapped_bidi_write_object
+    assert stream.metadata == (
+        ("x-goog-request-params", f"bucket=projects/_/buckets/{BUCKET}"),
+    )
+    assert stream.socket_like_rpc is None
+    assert not stream._is_stream_open
+    assert stream.first_bidi_write_req is None
+    assert stream.persisted_size == 0
+    assert stream.object_resource is None
+
+
+def test_async_write_object_stream_init_with_generation_and_handle(mock_client):
+    """Test the constructor with optional arguments."""
+    generation = 12345
+    write_handle = b"test-handle"
+    stream = _AsyncWriteObjectStream(
+        mock_client,
+        BUCKET,
+        OBJECT,
+        generation_number=generation,
+        write_handle=write_handle,
+    )
+
+    assert stream.generation_number == generation
+    assert stream.write_handle == write_handle
+
+
+def test_async_write_object_stream_init_raises_value_error():
+    """Test that the constructor raises ValueError for missing arguments."""
+    with pytest.raises(ValueError, match="client must be provided"):
+        _AsyncWriteObjectStream(None, BUCKET, OBJECT)
+
+    with pytest.raises(ValueError, match="bucket_name must be provided"):
+        _AsyncWriteObjectStream(mock.Mock(), None, OBJECT)
+
+    with pytest.raises(ValueError, match="object_name must be provided"):
+        _AsyncWriteObjectStream(mock.Mock(), BUCKET, None)
+
+
+@pytest.mark.asyncio
+async def test_unimplemented_methods_raise_error(mock_client):
+    """Test that unimplemented methods raise NotImplementedError."""
+    stream = _AsyncWriteObjectStream(mock_client, BUCKET, OBJECT)
+
+    with pytest.raises(NotImplementedError):
+        await stream.open()
+
+    with pytest.raises(NotImplementedError):
+        await stream.close()
+
+    with pytest.raises(NotImplementedError):
+        await stream.send(_storage_v2.BidiWriteObjectRequest())
+
+    with pytest.raises(NotImplementedError):
+        await stream.recv()

tests/unit/test_client.py

@@ -3145,6 +3145,7 @@ def test_list_buckets_w_partial_success(self):
             page_start=_buckets_page_start,
         )
 
+
 class Test__item_to_bucket(unittest.TestCase):
     def _call_fut(self, iterator, item):
         from google.cloud.storage.client import _item_to_bucket