Visitar URL original
feat(storage): support storage_class via Blob in AsyncAppendableObjec… · googleapis/google-cloud-python@b0872f2 · GitHub
Skip to content

Commit b0872f2

Browse files
chandra-sirichandra-siri
authored
feat(storage): support storage_class via Blob in AsyncAppendableObjectWriter (#18318)
## Description Adds support for propagating `storage_class` when creating an `AsyncAppendableObjectWriter` via `AsyncAppendableObjectWriter.from_blob(client, blob)`, mapping `storage_class` in `_grpc_conversions.blob_to_proto`, and updating system and unit tests accordingly. ### Changes - **Conversions & Writer**: - Added `"storage_class": "storage_class"` to `_BLOB_ATTR_TO_PROTO_FIELD` in `_grpc_conversions.py`. - Added `storage_class=blob.storage_class` in `AsyncAppendableObjectWriter.from_blob`. - **Unit Tests**: - Added unit tests for `storage_class` mapping in `test__grpc_conversions.py`. - Added unit tests for `storage_class` handling in `AsyncAppendableObjectWriter.from_blob` in `test_async_appendable_object_writer.py`. - **System Tests**: - Enabled `test_write_from_blob` and `test_write_blob_with_contexts` under RCU by setting `blob.storage_class = "RAPID"` when `RCU_SYSTEM_TESTS` is active.
1 parent dc30df2 commit b0872f2

6 files changed

Lines changed: 167 additions & 62 deletions

File tree

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
"content_language": "content_language",
2828
"temporary_hold": "temporary_hold",
2929
"event_based_hold": "event_based_hold",
30+
"storage_class": "storage_class",
3031
}
3132

3233

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -269,6 +269,7 @@ def from_blob(
269269
generation=blob.generation,
270270
write_handle=write_handle,
271271
writer_options=writer_options,
272+
storage_class=blob.storage_class,
272273
)
273274
instance.blob = blob
274275
return instance

‎packages/google-cloud-storage/tests/system/test_zonal.py‎

Lines changed: 39 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -422,6 +422,8 @@ async def _run():
422422
blob.content_encoding = content_encoding
423423
blob.content_language = content_language
424424
blob.custom_time = custom_time
425+
if RCU_SYSTEM_TESTS:
426+
blob.storage_class = "RAPID"
425427

426428
# 2. Use from_blob to create the writer
427429
writer = AsyncAppendableObjectWriter.from_blob(grpc_client, blob)
@@ -492,41 +494,45 @@ async def _run():
492494
event_loop.run_until_complete(_run())
493495

494496

495-
@pytest.mark.skipif(
496-
RCU_SYSTEM_TESTS, reason="Write blob with contexts for RCU not supported in SDK yet"
497-
)
498-
@pytest.mark.asyncio
499-
async def test_write_blob_with_contexts(storage_client, blobs_to_delete):
500-
async_client = await create_async_grpc_client()
501-
blob_name = f"ObjectContextsGrpc-{uuid.uuid4().hex}"
502-
503-
bucket = storage_client.bucket(_BUCKET_UNDER_TEST)
504-
blob = bucket.blob(blob_name)
505-
blob.contexts = ObjectContexts(
506-
blob, custom={"foo": ObjectCustomContextPayload(value="bar")}
507-
)
508-
writer = AsyncAppendableObjectWriter.from_blob(async_client, blob)
509-
await writer.open()
510-
await writer.append(b"grpc-test")
511-
await writer.close(finalize_on_close=True)
497+
def test_write_blob_with_contexts(
498+
event_loop, grpc_client_direct, storage_client, blobs_to_delete
499+
):
500+
async def _run():
501+
bucket = storage_client.bucket(_BUCKET_UNDER_TEST)
502+
blob_name = f"ObjectContextsGrpc-{uuid.uuid4().hex}"
503+
blob = bucket.blob(blob_name)
504+
blob.contexts = ObjectContexts(
505+
blob, custom={"foo": ObjectCustomContextPayload(value="bar")}
506+
)
507+
if RCU_SYSTEM_TESTS:
508+
blob.storage_class = "RAPID"
512509

