Add KMS encrypt on S3 and bumping version to 0b32

Author: igorborgest

.gitignore

@@ -67,6 +67,7 @@ instance/
 
 # Sphinx documentation
 docs/_build/
+docs/source/api/
 
 # PyBuilder
 target/

README.md

@@ -46,6 +46,19 @@ session.pandas.to_parquet(
 
 If a Glue Database name is passed, all the metadata will be created in the Glue Catalog. If not, only the s3 data write will be done.
 
+### Writing Pandas Dataframe to S3 as Parquet encrypting with a KMS key
+
+```py3
+extra_args = {
+    "ServerSideEncryption": "aws:kms",
+    "SSEKMSKeyId": "YOUR_KMY_KEY_ARN"
+}
+session = awswrangler.Session(s3_additional_kwargs=extra_args)
+session.pandas.to_parquet(
+    path="s3://..."
+)
+```
+
 ### Reading from AWS Athena to Pandas
 
 ```py3

awswrangler/__version__.py

@@ -1,4 +1,4 @@
 __title__ = "awswrangler"
 __description__ = "Utility belt to handle data on AWS."
-__version__ = "0.0b31"
+__version__ = "0.0b32"
 __license__ = "Apache License 2.0"

awswrangler/athena.py

@@ -54,6 +54,7 @@ def get_query_dtype(self, query_execution_id):
     def create_athena_bucket(self):
         """
         Creates the default Athena bucket if not exists
+
         :return: Bucket s3 path (E.g. s3://aws-athena-query-results-ACCOUNT-REGION/)
         """
         account_id = (self._session.boto3_session.client(
@@ -69,6 +70,7 @@ def create_athena_bucket(self):
     def run_query(self, query, database, s3_output=None):
         """
         Run a SQL Query against AWS Athena
+
         :param query: SQL query
         :param database: AWS Glue/Athena database name
         :param s3_output: AWS S3 path

awswrangler/glue.py

@@ -17,6 +17,7 @@ def __init__(self, session):
     def get_table_athena_types(self, database, table):
         """
         Get all columns names and the related data types
+
         :param database: Glue database's name
         :param table: Glue table's name
         :return: A dictionary as {"col name": "col dtype"}
@@ -34,6 +35,7 @@ def get_table_athena_types(self, database, table):
     def get_table_python_types(self, database, table):
         """
         Get all columns names and the related python types
+
         :param database: Glue database's name
         :param table: Glue table's name
         :return: A dictionary as {"col name": "col python type"}
@@ -178,7 +180,7 @@ def _build_schema(dataframe,
                       partition_cols,
                       preserve_index,
                       cast_columns=None):
-        print(f"dataframe.dtypes:\n{dataframe.dtypes}")
+        logger.debug(f"dataframe.dtypes:\n{dataframe.dtypes}")
         if not partition_cols:
             partition_cols = []
         schema_built = []

awswrangler/pandas.py

@@ -54,6 +54,7 @@ def read_csv(
         Try to mimic as most as possible pandas.read_csv()
         https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.read_csv.html
         P.S. max_result_size != None tries to mimic the chunksize behaviour in pandas.read_sql()
+
         :param path: AWS S3 path (E.g. S3://BUCKET_NAME/KEY_NAME)
         :param max_result_size: Max number of bytes on each request to S3
         :param header: Same as pandas.read_csv()
@@ -131,6 +132,7 @@ def _read_csv_iterator(
         Read CSV file from AWS S3 using optimized strategies.
         Try to mimic as most as possible pandas.read_csv()
         https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.read_csv.html
+
         :param client_s3: Boto3 S3 client object
         :param bucket_name: S3 bucket name
         :param key_path: S3 key path (W/o bucket)
@@ -235,6 +237,7 @@ def _extract_terminator_profile(body, sep, quotechar, lineterminator,
                                     last_index):
         """
         Backward parser for quoted CSV lines
+
         :param body: String
         :param sep: Same as pandas.read_csv()
         :param quotechar: Same as pandas.read_csv()
@@ -290,6 +293,7 @@ def _extract_terminator_profile(body, sep, quotechar, lineterminator,
     def _find_terminator(body, sep, quoting, quotechar, lineterminator):
         """
         Find for any suspicious of line terminator (From end to start)
+
         :param body: String
         :param sep: Same as pandas.read_csv()
         :param quoting: Same as pandas.read_csv()
@@ -345,6 +349,7 @@ def _read_csv_once(
         Read CSV file from AWS S3 using optimized strategies.
         Try to mimic as most as possible pandas.read_csv()
         https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.read_csv.html
+
         :param client_s3: Boto3 S3 client object
         :param bucket_name: S3 bucket name
         :param key_path: S3 key path (W/o bucket)
@@ -391,6 +396,7 @@ def read_sql_athena(self,
         """
         Executes any SQL query on AWS Athena and return a Dataframe of the result.
         P.S. If max_result_size is passed, then a iterator of Dataframes is returned.
+
         :param sql: SQL Query
         :param database: Glue/Athena Database
         :param s3_output: AWS S3 path
@@ -436,6 +442,7 @@ def to_csv(
         """
         Write a Pandas Dataframe as CSV files on S3
         Optionally writes metadata on AWS Glue.
+
         :param dataframe: Pandas Dataframe
         :param path: AWS S3 path (E.g. s3://bucket-name/folder_name/
         :param database: AWS Glue Database name
@@ -474,6 +481,7 @@ def to_parquet(self,
         """
         Write a Pandas Dataframe as parquet files on S3
         Optionally writes metadata on AWS Glue.
+
         :param dataframe: Pandas Dataframe
         :param path: AWS S3 path (E.g. s3://bucket-name/folder_name/
         :param database: AWS Glue Database name
@@ -483,8 +491,7 @@ def to_parquet(self,
         :param mode: "append", "overwrite", "overwrite_partitions"
         :param procs_cpu_bound: Number of cores used for CPU bound tasks
         :param procs_io_bound: Number of cores used for I/O bound tasks
-        :param cast_columns: Dictionary of columns names and Arrow types to be casted.
-               E.g. {"col name": "int64", "col2 name": "int32"}
+        :param cast_columns: Dictionary of columns names and Arrow types to be casted. (E.g. {"col name": "int64", "col2 name": "int32"})
         :return: List of objects written on S3
         """
         return self.to_s3(dataframe=dataframe,
@@ -514,6 +521,7 @@ def to_s3(self,
         """
         Write a Pandas Dataframe on S3
         Optionally writes metadata on AWS Glue.
+
         :param dataframe: Pandas Dataframe
         :param path: AWS S3 path (E.g. s3://bucket-name/folder_name/
         :param file_format: "csv" or "parquet"
@@ -524,9 +532,7 @@ def to_s3(self,
         :param mode: "append", "overwrite", "overwrite_partitions"
         :param procs_cpu_bound: Number of cores used for CPU bound tasks
         :param procs_io_bound: Number of cores used for I/O bound tasks
-        :param cast_columns: Dictionary of columns indexes and Arrow types to be casted.
-               E.g. {2: "int64", 5: "int32"}
-               Only for "parquet" file_format
+        :param cast_columns: Dictionary of columns indexes and Arrow types to be casted. (E.g. {2: "int64", 5: "int32"}) (Only for "parquet" file_format)
         :return: List of objects written on S3
         """
         if not partition_cols:
@@ -769,17 +775,16 @@ def to_redshift(
     ):
         """
         Load Pandas Dataframe as a Table on Amazon Redshift
+
         :param dataframe: Pandas Dataframe
         :param path: S3 path to write temporary files (E.g. s3://BUCKET_NAME/ANY_NAME/)
         :param connection: A PEP 249 compatible connection (Can be generated with Redshift.generate_connection())
         :param schema: The Redshift Schema for the table
         :param table: The name of the desired Redshift table
         :param iam_role: AWS IAM role with the related permissions
-        :param diststyle: Redshift distribution styles. Must be in ["AUTO", "EVEN", "ALL", "KEY"]
-               https://docs.aws.amazon.com/redshift/latest/dg/t_Distributing_data.html
+        :param diststyle: Redshift distribution styles. Must be in ["AUTO", "EVEN", "ALL", "KEY"] (https://docs.aws.amazon.com/redshift/latest/dg/t_Distributing_data.html)
         :param distkey: Specifies a column name or positional number for the distribution key
-        :param sortstyle: Sorting can be "COMPOUND" or "INTERLEAVED"
-               https://docs.aws.amazon.com/redshift/latest/dg/t_Sorting_data.html
+        :param sortstyle: Sorting can be "COMPOUND" or "INTERLEAVED" (https://docs.aws.amazon.com/redshift/latest/dg/t_Sorting_data.html)
         :param sortkey: List of columns to be sorted
         :param preserve_index: Should we preserve the Dataframe index?
         :param mode: append or overwrite

awswrangler/redshift.py

@@ -113,6 +113,7 @@ def load_table(
         """
         Load Parquet files into a Redshift table using a manifest file.
         Creates the table if necessary.
+
         :param dataframe: Pandas or Spark Dataframe
         :param dataframe_type: "pandas" or "spark"
         :param manifest_path: S3 path for manifest file (E.g. S3://...)
@@ -121,11 +122,9 @@ def load_table(
         :param redshift_conn: A PEP 249 compatible connection (Can be generated with Redshift.generate_connection())
         :param num_files: Number of files to be loaded
         :param iam_role: AWS IAM role with the related permissions
-        :param diststyle: Redshift distribution styles. Must be in ["AUTO", "EVEN", "ALL", "KEY"]
-               https://docs.aws.amazon.com/redshift/latest/dg/t_Distributing_data.html
+        :param diststyle: Redshift distribution styles. Must be in ["AUTO", "EVEN", "ALL", "KEY"] (https://docs.aws.amazon.com/redshift/latest/dg/t_Distributing_data.html)
         :param distkey: Specifies a column name or positional number for the distribution key
-        :param sortstyle: Sorting can be "COMPOUND" or "INTERLEAVED"
-               https://docs.aws.amazon.com/redshift/latest/dg/t_Sorting_data.html
+        :param sortstyle: Sorting can be "COMPOUND" or "INTERLEAVED" (https://docs.aws.amazon.com/redshift/latest/dg/t_Sorting_data.html)
         :param sortkey: List of columns to be sorted
         :param mode: append or overwrite
         :param preserve_index: Should we preserve the Dataframe index? (ONLY for Pandas Dataframe)
@@ -184,16 +183,15 @@ def _create_table(
     ):
         """
         Creates Redshift table.
+
         :param cursor: A PEP 249 compatible cursor
         :param dataframe: Pandas or Spark Dataframe
         :param dataframe_type: "pandas" or "spark"
         :param schema_name: Redshift schema
         :param table_name: Redshift table name
-        :param diststyle: Redshift distribution styles. Must be in ["AUTO", "EVEN", "ALL", "KEY"]
-               https://docs.aws.amazon.com/redshift/latest/dg/t_Distributing_data.html
+        :param diststyle: Redshift distribution styles. Must be in ["AUTO", "EVEN", "ALL", "KEY"] (https://docs.aws.amazon.com/redshift/latest/dg/t_Distributing_data.html)
         :param distkey: Specifies a column name or positional number for the distribution key
-        :param sortstyle: Sorting can be "COMPOUND" or "INTERLEAVED"
-               https://docs.aws.amazon.com/redshift/latest/dg/t_Sorting_data.html
+        :param sortstyle: Sorting can be "COMPOUND" or "INTERLEAVED" (https://docs.aws.amazon.com/redshift/latest/dg/t_Sorting_data.html)
         :param sortkey: List of columns to be sorted
         :param preserve_index: Should we preserve the Dataframe index? (ONLY for Pandas Dataframe)
         :return: None

awswrangler/s3.py

@@ -19,30 +19,36 @@ def mkdir_if_not_exists(fs, path):
             assert fs.exists(path)
 
 
-def get_fs(session_primitives):
-    aws_access_key_id, aws_secret_access_key, profile_name, config = None, None, None, None
+def get_fs(session_primitives=None):
+    aws_access_key_id, aws_secret_access_key, profile_name, config, s3_additional_kwargs = None, None, None, None, None
     if session_primitives:
-        aws_access_key_id = (session_primitives.aws_access_key_id
-                             if session_primitives.aws_access_key_id else None)
-        aws_secret_access_key = (session_primitives.aws_secret_access_key
-                                 if session_primitives.aws_secret_access_key
-                                 else None)
-        profile_name = (session_primitives.profile_name
-                        if session_primitives.profile_name else None)
-        config = {
-            "retries": {
-                "max_attempts": session_primitives.botocore_max_retries
+        if session_primitives.aws_access_key_id:
+            aws_access_key_id = session_primitives.aws_access_key_id
+        if session_primitives.aws_secret_access_key:
+            aws_secret_access_key = session_primitives.aws_secret_access_key
+        if session_primitives.profile_name:
+            profile_name = session_primitives.profile_name
+        if session_primitives.botocore_max_retries:
+            config = {
+                "retries": {
+                    "max_attempts": session_primitives.botocore_max_retries
+                }
             }
-        }
+        if session_primitives.s3_additional_kwargs:
+            s3_additional_kwargs = session_primitives.s3_additional_kwargs
     if profile_name:
-        return s3fs.S3FileSystem(profile_name=profile_name,
-                                 config_kwargs=config)
+        fs = s3fs.S3FileSystem(profile_name=profile_name,
+                               config_kwargs=config,
+                               s3_additional_kwargs=s3_additional_kwargs)
     elif aws_access_key_id and aws_secret_access_key:
-        return s3fs.S3FileSystem(key=aws_access_key_id,
-                                 secret=aws_secret_access_key,
-                                 config_kwargs=config)
+        fs = s3fs.S3FileSystem(key=aws_access_key_id,
+                               secret=aws_secret_access_key,
+                               config_kwargs=config,
+                               s3_additional_kwargs=s3_additional_kwargs)
     else:
-        return s3fs.S3FileSystem(config_kwargs=config)
+        fs = s3fs.S3FileSystem(config_kwargs=config,
+                               s3_additional_kwargs=s3_additional_kwargs)
+    return fs
 
 
 class S3:

awswrangler/session.py

@@ -37,23 +37,25 @@ def __init__(
             aws_session_token=None,
             region_name=None,
             botocore_max_retries=40,
+            s3_additional_kwargs=None,
             spark_context=None,
             spark_session=None,
             procs_cpu_bound=os.cpu_count(),
             procs_io_bound=os.cpu_count() * PROCS_IO_BOUND_FACTOR,
     ):
         """
-        Most parameters inherit from Boto3 ou Pyspark.
+        Most parameters inherit from Boto3 or Pyspark.
         https://boto3.amazonaws.com/v1/documentation/api/latest/guide/configuration.html
         https://spark.apache.org/docs/latest/api/python/index.html
 
-        :param boto3_session: Boto3.Session (Overwrite the others Boto3 parameters)
+        :param boto3_session: Boto3.Session (Overwrite others Boto3 parameters)
         :param profile_name: Boto3 profile_name
         :param aws_access_key_id: Boto3 aws_access_key_id
         :param aws_secret_access_key: Boto3 aws_secret_access_key
         :param aws_session_token: Boto3 aws_session_token
         :param region_name: Boto3 region_name
         :param botocore_max_retries: Botocore max retries
+        :param s3_additional_kwargs: Passed on to s3fs (https://s3fs.readthedocs.io/en/latest/#serverside-encryption)
         :param spark_context: Spark Context (pyspark.SparkContext)
         :param spark_session: Spark Session (pyspark.sql.SparkSession)
         :param procs_cpu_bound: number of processes that can be used in single
@@ -73,6 +75,7 @@ def __init__(
             retries={"max_attempts": self._botocore_max_retries})
         self._aws_session_token = aws_session_token
         self._region_name = boto3_session.region_name if boto3_session else region_name
+        self._s3_additional_kwargs = s3_additional_kwargs
         self._spark_context = spark_context
         self._spark_session = spark_session
         self._procs_cpu_bound = procs_cpu_bound
@@ -125,6 +128,7 @@ def _load_new_primitives(self):
             aws_session_token=self._aws_session_token,
             region_name=self._region_name,
             botocore_max_retries=self._botocore_max_retries,
+            s3_additional_kwargs=self._s3_additional_kwargs,
             botocore_config=self._botocore_config,
             procs_cpu_bound=self._procs_cpu_bound,
             procs_io_bound=self._procs_io_bound,
@@ -158,6 +162,10 @@ def botocore_max_retries(self):
     def botocore_config(self):
         return self._botocore_config
 
+    @property
+    def s3_additional_kwargs(self):
+        return self._s3_additional_kwargs
+
     @property
     def spark_context(self):
         return self._spark_context
@@ -235,6 +243,7 @@ def __init__(
             aws_session_token=None,
             region_name=None,
             botocore_max_retries=None,
+            s3_additional_kwargs=None,
             botocore_config=None,
             procs_cpu_bound=None,
             procs_io_bound=None,
@@ -249,6 +258,7 @@ def __init__(
         :param aws_session_token: Boto3 aws_session_token
         :param region_name: Boto3 region_name
         :param botocore_max_retries: Botocore max retries
+        :param s3_additional_kwargs: Passed on to s3fs (https://s3fs.readthedocs.io/en/latest/#serverside-encryption)
         :param botocore_config: Botocore configurations
         :param procs_cpu_bound: number of processes that can be used in single
         node applications for CPU bound case (Default: os.cpu_count())
@@ -261,6 +271,7 @@ def __init__(
         self._aws_session_token = aws_session_token
         self._region_name = region_name
         self._botocore_max_retries = botocore_max_retries
+        self._s3_additional_kwargs = s3_additional_kwargs
         self._botocore_config = botocore_config
         self._procs_cpu_bound = procs_cpu_bound
         self._procs_io_bound = procs_io_bound
@@ -293,6 +304,10 @@ def botocore_max_retries(self):
     def botocore_config(self):
         return self._botocore_config
 
+    @property
+    def s3_additional_kwargs(self):
+        return self._s3_additional_kwargs
+
     @property
     def procs_cpu_bound(self):
         return self._procs_cpu_bound
@@ -314,6 +329,7 @@ def session(self):
             aws_session_token=self._aws_session_token,
             region_name=self._region_name,
             botocore_max_retries=self._botocore_max_retries,
+            s3_additional_kwargs=self._s3_additional_kwargs,
             procs_cpu_bound=self._procs_cpu_bound,
             procs_io_bound=self._procs_io_bound,
         )