feat(storage): support storage_class in AsyncAppendableObjectWriter

packages/google-cloud-storage/google/cloud/storage/asyncio/async_appendable_object_writer.py

@@ -49,6 +49,7 @@
 _BIDI_WRITE_REDIRECTED_TYPE_URL = (
     "type.googleapis.com/google.storage.v2.BidiWriteObjectRedirectedError"
 )
+_SUPPORTED_STORAGE_CLASSES = ("STANDARD", "RAPID")
 logger = logging.getLogger(__name__)
 
 
@@ -111,6 +112,7 @@ def __init__(
         generation: Optional[int] = None,
         write_handle: Optional[_storage_v2.BidiWriteHandle] = None,
         writer_options: Optional[dict] = None,
+        storage_class: Optional[str] = None,
     ):
         """
         Class for appending data to a GCS Appendable Object.
@@ -179,13 +181,26 @@ def __init__(
                 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`.
+        :type storage_class: Optional[str]
+        :param storage_class: (Optional) Storage class of the object bytes.
+            Possible values are STANDARD | RAPID. If specified,
+            it overrides the bucket's `storage_class`. If not, object storage class
+            will be the same as bucket's storage_class.
         """
         _utils.raise_if_no_fast_crc32c()
+        if (
+            storage_class is not None
+            and storage_class not in _SUPPORTED_STORAGE_CLASSES
+        ):
+            raise ValueError(
+                f"storage_class must be either 'STANDARD' or 'RAPID', got '{storage_class}'"
+            )
         self.client = client
         self.bucket_name = bucket_name
         self.object_name = object_name
         self.write_handle = write_handle
         self.generation = generation
+        self.storage_class = storage_class
 
         self.write_obj_stream: Optional[_AsyncWriteObjectStream] = None
         self._is_stream_open: bool = False
@@ -361,6 +376,7 @@ async def _do_open():
                 generation_number=self.generation,
                 write_handle=self.write_handle,
                 routing_token=self._routing_token,
+                storage_class=self.storage_class,
             )
 
             if self._routing_token:

packages/google-cloud-storage/google/cloud/storage/asyncio/async_write_object_stream.py

@@ -25,6 +25,8 @@
 )
 from google.cloud.storage.asyncio.async_grpc_client import AsyncGrpcClient
 
