Add support for requester pays on s3. #430

Author: igorborgest

awswrangler/athena/_read.py

@@ -72,11 +72,17 @@ def _fix_csv_types(df: pd.DataFrame, parse_dates: List[str], binaries: List[str]
 
 
 def _delete_after_iterate(
-    dfs: Iterator[pd.DataFrame], paths: List[str], use_threads: bool, boto3_session: boto3.Session
+    dfs: Iterator[pd.DataFrame],
+    paths: List[str],
+    use_threads: bool,
+    boto3_session: boto3.Session,
+    s3_additional_kwargs: Optional[Dict[str, str]],
 ) -> Iterator[pd.DataFrame]:
     for df in dfs:
         yield df
-    s3.delete_objects(path=paths, use_threads=use_threads, boto3_session=boto3_session)
+    s3.delete_objects(
+        path=paths, use_threads=use_threads, boto3_session=boto3_session, s3_additional_kwargs=s3_additional_kwargs
+    )
 
 
 def _prepare_query_string_for_comparison(query_string: str) -> str:
@@ -213,6 +219,7 @@ def _fetch_parquet_result(
     chunksize: Optional[int],
     use_threads: bool,
     boto3_session: boto3.Session,
+    s3_additional_kwargs: Optional[Dict[str, Any]],
 ) -> Union[pd.DataFrame, Iterator[pd.DataFrame]]:
     ret: Union[pd.DataFrame, Iterator[pd.DataFrame]]
     chunked: Union[bool, int] = False if chunksize is None else chunksize
@@ -242,10 +249,21 @@ def _fetch_parquet_result(
     _logger.debug("type(ret): %s", type(ret))
     if chunked is False:
         if keep_files is False:
-            s3.delete_objects(path=paths_delete, use_threads=use_threads, boto3_session=boto3_session)
+            s3.delete_objects(
+                path=paths_delete,
+                use_threads=use_threads,
+                boto3_session=boto3_session,
+                s3_additional_kwargs=s3_additional_kwargs,
+            )
         return ret
     if keep_files is False:
-        return _delete_after_iterate(dfs=ret, paths=paths_delete, use_threads=use_threads, boto3_session=boto3_session)
+        return _delete_after_iterate(
+            dfs=ret,
+            paths=paths_delete,
+            use_threads=use_threads,
+            boto3_session=boto3_session,
+            s3_additional_kwargs=s3_additional_kwargs,
+        )
     return ret
 
 
@@ -255,6 +273,7 @@ def _fetch_csv_result(
     chunksize: Optional[int],
     use_threads: bool,
     boto3_session: boto3.Session,
+    s3_additional_kwargs: Optional[Dict[str, Any]],
 ) -> Union[pd.DataFrame, Iterator[pd.DataFrame]]:
     _chunksize: Optional[int] = chunksize if isinstance(chunksize, int) else None
     _logger.debug("_chunksize: %s", _chunksize)
@@ -282,13 +301,22 @@ def _fetch_csv_result(
         df = _fix_csv_types(df=ret, parse_dates=query_metadata.parse_dates, binaries=query_metadata.binaries)
         df = _apply_query_metadata(df=df, query_metadata=query_metadata)
         if keep_files is False:
-            s3.delete_objects(path=[path, f"{path}.metadata"], use_threads=use_threads, boto3_session=boto3_session)
+            s3.delete_objects(
+                path=[path, f"{path}.metadata"],
+                use_threads=use_threads,
+                boto3_session=boto3_session,
+                s3_additional_kwargs=s3_additional_kwargs,
+            )
         return df
     dfs = _fix_csv_types_generator(dfs=ret, parse_dates=query_metadata.parse_dates, binaries=query_metadata.binaries)
     dfs = _add_query_metadata_generator(dfs=dfs, query_metadata=query_metadata)
     if keep_files is False:
         return _delete_after_iterate(
-            dfs=dfs, paths=[path, f"{path}.metadata"], use_threads=use_threads, boto3_session=boto3_session
+            dfs=dfs,
+            paths=[path, f"{path}.metadata"],
+            use_threads=use_threads,
+            boto3_session=boto3_session,
+            s3_additional_kwargs=s3_additional_kwargs,
         )
     return dfs
 
@@ -299,6 +327,7 @@ def _resolve_query_with_cache(
     chunksize: Optional[Union[int, bool]],
     use_threads: bool,
     session: Optional[boto3.Session],
+    s3_additional_kwargs: Optional[Dict[str, Any]],
 ) -> Union[pd.DataFrame, Iterator[pd.DataFrame]]:
     """Fetch cached data and return it as a pandas DataFrame (or list of DataFrames)."""
     _logger.debug("cache_info:\n%s", cache_info)
@@ -319,6 +348,7 @@ def _resolve_query_with_cache(
             chunksize=chunksize,
             use_threads=use_threads,
             boto3_session=session,
+            s3_additional_kwargs=s3_additional_kwargs,
         )
     if cache_info.file_format == "csv":
         return _fetch_csv_result(
@@ -327,6 +357,7 @@ def _resolve_query_with_cache(
             chunksize=chunksize,
             use_threads=use_threads,
             boto3_session=session,
+            s3_additional_kwargs=s3_additional_kwargs,
         )
     raise exceptions.InvalidArgumentValue(f"Invalid data type: {cache_info.file_format}.")
 
@@ -345,6 +376,7 @@ def _resolve_query_without_cache_ctas(
     wg_config: _WorkGroupConfig,
     name: Optional[str],
     use_threads: bool,
+    s3_additional_kwargs: Optional[Dict[str, Any]],
     boto3_session: boto3.Session,
 ) -> Union[pd.DataFrame, Iterator[pd.DataFrame]]:
     path: str = f"{s3_output}/{name}"
@@ -412,6 +444,7 @@ def _resolve_query_without_cache_ctas(
         categories=categories,
         chunksize=chunksize,
         use_threads=use_threads,
+        s3_additional_kwargs=s3_additional_kwargs,
         boto3_session=boto3_session,
     )
 
@@ -429,6 +462,7 @@ def _resolve_query_without_cache_regular(
     kms_key: Optional[str],
     wg_config: _WorkGroupConfig,
     use_threads: bool,
+    s3_additional_kwargs: Optional[Dict[str, Any]],
     boto3_session: boto3.Session,
 ) -> Union[pd.DataFrame, Iterator[pd.DataFrame]]:
     _logger.debug("sql: %s", sql)
@@ -456,6 +490,7 @@ def _resolve_query_without_cache_regular(
         chunksize=chunksize,
         use_threads=use_threads,
         boto3_session=boto3_session,
+        s3_additional_kwargs=s3_additional_kwargs,
     )
 
 
@@ -474,6 +509,7 @@ def _resolve_query_without_cache(
     keep_files: bool,
     ctas_temp_table_name: Optional[str],
     use_threads: bool,
+    s3_additional_kwargs: Optional[Dict[str, Any]],
     boto3_session: boto3.Session,
 ) -> Union[pd.DataFrame, Iterator[pd.DataFrame]]:
     """
@@ -504,6 +540,7 @@ def _resolve_query_without_cache(
                 wg_config=wg_config,
                 name=name,
                 use_threads=use_threads,
+                s3_additional_kwargs=s3_additional_kwargs,
                 boto3_session=boto3_session,
             )
         finally:
@@ -521,6 +558,7 @@ def _resolve_query_without_cache(
         kms_key=kms_key,
         wg_config=wg_config,
         use_threads=use_threads,
+        s3_additional_kwargs=s3_additional_kwargs,
         boto3_session=boto3_session,
     )
 
@@ -546,6 +584,7 @@ def read_sql_query(
     max_local_cache_entries: int = 100,
     data_source: Optional[str] = None,
     params: Optional[Dict[str, Any]] = None,
+    s3_additional_kwargs: Optional[Dict[str, Any]] = None,
 ) -> Union[pd.DataFrame, Iterator[pd.DataFrame]]:
     """Execute any SQL query on AWS Athena and return the results as a Pandas DataFrame.
 
@@ -705,6 +744,9 @@ def read_sql_query(
         Dict of parameters that will be used for constructing the SQL query. Only named parameters are supported.
         The dict needs to contain the information in the form {'name': 'value'} and the SQL query needs to contain
         `:name;`.
+    s3_additional_kwargs : Optional[Dict[str, Any]]
+        Forward to botocore requests. Valid parameters: "RequestPayer".
+        e.g. s3_additional_kwargs={'RequestPayer': 'requester'}
 
     Returns
     -------
@@ -761,6 +803,7 @@ def read_sql_query(
                 chunksize=chunksize,
                 use_threads=use_threads,
                 session=session,
+                s3_additional_kwargs=s3_additional_kwargs,
             )
         except Exception as e:  # pylint: disable=broad-except
             _logger.error(e)  # if there is anything wrong with the cache, just fallback to the usual path
@@ -779,6 +822,7 @@ def read_sql_query(
         keep_files=keep_files,
         ctas_temp_table_name=ctas_temp_table_name,
         use_threads=use_threads,
+        s3_additional_kwargs=s3_additional_kwargs,
         boto3_session=session,
     )
 
@@ -803,6 +847,7 @@ def read_sql_table(
     max_remote_cache_entries: int = 50,
     max_local_cache_entries: int = 100,
     data_source: Optional[str] = None,
+    s3_additional_kwargs: Optional[Dict[str, Any]] = None,
 ) -> Union[pd.DataFrame, Iterator[pd.DataFrame]]:
     """Extract the full table AWS Athena and return the results as a Pandas DataFrame.
 
@@ -953,6 +998,9 @@ def read_sql_table(
         Only takes effect if max_cache_seconds > 0 and default value is 100.
     data_source : str, optional
         Data Source / Catalog name. If None, 'AwsDataCatalog' will be used by default.
+    s3_additional_kwargs : Optional[Dict[str, Any]]
+        Forward to botocore requests. Valid parameters: "RequestPayer".
+        e.g. s3_additional_kwargs={'RequestPayer': 'requester'}
 
     Returns
     -------
@@ -986,6 +1034,7 @@ def read_sql_table(
         max_cache_query_inspections=max_cache_query_inspections,
         max_remote_cache_entries=max_remote_cache_entries,
         max_local_cache_entries=max_local_cache_entries,
+        s3_additional_kwargs=s3_additional_kwargs,
     )
 
 

awswrangler/athena/_utils.py

@@ -191,6 +191,7 @@ def _fetch_txt_result(
     query_metadata: _QueryMetadata,
     keep_files: bool,
     boto3_session: boto3.Session,
+    s3_additional_kwargs: Optional[Dict[str, str]],
 ) -> pd.DataFrame:
     if query_metadata.output_location is None or query_metadata.output_location.endswith(".txt") is False:
         return pd.DataFrame()
@@ -211,7 +212,12 @@ def _fetch_txt_result(
         sep="\t",
     )
     if keep_files is False:
-        s3.delete_objects(path=[path, f"{path}.metadata"], use_threads=False, boto3_session=boto3_session)
+        s3.delete_objects(
+            path=[path, f"{path}.metadata"],
+            use_threads=False,
+            boto3_session=boto3_session,
+            s3_additional_kwargs=s3_additional_kwargs,
+        )
     return df
 
 
@@ -532,6 +538,7 @@ def describe_table(
     workgroup: Optional[str] = None,
     encryption: Optional[str] = None,
     kms_key: Optional[str] = None,
+    s3_additional_kwargs: Optional[Dict[str, Any]] = None,
     boto3_session: Optional[boto3.Session] = None,
 ) -> pd.DataFrame:
     """Show the list of columns, including partition columns: 'DESCRIBE table;'.
@@ -558,6 +565,9 @@ def describe_table(
         None, 'SSE_S3', 'SSE_KMS', 'CSE_KMS'.
     kms_key : str, optional
         For SSE-KMS and CSE-KMS , this is the KMS key ARN or ID.
+    s3_additional_kwargs : Optional[Dict[str, Any]]
+        Forward to botocore requests. Valid parameters: "RequestPayer".
+        e.g. s3_additional_kwargs={'RequestPayer': 'requester'}
     boto3_session : boto3.Session(), optional
         Boto3 Session. The default boto3 session will be used if boto3_session receive None.
 
@@ -587,9 +597,7 @@ def describe_table(
     )
     query_metadata: _QueryMetadata = _get_query_metadata(query_execution_id=query_id, boto3_session=session)
     raw_result = _fetch_txt_result(
-        query_metadata=query_metadata,
-        keep_files=True,
-        boto3_session=session,
+        query_metadata=query_metadata, keep_files=True, boto3_session=session, s3_additional_kwargs=s3_additional_kwargs
     )
     return _parse_describe_table(raw_result)
 
@@ -602,6 +610,7 @@ def show_create_table(
     workgroup: Optional[str] = None,
     encryption: Optional[str] = None,
     kms_key: Optional[str] = None,
+    s3_additional_kwargs: Optional[Dict[str, Any]] = None,
     boto3_session: Optional[boto3.Session] = None,
 ) -> str:
     """Generate the query that created it: 'SHOW CREATE TABLE table;'.
@@ -627,6 +636,9 @@ def show_create_table(
         None, 'SSE_S3', 'SSE_KMS', 'CSE_KMS'.
     kms_key : str, optional
         For SSE-KMS and CSE-KMS , this is the KMS key ARN or ID.
+    s3_additional_kwargs : Optional[Dict[str, Any]]
+        Forward to botocore requests. Valid parameters: "RequestPayer".
+        e.g. s3_additional_kwargs={'RequestPayer': 'requester'}
     boto3_session : boto3.Session(), optional
         Boto3 Session. The default boto3 session will be used if boto3_session receive None.
 
@@ -656,9 +668,7 @@ def show_create_table(
     )
     query_metadata: _QueryMetadata = _get_query_metadata(query_execution_id=query_id, boto3_session=session)
     raw_result = _fetch_txt_result(
-        query_metadata=query_metadata,
-        keep_files=True,
-        boto3_session=session,
+        query_metadata=query_metadata, keep_files=True, boto3_session=session, s3_additional_kwargs=s3_additional_kwargs
     )
     return cast(str, raw_result.createtab_stmt.str.strip().str.cat(sep=" "))
 

awswrangler/redshift.py

@@ -304,7 +304,9 @@ def _read_parquet_iterator(
     )
     yield from dfs
     if keep_files is False:
-        s3.delete_objects(path=path, use_threads=use_threads, boto3_session=boto3_session)
+        s3.delete_objects(
+            path=path, use_threads=use_threads, boto3_session=boto3_session, s3_additional_kwargs=s3_additional_kwargs
+        )
 
 
 def connect(
@@ -1012,7 +1014,9 @@ def unload(
             s3_additional_kwargs=s3_additional_kwargs,
         )
         if keep_files is False:
-            s3.delete_objects(path=path, use_threads=use_threads, boto3_session=session)
+            s3.delete_objects(
+                path=path, use_threads=use_threads, boto3_session=session, s3_additional_kwargs=s3_additional_kwargs
+            )
         return df
     return _read_parquet_iterator(
         path=path,
@@ -1129,7 +1133,7 @@ def copy_from_files(  # pylint: disable=too-many-locals,too-many-arguments
         Boto3 Session. The default boto3 session will be used if boto3_session receive None.
     s3_additional_kwargs:
         Forward to botocore requests. Valid parameters: "ACL", "Metadata", "ServerSideEncryption", "StorageClass",
-        "SSECustomerAlgorithm", "SSECustomerKey", "SSEKMSKeyId", "SSEKMSEncryptionContext", "Tagging".
+        "SSECustomerAlgorithm", "SSECustomerKey", "SSEKMSKeyId", "SSEKMSEncryptionContext", "Tagging", "RequestPayer".
         e.g. s3_additional_kwargs={'ServerSideEncryption': 'aws:kms', 'SSEKMSKeyId': 'YOUR_KMS_KEY_ARN'}
 
     Returns
@@ -1307,7 +1311,7 @@ def copy(  # pylint: disable=too-many-arguments
         Boto3 Session. The default boto3 session will be used if boto3_session receive None.
     s3_additional_kwargs:
         Forward to botocore requests. Valid parameters: "ACL", "Metadata", "ServerSideEncryption", "StorageClass",
-        "SSECustomerAlgorithm", "SSECustomerKey", "SSEKMSKeyId", "SSEKMSEncryptionContext", "Tagging".
+        "SSECustomerAlgorithm", "SSECustomerKey", "SSEKMSKeyId", "SSEKMSEncryptionContext", "Tagging", "RequestPayer".
         e.g. s3_additional_kwargs={'ServerSideEncryption': 'aws:kms', 'SSEKMSKeyId': 'YOUR_KMS_KEY_ARN'}
     max_rows_by_file : int
         Max number of rows in each file.
@@ -1338,7 +1342,7 @@ def copy(  # pylint: disable=too-many-arguments
     path = path[:-1] if path.endswith("*") else path
     path = path if path.endswith("/") else f"{path}/"
     session: boto3.Session = _utils.ensure_session(session=boto3_session)
-    if s3.list_objects(path=path, boto3_session=session):
+    if s3.list_objects(path=path, boto3_session=session, s3_additional_kwargs=s3_additional_kwargs):
         raise exceptions.InvalidArgument(
             f"The received S3 path ({path}) is not empty. "
             "Please, provide a different path or use wr.s3.delete_objects() to clean up the current one."
@@ -1380,4 +1384,6 @@ def copy(  # pylint: disable=too-many-arguments
         )
     finally:
         if keep_files is False:
-            s3.delete_objects(path=path, use_threads=use_threads, boto3_session=session)
+            s3.delete_objects(
+                path=path, use_threads=use_threads, boto3_session=session, s3_additional_kwargs=s3_additional_kwargs
+            )