diff --git a/openfga_sdk/oauth2.py b/openfga_sdk/oauth2.py index 6b065b2..084b497 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 8c55746..afc0800 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 ffe2bbd..f677d36 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 d6ebe72..83e15a7 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): """