feat: Support timeout as aiohttp.ClientTimeout and total_attempts (max retries) in AsyncAuthorizedSession (#1961)

b/485304839

---------

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Author: yinghsienwu

packages/google-auth/google/auth/aio/transport/aiohttp.py

@@ -12,12 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Transport adapter for Asynchronous HTTP Requests based on aiohttp.
-"""
+"""Transport adapter for Asynchronous HTTP Requests based on aiohttp."""
 
 import asyncio
 import logging
-from typing import AsyncGenerator, Mapping, Optional
+from typing import AsyncGenerator, Mapping, Optional, TYPE_CHECKING, Union
 
 try:
     import aiohttp  # type: ignore
@@ -31,6 +30,15 @@
 from google.auth.aio import _helpers as _helpers_async
 from google.auth.aio import transport
 
+if TYPE_CHECKING:  # pragma: NO COVER
+    from aiohttp import ClientTimeout  # type: ignore
+
+else:
+    try:
+        from aiohttp import ClientTimeout
+    except (ImportError, AttributeError):
+        ClientTimeout = None
+
 _LOGGER = logging.getLogger(__name__)
 
 
@@ -123,7 +131,7 @@ async def __call__(
         method: str = "GET",
         body: Optional[bytes] = None,
         headers: Optional[Mapping[str, str]] = None,
-        timeout: float = transport._DEFAULT_TIMEOUT_SECONDS,
+        timeout: Union[float, ClientTimeout] = transport._DEFAULT_TIMEOUT_SECONDS,
         **kwargs,
     ) -> transport.Response:
         """
@@ -158,7 +166,10 @@ async def __call__(
             if not self._session:
                 self._session = aiohttp.ClientSession()
 
-            client_timeout = aiohttp.ClientTimeout(total=timeout)
+            if isinstance(timeout, aiohttp.ClientTimeout):
+                client_timeout = timeout
+            else:
+                client_timeout = aiohttp.ClientTimeout(total=timeout)
             _helpers.request_log(_LOGGER, method, url, body, headers)
             response = await self._session.request(
                 method,
@@ -176,8 +187,12 @@ async def __call__(
             raise client_exc from caught_exc
 
         except asyncio.TimeoutError as caught_exc:
+            if isinstance(timeout, aiohttp.ClientTimeout):
+                timeout_seconds = timeout.total
+            else:
+                timeout_seconds = timeout
             timeout_exc = exceptions.TimeoutError(
-                f"Request timed out after {timeout} seconds."
+                f"Request timed out after {timeout_seconds} seconds."
             )
             raise timeout_exc from caught_exc
 

packages/google-auth/google/auth/aio/transport/sessions.py

@@ -16,13 +16,22 @@
 from contextlib import asynccontextmanager
 import functools
 import time
-from typing import Mapping, Optional
+from typing import Mapping, Optional, TYPE_CHECKING, Union
 
 from google.auth import _exponential_backoff, exceptions
 from google.auth.aio import transport
 from google.auth.aio.credentials import Credentials
 from google.auth.exceptions import TimeoutError
 
+if TYPE_CHECKING:  # pragma: NO COVER
+    from aiohttp import ClientTimeout  # type: ignore
+
+else:
+    try:
+        from aiohttp import ClientTimeout
+    except (ImportError, AttributeError):
+        ClientTimeout = None
+
 try:
     from google.auth.aio.transport.aiohttp import Request as AiohttpRequest
 
@@ -137,7 +146,8 @@ async def request(
         data: Optional[bytes] = None,
         headers: Optional[Mapping[str, str]] = None,
         max_allowed_time: float = transport._DEFAULT_TIMEOUT_SECONDS,
-        timeout: float = transport._DEFAULT_TIMEOUT_SECONDS,
+        timeout: Union[float, ClientTimeout] = transport._DEFAULT_TIMEOUT_SECONDS,
+        total_attempts: Optional[int] = transport.DEFAULT_MAX_RETRY_ATTEMPTS,
         **kwargs,
     ) -> transport.Response:
         """
@@ -146,14 +156,16 @@ async def request(
                 url (str): The URI to be requested.
                 data (Optional[bytes]): The payload or body in HTTP request.
                 headers (Optional[Mapping[str, str]]): Request headers.
-                timeout (float):
+                timeout (float, aiohttp.ClientTimeout):
                 The amount of time in seconds to wait for the server response
                 with each individual request.
                 max_allowed_time (float):
                 If the method runs longer than this, a ``Timeout`` exception is
                 automatically raised. Unlike the ``timeout`` parameter, this
                 value applies to the total method execution time, even if
                 multiple requests are made under the hood.
+                total_attempts (int):
+                The total number of retry attempts.
 
                 Mind that it is not guaranteed that the timeout error is raised
                 at ``max_allowed_time``. It might take longer, for example, if
@@ -172,7 +184,7 @@ async def request(
         """
 
         retries = _exponential_backoff.AsyncExponentialBackoff(
-            total_attempts=transport.DEFAULT_MAX_RETRY_ATTEMPTS
+            total_attempts=total_attempts,
         )
         async with timeout_guard(max_allowed_time) as with_timeout:
             await with_timeout(
@@ -198,11 +210,50 @@ async def get(
         data: Optional[bytes] = None,
         headers: Optional[Mapping[str, str]] = None,
         max_allowed_time: float = transport._DEFAULT_TIMEOUT_SECONDS,
-        timeout: float = transport._DEFAULT_TIMEOUT_SECONDS,
+        timeout: Union[float, ClientTimeout] = transport._DEFAULT_TIMEOUT_SECONDS,
+        total_attempts: Optional[int] = transport.DEFAULT_MAX_RETRY_ATTEMPTS,
         **kwargs,
     ) -> transport.Response:
+        """
+        Args:
+                url (str): The URI to be requested.
+                data (Optional[bytes]): The payload or body in HTTP request.
+                headers (Optional[Mapping[str, str]]): Request headers.
+                max_allowed_time (float):
+                If the method runs longer than this, a ``Timeout`` exception is
+                automatically raised. Unlike the ``timeout`` parameter, this
+                value applies to the total method execution time, even if
+                multiple requests are made under the hood.
+                timeout (float, aiohttp.ClientTimeout):
+                The amount of time in seconds to wait for the server response
+                with each individual request.
+                total_attempts (int):
+                The total number of retry attempts.
+
+                Mind that it is not guaranteed that the timeout error is raised
+                at ``max_allowed_time``. It might take longer, for example, if
+                an underlying request takes a lot of time, but the request
+                itself does not timeout, e.g. if a large file is being
+                transmitted. The timeout error will be raised after such
+                request completes.
+
+        Returns:
+                google.auth.aio.transport.Response: The HTTP response.
+
+        Raises:
+                google.auth.exceptions.TimeoutError: If the method does not complete within
+                the configured `max_allowed_time` or the request exceeds the configured
+                `timeout`.
+        """
         return await self.request(
-            "GET", url, data, headers, max_allowed_time, timeout, **kwargs
+            "GET",
+            url,
+            data,
+            headers,
+            max_allowed_time,
+            timeout,
+            total_attempts,
+            **kwargs,
         )
 
     @functools.wraps(request)
@@ -212,11 +263,50 @@ async def post(
         data: Optional[bytes] = None,
         headers: Optional[Mapping[str, str]] = None,
         max_allowed_time: float = transport._DEFAULT_TIMEOUT_SECONDS,
-        timeout: float = transport._DEFAULT_TIMEOUT_SECONDS,
+        timeout: Union[float, ClientTimeout] = transport._DEFAULT_TIMEOUT_SECONDS,
+        total_attempts: Optional[int] = transport.DEFAULT_MAX_RETRY_ATTEMPTS,
         **kwargs,
     ) -> transport.Response:
+        """
+        Args:
+                url (str): The URI to be requested.
+                data (Optional[bytes]): The payload or body in HTTP request.
+                headers (Optional[Mapping[str, str]]): Request headers.
+                max_allowed_time (float):
+                If the method runs longer than this, a ``Timeout`` exception is
+                automatically raised. Unlike the ``timeout`` parameter, this
+                value applies to the total method execution time, even if
+                multiple requests are made under the hood.
+                timeout (float, aiohttp.ClientTimeout):
+                The amount of time in seconds to wait for the server response
+                with each individual request.
+                total_attempts (int):
+                The total number of retry attempts.
+
+                Mind that it is not guaranteed that the timeout error is raised
+                at ``max_allowed_time``. It might take longer, for example, if
+                an underlying request takes a lot of time, but the request
+                itself does not timeout, e.g. if a large file is being
+                transmitted. The timeout error will be raised after such
+                request completes.
+
+        Returns:
+                google.auth.aio.transport.Response: The HTTP response.
+
+        Raises:
+                google.auth.exceptions.TimeoutError: If the method does not complete within
+                the configured `max_allowed_time` or the request exceeds the configured
+                `timeout`.
+        """
         return await self.request(
-            "POST", url, data, headers, max_allowed_time, timeout, **kwargs
+            "POST",
+            url,
+            data,
+            headers,
+            max_allowed_time,
+            timeout,
+            total_attempts,
+            **kwargs,
         )
 
     @functools.wraps(request)
@@ -226,11 +316,50 @@ async def put(
         data: Optional[bytes] = None,
         headers: Optional[Mapping[str, str]] = None,
         max_allowed_time: float = transport._DEFAULT_TIMEOUT_SECONDS,
-        timeout: float = transport._DEFAULT_TIMEOUT_SECONDS,
+        timeout: Union[float, ClientTimeout] = transport._DEFAULT_TIMEOUT_SECONDS,
+        total_attempts: Optional[int] = transport.DEFAULT_MAX_RETRY_ATTEMPTS,
         **kwargs,
     ) -> transport.Response:
+        """
+        Args:
+                url (str): The URI to be requested.
+                data (Optional[bytes]): The payload or body in HTTP request.
+                headers (Optional[Mapping[str, str]]): Request headers.
+                max_allowed_time (float):
+                If the method runs longer than this, a ``Timeout`` exception is
+                automatically raised. Unlike the ``timeout`` parameter, this
+                value applies to the total method execution time, even if
+                multiple requests are made under the hood.
+                timeout (float, aiohttp.ClientTimeout):
+                The amount of time in seconds to wait for the server response
+                with each individual request.
+                total_attempts (int):
+                The total number of retry attempts.
+
+                Mind that it is not guaranteed that the timeout error is raised
+                at ``max_allowed_time``. It might take longer, for example, if
+                an underlying request takes a lot of time, but the request
+                itself does not timeout, e.g. if a large file is being
+                transmitted. The timeout error will be raised after such
+                request completes.
+
+        Returns:
+                google.auth.aio.transport.Response: The HTTP response.
+
+        Raises:
+                google.auth.exceptions.TimeoutError: If the method does not complete within
+                the configured `max_allowed_time` or the request exceeds the configured
+                `timeout`.
+        """
         return await self.request(
-            "PUT", url, data, headers, max_allowed_time, timeout, **kwargs
+            "PUT",
+            url,
+            data,
+            headers,
+            max_allowed_time,
+            timeout,
+            total_attempts,
+            **kwargs,
         )
 
     @functools.wraps(request)
@@ -240,11 +369,50 @@ async def patch(
         data: Optional[bytes] = None,
         headers: Optional[Mapping[str, str]] = None,
         max_allowed_time: float = transport._DEFAULT_TIMEOUT_SECONDS,
-        timeout: float = transport._DEFAULT_TIMEOUT_SECONDS,
+        timeout: Union[float, ClientTimeout] = transport._DEFAULT_TIMEOUT_SECONDS,
+        total_attempts: Optional[int] = transport.DEFAULT_MAX_RETRY_ATTEMPTS,
         **kwargs,
     ) -> transport.Response:
+        """
+        Args:
+                url (str): The URI to be requested.
+                data (Optional[bytes]): The payload or body in HTTP request.
+                headers (Optional[Mapping[str, str]]): Request headers.
+                max_allowed_time (float):
+                If the method runs longer than this, a ``Timeout`` exception is
+                automatically raised. Unlike the ``timeout`` parameter, this
+                value applies to the total method execution time, even if
+                multiple requests are made under the hood.
+                timeout (float, aiohttp.ClientTimeout):
+                The amount of time in seconds to wait for the server response
+                with each individual request.
+                total_attempts (int):
+                The total number of retry attempts.
+
+                Mind that it is not guaranteed that the timeout error is raised
+                at ``max_allowed_time``. It might take longer, for example, if
+                an underlying request takes a lot of time, but the request
+                itself does not timeout, e.g. if a large file is being
+                transmitted. The timeout error will be raised after such
+                request completes.
+
+        Returns:
+                google.auth.aio.transport.Response: The HTTP response.
+
+        Raises:
+                google.auth.exceptions.TimeoutError: If the method does not complete within
+                the configured `max_allowed_time` or the request exceeds the configured
+                `timeout`.
+        """
         return await self.request(
-            "PATCH", url, data, headers, max_allowed_time, timeout, **kwargs
+            "PATCH",
+            url,
+            data,
+            headers,
+            max_allowed_time,
+            timeout,
+            total_attempts,
+            **kwargs,
         )
 
     @functools.wraps(request)
@@ -254,11 +422,50 @@ async def delete(
         data: Optional[bytes] = None,
         headers: Optional[Mapping[str, str]] = None,
         max_allowed_time: float = transport._DEFAULT_TIMEOUT_SECONDS,
-        timeout: float = transport._DEFAULT_TIMEOUT_SECONDS,
+        timeout: Union[float, ClientTimeout] = transport._DEFAULT_TIMEOUT_SECONDS,
+        total_attempts: Optional[int] = transport.DEFAULT_MAX_RETRY_ATTEMPTS,
         **kwargs,
     ) -> transport.Response:
+        """
+        Args:
+                url (str): The URI to be requested.
+                data (Optional[bytes]): The payload or body in HTTP request.
+                headers (Optional[Mapping[str, str]]): Request headers.
+                max_allowed_time (float):
+                If the method runs longer than this, a ``Timeout`` exception is
+                automatically raised. Unlike the ``timeout`` parameter, this
+                value applies to the total method execution time, even if
+                multiple requests are made under the hood.
+                timeout (float, aiohttp.ClientTimeout):
+                The amount of time in seconds to wait for the server response
+                with each individual request.
+                total_attempts (int):
+                The total number of retry attempts.
+
+                Mind that it is not guaranteed that the timeout error is raised
+                at ``max_allowed_time``. It might take longer, for example, if
+                an underlying request takes a lot of time, but the request
+                itself does not timeout, e.g. if a large file is being
+                transmitted. The timeout error will be raised after such
+                request completes.
+
+        Returns:
+                google.auth.aio.transport.Response: The HTTP response.
+
+        Raises:
+                google.auth.exceptions.TimeoutError: If the method does not complete within
+                the configured `max_allowed_time` or the request exceeds the configured
+                `timeout`.
+        """
         return await self.request(
-            "DELETE", url, data, headers, max_allowed_time, timeout, **kwargs
+            "DELETE",
+            url,
+            data,
+            headers,
+            max_allowed_time,
+            timeout,
+            total_attempts,
+            **kwargs,
         )
 
     async def close(self) -> None: