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

@@ -24,20 +24,50 @@ export STAX_SECRET_KEY=<your_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
 

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))
+

requirements.txt

@@ -8,6 +8,7 @@ pylint
 pytest
 pytest-cov
 responses
+requests
 pyjwt==2.4.0
 boto3
 aws_requests_auth

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,19 +26,32 @@ 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
     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,

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)

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)

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:

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()):