Add more Athena defaults to Session()

Author: igorborgest

awswrangler/athena.py

@@ -1,4 +1,4 @@
-from typing import Dict, List, Tuple, Optional, Any, Iterator
+from typing import Dict, List, Tuple, Optional, Any, Iterator, Union
 from time import sleep
 import logging
 import re
@@ -41,25 +41,53 @@ def create_athena_bucket(self):
         s3_resource.Bucket(s3_output)
         return s3_output
 
-    def run_query(self, query, database, s3_output=None, workgroup=None):
+    def run_query(self, query: str, database: Optional[str] = None, s3_output: Optional[str] = None, workgroup: Optional[str] = None, encryption: Optional[str] = None, kms_key: Optional[str] = None) -> str:
         """
         Run a SQL Query against AWS Athena
+        P.S All default values will be inherited from the Session()
 
         :param query: SQL query
         :param database: AWS Glue/Athena database name
         :param s3_output: AWS S3 path
         :param workgroup: Athena workgroup (By default uses de Session() workgroup)
+        :param encryption: None|'SSE_S3'|'SSE_KMS'|'CSE_KMS'
+        :param kms_key: For SSE-KMS and CSE-KMS , this is the KMS key ARN or ID.
         :return: Query execution ID
         """
+        args: Dict[str, Union[str, Dict[str, Union[str, Dict[str, str]]]]] = {"QueryString": query}
+
+        # s3_output
         if s3_output is None:
-            s3_output = self.create_athena_bucket()
-        if workgroup is None:
-            workgroup = self._session.athena_workgroup
-        logger.debug(f"Workgroup: {workgroup}")
-        response = self._client_athena.start_query_execution(QueryString=query,
-                                                             QueryExecutionContext={"Database": database},
-                                                             ResultConfiguration={"OutputLocation": s3_output},
-                                                             WorkGroup=workgroup)
+            if self._session.athena_s3_output is not None:
+                s3_output = self._session.athena_s3_output
+            else:
+                s3_output = self.create_athena_bucket()
+        args["ResultConfiguration"] = {"OutputLocation": s3_output}
+
+        # encryption
+        if encryption is not None:
+            args["ResultConfiguration"]["EncryptionConfiguration"] = {"EncryptionOption": encryption}
+            if kms_key is not None:
+                args["ResultConfiguration"]["EncryptionConfiguration"]["KmsKey"] = kms_key
+        elif self._session.athena_encryption is not None:
+            args["ResultConfiguration"]["EncryptionConfiguration"] = {"EncryptionOption": self._session.athena_encryption}
+            if self._session.athena_kms_key is not None:
+                args["ResultConfiguration"]["EncryptionConfiguration"]["KmsKey"] = self._session.athena_kms_key
+
+        # database
+        if database is not None:
+            args["QueryExecutionContext"] = {"Database": database}
+        elif self._session.athena_database is not None:
+            args["QueryExecutionContext"] = {"Database": self._session.athena_database}
+
+        # workgroup
+        if workgroup is not None:
+            args["WorkGroup"] = workgroup
+        elif self._session.athena_workgroup is not None:
+            args["WorkGroup"] = self._session.athena_workgroup
+
+        logger.debug(f"args: {args}")
+        response = self._client_athena.start_query_execution(**args)
         return response["QueryExecutionId"]
 
     def wait_query(self, query_execution_id):
@@ -84,7 +112,7 @@ def wait_query(self, query_execution_id):
             raise QueryCancelled(response["QueryExecution"]["Status"].get("StateChangeReason"))
         return response
 
-    def repair_table(self, database, table, s3_output=None, workgroup=None):
+    def repair_table(self, table: str, database: Optional[str] = None, s3_output: Optional[str] = None, workgroup: Optional[str] = None, encryption: Optional[str] = None, kms_key: Optional[str] = None):
         """
         Hive's metastore consistency check
         "MSCK REPAIR TABLE table;"
@@ -93,15 +121,18 @@ def repair_table(self, database, table, s3_output=None, workgroup=None):
         It is possible it will take some time to add all partitions.
         If this operation times out, it will be in an incomplete state
         where only a few partitions are added to the catalog.
+        P.S All default values will be inherited from the Session()
 
         :param database: Glue database name
         :param table: Glue table name
         :param s3_output: AWS S3 path
         :param workgroup: Athena workgroup (By default uses de Session() workgroup)
+        :param encryption: None|'SSE_S3'|'SSE_KMS'|'CSE_KMS'
+        :param kms_key: For SSE-KMS and CSE-KMS , this is the KMS key ARN or ID.
         :return: Query execution ID
         """
         query = f"MSCK REPAIR TABLE {table};"
-        query_id = self.run_query(query=query, database=database, s3_output=s3_output, workgroup=workgroup)
+        query_id = self.run_query(query=query, database=database, s3_output=s3_output, workgroup=workgroup, encryption=encryption, kms_key=kms_key)
         self.wait_query(query_execution_id=query_id)
         return query_id
 
@@ -142,18 +173,20 @@ def get_results(self, query_execution_id: str) -> Iterator[Dict[str, Any]]:
                 yield row
             next_token = res.get("NextToken")
 
