Skip to content

fix: retry token requests that fail with 429 or 5xx #331

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
cvas22 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix: retry token requests that fail with 429 or 5xx #331

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 23 additions & 11 deletions openfga_sdk/oauth2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down
34 changes: 23 additions & 11 deletions openfga_sdk/sync/oauth2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down
95 changes: 94 additions & 1 deletion test/oauth2_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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
Comment thread
coderabbitai[bot] marked this conversation as resolved.
Resolved

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):
"""
Expand Down
91 changes: 90 additions & 1 deletion test/sync/oauth2_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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):
"""
Expand Down