feat: expose persisted size in mrd (#1671)

feat: expose persisted size in MRD (MultiRangeReader)

Author: chandra-siri

google/cloud/storage/_experimental/asyncio/async_multi_range_downloader.py

@@ -180,6 +180,7 @@ def __init__(
         self._read_id_to_writable_buffer_dict = {}
         self._read_id_to_download_ranges_id = {}
         self._download_ranges_id_to_pending_read_ids = {}
+        self.persisted_size: Optional[int] = None  # updated after opening the stream
 
     async def open(self) -> None:
         """Opens the bidi-gRPC connection to read from the object.
@@ -206,6 +207,8 @@ async def open(self) -> None:
         if self.generation_number is None:
             self.generation_number = self.read_obj_str.generation_number
         self.read_handle = self.read_obj_str.read_handle
+        if self.read_obj_str.persisted_size is not None:
+            self.persisted_size = self.read_obj_str.persisted_size
         return
 
     async def download_ranges(

google/cloud/storage/_experimental/asyncio/async_read_object_stream.py

@@ -92,6 +92,7 @@ def __init__(
         self.metadata = (("x-goog-request-params", f"bucket={self._full_bucket_name}"),)
         self.socket_like_rpc: Optional[AsyncBidiRpc] = None
         self._is_stream_open: bool = False
+        self.persisted_size: Optional[int] = None
 
     async def open(self) -> None:
         """Opens the bidi-gRPC connection to read from the object.
@@ -106,8 +107,13 @@ async def open(self) -> None:
         )
         await self.socket_like_rpc.open()  # this is actually 1 send
         response = await self.socket_like_rpc.recv()
-        if self.generation_number is None:
-            self.generation_number = response.metadata.generation
+        # populated only in the first response of bidi-stream and when opened
+        # without using `read_handle`
+        if response.metadata:
+            if self.generation_number is None:
+                self.generation_number = response.metadata.generation
+            # update persisted size
+            self.persisted_size = response.metadata.size
 
         self.read_handle = response.read_handle
 

tests/system/test_zonal.py

@@ -58,6 +58,30 @@ async def test_basic_wrd(storage_client, blobs_to_delete, attempt_direct_path):
     await mrd.download_ranges([(0, 0, buffer)])
     await mrd.close()
     assert buffer.getvalue() == _BYTES_TO_UPLOAD
+    assert mrd.persisted_size == len(_BYTES_TO_UPLOAD)
+
+    # Clean up; use json client (i.e. `storage_client` fixture) to delete.
+    blobs_to_delete.append(storage_client.bucket(_ZONAL_BUCKET).blob(object_name))
+
+
+@pytest.mark.asyncio
+async def test_read_unfinalized_appendable_object(storage_client, blobs_to_delete):
+    object_name = f"read_unfinalized_appendable_object-{str(uuid.uuid4())[:4]}"
+    grpc_client = AsyncGrpcClient(attempt_direct_path=True).grpc_client
+
+    writer = AsyncAppendableObjectWriter(grpc_client, _ZONAL_BUCKET, object_name)
+    await writer.open()
+    await writer.append(_BYTES_TO_UPLOAD)
+    await writer.flush()
+
+    mrd = AsyncMultiRangeDownloader(grpc_client, _ZONAL_BUCKET, object_name)
+    buffer = BytesIO()
+    await mrd.open()
+    assert mrd.persisted_size == len(_BYTES_TO_UPLOAD)
+    # (0, 0) means read the whole object
+    await mrd.download_ranges([(0, 0, buffer)])
+    await mrd.close()
+    assert buffer.getvalue() == _BYTES_TO_UPLOAD
 
     # Clean up; use json client (i.e. `storage_client` fixture) to delete.
     blobs_to_delete.append(storage_client.bucket(_ZONAL_BUCKET).blob(object_name))

tests/unit/asyncio/test_async_multi_range_downloader.py

@@ -30,6 +30,7 @@
 
 _TEST_BUCKET_NAME = "test-bucket"
 _TEST_OBJECT_NAME = "test-object"
+_TEST_OBJECT_SIZE = 1024 * 1024  # 1 MiB
 _TEST_GENERATION_NUMBER = 123456789
 _TEST_READ_HANDLE = b"test-handle"
 
@@ -57,6 +58,7 @@ async def _make_mock_mrd(
         mock_stream = mock_cls_async_read_object_stream.return_value
         mock_stream.open = AsyncMock()
         mock_stream.generation_number = _TEST_GENERATION_NUMBER
+        mock_stream.persisted_size = _TEST_OBJECT_SIZE
         mock_stream.read_handle = _TEST_READ_HANDLE
 
         mrd = await AsyncMultiRangeDownloader.create_mrd(
@@ -106,6 +108,7 @@ async def test_create_mrd(
         assert mrd.object_name == _TEST_OBJECT_NAME
         assert mrd.generation_number == _TEST_GENERATION_NUMBER
         assert mrd.read_handle == _TEST_READ_HANDLE
+        assert mrd.persisted_size == _TEST_OBJECT_SIZE
         assert mrd.is_stream_open
 
     @mock.patch(

tests/unit/asyncio/test_async_read_object_stream.py

@@ -25,6 +25,7 @@
 _TEST_BUCKET_NAME = "test-bucket"
 _TEST_OBJECT_NAME = "test-object"
 _TEST_GENERATION_NUMBER = 12345
+_TEST_OBJECT_SIZE = 1024 * 1024  # 1 MiB
 _TEST_READ_HANDLE = b"test-read-handle"
 
 
@@ -37,6 +38,7 @@ async def instantiate_read_obj_stream(mock_client, mock_cls_async_bidi_rpc, open
     recv_response = mock.MagicMock(spec=_storage_v2.BidiReadObjectResponse)
     recv_response.metadata = mock.MagicMock(spec=_storage_v2.Object)
     recv_response.metadata.generation = _TEST_GENERATION_NUMBER
+    recv_response.metadata.size = _TEST_OBJECT_SIZE
     recv_response.read_handle = _TEST_READ_HANDLE
     socket_like_rpc.recv = AsyncMock(return_value=recv_response)
 
@@ -112,6 +114,7 @@ async def test_open(mock_client, mock_cls_async_bidi_rpc):
 
     assert read_obj_stream.generation_number == _TEST_GENERATION_NUMBER
     assert read_obj_stream.read_handle == _TEST_READ_HANDLE
+    assert read_obj_stream.persisted_size == _TEST_OBJECT_SIZE
     assert read_obj_stream.is_stream_open