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..77cba17c4335 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[ @@ -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( 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..09a0a71daa39 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 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