RedshiftDataAPI serverless support #1530

Author: cnfait

awswrangler/data_api/redshift.py

@@ -16,7 +16,9 @@ class RedshiftDataApi(connector.DataApiConnector):
     Parameters
     ----------
     cluster_id: str
-        Id for the target Redshift cluster.
+        Id for the target Redshift cluster - only required if `workgroup_name` not provided.
+    workgroup_name: str
+        Name for the target serverless Redshift workgroup - only required if `cluster_id` not provided.
     database: str
         Target database name.
     secret_arn: str
@@ -35,8 +37,9 @@ class RedshiftDataApi(connector.DataApiConnector):
 
     def __init__(
         self,
-        cluster_id: str,
-        database: str,
+        cluster_id: str = "",
+        workgroup_name: str = "",
+        database: str = "",
         secret_arn: str = "",
         db_user: str = "",
         sleep: float = 0.25,
@@ -45,6 +48,7 @@ def __init__(
         boto3_session: Optional[boto3.Session] = None,
     ) -> None:
         self.cluster_id = cluster_id
+        self.workgroup_name = workgroup_name
         self.database = database
         self.secret_arn = secret_arn
         self.db_user = db_user
@@ -53,22 +57,36 @@ def __init__(
         logger: logging.Logger = logging.getLogger(__name__)
         super().__init__(self.client, logger)
 
+    def _validate_redshift_target(self) -> None:
+        if self.database == "":
+            raise ValueError("`database` must be set for connection")
+        if self.cluster_id == "" and self.workgroup_name == "":
+            raise ValueError("Either `cluster_id` or `workgroup_name`(Redshift Serverless) must be set for connection")
+
     def _validate_auth_method(self) -> None:
-        if self.secret_arn == "" and self.db_user == "":
+        if self.workgroup_name == "" and self.secret_arn == "" and self.db_user == "":
             raise ValueError("Either `secret_arn` or `db_user` must be set for authentication")
 
     def _execute_statement(self, sql: str, database: Optional[str] = None) -> str:
+        self._validate_redshift_target()
         self._validate_auth_method()
-        credentials = {"SecretArn": self.secret_arn}
-        if self.db_user:
+        credentials = {}
+        if self.secret_arn:
+            credentials = {"SecretArn": self.secret_arn}
+        elif self.db_user:
             credentials = {"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}
+
         self.logger.debug("Executing %s", sql)
         response: Dict[str, Any] = self.client.execute_statement(
-            ClusterIdentifier=self.cluster_id,
+            **redshift_target,
             Database=database,
             Sql=sql,
             **credentials,
@@ -167,8 +185,9 @@ class RedshiftDataApiTimeoutException(Exception):
 
 
 def connect(
-    cluster_id: str,
-    database: str,
+    cluster_id: str = "",
+    workgroup_name: str = "",
+    database: str = "",
     secret_arn: str = "",
     db_user: str = "",
     boto3_session: Optional[boto3.Session] = None,
@@ -179,7 +198,9 @@ def connect(
     Parameters
     ----------
     cluster_id: str
-        Id for the target Redshift cluster.
+        Id for the target Redshift cluster - only required if `workgroup_name` not provided.
+    workgroup_name: str
+        Name for the target serverless Redshift workgroup - only required if `cluster_id` not provided.
     database: str
         Target database name.
     secret_arn: str
@@ -196,7 +217,13 @@ def connect(
     A RedshiftDataApi connection instance that can be used with `wr.redshift.data_api.read_sql_query`.
     """
     return RedshiftDataApi(
-        cluster_id, database, secret_arn=secret_arn, db_user=db_user, boto3_session=boto3_session, **kwargs
+        cluster_id=cluster_id,
+        workgroup_name=workgroup_name,
+        database=database,
+        secret_arn=secret_arn,
+        db_user=db_user,
+        boto3_session=boto3_session,
+        **kwargs,
     )
 
 

poetry.lock

@@ -27,8 +27,8 @@ classifiers = [
 [tool.poetry.dependencies]
 python = ">=3.7.1, <3.11"
 
-boto3 = "^1.20.17"
-botocore = "^1.23.17"
+boto3 = "^1.24.11"
+botocore = "^1.27.11"
 pandas = "^1.2.0"
 numpy = "^1.21.0"
 pyarrow = ">=2.0.0, <8.1.0"

pyproject.toml

@@ -130,7 +130,9 @@ def workgroup3(bucket, kms_key):
 
 @pytest.fixture(scope="session")
 def databases_parameters(cloudformation_outputs, db_password):
-    parameters = dict(postgresql={}, mysql={}, redshift={}, sqlserver={}, mysql_serverless={}, oracle={})
+    parameters = dict(
+        postgresql={}, mysql={}, redshift={}, sqlserver={}, mysql_serverless={}, oracle={}, redshift_serverless={}
+    )
     parameters["postgresql"]["host"] = cloudformation_outputs.get("PostgresqlAddress")
     parameters["postgresql"]["port"] = 3306
     parameters["postgresql"]["schema"] = "public"
@@ -160,6 +162,9 @@ def databases_parameters(cloudformation_outputs, db_password):
     parameters["oracle"]["port"] = 1521
     parameters["oracle"]["schema"] = "TEST"
     parameters["oracle"]["database"] = "ORCL"
+    parameters["redshift_serverless"]["secret_arn"] = cloudformation_outputs.get("RedshiftServerlessSecretArn")
+    parameters["redshift_serverless"]["workgroup"] = cloudformation_outputs.get("RedshiftServerlessWorkgroup")
+    parameters["redshift_serverless"]["database"] = "test"
     return parameters
 
 

tests/conftest.py

@@ -12,23 +12,44 @@ def redshift_connector(databases_parameters):
     cluster_id = databases_parameters["redshift"]["identifier"]
     database = databases_parameters["redshift"]["database"]
     secret_arn = databases_parameters["redshift"]["secret_arn"]
-    conn = wr.data_api.redshift.connect(cluster_id, database, secret_arn=secret_arn, boto3_session=None)
-    return conn
+    con = wr.data_api.redshift.connect(
+        cluster_id=cluster_id, database=database, secret_arn=secret_arn, boto3_session=None
+    )
+    return con
 
 
 def create_rds_connector(rds_type, parameters):
     cluster_id = parameters[rds_type]["arn"]
     database = parameters[rds_type]["database"]
     secret_arn = parameters[rds_type]["secret_arn"]
-    conn = wr.data_api.rds.connect(cluster_id, database, secret_arn=secret_arn, boto3_session=boto3.DEFAULT_SESSION)
-    return conn
+    con = wr.data_api.rds.connect(cluster_id, database, secret_arn=secret_arn, boto3_session=boto3.DEFAULT_SESSION)
+    return con
 
 
 @pytest.fixture
 def mysql_serverless_connector(databases_parameters):
     return create_rds_connector("mysql_serverless", databases_parameters)
 
 
+def test_connect_redshift_serverless_iam_role(databases_parameters):
+    workgroup_name = databases_parameters["redshift_serverless"]["workgroup"]
+    database = databases_parameters["redshift_serverless"]["database"]
+    con = wr.data_api.redshift.connect(workgroup_name=workgroup_name, 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_serverless_secrets_manager(databases_parameters):
+    workgroup_name = databases_parameters["redshift_serverless"]["workgroup"]
+    database = databases_parameters["redshift_serverless"]["database"]
+    secret_arn = databases_parameters["redshift_serverless"]["secret_arn"]
+    con = wr.data_api.redshift.connect(
+        workgroup_name=workgroup_name, database=database, secret_arn=secret_arn, boto3_session=None
+    )
+    df = wr.data_api.redshift.read_sql_query("SELECT 1", con=con)
+    assert df.shape == (1, 1)
+
+
 @pytest.fixture(scope="function")
 def mysql_serverless_table(mysql_serverless_connector):
     name = f"tbl_{get_time_str_with_random_suffix()}"

tests/test_data_api.py

@@ -573,21 +573,21 @@ def mock_data_api_connector(connector, has_result_set=True):
 
 def test_data_api_redshift_create_connection():
     cluster_id = "cluster123"
-    conn = wr.data_api.redshift.connect(cluster_id, "db1", db_user="admin")
-    assert conn.cluster_id == cluster_id
+    con = wr.data_api.redshift.connect(cluster_id=cluster_id, database="db1", db_user="admin")
+    assert con.cluster_id == cluster_id
 
 
 def test_data_api_redshift_read_sql_results():
     cluster_id = "cluster123"
-    con = wr.data_api.redshift.connect(cluster_id, "db1", db_user="admin")
+    con = wr.data_api.redshift.connect(cluster_id=cluster_id, database="db1", db_user="admin")
     expected_dataframe = mock_data_api_connector(con)
     dataframe = wr.data_api.redshift.read_sql_query("SELECT * FROM test", con=con)
     pd.testing.assert_frame_equal(dataframe, expected_dataframe)
 
 
 def test_data_api_redshift_read_sql_no_results():
     cluster_id = "cluster123"
-    con = wr.data_api.redshift.connect(cluster_id, "db1", db_user="admin")
+    con = wr.data_api.redshift.connect(cluster_id=cluster_id, database="db1", db_user="admin")
     mock_data_api_connector(con, has_result_set=False)
     dataframe = wr.data_api.redshift.read_sql_query("DROP TABLE test", con=con)
     assert dataframe.empty is True

tests/test_moto.py

@@ -66,8 +66,14 @@
     "    secret_arn=\"arn:aws:secretsmanager:us-east-1:111111111111:secret:aws-sdk-pandas/redshift-ewn43d\"\n",
     ")\n",
     "\n",
+    "con_redshift_serverless = wr.data_api.redshift.connect(\n",
+    "    workgroup_name=\"aws-sdk-pandas\",\n",
+    "    database=\"test_redshift\",\n",
+    "    secret_arn=\"arn:aws:secretsmanager:us-east-1:111111111111:secret:aws-sdk-pandas/redshift-f3en4w\"\n",
+    ")\n",
+    "\n",
     "con_mysql = wr.data_api.rds.connect(\n",
-    "    cluster_id=\"arn:aws:rds:us-east-1:111111111111:cluster:mysql-serverless-cluster-wrangler\",\n",
+    "    resource_arn=\"arn:aws:rds:us-east-1:111111111111:cluster:mysql-serverless-cluster-wrangler\",\n",
     "    database=\"test_rds\",\n",
     "    secret_arn=\"arn:aws:secretsmanager:us-east-1:111111111111:secret:aws-sdk-pandas/mysql-23df3\"\n",
     ")"
@@ -102,4 +108,4 @@
    "metadata": {}
   }
  ]
-}
+}