feat(storage): Add async client and async page iterator class along with their tests

Author: googlyrahman

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

@@ -0,0 +1,240 @@
+# 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.
+
+"""Asynchronous client for interacting with Google Cloud Storage."""
+
+import functools
+
+from google.cloud.storage._experimental.asyncio.async_helpers import ASYNC_DEFAULT_TIMEOUT
+from google.cloud.storage._experimental.asyncio.async_helpers import ASYNC_DEFAULT_RETRY
+from google.cloud.storage._experimental.asyncio.async_helpers import AsyncHTTPIterator
+from google.cloud.storage._experimental.asyncio.async_helpers import _do_nothing_page_start
+from google.cloud.storage._opentelemetry_tracing import create_trace_span
+from google.cloud.storage._experimental.asyncio.async_creds import AsyncCredsWrapper
+from google.cloud.storage.abstracts.base_client import BaseClient
+from google.cloud.storage._experimental.asyncio.async_connection import AsyncConnection
+from google.cloud.storage.abstracts import base_client
+
+try:
+    from google.auth.aio.transport import sessions
+    AsyncSession = sessions.AsyncAuthorizedSession
+    _AIO_AVAILABLE = True
+except ImportError:
+    AsyncSession = object
+    _AIO_AVAILABLE = False
+
+_marker = base_client.marker
+
+
+class AsyncClient(BaseClient):
+    """Asynchronous client to interact with Google Cloud Storage."""
+
+    def __init__(
+        self,
+        project=_marker,
+        credentials=None,
+        _async_http=None,
+        client_info=None,
+        client_options=None,
+        extra_headers={},
+        *,
+        api_key=None,
+    ):
+        if not _AIO_AVAILABLE:
+            # Python 3.9 or less comes with an older version of google-auth library which doesn't support asyncio
+            raise ImportError(
+                "Failed to import 'google.auth.aio', Consider using a newer python version (>=3.10)"
+                " or newer version of google-auth library to mitigate this issue."
+            )
+
+        if self._use_client_cert:
+            # google.auth.aio.transports.sessions.AsyncAuthorizedSession currently doesn't support configuring mTLS.
+            # In future, we can monkey patch the above, and do provide mTLS support, but that is not a priority
+            # at the moment.
+            raise ValueError("Async Client currently do not support mTLS")
+
+        # We initialize everything as per synchronous client.
+        super().__init__(
+            project=project,
+            credentials=credentials,
+            client_info=client_info,
+            client_options=client_options,
+            extra_headers=extra_headers,
+            api_key=api_key
+        )
+        self.credentials = AsyncCredsWrapper(self._credentials) # self._credential is synchronous.
+        self._connection = AsyncConnection(self, **self.connection_kw_args) # adapter for async communication
+        self._async_http_internal = _async_http
+        self._async_http_passed_by_user = (_async_http is not None)
+
+    @property
+    def async_http(self):
+        """Returns the existing asynchronous session, or create one if it does not exists."""
+        if self._async_http_internal is None:
+            self._async_http_internal = AsyncSession(credentials=self.credentials)
+        return self._async_http_internal
+
+    async def close(self):
+        """Close the session, if it exists"""
+        if self._async_http_internal is not None and not self._async_http_passed_by_user:
+            await self._async_http_internal.close()
+
+    async def _get_resource(
+        self,
+        path,
+        query_params=None,
+        headers=None,
+        timeout=ASYNC_DEFAULT_TIMEOUT,
+        retry=ASYNC_DEFAULT_RETRY,
+        _target_object=None,
+    ):
+        """See super() class"""
+        return await self._connection.api_request(
+            method="GET",
+            path=path,
+            query_params=query_params,
+            headers=headers,
+            timeout=timeout,
+            retry=retry,
+            _target_object=_target_object,
+        )
+
+    def _list_resource(
+        self,
+        path,
+        item_to_value,
+        page_token=None,
+        max_results=None,
+        extra_params=None,
+        page_start=_do_nothing_page_start,
+        page_size=None,
+        timeout=ASYNC_DEFAULT_TIMEOUT,
+        retry=ASYNC_DEFAULT_RETRY,
+    ):
+        """See super() class"""
+        kwargs = {
+            "method": "GET",
+            "path": path,
+            "timeout": timeout,
+        }
+        with create_trace_span(
+            name="Storage.AsyncClient._list_resource_returns_iterator",
+            client=self,
+            api_request=kwargs,
+            retry=retry,
+        ):
+            api_request = functools.partial(
+                self._connection.api_request, timeout=timeout, retry=retry
+            )
+        return AsyncHTTPIterator(
+            client=self,
+            api_request=api_request,
+            path=path,
+            item_to_value=item_to_value,
+            page_token=page_token,
+            max_results=max_results,
+            extra_params=extra_params,
+            page_start=page_start,
+            page_size=page_size,
+        )
+
+    async def _patch_resource(
+        self,
+        path,
+        data,
+        query_params=None,
+        headers=None,
+        timeout=ASYNC_DEFAULT_TIMEOUT,
+        retry=None,
+        _target_object=None,
+    ):
+        """See super() class"""
+        return await self._connection.api_request(
+            method="PATCH",
+            path=path,
+            data=data,
+            query_params=query_params,
+            headers=headers,
+            timeout=timeout,
+            retry=retry,
+            _target_object=_target_object,
+        )
+
+    async def _put_resource(
+        self,
+        path,
+        data,
+        query_params=None,
+        headers=None,
+        timeout=ASYNC_DEFAULT_TIMEOUT,
+        retry=None,
+        _target_object=None,
+    ):
+        """See super() class"""
+        return await self._connection.api_request(
+            method="PUT",
+            path=path,
+            data=data,
+            query_params=query_params,
+            headers=headers,
+            timeout=timeout,
+            retry=retry,
+            _target_object=_target_object,
+        )
+
+    async def _post_resource(
+        self,
+        path,
+        data,
+        query_params=None,
+        headers=None,
+        timeout=ASYNC_DEFAULT_TIMEOUT,
+        retry=None,
+        _target_object=None,
+    ):
+        """See super() class"""
+        return await self._connection.api_request(
+            method="POST",
+            path=path,
+            data=data,
+            query_params=query_params,
+            headers=headers,
+            timeout=timeout,
+            retry=retry,
+            _target_object=_target_object,
+        )
+
+    async def _delete_resource(
+        self,
+        path,
+        query_params=None,
+        headers=None,
+        timeout=ASYNC_DEFAULT_TIMEOUT,
+        retry=ASYNC_DEFAULT_RETRY,
+        _target_object=None,
+    ):
+        """See super() class"""
+        return await self._connection.api_request(
+            method="DELETE",
+            path=path,
+            query_params=query_params,
+            headers=headers,
+            timeout=timeout,
+            retry=retry,
+            _target_object=_target_object,
+        )
+
+    def bucket(self, bucket_name, user_project=None, generation=None):
+        """See super() class"""
+        raise NotImplementedError("This bucket class needs to be implemented.")

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

