feat(storage): support storage_class in _grpc_conversions

packages/google-cloud-storage/google/cloud/storage/_grpc_conversions.py

@@ -27,6 +27,7 @@
     "content_language": "content_language",
     "temporary_hold": "temporary_hold",
     "event_based_hold": "event_based_hold",
+    "storage_class": "storage_class",
 }
 
 

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[
@@ -122,9 +137,13 @@ async def open(self, metadata: Optional[List[Tuple[str, str]]] = None) -> None:
         if self.generation_number is None or self.generation_number == 0:
             if self.blob:
                 resource = _grpc_conversions.blob_to_proto(self.blob)
+                if not resource.storage_class and self.storage_class:
+                    resource.storage_class = self.storage_class
             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/google/cloud/storage/blob.py

@@ -183,6 +183,11 @@ class Blob(_PropertyMixin):
     :type generation: long
     :param generation:
         (Optional) If present, selects a specific revision of this object.
+
+    :type storage_class: str
+    :param storage_class:
+        (Optional) The storage class for the blob. Default value is None. If
+        nothing specified, its value will be the same as bucket's storage_class.
     """
 
     _chunk_size = None  # Default value for each instance.
@@ -216,6 +221,7 @@ def __init__(
         encryption_key=None,
         kms_key_name=None,
         generation=None,
+        storage_class=None,
     ):
         """
         property :attr:`name`
@@ -239,6 +245,9 @@ def __init__(
         if generation is not None:
             self._properties["generation"] = generation
 
+        if storage_class is not None:
+            self._properties["storageClass"] = storage_class
+
     @property
     def bucket(self):
         """Bucket which contains the object.
@@ -4943,28 +4952,42 @@ def kms_key_name(self, value):
         """
         self._patch_property("kmsKeyName", value)
 
-    storage_class = _scalar_property("storageClass")
-    """Retrieve the storage class for the object.
+    @property
+    def storage_class(self):
+        """Retrieve the storage class for the object.
 
-    This can only be set at blob / object **creation** time. If you'd
-    like to change the storage class **after** the blob / object already
-    exists in a bucket, call :meth:`update_storage_class` (which uses
-    :meth:`rewrite`).
+        Default value is None. If nothing specified, its value will be the
+        same as bucket's storage_class.
 
-    See https://cloud.google.com/storage/docs/storage-classes
+        This can only be set at blob / object **creation** time. If you'd
+        like to change the storage class **after** the blob / object already
+        exists in a bucket, call :meth:`update_storage_class` (which uses
+        :meth:`rewrite`).
 
-    :rtype: str or ``NoneType``
-    :returns:
-        If set, one of
-        :attr:`~google.cloud.storage.constants.STANDARD_STORAGE_CLASS`,
-        :attr:`~google.cloud.storage.constants.NEARLINE_STORAGE_CLASS`,
-        :attr:`~google.cloud.storage.constants.COLDLINE_STORAGE_CLASS`,
-        :attr:`~google.cloud.storage.constants.ARCHIVE_STORAGE_CLASS`,
-        :attr:`~google.cloud.storage.constants.MULTI_REGIONAL_LEGACY_STORAGE_CLASS`,
-        :attr:`~google.cloud.storage.constants.REGIONAL_LEGACY_STORAGE_CLASS`,
-        :attr:`~google.cloud.storage.constants.DURABLE_REDUCED_AVAILABILITY_STORAGE_CLASS`,
-        else ``None``.
-    """
+        See https://cloud.google.com/storage/docs/storage-classes
+
+        :rtype: str or ``NoneType``
+        :returns:
+            If set, one of
+            :attr:`~google.cloud.storage.constants.STANDARD_STORAGE_CLASS`,
+            :attr:`~google.cloud.storage.constants.NEARLINE_STORAGE_CLASS`,
+            :attr:`~google.cloud.storage.constants.COLDLINE_STORAGE_CLASS`,
+            :attr:`~google.cloud.storage.constants.ARCHIVE_STORAGE_CLASS`,
+            :attr:`~google.cloud.storage.constants.MULTI_REGIONAL_LEGACY_STORAGE_CLASS`,
+            :attr:`~google.cloud.storage.constants.REGIONAL_LEGACY_STORAGE_CLASS`,
+            :attr:`~google.cloud.storage.constants.DURABLE_REDUCED_AVAILABILITY_STORAGE_CLASS`,
+            else ``None``.
+        """
+        return self._properties.get("storageClass")
+
+    @storage_class.setter
+    def storage_class(self, value):
+        """Set the storage class for the object.
+
+        :type value: str or ``NoneType``
+        :param value: new storage class name (None to clear any existing storage class).
+        """
+        self._patch_property("storageClass", value)
 
     temporary_hold = _scalar_property("temporaryHold")
     """Is a temporary hold active on the object?

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."""