Add auto_create and db_groups to get_redshift_temp_engine(). #288

Author: igorborgest

awswrangler/db.py

@@ -315,6 +315,8 @@ def get_redshift_temp_engine(
     user: str,
     database: Optional[str] = None,
     duration: int = 900,
+    auto_create: bool = True,
+    db_groups: Optional[List[str]] = None,
     boto3_session: Optional[boto3.Session] = None,
 ) -> sqlalchemy.engine.Engine:
     """Get Glue connection details.
@@ -332,6 +334,12 @@ def get_redshift_temp_engine(
         The number of seconds until the returned temporary password expires.
         Constraint: minimum 900, maximum 3600.
         Default: 900
+    auto_create : bool
+        Create a database user with the name specified for the user named in user if one does not exist.
+    db_groups: List[str], optinal
+        A list of the names of existing database groups that the user named in DbUser will join for the current session,
+        in addition to any group memberships for an existing user.
+        If not specified, a new user is added only to PUBLIC.
     boto3_session : boto3.Session(), optional
         Boto3 Session. The default boto3 session will be used if boto3_session receive None.
 
@@ -347,9 +355,15 @@ def get_redshift_temp_engine(
 
     """
     client_redshift: boto3.client = _utils.client(service_name="redshift", session=boto3_session)
-    res: Dict[str, Any] = client_redshift.get_cluster_credentials(
-        DbUser=user, ClusterIdentifier=cluster_identifier, DurationSeconds=duration, AutoCreate=False
-    )
+    args: Dict[str, Any] = {
+        "DbUser": user,
+        "ClusterIdentifier": cluster_identifier,
+        "DurationSeconds": duration,
+        "AutoCreate": auto_create,
+    }
+    if db_groups is not None:
+        args["DbGroups"] = db_groups
+    res: Dict[str, Any] = client_redshift.get_cluster_credentials(**args)
     _user: str = _quote_plus(res["DbUser"])
     password: str = _quote_plus(res["DbPassword"])
     cluster: Dict[str, Any] = client_redshift.describe_clusters(ClusterIdentifier=cluster_identifier)["Clusters"][0]

tests/test_db.py

@@ -178,6 +178,15 @@ def test_redshift_temp_engine(parameters):
         assert cursor.fetchall()[0][0] == 1
 
 
+def test_redshift_temp_engine2(parameters):
+    engine = wr.db.get_redshift_temp_engine(
+        cluster_identifier=parameters["redshift"]["identifier"], user="john_doe", duration=900, db_groups=[]
+    )
+    with engine.connect() as con:
+        cursor = con.execute("SELECT 1")
+        assert cursor.fetchall()[0][0] == 1
+
+
 def test_postgresql_param():
     engine = wr.catalog.get_engine(connection="aws-data-wrangler-postgresql")
     df = wr.db.read_sql_query(sql="SELECT %(value)s as col0", con=engine, params={"value": 1})