@@ -0,0 +1,142 @@
+"""Helper utility for async client."""
+
+from google.api_core import retry_async
+from google.cloud.storage import retry as storage_retry
+from google.cloud.storage import _helpers
+from google.api_core.page_iterator_async import AsyncIterator
+from google.api_core.page_iterator import Page
+
+
+ASYNC_DEFAULT_RETRY = retry_async.AsyncRetry(predicate=storage_retry._should_retry)
+ASYNC_DEFAULT_TIMEOUT = _helpers._DEFAULT_TIMEOUT
+
+
+async def _do_nothing_page_start(iterator, page, response):
+    """Async Helper to provide custom behavior after a :class:`Page` is started.
+
+    This is a do-nothing stand-in as the default value.
+
+    Args:
+        iterator (Iterator): An iterator that holds some request info.
+        page (Page): The page that was just created.
+        response (Any): The API response for a page.
+    """
+    # pylint: disable=unused-argument
+    pass
+
+class AsyncHTTPIterator(AsyncIterator):
+    """A generic class for iterating through HTTP/JSON API list responses asynchronously.
+
+    Args:
+        client (google.cloud.storage._experimental.asyncio.async_client.AsyncClient): The API client.
+        api_request (Callable): The **async** function to use to make API requests.
+            This must be an awaitable.
+        path (str): The method path to query for the list of items.
+        item_to_value (Callable[AsyncIterator, Any]): Callable to convert an item 
+            from the type in the JSON response into a native object.
+        items_key (str): The key in the API response where the list of items
+            can be found.
+        page_token (str): A token identifying a page in a result set.
+        page_size (int): The maximum number of results to fetch per page.
+        max_results (int): The maximum number of results to fetch.
+        extra_params (dict): Extra query string parameters for the API call.
+        page_start (Callable): Callable to provide special behavior after a new page 
+            is created.
+        next_token (str): The name of the field used in the response for page tokens.
+    """
+
+    _DEFAULT_ITEMS_KEY = "items"
+    _PAGE_TOKEN = "pageToken"
+    _MAX_RESULTS = "maxResults"
+    _NEXT_TOKEN = "nextPageToken"
+    _RESERVED_PARAMS = frozenset([_PAGE_TOKEN])
+
+    def __init__(
+        self,
+        client,
+        api_request,
+        path,
+        item_to_value,
+        items_key=_DEFAULT_ITEMS_KEY,
+        page_token=None,
+        page_size=None,
+        max_results=None,
+        extra_params=None,
+        page_start=_do_nothing_page_start,
+        next_token=_NEXT_TOKEN,
+    ):
+        super().__init__(
+            client, item_to_value, page_token=page_token, max_results=max_results
+        )
+        self.api_request = api_request
+        self.path = path
+        self._items_key = items_key
+        self._page_size = page_size
+        self._page_start = page_start
+        self._next_token = next_token
+        self.extra_params = extra_params.copy() if extra_params else {}
+        self._verify_params()
+
+    def _verify_params(self):
+        """Verifies the parameters don't use any reserved parameter."""
+        reserved_in_use = self._RESERVED_PARAMS.intersection(self.extra_params)
+        if reserved_in_use:
+            raise ValueError("Using a reserved parameter", reserved_in_use)
+
+    async def _next_page(self):
+        """Get the next page in the iterator asynchronously.
+
+        Returns:
+            Optional[Page]: The next page in the iterator or None if
+                there are no pages left.
+        """
+        if self._has_next_page():
+            response = await self._get_next_page_response()
+            items = response.get(self._items_key, ())
+
+            # We reuse the synchronous Page class as it is just a container
+            page = Page(self, items, self.item_to_value, raw_page=response)
+
+            await self._page_start(self, page, response)
+            self.next_page_token = response.get(self._next_token)
+            return page
+        else:
+            return None
+
+    def _has_next_page(self):
+        """Determines whether or not there are more pages with results."""
+        if self.page_number == 0:
+            return True
+
+        if self.max_results is not None:
+            if self.num_results >= self.max_results:
+                return False
+
+        return self.next_page_token is not None
+
+    def _get_query_params(self):
+        """Getter for query parameters for the next request."""
+        result = {}
+        if self.next_page_token is not None:
+            result[self._PAGE_TOKEN] = self.next_page_token
+
+        page_size = None
+        if self.max_results is not None:
+            page_size = self.max_results - self.num_results
+            if self._page_size is not None:
+                page_size = min(page_size, self._page_size)
+        elif self._page_size is not None:
+            page_size = self._page_size
+
+        if page_size is not None:
+            result[self._MAX_RESULTS] = page_size
+
+        result.update(self.extra_params)
+        return result
+
+    async def _get_next_page_response(self):
+        """Requests the next page from the path provided asynchronously."""
+        params = self._get_query_params()
+        return await self.api_request(
+            method="GET", path=self.path, query_params=params
+        )