diff --git a/packages/google-cloud-storage/google/cloud/storage/_grpc_conversions.py b/packages/google-cloud-storage/google/cloud/storage/_grpc_conversions.py index 2071b3cb07e0..84ae8e7679d8 100644 --- a/packages/google-cloud-storage/google/cloud/storage/_grpc_conversions.py +++ b/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", } diff --git a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_appendable_object_writer.py b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_appendable_object_writer.py index 7ef8ef099e6a..c1067e688543 100644 --- a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_appendable_object_writer.py +++ b/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: diff --git a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_write_object_stream.py b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_write_object_stream.py index 48f786c7654d..fb194f713825 100644 --- a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_write_object_stream.py +++ b/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,6 +75,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") @@ -76,6 +83,13 @@ def __init__( 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( diff --git a/packages/google-cloud-storage/google/cloud/storage/blob.py b/packages/google-cloud-storage/google/cloud/storage/blob.py index 87493a30dfc0..7013b1caf58c 100644 --- a/packages/google-cloud-storage/google/cloud/storage/blob.py +++ b/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? diff --git a/packages/google-cloud-storage/tests/unit/asyncio/test_async_appendable_object_writer.py b/packages/google-cloud-storage/tests/unit/asyncio/test_async_appendable_object_writer.py index 3775343f8818..3a044abdc29d 100644 --- a/packages/google-cloud-storage/tests/unit/asyncio/test_async_appendable_object_writer.py +++ b/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.""" diff --git a/packages/google-cloud-storage/tests/unit/asyncio/test_async_write_object_stream.py b/packages/google-cloud-storage/tests/unit/asyncio/test_async_write_object_stream.py index 2fcca43b0882..67dbe21cfc5f 100644 --- a/packages/google-cloud-storage/tests/unit/asyncio/test_async_write_object_stream.py +++ b/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 @@ -178,6 +230,7 @@ async def test_open_new_object_with_blob_sync_attrs( mock_blob.content_language = "content-language" mock_blob.temporary_hold = True mock_blob.event_based_hold = True + mock_blob.storage_class = "RAPID" custom_time = datetime.datetime( 2025, 1, 1, 12, 0, 0, tzinfo=datetime.timezone.utc @@ -217,6 +270,7 @@ async def test_open_new_object_with_blob_sync_attrs( assert resource.content_language == "content-language" assert resource.temporary_hold is True assert resource.event_based_hold is True + assert resource.storage_class == "RAPID" assert int(resource.custom_time.timestamp()) == int(custom_time.timestamp()) @@ -231,6 +285,50 @@ async def test_open_new_object_with_blob_sync_attrs( assert "context-key" in resource.contexts.custom assert resource.contexts.custom["context-key"].value == "context-value" + @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_real_blob_storage_class( + self, mock_rpc_cls, mock_client, storage_class + ): + mock_rpc = mock_rpc_cls.return_value + mock_rpc.open = AsyncMock() + mock_rpc.recv = AsyncMock(return_value=MagicMock(resource=None)) + + mock_bucket = mock.Mock(spec=Bucket) + mock_bucket.name = BUCKET + + blob = Blob(name=OBJECT, bucket=mock_bucket, storage_class=storage_class) + stream = _AsyncWriteObjectStream(mock_client, BUCKET, OBJECT, blob=blob) + await stream.open() + + initial_request = mock_rpc_cls.call_args.kwargs["initial_request"] + resource = initial_request.write_object_spec.resource + assert resource.storage_class == storage_class + assert resource.name == OBJECT + + @mock.patch("google.cloud.storage.asyncio.async_write_object_stream.AsyncBidiRpc") + @pytest.mark.asyncio + async def test_open_new_object_with_blob_and_stream_storage_class_fallback( + self, mock_rpc_cls, mock_client + ): + mock_rpc = mock_rpc_cls.return_value + mock_rpc.open = AsyncMock() + mock_rpc.recv = AsyncMock(return_value=MagicMock(resource=None)) + + mock_bucket = mock.Mock(spec=Bucket) + mock_bucket.name = BUCKET + + blob = Blob(name=OBJECT, bucket=mock_bucket) # storage_class is None + stream = _AsyncWriteObjectStream( + mock_client, BUCKET, OBJECT, blob=blob, storage_class="RAPID" + ) + await stream.open() + + initial_request = mock_rpc_cls.call_args.kwargs["initial_request"] + resource = initial_request.write_object_spec.resource + assert resource.storage_class == "RAPID" + @pytest.mark.asyncio async def test_open_already_open_raises(self, mock_client): stream = _AsyncWriteObjectStream(mock_client, BUCKET, OBJECT) diff --git a/packages/google-cloud-storage/tests/unit/test__grpc_conversions.py b/packages/google-cloud-storage/tests/unit/test__grpc_conversions.py index 2620ce3a5897..9c2e049fd543 100644 --- a/packages/google-cloud-storage/tests/unit/test__grpc_conversions.py +++ b/packages/google-cloud-storage/tests/unit/test__grpc_conversions.py @@ -15,6 +15,7 @@ import datetime from unittest import mock +import pytest from google.cloud import _storage_v2 from google.cloud.storage import _grpc_conversions @@ -33,6 +34,7 @@ def test_blob_to_proto_simple_fields(): "content_language", "temporary_hold", "event_based_hold", + "storage_class", "custom_time", "acl", "retention", @@ -49,6 +51,7 @@ def test_blob_to_proto_simple_fields(): blob.content_language = "en" blob.temporary_hold = True blob.event_based_hold = False + blob.storage_class = "STANDARD" blob.custom_time = None blob.acl = None blob.retention = None @@ -66,6 +69,7 @@ def test_blob_to_proto_simple_fields(): assert proto.content_language == "en" assert proto.temporary_hold is True assert proto.event_based_hold is False + assert proto.storage_class == "STANDARD" def test_blob_to_proto_custom_time(): @@ -158,3 +162,49 @@ def test_blob_to_proto_contexts(): assert "key" in proto.contexts.custom assert proto.contexts.custom["key"].value == "val" + + +@pytest.mark.parametrize("storage_class", ["STANDARD", "RAPID", "NEARLINE"]) +def test_blob_to_proto_storage_class(storage_class): + blob = mock.Mock(spec=["name", "bucket", "storage_class"]) + blob.name = "blob-name" + blob.bucket.name = "bucket-name" + blob.storage_class = storage_class + for attr in _grpc_conversions._BLOB_ATTR_TO_PROTO_FIELD: + if attr != "storage_class": + setattr(blob, attr, None) + blob.custom_time = None + blob.acl = None + blob.retention = None + blob.contexts = None + + proto = _grpc_conversions.blob_to_proto(blob) + assert proto.storage_class == storage_class + + +def test_blob_to_proto_storage_class_default(): + blob = mock.Mock(spec=["name", "bucket"]) + blob.name = "blob-name" + blob.bucket.name = "bucket-name" + for attr in _grpc_conversions._BLOB_ATTR_TO_PROTO_FIELD: + setattr(blob, attr, None) + blob.custom_time = None + blob.acl = None + blob.retention = None + blob.contexts = None + + proto = _grpc_conversions.blob_to_proto(blob) + assert proto.storage_class == "" + + +def test_blob_to_proto_real_blob(): + from google.cloud.storage.blob import Blob + from google.cloud.storage.bucket import Bucket + + bucket = mock.Mock(spec=Bucket) + bucket.name = "my-bucket" + blob = Blob(name="my-blob", bucket=bucket, storage_class="RAPID") + proto = _grpc_conversions.blob_to_proto(blob) + assert proto.name == "my-blob" + assert proto.bucket == "projects/_/buckets/my-bucket" + assert proto.storage_class == "RAPID" diff --git a/packages/google-cloud-storage/tests/unit/test_blob.py b/packages/google-cloud-storage/tests/unit/test_blob.py index c5609e3dd5c6..752a6f2bfd56 100644 --- a/packages/google-cloud-storage/tests/unit/test_blob.py +++ b/packages/google-cloud-storage/tests/unit/test_blob.py @@ -86,6 +86,7 @@ def test_ctor_wo_encryption_key(self): self.assertIs(blob._acl.blob, blob) self.assertEqual(blob._encryption_key, None) self.assertEqual(blob.kms_key_name, None) + self.assertIsNone(blob.storage_class) def test_ctor_with_encoded_unicode(self): blob_name = b"wet \xe2\x9b\xb5" @@ -139,6 +140,43 @@ def test_ctor_with_generation(self): blob = self._make_one(BLOB_NAME, bucket=bucket, generation=GENERATION) self.assertEqual(blob.generation, GENERATION) + def test_ctor_with_storage_class(self): + BLOB_NAME = "blob-name" + STORAGE_CLASS = "STANDARD" + bucket = _Bucket() + blob = self._make_one(BLOB_NAME, bucket=bucket, storage_class=STORAGE_CLASS) + self.assertEqual(blob.storage_class, STORAGE_CLASS) + self.assertEqual(blob._properties.get("storageClass"), STORAGE_CLASS) + + def test_ctor_with_storage_class_rapid(self): + BLOB_NAME = "blob-name" + STORAGE_CLASS = "RAPID" + bucket = _Bucket() + blob = self._make_one(BLOB_NAME, bucket=bucket, storage_class=STORAGE_CLASS) + self.assertEqual(blob.storage_class, STORAGE_CLASS) + self.assertEqual(blob._properties.get("storageClass"), STORAGE_CLASS) + + def test_ctor_with_storage_class_default(self): + BLOB_NAME = "blob-name" + bucket = _Bucket() + blob = self._make_one(BLOB_NAME, bucket=bucket) + self.assertIsNone(blob.storage_class) + self.assertNotIn("storageClass", blob._properties) + + def test_storage_class_property(self): + BLOB_NAME = "blob-name" + bucket = _Bucket() + blob = self._make_one(BLOB_NAME, bucket=bucket) + self.assertIsNone(blob.storage_class) + blob.storage_class = "STANDARD" + self.assertEqual(blob.storage_class, "STANDARD") + self.assertEqual(blob._properties.get("storageClass"), "STANDARD") + self.assertIn("storageClass", blob._changes) + + blob.storage_class = None + self.assertIsNone(blob.storage_class) + self.assertIsNone(blob._properties.get("storageClass")) + def _set_properties_helper(self, kms_key_name=None): from google.cloud._helpers import _RFC3339_MICROS