Merge branch 'je-httpx-retry' into je-http2-client

Author: jonathanedey

firebase_admin/_http_client.py

@@ -43,7 +43,8 @@
     connect=1, read=1, status=4, status_forcelist=[500, 503],
     raise_on_status=False, backoff_factor=0.5, **_ANY_METHOD)
 
-DEFAULT_HTTPX_RETRY_CONFIG = HttpxRetry(status=4, status_forcelist=[500, 503], backoff_factor=0.5)
+DEFAULT_HTTPX_RETRY_CONFIG = HttpxRetry(
+    max_retries=4, status_forcelist=[500, 503], backoff_factor=0.5)
 
 
 DEFAULT_TIMEOUT_SECONDS = 120
@@ -166,16 +167,6 @@ def parse_body(self, resp):
         return resp.json()
 
 
-# Auth Flow
-# TODO: Remove comments
-# The aim here is to be able to get auth credentials right before the request is sent.
-# This is similar to what is done in transport.requests.AuthorizedSession().
-# We can then pass this in at the client level.
-
-# Notes:
-# - This implementations does not cover timeouts on requests sent to refresh credentials.
-# - Uses HTTP/1 and a blocking credential for refreshing.
-# - Network error retries for refreshing credentials.
 class GoogleAuthCredentialFlow(httpx.Auth):
     """Google Auth Credential Auth Flow"""
     def __init__(self, credential: credentials.Credentials):

firebase_admin/_retry.py

@@ -34,23 +34,18 @@
 
 class HttpxRetry:
     """HTTPX based retry config"""
-    # TODO: Decide
-    # urllib3.Retry ignores the status_forcelist when respecting Retry-After header
-    # Only 413, 429 and 503 errors are retried with the Retry-After header.
-    # Should we do the same?
-    # Default status codes to be used for ``status_forcelist``
+    # Status codes to be used for respecting `Retry-After` header
     RETRY_AFTER_STATUS_CODES = frozenset([413, 429, 503])
 