-    def query(self, query: str, database: str, s3_output: str = None,
-              workgroup: str = None) -> Iterator[Dict[str, Any]]:
+    def query(self, query: str, database: Optional[str] = None, s3_output: Optional[str] = None, workgroup: Optional[str] = None, encryption: Optional[str] = None, kms_key: Optional[str] = None) -> Iterator[Dict[str, Any]]:
         """
         Run a SQL Query against AWS Athena and return the result as a Iterator of lists
+        P.S All default values will be inherited from the Session()
 
         :param query: SQL query
         :param database: Glue database name
         :param s3_output: AWS S3 path
         :param workgroup: Athena workgroup (By default uses de Session() workgroup)
+        :param encryption: None|'SSE_S3'|'SSE_KMS'|'CSE_KMS'
+        :param kms_key: For SSE-KMS and CSE-KMS , this is the KMS key ARN or ID.
         :return: Query execution ID
         """
-        query_id: str = self.run_query(query=query, database=database, s3_output=s3_output, workgroup=workgroup)
+        query_id: str = self.run_query(query=query, database=database, s3_output=s3_output, workgroup=workgroup, encryption=encryption, kms_key=kms_key)
         self.wait_query(query_execution_id=query_id)
         return self.get_results(query_execution_id=query_id)
 

awswrangler/pandas.py

@@ -1,4 +1,4 @@
-from typing import Dict, List, Tuple, Optional, Any
+from typing import Dict, List, Tuple, Optional, Any, Union
 from io import BytesIO, StringIO
 import multiprocessing as mp
 import logging
@@ -32,7 +32,6 @@ def _get_bounders(dataframe, num_partitions):
 
 
 class Pandas:
-
     VALID_CSV_SERDES = ["OpenCSVSerDe", "LazySimpleSerDe"]
     VALID_CSV_COMPRESSIONS = [None]
     VALID_PARQUET_COMPRESSIONS = [None, "snappy", "gzip"]
@@ -61,7 +60,7 @@ def read_csv(
             quotechar='"',
             quoting=csv.QUOTE_MINIMAL,
             escapechar=None,
-            parse_dates=False,
+            parse_dates: Union[bool, Dict, List] = False,
             infer_datetime_format=False,
             encoding="utf-8",
             converters=None,
@@ -153,7 +152,7 @@ def _read_csv_iterator(
             quotechar='"',
             quoting=csv.QUOTE_MINIMAL,
             escapechar=None,
-            parse_dates=False,
+            parse_dates: Union[bool, Dict, List] = False,
             infer_datetime_format=False,
             encoding="utf-8",
             converters=None,
@@ -365,7 +364,7 @@ def _read_csv_once(
             quotechar='"',
             quoting=0,
             escapechar=None,
-            parse_dates=False,
+            parse_dates: Union[bool, Dict, List] = False,
             infer_datetime_format=False,
             encoding=None,
             converters=None,
@@ -446,20 +445,25 @@ def _get_query_dtype(self, query_execution_id: str) -> Tuple[Dict[str, str], Lis
         logger.debug(f"parse_dates: {parse_dates}")
         return dtype, parse_timestamps, parse_dates, converters
 
-    def read_sql_athena(self, sql, database, s3_output=None, max_result_size=None):
+    def read_sql_athena(self, sql, database=None, s3_output=None, max_result_size=None, workgroup=None,
+                        encryption=None, kms_key=None):
         """
         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.
+        P.S.S. All default values will be inherited from the Session()
 
         :param sql: SQL Query
         :param database: Glue/Athena Database
         :param s3_output: AWS S3 path
         :param max_result_size: Max number of bytes on each request to S3
+        :param workgroup: The name of the workgroup in which the query is being started. (By default uses de Session() workgroup)
+        :param encryption: None|'SSE_S3'|'SSE_KMS'|'CSE_KMS'
+        :param kms_key: For SSE-KMS and CSE-KMS , this is the KMS key ARN or ID.
         :return: Pandas Dataframe or Iterator of Pandas Dataframes if max_result_size != None
         """
         if not s3_output:
             s3_output = self._session.athena.create_athena_bucket()
-        query_execution_id = self._session.athena.run_query(query=sql, database=database, s3_output=s3_output)
+        query_execution_id = self._session.athena.run_query(query=sql, database=database, s3_output=s3_output, workgroup=workgroup, encryption=encryption, kms_key=kms_key)
         query_response = self._session.athena.wait_query(query_execution_id=query_execution_id)
         if query_response["QueryExecution"]["Status"]["State"] in ["FAILED", "CANCELLED"]:
             reason = query_response["QueryExecution"]["Status"]["StateChangeReason"]
@@ -497,7 +501,7 @@ def to_csv(
             path,
             sep=",",
             serde="OpenCSVSerDe",
-            database=None,
+            database: Optional[str] = None,
             table=None,
             partition_cols=None,
             preserve_index=True,
@@ -544,7 +548,7 @@ def to_csv(
     def to_parquet(self,
                    dataframe,
                    path,
-                   database=None,
+                   database: Optional[str] = None,
                    table=None,
                    partition_cols=None,
                    preserve_index=True,
@@ -766,7 +770,7 @@ def _data_to_s3_dataset_writer(dataframe,
             for keys, subgroup in dataframe.groupby(partition_cols):
                 subgroup = subgroup.drop(partition_cols, axis="columns")
                 if not isinstance(keys, tuple):
-                    keys = (keys, )
+                    keys = (keys,)
                 subdir = "/".join([f"{name}={val}" for name, val in zip(partition_cols, keys)])
                 prefix = "/".join([path, subdir])
                 object_path = Pandas._data_to_s3_object_writer(dataframe=subgroup,