feat(zb-experimental): Add AsyncMultiRangeDownloader (#1550)

feat(zb-experimental): Add AsyncMultiRangeDownloader and it's init method.

Author: chandra-siri

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

@@ -0,0 +1,165 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from typing import Any, List, Optional, Tuple
+
+from google.cloud.storage._experimental.asyncio.async_read_object_stream import (
+ _AsyncReadObjectStream,
+)
+from google.cloud.storage._experimental.asyncio.async_grpc_client import (
+ AsyncGrpcClient,
+)
+
+from io import BytesIO
+
+
+class AsyncMultiRangeDownloader:
+ """Provides an interface for downloading multiple ranges of a GCS ``Object``
+ concurrently.
+
+ Example usage:
+
+ .. code-block:: python
+
+ client = AsyncGrpcClient().grpc_client
+ mrd = await AsyncMultiRangeDownloader.create_mrd(
+ client, bucket_name="chandrasiri-rs", object_name="test_open9"
+ )
+ my_buff1 = BytesIO()
+ my_buff2 = BytesIO()
+ my_buff3 = BytesIO()
+ my_buff4 = BytesIO()
+ buffers = [my_buff1, my_buff2, my_buff3, my_buff4]
+ await mrd.download_ranges(
+ [
+ (0, 100, my_buff1),
+ (100, 200, my_buff2),
+ (200, 300, my_buff3),
+ (300, 400, my_buff4),
+ ]
+ )
+ for buff in buffers:
+ print("downloaded bytes", buff.getbuffer().nbytes)
+
+ """
+
+ @classmethod
+ async def create_mrd(
+ cls,
+ client: AsyncGrpcClient.grpc_client,
+ bucket_name: str,
+ object_name: str,
+ generation_number: Optional[int] = None,
+ read_handle: Optional[bytes] = None,
+ ) -> AsyncMultiRangeDownloader:
+ """Initializes a MultiRangeDownloader and opens the underlying bidi-gRPC
+ object for reading.
+
+ :type client: :class:`~google.cloud.storage._experimental.asyncio.async_grpc_client.AsyncGrpcClient.grpc_client`
+ :param client: The asynchronous client to use for making API requests.
+
+ :type bucket_name: str
+ :param bucket_name: The name of the bucket containing the object.
+
+ :type object_name: str
+ :param object_name: The name of the object to be read.
+
+ :type generation_number: int
+ :param generation_number: (Optional) If present, selects a specific
+ revision of this object.
+
+ :type read_handle: bytes
+ :param read_handle: (Optional) An existing handle for reading the object.
+ If provided, opening the bidi-gRPC connection will be faster.
+
+ :rtype: :class:`~google.cloud.storage._experimental.asyncio.async_multi_range_downloader.AsyncMultiRangeDownloader`
+ :returns: An initialized AsyncMultiRangeDownloader instance for reading.
+ """
+ mrd = cls(client, bucket_name, object_name, generation_number, read_handle)
+ await mrd.open()
+ return mrd
+
+ def __init__(
+ self,
+ client: AsyncGrpcClient.grpc_client,
+ bucket_name: str,
+ object_name: str,
+ generation_number: Optional[int] = None,
+ read_handle: Optional[bytes] = None,
+ ) -> None:
+ """Constructor for AsyncMultiRangeDownloader, clients are not adviced to
+ use it directly. Instead it's adviced to use the classmethod `create_mrd`.
+
+ :type client: :class:`~google.cloud.storage._experimental.asyncio.async_grpc_client.AsyncGrpcClient.grpc_client`
+ :param client: The asynchronous client to use for making API requests.
+
+ :type bucket_name: str
+ :param bucket_name: The name of the bucket containing the object.
+
+ :type object_name: str
+ :param object_name: The name of the object to be read.
+
+ :type generation_number: int
+ :param generation_number: (Optional) If present, selects a specific revision of
+ this object.
+
+ :type read_handle: bytes
+ :param read_handle: (Optional) An existing read handle.
+ """
+ self.client = client
+ self.bucket_name = bucket_name
+ self.object_name = object_name
+ self.generation_number = generation_number
+ self.read_handle = read_handle
+ self.read_obj_str: _AsyncReadObjectStream = None
+
+ async def open(self) -> None:
+ """Opens the bidi-gRPC connection to read from the object.
+
+ This method initializes and opens an `_AsyncReadObjectStream` (bidi-gRPC stream) to
+ for downloading ranges of data from GCS ``Object``.
+
+ "Opening" constitutes fetching object metadata such as generation number
+ and read handle and sets them as attributes if not already set.
+ """
+ self.read_obj_str = _AsyncReadObjectStream(
+ client=self.client,
+ bucket_name=self.bucket_name,
+ object_name=self.object_name,
+ generation_number=self.generation_number,
+ read_handle=self.read_handle,
+ )
+ await self.read_obj_str.open()
+ if self.generation_number is None:
+ self.generation_number = self.read_obj_str.generation_number
+ self.read_handle = self.read_obj_str.read_handle
+ return
+
+ async def download_ranges(self, read_ranges: List[Tuple[int, int, BytesIO]]) -> Any:
+ """Downloads multiple byte ranges from the object into the buffers
+ provided by user.
+
+ :type read_ranges: List[Tuple[int, int, "BytesIO"]]
+ :param read_ranges: A list of tuples, where each tuple represents a
+ byte range (start_byte, end_byte, buffer) to download. Buffer has to
+ be provided by the user, and user has to make sure appropriate
+ memory is available in the application to avoid out-of-memory crash.
+
+
+ Raises:
+ NotImplementedError: This method is not yet implemented.
+ """
+ raise NotImplementedError("TODO")

tests/unit/asyncio/test_async_multi_range_downloader.py

@@ -0,0 +1,79 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+from unittest import mock
+from unittest.mock import AsyncMock
+
+from google.cloud.storage._experimental.asyncio.async_multi_range_downloader import (
+ AsyncMultiRangeDownloader,
+)
+from io import BytesIO
+
+
+_TEST_BUCKET_NAME = "test-bucket"
+_TEST_OBJECT_NAME = "test-object"
+_TEST_GENERATION_NUMBER = 123456789
+_TEST_READ_HANDLE = b"test-handle"
+
+
+@mock.patch(
+ "google.cloud.storage._experimental.asyncio.async_multi_range_downloader._AsyncReadObjectStream"
+)
+@mock.patch(
+ "google.cloud.storage._experimental.asyncio.async_grpc_client.AsyncGrpcClient.grpc_client"
+)
+@pytest.mark.asyncio
+async def test_create_mrd(mock_async_grpc_client, async_read_object_stream):
+ # Arrange
+ mock_stream_instance = async_read_object_stream.return_value
+ mock_stream_instance.open = AsyncMock()
+ mock_stream_instance.generation_number = _TEST_GENERATION_NUMBER
+ mock_stream_instance.read_handle = _TEST_READ_HANDLE
+
+ # act
+ mrd = await AsyncMultiRangeDownloader.create_mrd(
+ mock_async_grpc_client, _TEST_BUCKET_NAME, _TEST_OBJECT_NAME
+ )
+
+ # Assert
+ async_read_object_stream.assert_called_once_with(
+ client=mock_async_grpc_client,
+ bucket_name=_TEST_BUCKET_NAME,
+ object_name=_TEST_OBJECT_NAME,
+ generation_number=None,
+ read_handle=None,
+ )
+ mock_stream_instance.open.assert_called_once()
+
+ assert mrd.client == mock_async_grpc_client
+ assert mrd.bucket_name == _TEST_BUCKET_NAME
+ assert mrd.object_name == _TEST_OBJECT_NAME
+ assert mrd.generation_number == _TEST_GENERATION_NUMBER
+ assert mrd.read_handle == _TEST_READ_HANDLE
+ assert mrd.read_obj_str is mock_stream_instance
+
+
+@mock.patch(
+ "google.cloud.storage._experimental.asyncio.async_grpc_client.AsyncGrpcClient.grpc_client"
+)
+@pytest.mark.asyncio
+async def test_download_ranges(mock_async_grpc_client):
+ """Test that download_ranges() raises NotImplementedError."""
+ mrd = AsyncMultiRangeDownloader(
+ mock_async_grpc_client, _TEST_BUCKET_NAME, _TEST_OBJECT_NAME
+ )
+
+ with pytest.raises(NotImplementedError):
+ await mrd.download_ranges([(0, 100, BytesIO())])