feat: Add support for asynchronous AuthorizedSession api (#1577)

* chore: initial setup for async auth sessions api (#1571) * chore: initial setup for async auth sessions api * fix whitespace * add init file * update file names to aiohttp * update import statement * feat: Implement asynchronous timeout context manager (#1569) * feat: implement async timeout guard * add docstring * clean whitespace * update import file name * add missing return statement * update test cases * update test cases * include underlying timeout exception in trace * avoid the cost of actual time * feat: Implement asynchronous `AuthorizedSession` api response class (#1575) * feat: implement asynchronous response class for AuthorizedSessions API * check if aiohttp is installed and avoid tests dependency * update content to be async * update docstring to be specific to aiohttp * add type checking and avoid leaking underlying API responses * add test case for iterating chunks * add read method to response interface * address PR comments * fix lint issues * feat: Implement asynchronous `AuthorizedSession` api request class (#1579) * feat: implement request class for asynchoronous AuthorizedSession API * add type checking and address TODOs * remove default values from interface methods * aiohttp reponse close method must not be awaited * cleanup * update Request class docstring * feat: Implement asynchronous `AuthorizedSession` class (#1580) * feat: Implement Asynchronous AuthorizedSession class * add comment for implementing locks within refresh * move timeout guard to sessions * add unit tests and code cleanup * implement async exponential backoff iterator * cleanup * add testing for http methods and cleanup * update number of retries to 3 * refactor test cases * fix linter and mypy issues * fix pytest code coverage * fix: avoid leaking api error for closed session * add error handling for response * cleanup default values and add test coverage * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * cleanup: minor code cleanup (#1589) * chore: Add aiohttp requirements test constraint. (#1566) See #1565 for more information. * chore(main): release 2.33.0 (#1560) * chore(main): release 2.33.0 * fix: retry token request on retryable status code (#1563) * fix: retry token request on retryable status code * feat(auth): Update get_client_ssl_credentials to support X.509 workload certs (#1558) * feat(auth): Update get_client_ssl_credentials to support X.509 workload certs * feat(auth): Update has_default_client_cert_source * feat(auth): Fix formatting * feat(auth): Fix test__mtls_helper.py * feat(auth): Fix function name in tests * chore: Refresh system test creds. * feat(auth): Fix style * feat(auth): Fix casing * feat(auth): Fix linter issue * feat(auth): Fix coverage issue --------- Co-authored-by: Carl Lundin <[email protected]> Co-authored-by: Carl Lundin <[email protected]> * chore: Update ECP deps. (#1583) * chore(main): release 2.34.0 (#1574) * cleanup: minor code cleanup * fix lint issues --------- Co-authored-by: Carl Lundin <[email protected]> Co-authored-by: release-please[bot] <55107282+release-please[bot]@users.noreply.github.com> Co-authored-by: Andy Zhao <[email protected]> Co-authored-by: Carl Lundin <[email protected]> * update secrets from forked repo --------- Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com> Co-authored-by: arithmetic1728 <[email protected]> Co-authored-by: Carl Lundin <[email protected]> Co-authored-by: release-please[bot] <55107282+release-please[bot]@users.noreply.github.com> Co-authored-by: Andy Zhao <[email protected]> Co-authored-by: Carl Lundin <[email protected]>
googleapis · Sep 16, 2024 · 2910b6b · 2910b6b
1 parent 6f75dd5
commit 2910b6b
Show file tree

Hide file tree

Showing 9 changed files with 1,187 additions and 16 deletions.
diff --git a/google/auth/_exponential_backoff.py b/google/auth/_exponential_backoff.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import asyncio
 import random
 import time
 
@@ -38,9 +39,8 @@
 """
 
 
-class ExponentialBackoff:
-    """An exponential backoff iterator. This can be used in a for loop to
-    perform requests with exponential backoff.
+class _BaseExponentialBackoff:
+    """An exponential backoff iterator base class.
 
     Args:
         total_attempts Optional[int]:
@@ -84,9 +84,40 @@ def __init__(
         self._multiplier = multiplier
         self._backoff_count = 0
 
-    def __iter__(self):
+    @property
+    def total_attempts(self):
+        """The total amount of backoff attempts that will be made."""
+        return self._total_attempts
+
+    @property
+    def backoff_count(self):
+        """The current amount of backoff attempts that have been made."""
+        return self._backoff_count
+
+    def _reset(self):
         self._backoff_count = 0
         self._current_wait_in_seconds = self._initial_wait_seconds
+
+    def _calculate_jitter(self):
+        jitter_variance = self._current_wait_in_seconds * self._randomization_factor
+        jitter = random.uniform(
+            self._current_wait_in_seconds - jitter_variance,
+            self._current_wait_in_seconds + jitter_variance,
+        )
+
+        return jitter
+
+
+class ExponentialBackoff(_BaseExponentialBackoff):
+    """An exponential backoff iterator. This can be used in a for loop to
+    perform requests with exponential backoff.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super(ExponentialBackoff, self).__init__(*args, **kwargs)
+
+    def __iter__(self):
+        self._reset()
         return self
 
     def __next__(self):
@@ -97,23 +128,37 @@ def __next__(self):
         if self._backoff_count <= 1:
             return self._backoff_count
 
-        jitter_variance = self._current_wait_in_seconds * self._randomization_factor
-        jitter = random.uniform(
-            self._current_wait_in_seconds - jitter_variance,
-            self._current_wait_in_seconds + jitter_variance,
-        )
+        jitter = self._calculate_jitter()
 
         time.sleep(jitter)
 
         self._current_wait_in_seconds *= self._multiplier
         return self._backoff_count
 
-    @property
-    def total_attempts(self):
-        """The total amount of backoff attempts that will be made."""
-        return self._total_attempts
 
-    @property
-    def backoff_count(self):
-        """The current amount of backoff attempts that have been made."""
+class AsyncExponentialBackoff(_BaseExponentialBackoff):
+    """An async exponential backoff iterator. This can be used in a for loop to
+    perform async requests with exponential backoff.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super(AsyncExponentialBackoff, self).__init__(*args, **kwargs)
+
+    def __aiter__(self):
+        self._reset()
+        return self
+
+    async def __anext__(self):
+        if self._backoff_count >= self._total_attempts:
+            raise StopAsyncIteration
+        self._backoff_count += 1
+
+        if self._backoff_count <= 1:
+            return self._backoff_count
+
+        jitter = self._calculate_jitter()
+
+        await asyncio.sleep(jitter)
+
+        self._current_wait_in_seconds *= self._multiplier
         return self._backoff_count
diff --git a/google/auth/aio/transport/__init__.py b/google/auth/aio/transport/__init__.py
@@ -0,0 +1,144 @@
+# Copyright 2024 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.
+
+"""Transport - Asynchronous HTTP client library support.
+
+:mod:`google.auth.aio` is designed to work with various asynchronous client libraries such
+as aiohttp. In order to work across these libraries with different
+interfaces some abstraction is needed.
+
+This module provides two interfaces that are implemented by transport adapters
+to support HTTP libraries. :class:`Request` defines the interface expected by
+:mod:`google.auth` to make asynchronous requests. :class:`Response` defines the interface
+for the return value of :class:`Request`.
+"""
+
+import abc
+from typing import AsyncGenerator, Mapping, Optional
+
+import google.auth.transport
+
+
+_DEFAULT_TIMEOUT_SECONDS = 180
+
+DEFAULT_RETRYABLE_STATUS_CODES = google.auth.transport.DEFAULT_RETRYABLE_STATUS_CODES
+"""Sequence[int]:  HTTP status codes indicating a request can be retried.
+"""
+
+
+DEFAULT_MAX_RETRY_ATTEMPTS = 3
+"""int: How many times to retry a request."""
+
+
+class Response(metaclass=abc.ABCMeta):
+    """Asynchronous HTTP Response Interface."""
+
+    @property
+    @abc.abstractmethod
+    def status_code(self) -> int:
+        """
+        The HTTP response status code.
+
+        Returns:
+            int: The HTTP response status code.
+
+        """
+        raise NotImplementedError("status_code must be implemented.")
+
+    @property
+    @abc.abstractmethod
+    def headers(self) -> Mapping[str, str]:
+        """The HTTP response headers.
+
+        Returns:
+            Mapping[str, str]: The HTTP response headers.
+        """
+        raise NotImplementedError("headers must be implemented.")
+
+    @abc.abstractmethod
+    async def content(self, chunk_size: int) -> AsyncGenerator[bytes, None]:
+        """The raw response content.
+
+        Args:
+            chunk_size (int): The size of each chunk.
+
+        Yields:
+            AsyncGenerator[bytes, None]: An asynchronous generator yielding
+            response chunks as bytes.
+        """
+        raise NotImplementedError("content must be implemented.")
+
+    @abc.abstractmethod
+    async def read(self) -> bytes:
+        """Read the entire response content as bytes.
+
+        Returns:
+            bytes: The entire response content.
+        """
+        raise NotImplementedError("read must be implemented.")
+
+    @abc.abstractmethod
+    async def close(self):
+        """Close the response after it is fully consumed to resource."""
+        raise NotImplementedError("close must be implemented.")
+
+
+class Request(metaclass=abc.ABCMeta):
+    """Interface for a callable that makes HTTP requests.
+
+    Specific transport implementations should provide an implementation of
+    this that adapts their specific request / response API.
+
+    .. automethod:: __call__
+    """
+
+    @abc.abstractmethod
+    async def __call__(
+        self,
+        url: str,
+        method: str,
+        body: Optional[bytes],
+        headers: Optional[Mapping[str, str]],
+        timeout: float,
+        **kwargs
+    ) -> Response:
+        """Make an HTTP request.
+
+        Args:
+            url (str): The URI to be requested.
+            method (str): The HTTP method to use for the request. Defaults
+                to 'GET'.
+            body (Optional[bytes]): The payload / body in HTTP request.
+            headers (Mapping[str, str]): Request headers.
+            timeout (float): The number of seconds to wait for a
+                response from the server. If not specified or if None, the
+                transport-specific default timeout will be used.
+            kwargs: Additional arguments passed on to the transport's
+                request method.
+
+        Returns:
+            google.auth.aio.transport.Response: The HTTP response.
+
+        Raises:
+            google.auth.exceptions.TransportError: If any exception occurred.
+        """
+        # pylint: disable=redundant-returns-doc, missing-raises-doc
+        # (pylint doesn't play well with abstract docstrings.)
+        raise NotImplementedError("__call__ must be implemented.")
+
+    async def close(self) -> None:
+        """
+        Close the underlying session.
+        """
+        raise NotImplementedError("close must be implemented.")