Feature/copy creds (#485)

* [feat] added functionality to use boto3_session to unload to s3.

* [refactor] Refactored authorization string creation and getting credentials from boto3 session.

* [feat/doc] Added functionality to use boto3 session credentials for unload function. Also updated doc for db_users.

* [doc] Updated docstring

* [refactor] reformatted using black

* [refactor] fixed black string reformat

Author: danielwo

awswrangler/_utils.py

@@ -232,6 +232,16 @@ def get_region_from_session(boto3_session: Optional[boto3.Session] = None, defau
     raise exceptions.InvalidArgument("There is no region_name defined on boto3, please configure it.")
 
 
+def get_credentials_from_session(
+    boto3_session: Optional[boto3.Session] = None,
+) -> botocore.credentials.ReadOnlyCredentials:
+    """Get AWS credentials from boto3 session."""
+    session: boto3.Session = ensure_session(session=boto3_session)
+    credentials: botocore.credentials.Credentials = session.get_credentials()
+    frozen_credentials: botocore.credentials.ReadOnlyCredentials = credentials.get_frozen_credentials()
+    return frozen_credentials
+
+
 def list_sampling(lst: List[Any], sampling: float) -> List[Any]:
     """Random List sampling."""
     if sampling == 1.0:

awswrangler/redshift.py

@@ -5,6 +5,7 @@
 from typing import Any, Dict, Iterator, List, Optional, Tuple, Union
 
 import boto3
+import botocore
 import pandas as pd
 import pyarrow as pa
 import redshift_connector
@@ -54,18 +55,52 @@ def _does_table_exist(cursor: redshift_connector.Cursor, schema: Optional[str],
     return len(cursor.fetchall()) > 0
 
 
+def _make_s3_auth_string(
+    aws_access_key_id: Optional[str] = None,
+    aws_secret_access_key: Optional[str] = None,
+    aws_session_token: Optional[str] = None,
+    iam_role: Optional[str] = None,
+    boto3_session: Optional[boto3.Session] = None,
+) -> str:
+    if aws_access_key_id is not None and aws_secret_access_key is not None:
+        auth_str: str = f"ACCESS_KEY_ID '{aws_access_key_id}'\nSECRET_ACCESS_KEY '{aws_secret_access_key}'\n"
+        if aws_session_token is not None:
+            auth_str += f"SESSION_TOKEN '{aws_session_token}'\n"
+    elif iam_role is not None:
+        auth_str = f"IAM_ROLE '{iam_role}'\n"
+    else:
+        _logger.debug("Attempting to get S3 authorization credentials from boto3 session.")
+        credentials: botocore.credentials.ReadOnlyCredentials
+        credentials = _utils.get_credentials_from_session(boto3_session=boto3_session)
+        if credentials.access_key is None or credentials.secret_key is None:
+            raise exceptions.InvalidArgument(
+                "One of IAM Role or AWS ACCESS_KEY_ID and SECRET_ACCESS_KEY must be "
+                "given. Unable to find ACCESS_KEY_ID and SECRET_ACCESS_KEY in boto3 "
+                "session."
+            )
+
+        auth_str = f"ACCESS_KEY_ID '{credentials.access_key}'\nSECRET_ACCESS_KEY '{credentials.secret_key}'\n"
+        if credentials.token is not None:
+            auth_str += f"SESSION_TOKEN '{credentials.token}'\n"
+
+    return auth_str
+
+
 def _copy(
     cursor: redshift_connector.Cursor,
     path: str,
     table: str,
-    iam_role: str,
+    iam_role: Optional[str] = None,
+    boto3_session: Optional[str] = None,
     schema: Optional[str] = None,
 ) -> None:
     if schema is None:
         table_name: str = f'"{table}"'
     else:
         table_name = f'"{schema}"."{table}"'
-    sql: str = f"COPY {table_name} FROM '{path}'\nIAM_ROLE '{iam_role}'\nFORMAT AS PARQUET"
+
+    auth_str: str = _make_s3_auth_string(iam_role=iam_role, boto3_session=boto3_session)
+    sql: str = f"COPY {table_name} FROM '{path}'{auth_str}\nFORMAT AS PARQUET"
     _logger.debug("copy query:\n%s", sql)
     cursor.execute(sql)
 
@@ -388,6 +423,9 @@ def connect_temp(
         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], optional
+        A list of the names of existing database groups that the user named in user 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.
     ssl: bool
@@ -703,7 +741,7 @@ def unload_to_files(
     sql: str,
     path: str,
     con: redshift_connector.Connection,
-    iam_role: str,
+    iam_role: Optional[str] = None,
     region: Optional[str] = None,
     max_file_size: Optional[float] = None,
     kms_key_id: Optional[str] = None,
@@ -730,7 +768,7 @@ def unload_to_files(
     con : redshift_connector.Connection
         Use redshift_connector.connect() to use "
         "credentials directly or wr.redshift.connect() to fetch it from the Glue Catalog.
-    iam_role : str
+    iam_role : str, optional
         AWS IAM role with the related permissions.
     region : str, optional
         Specifies the AWS Region where the target Amazon S3 bucket is located.
@@ -782,10 +820,13 @@ def unload_to_files(
         region_str: str = f"\nREGION AS '{region}'" if region is not None else ""
         max_file_size_str: str = f"\nMAXFILESIZE AS {max_file_size} MB" if max_file_size is not None else ""
         kms_key_id_str: str = f"\nKMS_KEY_ID '{kms_key_id}'" if kms_key_id is not None else ""
+
+        auth_str: str = _make_s3_auth_string(iam_role=iam_role, boto3_session=boto3_session)
+
         sql = (
             f"UNLOAD ('{sql}')\n"
             f"TO '{path}'\n"
-            f"IAM_ROLE '{iam_role}'\n"
+            f"{auth_str}"
             "ALLOWOVERWRITE\n"
             "PARALLEL ON\n"
             "FORMAT PARQUET\n"
@@ -804,7 +845,7 @@ def unload(
     sql: str,
     path: str,
     con: redshift_connector.Connection,
-    iam_role: str,
+    iam_role: Optional[str],
     region: Optional[str] = None,
     max_file_size: Optional[float] = None,
     kms_key_id: Optional[str] = None,
@@ -857,7 +898,7 @@ def unload(
     con : redshift_connector.Connection
         Use redshift_connector.connect() to use "
         "credentials directly or wr.redshift.connect() to fetch it from the Glue Catalog.
-    iam_role : str
+    iam_role : str, optional
         AWS IAM role with the related permissions.
     region : str, optional
         Specifies the AWS Region where the target Amazon S3 bucket is located.
@@ -949,7 +990,7 @@ def copy_from_files(  # pylint: disable=too-many-locals,too-many-arguments
     con: redshift_connector.Connection,
     table: str,
     schema: str,
-    iam_role: str,
+    iam_role: Optional[str] = None,
     parquet_infer_sampling: float = 1.0,
     mode: str = "append",
     diststyle: str = "AUTO",
@@ -992,7 +1033,7 @@ def copy_from_files(  # pylint: disable=too-many-locals,too-many-arguments
         Table name
     schema : str
         Schema name
-    iam_role : str
+    iam_role : str, optional
         AWS IAM role with the related permissions.
     parquet_infer_sampling : float
         Random sample ratio of files that will have the metadata inspected.
@@ -1089,6 +1130,7 @@ def copy_from_files(  # pylint: disable=too-many-locals,too-many-arguments
                 table=created_table,
                 schema=created_schema,
                 iam_role=iam_role,
+                boto3_session=boto3_session,
             )
             if table != created_table:  # upsert
                 _upsert(cursor=cursor, schema=schema, table=table, temp_table=created_table, primary_keys=primary_keys)
@@ -1107,7 +1149,7 @@ def copy(  # pylint: disable=too-many-arguments
     con: redshift_connector.Connection,
     table: str,
     schema: str,
-    iam_role: str,
+    iam_role: Optional[str] = None,
     index: bool = False,
     dtype: Optional[Dict[str, str]] = None,
     mode: str = "append",
@@ -1161,7 +1203,7 @@ def copy(  # pylint: disable=too-many-arguments
         Table name
     schema : str
         Schema name
-    iam_role : str
+    iam_role : str, optional
         AWS IAM role with the related permissions.
     index : bool
         True to store the DataFrame index in file, otherwise False to ignore it.