From 96fd9ce445dd1d15ec08103087457805a833d6cd Mon Sep 17 00:00:00 2001 From: Sam Dammers Date: Tue, 27 Sep 2022 10:08:39 +1000 Subject: [PATCH] feat(auth): add retries to aws boto calls (#88) * feat(auth): add retries to aws boto calls the standard boto client defaults to 'legacy' which retries certain HTTP status codes and a limited set of service errors. update retry config to 'standard' which has a broader set of errors/exceptions. in particular it will retry 'TooManyRequestsException' which can be throttled in high burst situations. max retries can be configured by setting the 'STAX_API_AUTH_MAX_RETRIES' environment var. * feat(api): add default retry configuration to stax api calls add default retries for stax api calls. allow these configuration to be modified by consumers. amended prior boto3 configuration to opt for storing in the Stax config object. marked existing token threshold ENV configuration as deprecated. * chore(sdk): update comment with deprecated notice --- README.md | 42 +++++++++++++++++++++++++++++++++------ examples/retry.py | 27 +++++++++++++++++++++++++ requirements.txt | 1 + staxapp/api.py | 29 +++++++++++++++++++++------ staxapp/auth.py | 27 ++++++++++++++++--------- staxapp/aws_srp.py | 9 ++++++++- staxapp/config.py | 49 ++++++++++++++++++++++++++++++++++++++++++---- staxapp/openapi.py | 6 ++++-- staxapp/retry.py | 37 ++++++++++++++++++++++++++++++++++ tests/test_api.py | 7 +++++-- tests/test_auth.py | 30 ++++++++++------------------ 11 files changed, 214 insertions(+), 50 deletions(-) create mode 100644 examples/retry.py create mode 100644 staxapp/retry.py diff --git a/README.md b/README.md index d5e8464..eaaa5e8 100644 --- a/README.md +++ b/README.md @@ -24,20 +24,50 @@ export STAX_SECRET_KEY= ``` ##### Client Auth Configuration -You can configure each client individually by passing in a config on init. -When a client is created it's configuration will be locked in and any change to the configurations will not affect the client. +The Stax SDK can configure each client individually by passing in a config on init. +When a client is created its configuration will be locked in and any change to the configurations will not affect the client. This can be seen in our [guide](https://github.com/stax-labs/lib-stax-python-sdk/blob/master/examples/auth.py). *Optional configuration:* -##### Authentication token expiry +##### Token expiry -Allows configuration of the threshold to when the Auth library should re-cache the credentials +The Stax SDK can be configured to refresh the API Token prior to expiry. *Suggested use when running within CI/CD tools to reduce overall auth calls* -~~~bash + +```python +from staxapp.config import Config, StaxAuthRetryConfig + +auth_retry_config = StaxAuthRetryConfig +auth_retry_config.token_expiry_threshold = 2 +Config.api_auth_retry_config = auth_retry_config +``` + +(Deprecated): This value can also be set via the following Environment Var `TOKEN_EXPIRY_THRESHOLD_IN_MINS` +```bash export TOKEN_EXPIRY_THRESHOLD_IN_MINS=2 # Type: Integer representing minutes -~~~ +``` + +##### Retries + +The Stax SDK has configured safe defaults for Auth and API retries. +This behaviour can be adjusted via the SDK config: [example](https://github.com/stax-labs/lib-stax-python-sdk/blob/master/examples/retry.py). + +```python +from staxapp.config import Config, StaxAPIRetryConfig, StaxAuthRetryConfig + +retry_config = StaxAPIRetryConfig +retry_config.retry_methods = ('GET', 'POST', 'PUT', 'DELETE', 'OPTIONS') +retry_config.status_codes = (429, 500, 502, 504) +retry_config.backoff_factor = 1.2 +retry_config.max_attempts = 3 +Config.api_retry_config = retry_config + +auth_retry_config = StaxAuthRetryConfig +auth_retry_config.max_attempts = 3 +Config.api_auth_retry_config = auth_retry_config +``` ##### Logging levels diff --git a/examples/retry.py b/examples/retry.py new file mode 100644 index 0000000..f5642ea --- /dev/null +++ b/examples/retry.py @@ -0,0 +1,27 @@ +import json +import os + +from staxapp.config import Config, StaxAPIRetryConfig, StaxAuthRetryConfig +from staxapp.openapi import StaxClient + +Config.access_key = os.getenv("STAX_ACCESS_KEY") +Config.secret_key = os.getenv("STAX_SECRET_KEY") + +# Retry Config for Stax API calls +retry_config = StaxAPIRetryConfig +retry_config.retry_methods = ('GET', 'POST', 'PUT', 'DELETE', 'OPTIONS') +retry_config.status_codes = (429, 500) + +Config.api_retry_config = retry_config + +# Retry config for Stax Authentication calls +auth_retry_config = StaxAuthRetryConfig +auth_retry_config.max_attempts = 3 + +Config.api_auth_retry_config = auth_retry_config + +# Read all accounts within your Stax Organisation +accounts = StaxClient("accounts") +response = accounts.ReadAccounts() +print(json.dumps(response, indent=4, sort_keys=True)) + diff --git a/requirements.txt b/requirements.txt index 5904010..658a07a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -8,6 +8,7 @@ pylint pytest pytest-cov responses +requests pyjwt==2.4.0 boto3 aws_requests_auth diff --git a/staxapp/api.py b/staxapp/api.py index 8518927..be7f8e3 100644 --- a/staxapp/api.py +++ b/staxapp/api.py @@ -1,9 +1,13 @@ -import json +""" +This module contains the http api handlers. +""" +import logging import requests -from staxapp.config import Config +from staxapp.config import Config, StaxAPIRetryConfig from staxapp.exceptions import ApiException +from staxapp.retry import requests_retry_session class Api: @@ -22,11 +26,23 @@ def _headers(cls, custom_headers) -> dict: } return headers + @classmethod + def request_session(cls, config: StaxAPIRetryConfig): + """Requests retry session with backoff""" + print(config.retry_methods) + return requests_retry_session( + retries=config.max_attempts, + status_list=config.status_codes, + allowed_methods=config.retry_methods, + backoff_factor=config.backoff_factor, + ) + @staticmethod def handle_api_response(response): try: response.raise_for_status() except requests.exceptions.HTTPError as e: + # logging.debug(f"request retried {len(response.raw.retries.history)} times") ## Useful to prove working raise ApiException(str(e), response) @classmethod @@ -34,7 +50,8 @@ def get(cls, url_frag, params={}, config=None, **kwargs): config = cls.get_config(config) url_frag = url_frag.replace(f"/{config.API_VERSION}", "") url = f"{config.api_base_url()}/{url_frag.lstrip('/')}" - response = requests.get( + + response = cls.request_session(config.api_retry_config).get( url, auth=config._auth(), params=params, @@ -50,7 +67,7 @@ def post(cls, url_frag, payload={}, config=None, **kwargs): url_frag = url_frag.replace(f"/{config.API_VERSION}", "") url = f"{config.api_base_url()}/{url_frag.lstrip('/')}" - response = requests.post( + response = cls.request_session(config.api_retry_config).post( url, json=payload, auth=config._auth(), @@ -66,7 +83,7 @@ def put(cls, url_frag, payload={}, config=None, **kwargs): url_frag = url_frag.replace(f"/{config.API_VERSION}", "") url = f"{config.api_base_url()}/{url_frag.lstrip('/')}" - response = requests.put( + response = cls.request_session(config.api_retry_config).put( url, json=payload, auth=config._auth(), @@ -82,7 +99,7 @@ def delete(cls, url_frag, params={}, config=None, **kwargs): url_frag = url_frag.replace(f"/{config.API_VERSION}", "") url = f"{config.api_base_url()}/{url_frag.lstrip('/')}" - response = requests.delete( + response = cls.request_session(config.api_retry_config).delete( url, auth=config._auth(), params=params, diff --git a/staxapp/auth.py b/staxapp/auth.py index bb45b7f..44183a4 100644 --- a/staxapp/auth.py +++ b/staxapp/auth.py @@ -1,6 +1,5 @@ #!/usr/local/bin/python3 from datetime import datetime, timedelta, timezone -from os import environ import boto3 from aws_requests_auth.aws_auth import AWSRequestsAuth @@ -14,7 +13,7 @@ class StaxAuth: - def __init__(self, config_branch: str, config: StaxConfig, max_retries: int = 3): + def __init__(self, config_branch: str, config: StaxConfig, max_retries: int = 5): self.config = config api_config = self.config.api_config self.identity_pool = api_config.get(config_branch).get("identityPoolId") @@ -52,7 +51,10 @@ def id_token_from_cognito( srp_client = boto3.client( "cognito-idp", region_name=self.aws_region, - config=BotoConfig(signature_version=UNSIGNED), + config=BotoConfig( + signature_version=UNSIGNED, + retries={"max_attempts": self.max_retries, "mode": "standard"}, + ), ) aws = AWSSRP( username=username, @@ -86,7 +88,10 @@ def sts_from_cognito_identity_pool(self, token, cognito_client=None, **kwargs): cognito_client = boto3.client( "cognito-identity", region_name=self.aws_region, - config=BotoConfig(signature_version=UNSIGNED), + config=BotoConfig( + signature_version=UNSIGNED, + retries={"max_attempts": self.max_retries, "mode": "standard"}, + ), ) for i in range(self.max_retries): @@ -105,7 +110,7 @@ def sts_from_cognito_identity_pool(self, token, cognito_client=None, **kwargs): ) break except ClientError as e: - # AWS eventual consistency, attempt to retry up to 3 times + # AWS eventual consistency, attempt to retry up to n (max_retries) times if "Couldn't verify signed token" in str(e): continue else: @@ -146,10 +151,14 @@ def requests_auth(username, password, **kwargs): class ApiTokenAuth: @staticmethod def requests_auth(config: StaxConfig, **kwargs): - # Minimize the potentical for token to expire while still being used for auth (say within a lambda function) + # Minimize the potential for token to expire while still being used for auth (say within a lambda function) + print(config.api_auth_retry_config.token_expiry_threshold) if config.expiration and config.expiration - timedelta( - minutes=int(environ.get("TOKEN_EXPIRY_THRESHOLD_IN_MINS", 1)) + minutes=config.api_auth_retry_config.token_expiry_threshold ) > datetime.now(timezone.utc): return config.auth - - return StaxAuth("ApiAuth", config).requests_auth(**kwargs) + return StaxAuth( + "ApiAuth", + config, + max_retries=config.api_auth_retry_config.max_attempts, + ).requests_auth(**kwargs) diff --git a/staxapp/aws_srp.py b/staxapp/aws_srp.py index 0bbc11a..096e4b4 100644 --- a/staxapp/aws_srp.py +++ b/staxapp/aws_srp.py @@ -27,6 +27,7 @@ import boto3 import six +from botocore.config import Config as BotoConfig class WarrantException(Exception): @@ -153,7 +154,13 @@ def __init__( self.client_id = client_id self.client_secret = client_secret self.client = ( - client if client else boto3.client("cognito-idp", region_name=pool_region) + client + if client + else boto3.client( + "cognito-idp", + region_name=pool_region, + config=BotoConfig(retries={"max_attempts": 5, "mode": "standard"}), + ) ) self.big_n = hex_to_long(n_hex) self.g = hex_to_long(g_hex) diff --git a/staxapp/config.py b/staxapp/config.py index 1a1e020..d0f7520 100644 --- a/staxapp/config.py +++ b/staxapp/config.py @@ -1,8 +1,7 @@ import logging import os import platform as sysinfo -from distutils.command.config import config -from email.policy import default +from typing import NamedTuple import requests @@ -12,9 +11,39 @@ logging.getLogger().setLevel(os.environ.get("LOG_LEVEL", logging.INFO)) +class StaxAuthRetryConfig(NamedTuple): + """ + Configuration options for the Stax API Auth retry + max_attempts: int: number of attempts to make + token_expiry_threshold: int: number of minutes before expiry to refresh + """ + + max_attempts = 5 + token_expiry_threshold = int( + # Env Var for backwards compatability, deprecated since 1.3.0 + os.getenv("TOKEN_EXPIRY_THRESHOLD_IN_MINS", 1) + ) + + +class StaxAPIRetryConfig(NamedTuple): + """ + Configuration options for the Stax API Auth retry + + max_attempts: int: number of attempts to make + backoff_factor: float: exponential backoff factor + status_codes: Tuple[int]: number of attempts to make + retry_methods: Tuple[str]: http methods to perform retries on + """ + + max_attempts = 5 + backoff_factor = 1.0 + status_codes = (429, 500, 502, 504) + retry_methods = ("GET", "PUT", "DELETE", "OPTIONS") + + class Config: """ - Insert doco here + Stax SDK Config """ STAX_REGION = os.getenv("STAX_REGION", "au1.staxapp.cloud") @@ -38,6 +67,9 @@ class Config: python_version = sysinfo.python_version() sdk_version = staxapp.__version__ + api_auth_retry_config = StaxAuthRetryConfig + api_retry_config = StaxAPIRetryConfig + def set_config(self): self.base_url = f"https://{self.hostname}/{self.API_VERSION}" config_url = f"{self.api_base_url()}/public/config" @@ -60,11 +92,20 @@ def get_api_config(cls, config_url): cls.cached_api_config["caching"] = config_url return config_response.json() - def __init__(self, hostname=None, access_key=None, secret_key=None): + def __init__( + self, + hostname=None, + access_key=None, + secret_key=None, + api_auth_retry_config=StaxAuthRetryConfig, + api_retry_config=StaxAPIRetryConfig, + ): if hostname is not None: self.hostname = hostname self.access_key = access_key self.secret_key = secret_key + self.api_auth_retry_config = api_auth_retry_config + self.api_retry_config = api_retry_config def init(self): if self._initialized: diff --git a/staxapp/openapi.py b/staxapp/openapi.py index 24e4e13..8da85c0 100644 --- a/staxapp/openapi.py +++ b/staxapp/openapi.py @@ -7,7 +7,7 @@ from staxapp.auth import ApiTokenAuth from staxapp.config import Config from staxapp.contract import StaxContract -from staxapp.exceptions import ApiException, ValidationException +from staxapp.exceptions import ValidationException class StaxClient: @@ -24,6 +24,8 @@ def __init__(self, classname, force=False, config=None): hostname=config.hostname, access_key=config.access_key, secret_key=config.secret_key, + api_auth_retry_config=config.api_auth_retry_config, + api_retry_config=config.api_retry_config, ) if not self._config._initialized: self._config.init() @@ -111,7 +113,7 @@ def stax_wrapper(*args, **kwargs): ] # Sort the operation map parameters parameter_index = -1 - # Check if the any of the parameter schemas match parameters provided + # Check if any of the parameter schemas match parameters provided for index in range(0, len(operation_parameters)): # Get any parameters from the keyword args and remove them from the payload if set(operation_parameters[index]).issubset(payload.keys()): diff --git a/staxapp/retry.py b/staxapp/retry.py new file mode 100644 index 0000000..69e5014 --- /dev/null +++ b/staxapp/retry.py @@ -0,0 +1,37 @@ +import requests +from requests.adapters import HTTPAdapter +from requests.packages.urllib3.util.retry import Retry + +N_RETRIES = 5 +BACKOFF_FACTOR = 1.0 +RETRY_STATUSES = (429, 500, 502, 504) +DEFAULT_ALLOWED_METHODS = ("GET", "PUT", "DELETE", "OPTIONS") + + +def requests_retry_session( + retries=N_RETRIES, + backoff_factor=BACKOFF_FACTOR, + status_list=RETRY_STATUSES, + allowed_methods=DEFAULT_ALLOWED_METHODS, + session=None, +): + """ + Initialises a retry requests session with configured exponential backoff + https://www.peterbe.com/plog/best-practice-with-retries-with-requests + """ + session = session or requests.Session() + retry = Retry( + total=retries, + read=retries, + connect=retries, + backoff_factor=backoff_factor, + status_forcelist=status_list, + allowed_methods=allowed_methods, + # Do not raise Retry Exception for backwards compatibility, return last response so ApiException raised + raise_on_status=False, + ) + adapter = HTTPAdapter(max_retries=retry) + session.mount("http://", adapter) + session.mount("https://", adapter) + + return session diff --git a/tests/test_api.py b/tests/test_api.py index 9dc682c..b3fb60e 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -4,13 +4,16 @@ To run: nose2 -v basics """ - +import os from datetime import datetime, timezone import unittest +from unittest.mock import patch import responses from staxapp.exceptions import ApiException -from staxapp.config import Config +# Due to os.getenv loading during import, need to patch in Env Var +with patch.dict(os.environ, {"TOKEN_EXPIRY_THRESHOLD_IN_MINS": "10"}, clear=True): + from staxapp.config import Config from staxapp.api import Api diff --git a/tests/test_auth.py b/tests/test_auth.py index e00ee52..a43fe33 100644 --- a/tests/test_auth.py +++ b/tests/test_auth.py @@ -4,7 +4,7 @@ To run: nose2 -v basics """ - +import os import unittest from unittest.mock import patch import jwt @@ -16,11 +16,11 @@ from botocore.client import Config as BotoConfig from botocore.stub import Stubber, ANY from datetime import datetime, timedelta, timezone -from os import environ -from staxapp.openapi import StaxClient +with patch.dict(os.environ, {"TOKEN_EXPIRY_THRESHOLD_IN_MINS": "10"}, clear=True): + from staxapp.config import Config + from staxapp.auth import StaxAuth, ApiTokenAuth, RootAuth -from staxapp.config import Config from staxapp.exceptions import InvalidCredentialsException @@ -28,7 +28,8 @@ class StaxAuthTests(unittest.TestCase): """ Inherited class to run all unit tests for this module """ - + + def setUp(self): self.cognito_client = botocore.session.get_session().create_client( "cognito-identity", @@ -328,6 +329,9 @@ def testApiTokenAuthExpiring(self, requests_auth_mock): """ sa = StaxAuth("ApiAuth", self.config) StaxConfig = self.config + ## Assert ENV Var used instead of default value of 1 + assert StaxConfig.api_auth_retry_config.token_expiry_threshold == 10 + ## expiration 20 minutes in the future, no need to refresh StaxConfig.expiration = datetime.now(timezone.utc) + timedelta(minutes=20) @@ -339,21 +343,7 @@ def testApiTokenAuthExpiring(self, requests_auth_mock): requests_auth_mock.assert_not_called() requests_auth_mock.reset_mock() - ## expiration in 5 seconds from now, refresh to avoid token becoming stale used - StaxConfig.expiration = datetime.now(timezone.utc) + timedelta(seconds=5) - - ApiTokenAuth.requests_auth( - config=self.config, - srp_client=self.aws_srp_client, - cognito_client=self.cognito_client, - ) - requests_auth_mock.assert_called_once() - - - requests_auth_mock.reset_mock() - ## expiration in 5 minutes from now, refresh to avoid token becoming stale used - ## override default triggering library to not refresh - environ["TOKEN_EXPIRY_THRESHOLD_IN_MINS"] = "10" + ## expiration in 2 minutes from now, refresh to avoid token becoming stale used StaxConfig.expiration = datetime.now(timezone.utc) + timedelta(minutes=2) ApiTokenAuth.requests_auth(