From 67bc469768bf8265d286e56aab73b3eb560a9fe7 Mon Sep 17 00:00:00 2001
From: Devanathan Sabapathy
 <123671658+DevanathanSabapathy1@users.noreply.github.com>
Date: Fri, 5 Jul 2024 10:11:44 -0700
Subject: [PATCH] fixing bug to allow replay to connect to multiple databases
 and lowering numpy version(#152)

numpy version

Co-authored-by: Devanathan Sabapathy <devsaba@dev-dsk-devsaba-2a-7e44d235.us-west-2.amazon.com>
---
 core/replay/prep.py     |  3 ++-
 core/tests/test_prep.py | 24 ++++++++++++++++++++++++
 requirements.txt        |  2 +-
 3 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/core/replay/prep.py b/core/replay/prep.py
index ee9a7c2f0..863a74d60 100644
--- a/core/replay/prep.py
+++ b/core/replay/prep.py
@@ -139,7 +139,8 @@ def get_connection_credentials(self, username, database=None, skip_cache=False):
             cluster_host = cluster_endpoint.split(":")[0]
 
         cluster_port = cluster_endpoint_split[5].split("/")[0][4:]
-        database = cluster_endpoint_split[5].split("/")[1]
+        if database == None:
+            database = cluster_endpoint_split[5].split("/")[1]
 
         additional_args = {}
         if os.environ.get("ENDPOINT_URL"):
diff --git a/core/tests/test_prep.py b/core/tests/test_prep.py
index 2b4cf7e9b..d2eab1011 100644
--- a/core/tests/test_prep.py
+++ b/core/tests/test_prep.py
@@ -266,6 +266,30 @@ def test_get_connection_credentials_provisioned_from_redshift_api(
                 "database": "dev",
             },
         )
+    
+    @patch("core.replay.prep.redshift_get_cluster_credentials")
+    def test_get_connection_credentials_for_database_not_being_none(self, mock_get_cluster_creds):
+        mock_get_cluster_creds.return_value = {"DbUser": "testAdmin", "DbPassword": "testPassword"}
+        prep = ReplayPrep(
+            {
+                "target_cluster_endpoint": "test.111222333222.us-east-1.redshift.amazonaws.com:5439/dev",
+                "odbc_driver": "test",
+                "nlb_nat_dns": None,
+                "target_cluster_region": "us-east-1",
+            }
+        )
+
+        result = prep.get_connection_credentials("testUsername","tcpds_100g")
+        self.assertEqual(
+            result["psql"],
+            {
+                "username": "testAdmin",
+                "password": "testPassword",
+                "host": "test.111222333222.us-east-1.redshift.amazonaws.com",
+                "port": "5439",
+                "database": "tcpds_100g",
+            },
+        )
 
 
 class TestCorrelateTransactionsWithConnections(unittest.TestCase):
diff --git a/requirements.txt b/requirements.txt
index c99c35858..6dbbe5691 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -3,7 +3,7 @@ botocore==1.27.24
 Flask==2.2.5
 matplotlib==3.5.2
 moto==3.1.16
-numpy
+numpy==1.26.4
 pandas
 python-dateutil==2.8.1
 PyYAML==6.0