From f298298c0038013124dc6a58e19130fd8ed58852 Mon Sep 17 00:00:00 2001
From: Sathiish Kumar <sathiish@amazon.com>
Date: Tue, 18 Jul 2023 09:00:13 -0700
Subject: [PATCH] Add retries for get cluster credentials API call and remove
 unused variables

---
 common/aws_service.py | 8 ++++++--
 core/replay/prep.py   | 8 ++------
 2 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/common/aws_service.py b/common/aws_service.py
index f75f0f394..f3534c181 100644
--- a/common/aws_service.py
+++ b/common/aws_service.py
@@ -3,6 +3,7 @@
 import json
 import logging
 import boto3
+from botocore.config import Config
 from botocore.exceptions import ClientError
 import asyncio
 import functools
@@ -34,9 +35,11 @@ def redshift_describe_logging_status(source_cluster_endpoint):
 
 
 def redshift_get_cluster_credentials(
-    region, user, database_name, cluster_id, duration=900, auto_create=False
+    region, user, database_name, cluster_id, duration=900, auto_create=False, additional_args={}
 ):
-    rs_client = boto3.client("redshift", region)
+    # see https://boto3.amazonaws.com/v1/documentation/api/latest/guide/retries.html#standard-retry-mode
+    config = Config(retries={"max_attempts": 10, "mode": "standard"})
+    rs_client = boto3.client("redshift", region, config=config)
     try:
         response = rs_client.get_cluster_credentials(
             DbUser=user,
@@ -44,6 +47,7 @@ def redshift_get_cluster_credentials(
             ClusterIdentifier=cluster_id,
             DurationSeconds=duration,
             AutoCreate=auto_create,
+            **additional_args,
         )
     except Exception as e:
         if e == rs_client.exceptions.ClusterNotFoundFault:
diff --git a/core/replay/prep.py b/core/replay/prep.py
index 455d77d2c..ee9a7c2f0 100644
--- a/core/replay/prep.py
+++ b/core/replay/prep.py
@@ -105,11 +105,8 @@ def correlate_transactions_with_connections(self, replay_id):
             total_connections,
         )
 
-    def get_connection_credentials(
-        self, username, database=None, max_attempts=10, skip_cache=False
-    ):
+    def get_connection_credentials(self, username, database=None, skip_cache=False):
         credentials_timeout_sec = 3600
-        retry_delay_sec = 10
 
         # how long to cache credentials per user
         cache_timeout_sec = 1800
@@ -155,8 +152,6 @@ def get_connection_credentials(
                 "verify": False,
             }
 
-        response = None
-
         if is_serverless(self.config) and self.config.get("secret_name", None) is not None:
             logger.info(f"Fetching secrets from: {self.config['secret_name']}")
             secret_keys = ["admin_username", "admin_password"]
@@ -179,6 +174,7 @@ def get_connection_credentials(
                 cluster_id=cluster_id,
                 duration=credentials_timeout_sec,
                 auto_create=False,
+                additional_args=additional_args,
             )
 
         if response is None: