From e27f2d0db836caf1cd8fc2c24f435f5fd766adc0 Mon Sep 17 00:00:00 2001
From: Blayne Chard <bchard@linz.govt.nz>
Date: Mon, 29 Aug 2022 16:18:57 +1200
Subject: [PATCH] fix: retry upto 3 times when InvalidIdentityTokenException
 happens (#115)

* fix: retry upto 3 times when InvalidIdentityTokenException happens

InvalidIdentityTokenException seems to intermitently cause workflows to fail when trying to get_credentials.

catch the failure and attempt to retry.

* fix: retry fetching credentials on failure
---
 scripts/aws/aws_helper.py   | 31 +++++++++++++++++++++++++++++++
 scripts/gdal/gdal_helper.py |  5 ++---
 2 files changed, 33 insertions(+), 3 deletions(-)

diff --git a/scripts/aws/aws_helper.py b/scripts/aws/aws_helper.py
index 534f64013..854361b66 100644
--- a/scripts/aws/aws_helper.py
+++ b/scripts/aws/aws_helper.py
@@ -1,5 +1,7 @@
 import json
+from dataclasses import dataclass
 from os import environ
+from time import sleep
 from typing import Any, Dict, List, NamedTuple, Optional
 from urllib.parse import urlparse
 
@@ -79,6 +81,35 @@ def get_session(prefix: str) -> boto3.Session:
     return current_session
 
 
+@dataclass
+class AwsFrozenCredentials:
+    """
+    work around as I couldn't find the type for get_frozen_credentials()
+    """
+
+    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 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)
+            last_error = e
+
+    raise last_error
+
+
 def _get_credential_config(prefix: str) -> Optional[CredentialSource]:
     get_log().debug("get_credentials_bucket_name", prefix=prefix)
     if not bucket_roles:
diff --git a/scripts/gdal/gdal_helper.py b/scripts/gdal/gdal_helper.py
index 5cfd5666b..277b0d78a 100644
--- a/scripts/gdal/gdal_helper.py
+++ b/scripts/gdal/gdal_helper.py
@@ -4,7 +4,7 @@
 
 from linz_logger import get_log
 
-from scripts.aws.aws_helper import get_session, is_s3
+from scripts.aws.aws_helper import get_session_credentials, is_s3
 from scripts.logging.time_helper import time_in_ms
 
 
@@ -60,8 +60,7 @@ def run_gdal(
     if input_file:
         if is_s3(input_file):
             # Set the credentials for GDAL to be able to read the input file
-            session = get_session(input_file)
-            credentials = session.get_credentials()
+            credentials = get_session_credentials(input_file)
             gdal_env["AWS_ACCESS_KEY_ID"] = credentials.access_key
             gdal_env["AWS_SECRET_ACCESS_KEY"] = credentials.secret_key
             gdal_env["AWS_SESSION_TOKEN"] = credentials.token