fixing bug to allow replay to connect to multiple databases and lower…

…ing numpy version(#152) numpy version Co-authored-by: Devanathan Sabapathy <[email protected]>
aws · Jul 5, 2024 · 67bc469 · 67bc469
1 parent 616b5cd
commit 67bc469
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 2 deletions.
diff --git 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
@@ -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
@@ -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