Merge branch 'aws:main' into main

aws · Sep 30, 2023 · 8169293 · 8169293
2 parents bcfa0ae + 5ee1081
commit 8169293
Show file tree

Hide file tree

Showing 4 changed files with 44 additions and 12 deletions.
diff --git a/awswrangler/data_api/redshift.py b/awswrangler/data_api/redshift.py
@@ -93,8 +93,12 @@ def _validate_redshift_target(self) -> None:
             raise ValueError("Either `cluster_id` or `workgroup_name`(Redshift Serverless) must be set for connection")
 
     def _validate_auth_method(self) -> None:
-        if not self.workgroup_name and not self.secret_arn and not self.db_user:
-            raise ValueError("Either `secret_arn` or `db_user` must be set for authentication")
+        if not self.workgroup_name and not self.secret_arn and not self.db_user and not self.cluster_id:
+            raise exceptions.InvalidArgumentCombination(
+                "Either `secret_arn`, `workgroup_name`, `db_user`, or `cluster_id` must be set for authentication."
+            )
+        if self.db_user and self.secret_arn:
+            raise exceptions.InvalidArgumentCombination("Only one of `secret_arn` or `db_user` is allowed.")
 
     def _execute_statement(
         self,
@@ -110,26 +114,25 @@ def _execute_statement(
 
         self._validate_redshift_target()
         self._validate_auth_method()
-        credentials = {}
+        args = {}
         if self.secret_arn:
-            credentials = {"SecretArn": self.secret_arn}
-        elif self.db_user:
-            credentials = {"DbUser": self.db_user}
+            args["SecretArn"] = self.secret_arn
+        if self.db_user:
+            args["DbUser"] = self.db_user
 
         if database is None:
             database = self.database
 
         if self.cluster_id:
-            redshift_target = {"ClusterIdentifier": self.cluster_id}
-        elif self.workgroup_name:
-            redshift_target = {"WorkgroupName": self.workgroup_name}
+            args["ClusterIdentifier"] = self.cluster_id
+        if self.workgroup_name:
+            args["WorkgroupName"] = self.workgroup_name
 
         _logger.debug("Executing %s", sql)
         response = self.client.execute_statement(
-            **redshift_target,  # type: ignore[arg-type]
             Database=database,
             Sql=sql,
-            **credentials,  # type: ignore[arg-type]
+            **args,  # type: ignore[arg-type]
         )
         return response["Id"]
 

diff --git a/docs/source/layers.rst b/docs/source/layers.rst
@@ -408,3 +408,15 @@ Version 3.4.0
 +----------------+--------+-------+-----------------------------------------------------------------------------------+
 | me-south-1     | 3.9    | arm64 | arn:aws:lambda:me-south-1:938046470361:layer:AWSSDKPandas-Python39-Arm64:1        |
 +----------------+--------+-------+-----------------------------------------------------------------------------------+
+| cn-north-1     | 3.10   | x86_64| arn:aws-cn:lambda:cn-north-1:406640652441:layer:AWSSDKPandas-Python310:1          |
++----------------+--------+-------+-----------------------------------------------------------------------------------+
+| cn-north-1     | 3.8    | x86_64| arn:aws-cn:lambda:cn-north-1:406640652441:layer:AWSSDKPandas-Python38:1           |
++----------------+--------+-------+-----------------------------------------------------------------------------------+
+| cn-north-1     | 3.9    | x86_64| arn:aws-cn:lambda:cn-north-1:406640652441:layer:AWSSDKPandas-Python39:1           |
++----------------+--------+-------+-----------------------------------------------------------------------------------+
+| cn-northwest-1 | 3.10   | x86_64| arn:aws-cn:lambda:cn-northwest-1:406640652441:layer:AWSSDKPandas-Python310:1      |
++----------------+--------+-------+-----------------------------------------------------------------------------------+
+| cn-northwest-1 | 3.8    | x86_64| arn:aws-cn:lambda:cn-northwest-1:406640652441:layer:AWSSDKPandas-Python38:1       |
++----------------+--------+-------+-----------------------------------------------------------------------------------+
+| cn-northwest-1 | 3.9    | x86_64| arn:aws-cn:lambda:cn-northwest-1:406640652441:layer:AWSSDKPandas-Python39:1       |
++----------------+--------+-------+-----------------------------------------------------------------------------------+
diff --git a/tests/unit/test_data_api.py b/tests/unit/test_data_api.py
@@ -54,6 +54,23 @@ def test_connect_redshift_serverless_iam_role(databases_parameters: Dict[str, An
     assert df.shape == (1, 1)
 
 
+def test_connect_redshift_cluster_iam_role(databases_parameters: Dict[str, Any]) -> None:
+    cluster_id = databases_parameters["redshift"]["identifier"]
+    database = databases_parameters["redshift"]["database"]
+    con = wr.data_api.redshift.connect(cluster_id=cluster_id, database=database, boto3_session=None)
+    df = wr.data_api.redshift.read_sql_query("SELECT 1", con=con)
+    assert df.shape == (1, 1)
+
+
+def test_connect_redshift_cluster_db_user(databases_parameters: Dict[str, Any]) -> None:
+    cluster_id = databases_parameters["redshift"]["identifier"]
+    database = databases_parameters["redshift"]["database"]
+    db_user = databases_parameters["user"]
+    con = wr.data_api.redshift.connect(cluster_id=cluster_id, database=database, db_user=db_user, boto3_session=None)
+    df = wr.data_api.redshift.read_sql_query("SELECT 1", con=con)
+    assert df.shape == (1, 1)
+
+
 def test_connect_redshift_serverless_secrets_manager(databases_parameters: Dict[str, Any]) -> None:
     workgroup_name = databases_parameters["redshift_serverless"]["workgroup"]
     database = databases_parameters["redshift_serverless"]["database"]

diff --git a/tox.ini b/tox.ini
@@ -10,7 +10,7 @@ passenv =
        AWS_SECRET_ACCESS_KEY
        AWS_SESSION_TOKEN
 setenv =
-       COV_FAIL_UNDER = 82.00
+       COV_FAIL_UNDER = 87.00
 allowlist_externals = poetry
 commands_pre =
        poetry install --no-root --sync --extras "deltalake gremlin mysql opencypher opensearch oracle postgres redshift sparql sqlserver geopandas"