513-
try:
514-
blobs = list(
515-
storage_client.list_blobs(
516-
_BUCKET_UNDER_TEST, filter_='contexts."foo"="bar"'
510+
writer = AsyncAppendableObjectWriter.from_blob(grpc_client_direct, blob)
511+
await writer.open()
512+
await writer.append(b"grpc-test")
513+
await writer.close(finalize_on_close=True)
514+
515+
try:
516+
blobs = list(
517+
storage_client.list_blobs(
518+
_BUCKET_UNDER_TEST, filter_='contexts."foo"="bar"'
519+
)
517520
)
518-
)
519-
names = [b.name for b in blobs]
520-
assert blob_name in names
521-
522-
# Assert contexts via gRPC GetObject
523-
obj_proto = await async_client.get_object(_BUCKET_UNDER_TEST, blob_name)
524-
assert "foo" in obj_proto.contexts.custom
525-
assert obj_proto.contexts.custom["foo"].value == "bar"
526-
finally:
527-
blobs_to_delete.append(
528-
storage_client.bucket(_BUCKET_UNDER_TEST).blob(blob_name)
529-
)
521+
names = [b.name for b in blobs]
522+
assert blob_name in names
523+
524+
# Assert contexts via gRPC GetObject
525+
obj_proto = await grpc_client_direct.get_object(
526+
_BUCKET_UNDER_TEST, blob_name
527+
)
528+
assert "foo" in obj_proto.contexts.custom
529+
assert obj_proto.contexts.custom["foo"].value == "bar"
530+
finally:
531+
blobs_to_delete.append(
532+
storage_client.bucket(_BUCKET_UNDER_TEST).blob(blob_name)
533+
)
534+
535+
event_loop.run_until_complete(_run())
530536

531537

532538
def test_read_unfinalized_appendable_object(

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

Lines changed: 82 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -113,6 +113,13 @@ def mock_appendable_writer():
113113
mock_stream.persisted_size = 0
114114
mock_stream.generation_number = GENERATION
115115
mock_stream.write_handle = WRITE_HANDLE
116+
mock_stream.storage_class = None
117+
118+
def _create_stream(*args, **kwargs):
119+
mock_stream.storage_class = kwargs.get("storage_class")
120+
return mock_stream
121+
122+
mock_stream_cls.side_effect = _create_stream
116123

117124
yield {
118125
"mock_client": mock_client,
@@ -140,14 +147,35 @@ def test_init_defaults(self, mock_appendable_writer):
140147
assert writer.flush_interval == _DEFAULT_FLUSH_INTERVAL_BYTES
141148
assert writer.storage_class is None
142149

150+
@pytest.mark.asyncio
143151
@pytest.mark.parametrize("storage_class", ["STANDARD", "RAPID"])
144-
def test_init_with_storage_class(self, mock_appendable_writer, storage_class):
152+
async def test_init_with_storage_class(self, mock_appendable_writer, storage_class):
145153
writer = self._make_one(
146154
mock_appendable_writer["mock_client"],
147155
storage_class=storage_class,
148156
)
149157
assert writer.storage_class == storage_class
150158

159+
mock_appendable_writer["mock_stream"].generation_number = 456
160+
mock_appendable_writer["mock_stream"].write_handle = b"new-h"
161+
mock_appendable_writer["mock_stream"].persisted_size = 0
162+
163+
await writer.open()
164+
165+
assert writer._is_stream_open
166+
mock_stream_cls = mock_appendable_writer["mock_stream_cls"]
167+
mock_stream_cls.assert_called_once_with(
168+
client=mock_appendable_writer["mock_client"].grpc_client,
169+
bucket_name=BUCKET,
170+
object_name=OBJECT,
171+
blob=None,
172+
generation_number=None,
173+
write_handle=None,
174+
routing_token=None,
175+
storage_class=storage_class,
176+
)
177+
assert writer.write_obj_stream.storage_class == storage_class
178+
151179
def test_init_with_writer_options(self, mock_appendable_writer):
152180
writer = self._make_one(
153181
mock_appendable_writer["mock_client"],
@@ -182,6 +210,7 @@ def test_from_blob(self, mock_appendable_writer):
182210
mock_blob.name = OBJECT
183211
mock_blob.bucket.name = BUCKET
184212
mock_blob.generation = GENERATION
213+
mock_blob.storage_class = "RAPID"
185214

186215
writer = AsyncAppendableObjectWriter.from_blob(
187216
mock_appendable_writer["mock_client"],
@@ -194,6 +223,57 @@ def test_from_blob(self, mock_appendable_writer):
194223
assert writer.generation == GENERATION
195224
assert writer.flush_interval == EIGHT_MIB
196225
assert writer.blob == mock_blob
226+
assert writer.storage_class == "RAPID"
227+
228+
@pytest.mark.parametrize("storage_class", ["STANDARD", "RAPID", None])
229+
def test_from_blob_storage_class(self, mock_appendable_writer, storage_class):
230+
mock_blob = mock.Mock(spec=Blob)
231+
mock_blob.name = OBJECT
232+
mock_blob.bucket.name = BUCKET
233+
mock_blob.generation = GENERATION
234+
mock_blob.storage_class = storage_class
235+
236+
writer = AsyncAppendableObjectWriter.from_blob(
237+
mock_appendable_writer["mock_client"],
238+
mock_blob,
239+
)
240+
241+
assert writer.storage_class == storage_class
242+
243+
@pytest.mark.asyncio
244+
@pytest.mark.parametrize("storage_class", ["STANDARD", "RAPID"])
245+
async def test_from_blob_open_passes_storage_class(
246+
self, mock_appendable_writer, storage_class
247+
):
248+
mock_blob = mock.Mock(spec=Blob)
249+
mock_blob.name = OBJECT
250+
mock_blob.bucket.name = BUCKET
251+
mock_blob.generation = GENERATION
252+
mock_blob.storage_class = storage_class
253+
254+
writer = AsyncAppendableObjectWriter.from_blob(
255+
mock_appendable_writer["mock_client"],
256+
mock_blob,
257+
)
258+
mock_appendable_writer["mock_stream"].generation_number = 456
259+
mock_appendable_writer["mock_stream"].write_handle = b"new-h"
260+
mock_appendable_writer["mock_stream"].persisted_size = 0
261+
262+
await writer.open()
263+
264+
assert writer._is_stream_open
265+
mock_stream_cls = mock_appendable_writer["mock_stream_cls"]
266+
mock_stream_cls.assert_called_once_with(
267+
client=mock_appendable_writer["mock_client"].grpc_client,
268+
bucket_name=BUCKET,
269+
object_name=OBJECT,
270+
blob=mock_blob,
271+
generation_number=GENERATION,
272+
write_handle=None,
273+
routing_token=None,
274+
storage_class=storage_class,
275+
)
276+
assert writer.write_obj_stream.storage_class == storage_class
197277

198278
# -------------------------------------------------------------------------
199279
# Stream Lifecycle Tests
@@ -230,34 +310,7 @@ async def test_open_success(self, mock_appendable_writer):
230310
mock_appendable_writer["mock_stream"].open.assert_awaited_once()
231311
mock_stream_cls = mock_appendable_writer["mock_stream_cls"]
232312
assert mock_stream_cls.call_args.kwargs["storage_class"] is None
233-
234-
@pytest.mark.asyncio
235-
@pytest.mark.parametrize("storage_class", ["STANDARD", "RAPID"])
236-
async def test_open_passes_storage_class(
237-
self, mock_appendable_writer, storage_class
238-
):
239-
writer = self._make_one(
240-
mock_appendable_writer["mock_client"],
241-
storage_class=storage_class,
242-
)
243-
mock_appendable_writer["mock_stream"].generation_number = 456
244-
mock_appendable_writer["mock_stream"].write_handle = b"new-h"
245-
mock_appendable_writer["mock_stream"].persisted_size = 0
246-
247-
await writer.open()
248-
249-
assert writer._is_stream_open
250-
mock_stream_cls = mock_appendable_writer["mock_stream_cls"]
251-
mock_stream_cls.assert_called_once_with(
252-
client=mock_appendable_writer["mock_client"].grpc_client,
253-
bucket_name=BUCKET,
254-
object_name=OBJECT,
255-
blob=None,
256-
generation_number=None,
257-
write_handle=None,
258-
routing_token=None,
259-
storage_class=storage_class,
260-
)
313+
assert writer.write_obj_stream.storage_class is None
261314

262315
def test_on_open_error_redirection(self, mock_appendable_writer):
263316
"""Verify redirect info is extracted from helper."""

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

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,17 @@ def test_init_with_storage_class(self, mock_client, storage_class):
7272
)
7373
assert stream.storage_class == storage_class
7474

75+
@pytest.mark.parametrize("storage_class", ["STANDARD", "RAPID"])
76+
def test_init_with_blob_storage_class(self, mock_client, storage_class):
77+
mock_blob = mock.Mock(spec=Blob)
78+
mock_blob.name = OBJECT
79+
mock_blob.bucket = mock.Mock(spec=Bucket)
80+
mock_blob.bucket.name = BUCKET
81+
mock_blob.storage_class = storage_class
82+
83+
stream = _AsyncWriteObjectStream(mock_client, BUCKET, OBJECT, blob=mock_blob)
84+
assert stream.blob.storage_class == storage_class
85+
7586
def test_init_raises_value_error(self, mock_client):
7687
with pytest.raises(ValueError, match="client must be provided"):
7788
_AsyncWriteObjectStream(None, BUCKET, OBJECT)
@@ -220,6 +231,7 @@ async def test_open_new_object_with_blob_sync_attrs(
220231
mock_blob.content_language = "content-language"
221232
mock_blob.temporary_hold = True
222233
mock_blob.event_based_hold = True
234+
mock_blob.storage_class = "RAPID"
223235

224236
custom_time = datetime.datetime(
225237
2025, 1, 1, 12, 0, 0, tzinfo=datetime.timezone.utc
@@ -244,6 +256,7 @@ async def test_open_new_object_with_blob_sync_attrs(
244256
mock_blob.contexts = ObjectContexts(mock_blob, custom={"context-key": payload})
245257

246258
stream = _AsyncWriteObjectStream(mock_client, BUCKET, OBJECT, blob=mock_blob)
259+
assert stream.blob.storage_class == "RAPID"
247260
await stream.open()
248261

249262
# Verify initial request contains synced attributes from blob
@@ -259,6 +272,7 @@ async def test_open_new_object_with_blob_sync_attrs(
259272
assert resource.content_language == "content-language"
260273
assert resource.temporary_hold is True
261274
assert resource.event_based_hold is True
275+
assert resource.storage_class == "RAPID"
262276

263277
assert int(resource.custom_time.timestamp()) == int(custom_time.timestamp())
264278

‎packages/google-cloud-storage/tests/unit/test__grpc_conversions.py‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,8 @@
1515
import datetime
1616
from unittest import mock
1717

18+
import pytest
19+
1820
from google.cloud import _storage_v2
1921
from google.cloud.storage import _grpc_conversions
2022

@@ -33,6 +35,7 @@ def test_blob_to_proto_simple_fields():
3335
"content_language",
3436
"temporary_hold",
3537
"event_based_hold",
38+
"storage_class",
3639
"custom_time",
3740
"acl",
3841
"retention",
@@ -49,6 +52,7 @@ def test_blob_to_proto_simple_fields():
4952
blob.content_language = "en"
5053
blob.temporary_hold = True
5154
blob.event_based_hold = False
55+
blob.storage_class = "RAPID"
5256
blob.custom_time = None
5357
blob.acl = None
5458
blob.retention = None
@@ -66,6 +70,32 @@ def test_blob_to_proto_simple_fields():
6670
assert proto.content_language == "en"
6771
assert proto.temporary_hold is True
6872
assert proto.event_based_hold is False
73+
assert proto.storage_class == "RAPID"
74+
75+
76+
@pytest.mark.parametrize(
77+
"storage_class", ["STANDARD", "RAPID", "NEARLINE", "COLDLINE", "ARCHIVE"]
78+
)
79+
def test_blob_to_proto_storage_class(storage_class):
80+
blob = mock.Mock(spec=["name", "bucket", "storage_class"])
81+
blob.name = "blob-name"
82+
blob.bucket.name = "bucket-name"
83+
blob.storage_class = storage_class
84+
85+
proto = _grpc_conversions.blob_to_proto(blob)
86+
87+
assert proto.storage_class == storage_class
88+
89+
90+
def test_blob_to_proto_storage_class_none():
91+
blob = mock.Mock(spec=["name", "bucket", "storage_class"])
92+
blob.name = "blob-name"
93+
blob.bucket.name = "bucket-name"
94+
blob.storage_class = None
95+
96+
proto = _grpc_conversions.blob_to_proto(blob)
97+
98+
assert proto.storage_class == ""
6999

70100

71101
def test_blob_to_proto_custom_time():

0 commit comments

Comments
 (0)