diff --git a/core/replay/connection_thread.py b/core/replay/connection_thread.py index 0db50aaf..96384aa9 100644 --- a/core/replay/connection_thread.py +++ b/core/replay/connection_thread.py @@ -104,9 +104,7 @@ def initiate_connection(self, username): else: interface = "psql" r = ReplayPrep(self.config) - credentials = r.get_connection_credentials( - username, database=self.connection_log.database_name - ) + credentials = r.get_connection_credentials(username, database=self.connection_log.database_name) try: try: diff --git a/core/replay/replay.py b/core/replay/replay.py index caa94f01..2fca2532 100644 --- a/core/replay/replay.py +++ b/core/replay/replay.py @@ -83,11 +83,7 @@ def main(): # test connection try: # use the first user as a test - prep.get_connection_credentials( - connection_logs[0].username, - database=connection_logs[0].database_name, - max_attempts=1, - ) + prep.get_connection_credentials(connection_logs[0].username, database=connection_logs[0].database_name) except CredentialsException as e: logger.error( f"Unable to retrieve credentials using GetClusterCredentials ({str(e)}). " diff --git a/core/replay/unload_sys_table.py b/core/replay/unload_sys_table.py index 9e492563..3001adfd 100644 --- a/core/replay/unload_sys_table.py +++ b/core/replay/unload_sys_table.py @@ -17,9 +17,7 @@ def __init__(self, config, replay_id): def unload_system_table(self): # TODO: wrap this in retries and proper logging prep = ReplayPrep(self.config) - credentials = prep.get_connection_credentials( - self.config["master_username"], max_attempts=3 - ) + credentials = prep.get_connection_credentials(self.config["master_username"]) try: conn = db_connect( self.default_interface, diff --git a/core/tests/test_unload_sys_table.py b/core/tests/test_unload_sys_table.py index 5ec18a36..25794854 100644 --- a/core/tests/test_unload_sys_table.py +++ b/core/tests/test_unload_sys_table.py @@ -38,7 +38,7 @@ conn.cursor.return_value = cursor -def mock_get_connection_cred(self, val, max_attempts): +def mock_get_connection_cred(self, val): return { "host": "somehost", "port": 5437, @@ -82,9 +82,7 @@ def test_unload_system_table_with_unload_query(self, mock_boto, mock_debug): unload_object.unload_system_table() - mock_debug.assert_called_once_with( - "Executed unload query: stl_unload" - ) + mock_debug.assert_called_once_with("Executed unload query: stl_unload") @patch.object(ReplayPrep, "get_connection_credentials", mock_get_connection_cred) @patch("core.replay.unload_sys_table.db_connect", mock_db_connect_error)