feat(storage): support storage_class in AsyncAppendableObjectWriter

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

@@ -111,6 +111,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 +180,18 @@ 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.
+            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()
         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 +367,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

@@ -58,6 +58,9 @@ 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.
     """
 
     def __init__(
@@ -69,6 +72,7 @@ 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")
@@ -86,6 +90,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 +129,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/system/conftest.py

@@ -30,13 +30,17 @@
 
 import pytest  # noqa: E402
 from google.api_core import exceptions  # noqa: E402
+from google.api_core.client_options import ClientOptions  # noqa: E402
 
 from google.cloud import kms  # noqa: E402
 from google.cloud.storage._helpers import _base64_md5hash  # noqa: E402
 from google.cloud.storage.retry import DEFAULT_RETRY  # noqa: E402
 
 from . import _helpers  # noqa: E402
 
+PREPROD_JSON_HOST = "https://storage-preprod-test-unified.googleusercontent.com"
+RUN_SYSTEM_TESTS_ON_PREPROD = os.getenv("RUN_SYSTEM_TESTS_ON_PREPROD") == "True"
+
 if trace_api is not None:
     _global_exporter = InMemorySpanExporter()
     _provider = TracerProvider()
@@ -112,10 +116,19 @@ def _kms_key_name(client, bucket, key_name):
 
 
 @pytest.fixture(scope="session")
-def storage_client():
+def run_system_tests_on_preprod():
+    return os.getenv("RUN_SYSTEM_TESTS_ON_PREPROD") == "True"
+
+
+@pytest.fixture(scope="session")
+def storage_client(run_system_tests_on_preprod):
     from google.cloud.storage import Client
 
-    client = Client()
+    client = Client(
+        client_options=ClientOptions(api_endpoint=PREPROD_JSON_HOST)
+        if run_system_tests_on_preprod
+        else None
+    )
     with contextlib.closing(client):
         yield client