+_SUPPORTED_STORAGE_CLASSES = ("STANDARD", "RAPID")
+
 
 class _AsyncWriteObjectStream(_AsyncAbstractObjectStream):
     """Class representing a gRPC bidi-stream for writing data from a GCS
@@ -58,6 +60,10 @@ class _AsyncWriteObjectStream(_AsyncAbstractObjectStream):
     :type write_handle: _storage_v2.BidiWriteHandle
     :param write_handle: (Optional) An existing handle for writing the object.
                         If provided, opening the bidi-gRPC connection will be faster.
+
+    :type storage_class: Optional[str]
+    :param storage_class: (Optional) The storage class of the object.
+                        Could be either STANDARD | RAPID.
     """
 
     def __init__(
@@ -69,13 +75,21 @@ def __init__(
         write_handle: Optional[_storage_v2.BidiWriteHandle] = None,
         routing_token: Optional[str] = None,
         blob: Optional[Blob] = None,
+        storage_class: Optional[str] = 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")
+        if (
+            storage_class is not None
+            and storage_class not in _SUPPORTED_STORAGE_CLASSES
+        ):
+            raise ValueError(
+                f"storage_class must be either 'STANDARD' or 'RAPID', got '{storage_class}'"
+            )
 
         super().__init__(
             bucket_name=bucket_name,
@@ -86,6 +100,7 @@ def __init__(
         self.write_handle: Optional[_storage_v2.BidiWriteHandle] = write_handle
         self.routing_token: Optional[str] = routing_token
         self.blob: Optional[Blob] = blob
+        self.storage_class: Optional[str] = storage_class
         self._full_bucket_name = f"projects/_/buckets/{self.bucket_name}"
 
         self.rpc = self.client._client._transport._wrapped_methods[
@@ -124,7 +139,9 @@ async def open(self, metadata: Optional[List[Tuple[str, str]]] = None) -> None:
                 resource = _grpc_conversions.blob_to_proto(self.blob)
             else:
                 resource = _storage_v2.Object(
-                    name=self.object_name, bucket=self._full_bucket_name
+                    name=self.object_name,
+                    bucket=self._full_bucket_name,
+                    storage_class=self.storage_class,
                 )
             self.first_bidi_write_req = _storage_v2.BidiWriteObjectRequest(
                 write_object_spec=_storage_v2.WriteObjectSpec(

packages/google-cloud-storage/tests/unit/asyncio/test_async_appendable_object_writer.py

@@ -117,6 +117,7 @@ def mock_appendable_writer():
     yield {
         "mock_client": mock_client,
         "mock_stream": mock_stream,
+        "mock_stream_cls": mock_stream_cls,
     }
 
     stream_patcher.stop()
@@ -137,6 +138,24 @@ def test_init_defaults(self, mock_appendable_writer):
         assert writer.persisted_size is None
         assert writer.bytes_appended_since_last_flush == 0
         assert writer.flush_interval == _DEFAULT_FLUSH_INTERVAL_BYTES
+        assert writer.storage_class is None
+
+    @pytest.mark.parametrize("storage_class", ["STANDARD", "RAPID"])
+    def test_init_with_storage_class(self, mock_appendable_writer, storage_class):
+        writer = self._make_one(
+            mock_appendable_writer["mock_client"],
+            storage_class=storage_class,
+        )
+        assert writer.storage_class == storage_class
+
+    def test_init_with_invalid_storage_class_raises(self, mock_appendable_writer):
+        with pytest.raises(
+            ValueError, match="storage_class must be either 'STANDARD' or 'RAPID'"
+        ):
+            self._make_one(
+                mock_appendable_writer["mock_client"],
+                storage_class="INVALID",
+            )
 
     def test_init_with_writer_options(self, mock_appendable_writer):
         writer = self._make_one(
@@ -218,6 +237,36 @@ async def test_open_success(self, mock_appendable_writer):
         assert writer.generation == 456
         assert writer.write_handle == b"new-h"
         mock_appendable_writer["mock_stream"].open.assert_awaited_once()
+        mock_stream_cls = mock_appendable_writer["mock_stream_cls"]
+        assert mock_stream_cls.call_args.kwargs["storage_class"] is None
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("storage_class", ["STANDARD", "RAPID"])
+    async def test_open_passes_storage_class(
+        self, mock_appendable_writer, storage_class
+    ):
+        writer = self._make_one(
+            mock_appendable_writer["mock_client"],
+            storage_class=storage_class,
+        )
+        mock_appendable_writer["mock_stream"].generation_number = 456
+        mock_appendable_writer["mock_stream"].write_handle = b"new-h"
+        mock_appendable_writer["mock_stream"].persisted_size = 0
+
+        await writer.open()
+
+        assert writer._is_stream_open
+        mock_stream_cls = mock_appendable_writer["mock_stream_cls"]
+        mock_stream_cls.assert_called_once_with(
+            client=mock_appendable_writer["mock_client"].grpc_client,
+            bucket_name=BUCKET,
+            object_name=OBJECT,
+            blob=None,
+            generation_number=None,
+            write_handle=None,
+            routing_token=None,
+            storage_class=storage_class,
+        )
 
     def test_on_open_error_redirection(self, mock_appendable_writer):
         """Verify redirect info is extracted from helper."""

packages/google-cloud-storage/tests/unit/asyncio/test_async_write_object_stream.py

@@ -63,6 +63,22 @@ def test_init_basic(self, mock_client):
             ("x-goog-request-params", f"bucket={FULL_BUCKET_PATH}"),
         )
         assert not stream.is_stream_open
+        assert stream.storage_class is None
+
+    @pytest.mark.parametrize("storage_class", ["STANDARD", "RAPID"])
+    def test_init_with_storage_class(self, mock_client, storage_class):
+        stream = _AsyncWriteObjectStream(
+            mock_client, BUCKET, OBJECT, storage_class=storage_class
+        )
+        assert stream.storage_class == storage_class
+
+    def test_init_with_invalid_storage_class_raises(self, mock_client):
+        with pytest.raises(
+            ValueError, match="storage_class must be either 'STANDARD' or 'RAPID'"
+        ):
+            _AsyncWriteObjectStream(
+                mock_client, BUCKET, OBJECT, storage_class="INVALID"
+            )
 
     def test_init_raises_value_error(self, mock_client):
         with pytest.raises(ValueError, match="client must be provided"):
@@ -94,10 +110,46 @@ async def test_open_new_object(self, mock_rpc_cls, mock_client):
         await stream.open()
 
         # Check if BidiRpc was initialized with WriteObjectSpec
+        call_args = mock_rpc_cls.call_args
+        initial_request = call_args.kwargs["initial_request"]
+        # In proto3, string fields default to "" rather than None
+        resource = initial_request.write_object_spec.resource
+        assert "storage_class" not in resource
+        assert resource.storage_class == ""
+        assert initial_request.write_object_spec.appendable
+
+        assert stream.is_stream_open
+        assert stream.write_handle == WRITE_HANDLE
+        assert stream.generation_number == GENERATION
+
+    @mock.patch("google.cloud.storage.asyncio.async_write_object_stream.AsyncBidiRpc")
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("storage_class", ["STANDARD", "RAPID"])
+    async def test_open_new_object_with_storage_class(
+        self, mock_rpc_cls, mock_client, storage_class
+    ):
+        mock_rpc = mock_rpc_cls.return_value
+        mock_rpc.open = AsyncMock()
+
+        mock_response = MagicMock()
+        mock_response.persisted_size = 0
+        mock_response.resource.generation = GENERATION
+        mock_response.resource.size = 0
+        mock_response.write_handle = WRITE_HANDLE
+        mock_rpc.recv = AsyncMock(return_value=mock_response)
+
+        stream = _AsyncWriteObjectStream(
+            mock_client, BUCKET, OBJECT, storage_class=storage_class
+        )
+        await stream.open()
+
         call_args = mock_rpc_cls.call_args
         initial_request = call_args.kwargs["initial_request"]
         assert initial_request.write_object_spec is not None
         assert initial_request.write_object_spec.resource.name == OBJECT
+        assert (
+            initial_request.write_object_spec.resource.storage_class == storage_class
+        )
         assert initial_request.write_object_spec.appendable
 
         assert stream.is_stream_open