From 632025cb65304ddf1e0c85d8cdb2ea3ccdfad133 Mon Sep 17 00:00:00 2001 From: Srinivas Chippagiri Date: Thu, 5 Nov 2026 02:03:23 -0800 Subject: [PATCH] fix: retry token requests that fail with 429 or 5xx The token request loop retried on 429 and 5xx by checking the status of the response, but the REST client raises RateLimitExceededError or ServiceException for those statuses instead of returning them. The first failed token request was therefore raised to the caller without any retry. Catch those errors in the async and sync OAuth2 clients and retry them, except 501, until max_retry is reached. Fixes #139 Signed-off-by: Srinivas Chippagiri --- openfga_sdk/oauth2.py | 34 +++++++++----- openfga_sdk/sync/oauth2.py | 34 +++++++++----- test/oauth2_test.py | 95 +++++++++++++++++++++++++++++++++++++- test/sync/oauth2_test.py | 91 +++++++++++++++++++++++++++++++++++- 4 files changed, 230 insertions(+), 24 deletions(-) diff --git a/openfga_sdk/oauth2.py b/openfga_sdk/oauth2.py index 6b065b2e..084b497d 100644 --- a/openfga_sdk/oauth2.py +++ b/openfga_sdk/oauth2.py @@ -13,7 +13,11 @@ TOKEN_EXPIRY_THRESHOLD_BUFFER_IN_SEC, ) from openfga_sdk.credentials import Credentials -from openfga_sdk.exceptions import AuthenticationError +from openfga_sdk.exceptions import ( + AuthenticationError, + RateLimitExceededError, + ServiceException, +) from openfga_sdk.oauth2_common import _TokenState, jitter from openfga_sdk.telemetry.attributes import TelemetryAttributes from openfga_sdk.telemetry.telemetry import Telemetry @@ -103,16 +107,24 @@ async def _obtain_token(self, client): ) for attempt in range(max_retry + 1): - raw_response = await client.request( - method="POST", - url=token_url, - headers=headers, - query_params=None, - body=None, - _preload_content=True, - _request_timeout=None, - post_params=post_params, - ) + try: + raw_response = await client.request( + method="POST", + url=token_url, + headers=headers, + query_params=None, + body=None, + _preload_content=True, + _request_timeout=None, + post_params=post_params, + ) + except (RateLimitExceededError, ServiceException) as err: + # The REST client raises on 429 and 5xx responses instead of returning them, + # so retry here, like the status check below intends. 501 isn't retried. + if attempt < max_retry and err.status != 501: + await asyncio.sleep(jitter(attempt, min_wait_in_ms)) + continue + raise if 500 <= raw_response.status <= 599 or raw_response.status == 429: if attempt < max_retry and raw_response.status != 501: diff --git a/openfga_sdk/sync/oauth2.py b/openfga_sdk/sync/oauth2.py index 8c55746f..afc0800d 100644 --- a/openfga_sdk/sync/oauth2.py +++ b/openfga_sdk/sync/oauth2.py @@ -14,7 +14,11 @@ TOKEN_EXPIRY_THRESHOLD_BUFFER_IN_SEC, ) from openfga_sdk.credentials import Credentials -from openfga_sdk.exceptions import AuthenticationError +from openfga_sdk.exceptions import ( + AuthenticationError, + RateLimitExceededError, + ServiceException, +) from openfga_sdk.oauth2_common import _TokenState, jitter from openfga_sdk.telemetry.attributes import TelemetryAttributes from openfga_sdk.telemetry.telemetry import Telemetry @@ -104,16 +108,24 @@ def _obtain_token(self, client): ) for attempt in range(max_retry + 1): - raw_response = client.request( - method="POST", - url=token_url, - headers=headers, - query_params=None, - body=None, - _preload_content=True, - _request_timeout=None, - post_params=post_params, - ) + try: + raw_response = client.request( + method="POST", + url=token_url, + headers=headers, + query_params=None, + body=None, + _preload_content=True, + _request_timeout=None, + post_params=post_params, + ) + except (RateLimitExceededError, ServiceException) as err: + # The REST client raises on 429 and 5xx responses instead of returning them, + # so retry here, like the status check below intends. 501 isn't retried. + if attempt < max_retry and err.status != 501: + time.sleep(jitter(attempt, min_wait_in_ms)) + continue + raise if 500 <= raw_response.status <= 599 or raw_response.status == 429: if attempt < max_retry and raw_response.status != 501: diff --git a/test/oauth2_test.py b/test/oauth2_test.py index ffe2bbd7..f677d36d 100644 --- a/test/oauth2_test.py +++ b/test/oauth2_test.py @@ -11,7 +11,11 @@ from openfga_sdk.configuration import Configuration from openfga_sdk.constants import TOKEN_EXPIRY_THRESHOLD_BUFFER_IN_SEC from openfga_sdk.credentials import CredentialConfiguration, Credentials -from openfga_sdk.exceptions import AuthenticationError +from openfga_sdk.exceptions import ( + AuthenticationError, + RateLimitExceededError, + ServiceException, +) from openfga_sdk.oauth2 import OAuth2Client from openfga_sdk.oauth2_common import _TokenState @@ -222,6 +226,95 @@ async def test_get_authentication_erroneous_response_fails(self, mock_request): await client.get_authentication_header(rest_client) await rest_client.close() + def _client_credentials(self): + return Credentials( + method="client_credentials", + configuration=CredentialConfiguration( + client_id="myclientid", + client_secret="mysecret", + api_issuer="issuer.fga.example", + api_audience="myaudience", + ), + ) + + @patch.object(rest.RESTClientObject, "request") + async def test_get_authentication_retries_raised_429_and_5xx_errors( + self, mock_request + ): + """ + The REST client raises on 429 and 5xx responses, and those token requests should be retried + """ + error_body = '{"code": "internal_error", "message": "Internal Server Error"}' + response_body = '{"expires_in": 120, "access_token": "AABBCCDD"}' + + mock_request.side_effect = [ + ServiceException(http_resp=mock_response(error_body, 500)), + RateLimitExceededError(http_resp=mock_response(error_body, 429)), + ServiceException(http_resp=mock_response(error_body, 503)), + mock_response(response_body, 200), + ] + + configuration = Configuration() + configuration.retry_params.max_retry = 5 + configuration.retry_params.min_wait_in_ms = 1 + + rest_client = rest.RESTClientObject(configuration) + client = OAuth2Client(self._client_credentials(), configuration) + + auth_header = await client.get_authentication_header(rest_client) + + self.assertEqual(mock_request.call_count, 4) + self.assertEqual(auth_header, {"Authorization": "Bearer AABBCCDD"}) + await rest_client.close() + + @patch.object(rest.RESTClientObject, "request") + async def test_get_authentication_stops_retrying_after_max_retry( + self, mock_request + ): + """ + After max_retry retries, the last error is raised + """ + error_body = '{"code": "internal_error", "message": "Internal Server Error"}' + mock_request.side_effect = [ + ServiceException(http_resp=mock_response(error_body, 500)) for _ in range(3) + ] + + configuration = Configuration() + configuration.retry_params.max_retry = 2 + configuration.retry_params.min_wait_in_ms = 1 + + rest_client = rest.RESTClientObject(configuration) + client = OAuth2Client(self._client_credentials(), configuration) + + with self.assertRaises(ServiceException): + await client.get_authentication_header(rest_client) + + self.assertEqual(mock_request.call_count, 3) + await rest_client.close() + + @patch.object(rest.RESTClientObject, "request") + async def test_get_authentication_does_not_retry_501(self, mock_request): + """ + A 501 (Not Implemented) response is not retried + """ + error_body = '{"code": "not_implemented", "message": "Not Implemented"}' + mock_request.side_effect = [ + ServiceException(http_resp=mock_response(error_body, 501)) + ] + + configuration = Configuration() + configuration.retry_params.max_retry = 5 + configuration.retry_params.min_wait_in_ms = 1 + + rest_client = rest.RESTClientObject(configuration) + client = OAuth2Client(self._client_credentials(), configuration) + + with self.assertRaises(ServiceException): + await client.get_authentication_header(rest_client) + + self.assertEqual(mock_request.call_count, 1) + await rest_client.close() + @patch.object(rest.RESTClientObject, "request") async def test_get_authentication_retries_5xx_responses(self, mock_request): """ diff --git a/test/sync/oauth2_test.py b/test/sync/oauth2_test.py index d6ebe727..83e15a7d 100644 --- a/test/sync/oauth2_test.py +++ b/test/sync/oauth2_test.py @@ -11,7 +11,11 @@ from openfga_sdk.configuration import Configuration from openfga_sdk.constants import TOKEN_EXPIRY_THRESHOLD_BUFFER_IN_SEC from openfga_sdk.credentials import CredentialConfiguration, Credentials -from openfga_sdk.exceptions import AuthenticationError +from openfga_sdk.exceptions import ( + AuthenticationError, + RateLimitExceededError, + ServiceException, +) from openfga_sdk.oauth2_common import _TokenState from openfga_sdk.sync import rest from openfga_sdk.sync.oauth2 import OAuth2Client @@ -339,6 +343,91 @@ async def test_get_authentication_erroneous_response_fails(self, mock_request): client.get_authentication_header(rest_client) rest_client.close() + def _client_credentials(self): + return Credentials( + method="client_credentials", + configuration=CredentialConfiguration( + client_id="myclientid", + client_secret="mysecret", + api_issuer="issuer.fga.example", + api_audience="myaudience", + ), + ) + + @patch.object(rest.RESTClientObject, "request") + def test_get_authentication_retries_raised_429_and_5xx_errors(self, mock_request): + """ + The REST client raises on 429 and 5xx responses, and those token requests should be retried + """ + error_body = '{"code": "internal_error", "message": "Internal Server Error"}' + response_body = '{"expires_in": 120, "access_token": "AABBCCDD"}' + + mock_request.side_effect = [ + ServiceException(http_resp=mock_response(error_body, 500)), + RateLimitExceededError(http_resp=mock_response(error_body, 429)), + ServiceException(http_resp=mock_response(error_body, 503)), + mock_response(response_body, 200), + ] + + configuration = Configuration() + configuration.retry_params.max_retry = 5 + configuration.retry_params.min_wait_in_ms = 1 + + rest_client = rest.RESTClientObject(configuration) + client = OAuth2Client(self._client_credentials(), configuration) + + auth_header = client.get_authentication_header(rest_client) + + self.assertEqual(mock_request.call_count, 4) + self.assertEqual(auth_header, {"Authorization": "Bearer AABBCCDD"}) + rest_client.close() + + @patch.object(rest.RESTClientObject, "request") + def test_get_authentication_stops_retrying_after_max_retry(self, mock_request): + """ + After max_retry retries, the last error is raised + """ + error_body = '{"code": "internal_error", "message": "Internal Server Error"}' + mock_request.side_effect = [ + ServiceException(http_resp=mock_response(error_body, 500)) for _ in range(3) + ] + + configuration = Configuration() + configuration.retry_params.max_retry = 2 + configuration.retry_params.min_wait_in_ms = 1 + + rest_client = rest.RESTClientObject(configuration) + client = OAuth2Client(self._client_credentials(), configuration) + + with self.assertRaises(ServiceException): + client.get_authentication_header(rest_client) + + self.assertEqual(mock_request.call_count, 3) + rest_client.close() + + @patch.object(rest.RESTClientObject, "request") + def test_get_authentication_does_not_retry_501(self, mock_request): + """ + A 501 (Not Implemented) response is not retried + """ + error_body = '{"code": "not_implemented", "message": "Not Implemented"}' + mock_request.side_effect = [ + ServiceException(http_resp=mock_response(error_body, 501)) + ] + + configuration = Configuration() + configuration.retry_params.max_retry = 5 + configuration.retry_params.min_wait_in_ms = 1 + + rest_client = rest.RESTClientObject(configuration) + client = OAuth2Client(self._client_credentials(), configuration) + + with self.assertRaises(ServiceException): + client.get_authentication_header(rest_client) + + self.assertEqual(mock_request.call_count, 1) + rest_client.close() + @patch.object(rest.RESTClientObject, "request") def test_get_authentication_retries_5xx_responses(self, mock_request): """