Add support for specifying extra s3 args for copy_objects()

Author: donghan

awswrangler/s3/_copy.py

@@ -13,7 +13,7 @@
 _logger: logging.Logger = logging.getLogger(__name__)
 
 
-def _copy_objects(batch: List[Tuple[str, str]], use_threads: bool, boto3_session: boto3.Session) -> None:
+def _copy_objects(batch: List[Tuple[str, str]], s3_additional_kwargs: Optional[Dict[str, str]], use_threads: bool, boto3_session: boto3.Session) -> None:
     _logger.debug("len(batch): %s", len(batch))
     client_s3: boto3.client = _utils.client(service_name="s3", session=boto3_session)
     resource_s3: boto3.resource = _utils.resource(service_name="s3", session=boto3_session)
@@ -26,6 +26,7 @@ def _copy_objects(batch: List[Tuple[str, str]], use_threads: bool, boto3_session
             Bucket=target_bucket,
             Key=target_key,
             SourceClient=client_s3,
+            ExtraArgs=s3_additional_kwargs,
             Config=TransferConfig(num_download_attempts=15, use_threads=use_threads),
         )
 
@@ -36,6 +37,7 @@ def merge_datasets(
     mode: str = "append",
     use_threads: bool = True,
     boto3_session: Optional[boto3.Session] = None,
+    s3_additional_kwargs: Optional[Dict[str, str]] = None,
 ) -> List[str]:
     """Merge a source dataset into a target dataset.
 
@@ -67,6 +69,10 @@ def merge_datasets(
         If enabled os.cpu_count() will be used as the max number of threads.
     boto3_session : boto3.Session(), optional
         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".
+        e.g. s3_additional_kwargs={'ServerSideEncryption': 'aws:kms', 'SSEKMSKeyId': 'YOUR_KMY_KEY_ARN'}
 
     Returns
     -------
@@ -108,7 +114,7 @@ def merge_datasets(
         raise exceptions.InvalidArgumentValue(f"{mode} is a invalid mode option.")
 
     new_objects: List[str] = copy_objects(
-        paths=paths, source_path=source_path, target_path=target_path, use_threads=use_threads, boto3_session=session
+        paths=paths, source_path=source_path, target_path=target_path, use_threads=use_threads, boto3_session=session, s3_additional_kwargs=s3_additional_kwargs
     )
     _logger.debug("len(new_objects): %s", len(new_objects))
     return new_objects
@@ -119,6 +125,7 @@ def copy_objects(
     source_path: str,
     target_path: str,
     replace_filenames: Optional[Dict[str, str]] = None,
+    s3_additional_kwargs: Optional[Dict[str, str]] = None,
     use_threads: bool = True,
     boto3_session: Optional[boto3.Session] = None,
 ) -> List[str]:
@@ -144,6 +151,10 @@ def copy_objects(
         If enabled os.cpu_count() will be used as the max number of threads.
     boto3_session : boto3.Session(), optional
         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".
+        e.g. s3_additional_kwargs={'ServerSideEncryption': 'aws:kms', 'SSEKMSKeyId': 'YOUR_KMY_KEY_ARN'}
 
     Returns
     -------
@@ -154,9 +165,13 @@ def copy_objects(
     --------
     >>> import awswrangler as wr
     >>> wr.s3.copy_objects(
-    ...     paths=["s3://bucket0/dir0/key0", "s3://bucket0/dir0/key1"])
+    ...     paths=["s3://bucket0/dir0/key0", "s3://bucket0/dir0/key1"],
     ...     source_path="s3://bucket0/dir0/",
     ...     target_path="s3://bucket1/dir1/",
+    ...     s3_additional_kwargs={
+    ...         'ServerSideEncryption': 'aws:kms',
+    ...         'SSEKMSKeyId': 'YOUR_KMY_KEY_ARN'
+    ...     }
     ... )
     ["s3://bucket1/dir1/key0", "s3://bucket1/dir1/key1"]
 
@@ -184,5 +199,5 @@ def copy_objects(
         new_objects.append(path_final)
         batch.append((path, path_final))
     _logger.debug("len(new_objects): %s", len(new_objects))
-    _copy_objects(batch=batch, use_threads=use_threads, boto3_session=session)
+    _copy_objects(batch=batch, use_threads=use_threads, boto3_session=session, s3_additional_kwargs=s3_additional_kwargs)
     return new_objects