-    #: Default maximum backoff time.
+    # Default maximum backoff time.
     DEFAULT_BACKOFF_MAX = 120
 
     def __init__(
             self,
-            status: int = 10,
+            max_retries: int = 10,
             status_forcelist: Optional[List[int]] = None,
             backoff_factor: float = 0,
             backoff_max: float = DEFAULT_BACKOFF_MAX,
-            raise_on_status: bool = False,
             backoff_jitter: float = 0,
             history: Optional[List[Tuple[
                 httpx.Request,
@@ -59,11 +54,10 @@ def __init__(
             ]]] = None,
             respect_retry_after_header: bool = False,
     ) -> None:
-        self.status = status
+        self.retries_left = max_retries
         self.status_forcelist = status_forcelist
         self.backoff_factor = backoff_factor
         self.backoff_max = backoff_max
-        self.raise_on_status = raise_on_status
         self.backoff_jitter = backoff_jitter
         if history:
             self.history = history
@@ -90,16 +84,10 @@ def is_retryable_response(self, response: httpx.Response) -> bool:
 
         return False
 
-    # Placeholder for exception retrying
-    def is_retryable_error(self, error: Exception):
-        """Determine if the error implies that the request should be retired if possible."""
-        logger.debug(error)
-        return False
-
     def is_exhausted(self) -> bool:
         """Determine if there are anymore more retires."""
-        # status count is negative
-        return self.status < 0
+        # retries_left is negative
+        return self.retries_left < 0
 
     # Identical implementation of `urllib3.Retry.parse_retry_after()`
     def _parse_retry_after(self, retry_after_header: str) -> float | None:
@@ -111,7 +99,6 @@ def _parse_retry_after(self, retry_after_header: str) -> float | None:
         else:
             retry_date_tuple = email.utils.parsedate_tz(retry_after_header)
             if retry_date_tuple is None:
-                # TODO: Verify if this is the appropriate way to handle this.
                 raise httpx.RemoteProtocolError(f"Invalid Retry-After header: {retry_after_header}")
 
             retry_date = email.utils.mktime_tz(retry_date_tuple)
@@ -167,53 +154,37 @@ def increment(
             error: Optional[Exception] = None
     ) -> None:
         """Update the retry state based on request attempt."""
-        if response and self.is_retryable_response(response):
-            self.status -= 1
+        self.retries_left -= 1
         self.history.append((request, response, error))
 
 
-# TODO: Remove comments
-# Note - This implementation currently covers:
-#   - basic retires for pre-defined status errors
-#   - applying retry backoff and backoff jitter
-#   - ability to respect a response's retry-after header
 class HttpxRetryTransport(httpx.AsyncBaseTransport):
     """HTTPX transport with retry logic."""
 
-    # DEFAULT_RETRY = HttpxRetry(
-    #     connect=1, read=1, status=4, status_forcelist=[500, 503],
-    #     raise_on_status=False, backoff_factor=0.5, allowed_methods=None
-    # )
-    DEFAULT_RETRY = HttpxRetry(status=4, status_forcelist=[500, 503], backoff_factor=0.5)
+    DEFAULT_RETRY = HttpxRetry(max_retries=4, status_forcelist=[500, 503], backoff_factor=0.5)
 
-    # We could also support passing kwargs here
     def __init__(self, retry: HttpxRetry = DEFAULT_RETRY, **kwargs) -> None:
         self._retry = retry
 
         transport_kwargs = kwargs.copy()
         transport_kwargs.update({'retries': 0, 'http2': True})
-        # We should use a full AsyncHTTPTransport under the hood since that is
-        # fully implemented. We could consider making this class extend a
-        # AsyncHTTPTransport instead and use the parent class's methods to handle
-        # requests. We sould also ensure that that transport's internal retry is
-        # not enabled.
+        # We use a full AsyncHTTPTransport under the hood that is already
+        # set up to handle requests. We also insure that that transport's internal
+        # retries are not allowed.
         self._wrapped_transport = httpx.AsyncHTTPTransport(**transport_kwargs)
 
     async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
         return await self._dispatch_with_retry(
             request, self._wrapped_transport.handle_async_request)
 
-    # Two types of retries
-    # - Status code (500s, redirect)
-    # - Error code (read, connect, other)
     async def _dispatch_with_retry(
             self,
             request: httpx.Request,
             dispatch_method: Callable[[httpx.Request], CoroutineType[Any, Any, httpx.Response]]
     ) -> httpx.Response:
         """Sends a request with retry logic using a provided dispatch method."""
         # This request config is used across all requests that use this transport and therefore
-        # needs to be copied to be used for just this request ans it's retries.
+        # needs to be copied to be used for just this request and it's retries.
         retry = self._retry.copy()
         # First request
         response, error = None, None
@@ -238,16 +209,16 @@ async def _dispatch_with_retry(
             if response and not retry.is_retryable_response(response):
                 return response
 
-            if error and not retry.is_retryable_error(error):
+            if error:
                 raise error
 
-            retry.increment(request, response)
+            retry.increment(request, response, error)
 
         if response:
             return response
         if error:
             raise error
-        raise Exception('_dispatch_with_retry() ended with no response or exception')
+        raise AssertionError('_dispatch_with_retry() ended with no response or exception')
 
     async def aclose(self) -> None:
         await self._wrapped_transport.aclose()

firebase_admin/messaging.py

@@ -670,11 +670,6 @@ def _handle_batch_error(self, error):
         return _gapic_utils.handle_platform_error_from_googleapiclient(
             error, _MessagingService._build_fcm_error_googleapiclient)
 
-    # TODO: Remove comments
-    # We should be careful to clean up the httpx clients.
-    # Since we are using an async client we must also close in async. However we can sync wrap this.
-    # The close method is called by the app on shutdown/clean-up of each service. We don't seem to
-    # make use of this much elsewhere.
     def close(self) -> None:
         asyncio.run(self._async_client.aclose())
 

tests/test_http_client.py

@@ -276,7 +276,7 @@ def test_init_with_custom_settings(self, mocker: MockerFixture):
 
         mock_credential = testutils.MockGoogleCredential()
         headers = {'X-Custom': 'Test'}
-        custom_retry = HttpxRetry(status=1, status_forcelist=[429], backoff_factor=0)
+        custom_retry = HttpxRetry(max_retries=1, status_forcelist=[429], backoff_factor=0)
         timeout = 60
         http2 = False
 

tests/test_retry.py

@@ -37,7 +37,7 @@ class TestHttpxRetryTransport():
     @respx.mock
     async def test_no_retry_on_success(self, base_url: str, mocker: MockerFixture):
         """Test that a successful response doesn't trigger retries."""
-        retry_config = HttpxRetry(status=3, status_forcelist=[500])
+        retry_config = HttpxRetry(max_retries=3, status_forcelist=[500])
         transport = HttpxRetryTransport(retry=retry_config)
         client = httpx.AsyncClient(transport=transport)
 
@@ -55,7 +55,7 @@ async def test_no_retry_on_success(self, base_url: str, mocker: MockerFixture):
     @respx.mock
     async def test_no_retry_on_non_retryable_status(self, base_url: str, mocker: MockerFixture):
         """Test that a non-retryable error status doesn't trigger retries."""
-        retry_config = HttpxRetry(status=3, status_forcelist=[500, 503])
+        retry_config = HttpxRetry(max_retries=3, status_forcelist=[500, 503])
         transport = HttpxRetryTransport(retry=retry_config)
         client = httpx.AsyncClient(transport=transport)
 
@@ -75,7 +75,7 @@ async def test_retry_on_status_code_success_on_last_retry(
             self, base_url: str, mocker: MockerFixture
     ):
         """Test retry on status code from status_forcelist, succeeding on the last attempt."""
-        retry_config = HttpxRetry(status=2, status_forcelist=[503, 500], backoff_factor=0.5)
+        retry_config = HttpxRetry(max_retries=2, status_forcelist=[503, 500], backoff_factor=0.5)
         transport = HttpxRetryTransport(retry=retry_config)
         client = httpx.AsyncClient(transport=transport)
 
@@ -101,7 +101,7 @@ async def test_retry_exhausted_returns_last_response(
             self, base_url: str, mocker: MockerFixture
     ):
         """Test that the last response is returned when retries are exhausted."""
-        retry_config = HttpxRetry(status=1, status_forcelist=[500], backoff_factor=0)
+        retry_config = HttpxRetry(max_retries=1, status_forcelist=[500], backoff_factor=0)
         transport = HttpxRetryTransport(retry=retry_config)
         client = httpx.AsyncClient(transport=transport)
 
@@ -124,7 +124,8 @@ async def test_retry_exhausted_returns_last_response(
     @respx.mock
     async def test_retry_after_header_seconds(self, base_url: str, mocker: MockerFixture):
         """Test respecting Retry-After header with seconds value."""
-        retry_config = HttpxRetry(status=1, respect_retry_after_header=True, backoff_factor=100)
+        retry_config = HttpxRetry(
+            max_retries=1, respect_retry_after_header=True, backoff_factor=100)
         transport = HttpxRetryTransport(retry=retry_config)
         client = httpx.AsyncClient(transport=transport)
 
@@ -146,7 +147,8 @@ async def test_retry_after_header_seconds(self, base_url: str, mocker: MockerFix
     @respx.mock
     async def test_retry_after_header_http_date(self, base_url: str, mocker: MockerFixture):
         """Test respecting Retry-After header with an HTTP-date value."""
-        retry_config = HttpxRetry(status=1, respect_retry_after_header=True, backoff_factor=100)
+        retry_config = HttpxRetry(
+            max_retries=1, respect_retry_after_header=True, backoff_factor=100)
         transport = HttpxRetryTransport(retry=retry_config)
         client = httpx.AsyncClient(transport=transport)
 
@@ -181,7 +183,7 @@ async def test_retry_after_header_http_date(self, base_url: str, mocker: MockerF
     async def test_retry_after_ignored_when_disabled(self, base_url: str, mocker: MockerFixture):
         """Test Retry-After header is ignored if `respect_retry_after_header` is `False`."""
         retry_config = HttpxRetry(
-            status=3, respect_retry_after_header=False, status_forcelist=[429],
+            max_retries=3, respect_retry_after_header=False, status_forcelist=[429],
             backoff_factor=0.5, backoff_max=10)
         transport = HttpxRetryTransport(retry=retry_config)
         client = httpx.AsyncClient(transport=transport)
@@ -215,7 +217,7 @@ async def test_retry_after_header_missing_backoff_fallback(
         """Test Retry-After header is ignored if `respect_retry_after_header`is `True` but header is
         not set."""
         retry_config = HttpxRetry(
-            status=3, respect_retry_after_header=True, status_forcelist=[429],
+            max_retries=3, respect_retry_after_header=True, status_forcelist=[429],
             backoff_factor=0.5, backoff_max=10)
         transport = HttpxRetryTransport(retry=retry_config)
         client = httpx.AsyncClient(transport=transport)
@@ -247,7 +249,7 @@ async def test_exponential_backoff(self, base_url: str, mocker: MockerFixture):
         """Test that sleep time increases exponentially with `backoff_factor`."""
         # status=3 allows 3 retries (attempts 2, 3, 4)
         retry_config = HttpxRetry(
-            status=3, status_forcelist=[500], backoff_factor=0.1, backoff_max=10.0)
+            max_retries=3, status_forcelist=[500], backoff_factor=0.1, backoff_max=10.0)
         transport = HttpxRetryTransport(retry=retry_config)
         client = httpx.AsyncClient(transport=transport)
 
@@ -278,7 +280,7 @@ async def test_backoff_max(self, base_url: str, mocker: MockerFixture):
         """Test that backoff time respects `backoff_max`."""
         # status=4 allows 4 retries. backoff_factor=1 causes rapid increase.
         retry_config = HttpxRetry(
-            status=4, status_forcelist=[500], backoff_factor=1, backoff_max=3.0)
+            max_retries=4, status_forcelist=[500], backoff_factor=1, backoff_max=3.0)
         transport = HttpxRetryTransport(retry=retry_config)
         client = httpx.AsyncClient(transport=transport)
 
@@ -310,7 +312,7 @@ async def test_backoff_max(self, base_url: str, mocker: MockerFixture):
     async def test_backoff_jitter(self, base_url: str, mocker: MockerFixture):
         """Test that `backoff_jitter` adds randomness within bounds."""
         retry_config = HttpxRetry(
-            status=3, status_forcelist=[500], backoff_factor=0.2, backoff_jitter=0.1)
+            max_retries=3, status_forcelist=[500], backoff_factor=0.2, backoff_jitter=0.1)
         transport = HttpxRetryTransport(retry=retry_config)
         client = httpx.AsyncClient(transport=transport)
 
@@ -343,7 +345,7 @@ async def test_backoff_jitter(self, base_url: str, mocker: MockerFixture):
     @respx.mock
     async def test_error_not_retryable(self, base_url):
         """Test that non-HTTP errors are raised immediately if not retryable."""
-        retry_config = HttpxRetry(status=3)
+        retry_config = HttpxRetry(max_retries=3)
         transport = HttpxRetryTransport(retry=retry_config)
         client = httpx.AsyncClient(transport=transport)
 
@@ -362,7 +364,7 @@ class TestHttpxRetry():
 
     def test_httpx_retry_copy(self, base_url):
         """Test that `HttpxRetry.copy()` creates a deep copy."""
-        original = HttpxRetry(status=5, status_forcelist=[500, 503], backoff_factor=0.5)
+        original = HttpxRetry(max_retries=5, status_forcelist=[500, 503], backoff_factor=0.5)
         original.history.append((base_url, None, None)) # Add something mutable
 
         copied = original.copy()
@@ -372,17 +374,17 @@ def test_httpx_retry_copy(self, base_url):
         assert original.history is not copied.history
 
         # Assert values are the same initially
-        assert copied.status == original.status
+        assert copied.retries_left == original.retries_left
         assert copied.status_forcelist == original.status_forcelist
         assert copied.backoff_factor == original.backoff_factor
         assert len(copied.history) == 1
 
         # Modify the copy and check original is unchanged
-        copied.status = 1
+        copied.retries_left = 1
         copied.status_forcelist = [404]
         copied.history.append((base_url, None, None))
 
-        assert original.status == 5
+        assert original.retries_left == 5
         assert original.status_forcelist == [500, 503]
         assert len(original.history) == 1
 
@@ -413,7 +415,8 @@ def test_parse_retry_after_invalid_date(self):
             retry._parse_retry_after('Invalid Date Format')
 
     def test_get_backoff_time_calculation(self):
-        retry = HttpxRetry(status=6, status_forcelist=[503], backoff_factor=0.5, backoff_max=10.0)
+        retry = HttpxRetry(
+            max_retries=6, status_forcelist=[503], backoff_factor=0.5, backoff_max=10.0)
         response = httpx.Response(503)
         # No history -> attempt 1 -> no backoff before first request
         # Note: get_backoff_time() is typically called *before* the *next* request,
@@ -447,5 +450,5 @@ def test_get_backoff_time_calculation(self):
 
         # Simulate attempt 6 completed
         retry.increment(self._TEST_REQUEST, response)
-        # History len 6, attempt 7 -> 0.5*(2^4) = 10.0
+        # History len 6, attempt 7 -> 0.5*(2^5) = 16.0 Clamped to 10
         assert retry.get_backoff_time() == pytest.approx(10.0)