fix: retry fetching credentials on failure

scripts/aws/aws_helper.py

@@ -1,4 +1,5 @@
 import json
+from dataclasses import dataclass
 from os import environ
 from time import sleep
 from typing import Any, Dict, List, NamedTuple, Optional
@@ -12,6 +13,7 @@
 from scripts.aws.aws_credential_source import CredentialSource
 
 S3Path = NamedTuple("S3Path", [("bucket", str), ("key", str)])
+AwsCredentials = NamedTuple("AwsCredentials", [("access_key", str), ("secret_key", str), ("token", str)])
 
 aws_profile = environ.get("AWS_PROFILE")
 session = boto3.Session(profile_name=aws_profile)
@@ -80,16 +82,34 @@ def get_session(prefix: str) -> boto3.Session:
     get_log().info("role_assume", prefix=prefix, bucket=cfg.bucket, role_arn=cfg.roleArn)
     return current_session
 
-def get_session_credentials(prefix: str, retry_count=3) -> boto3.Credentials:
+
+@dataclass
+class AwsFrozenCredentials:
     """
-    Attempt to get cretentials for a prefix, retrying upto retry_count amount of times
+    work around as I couldn't find the type for get_frozen_credentials()
     """
-    for retry in range(retry_count):
+
+    access_key: str
+    secret_key: str
+    token: str
+
+
+def get_session_credentials(prefix: str, retry_count: int = 3) -> AwsFrozenCredentials:
+    """
+    Attempt to get credentials for a prefix, retrying upto retry_count amount of times
+    """
+    last_error: Exception = Exception(f"Invalid retry count: {retry_count}")
+    for retry in range(1, retry_count + 1):
         try:
-            get_session(prefix).get_credentials()
-        except botocore.errorfactory.InvalidIdentityTokenException: 
+            # Get credentials may give differing access_key and secret_key
+            credentials: AwsFrozenCredentials = get_session(prefix).get_frozen_credentials()
+            return credentials
+        except client_sts.meta.client.exceptions.InvalidIdentityTokenException as e:
             get_log().warn("bucket_load_retry", retry_count=retry)
-            sleep(0.5 * (retry + 1))
+            sleep(0.5 * retry)
+            last_error = e
+
+    raise last_error
 
 
 def _get_credential_config(prefix: str) -> Optional[CredentialSource]:

scripts/gdal/gdal_helper.py

@@ -4,7 +4,7 @@
 
 from linz_logger import get_log
 
-from scripts.aws.aws_helper import is_s3
+from scripts.aws.aws_helper import get_session_credentials, is_s3
 from scripts.logging.time_helper import time_in_ms