From 029737bd5859ff4a91c899bb3c02484b2d3f0278 Mon Sep 17 00:00:00 2001
From: Darren Weber <dweber.consulting@gmail.com>
Date: Tue, 27 Jun 2023 19:27:48 -0700
Subject: [PATCH] Add S3 IO for file copy

---
 aio_aws/s3_io.py    | 37 +++++++++++++++++++++++++++++++++++++
 tests/test_s3_io.py | 23 +++++++++++++++++++++++
 2 files changed, 60 insertions(+)

diff --git a/aio_aws/s3_io.py b/aio_aws/s3_io.py
index fdf76d7..b88d8db 100644
--- a/aio_aws/s3_io.py
+++ b/aio_aws/s3_io.py
@@ -182,6 +182,43 @@ def s3_file_wait(
         raise err
 
 
+def s3_file_copy(
+    src_s3_uri: str, dst_s3_uri: str, *args, s3_client: BaseClient = None, **kwargs
+) -> bool:
+    """
+    Copy s3 URI for source to destination.
+
+    The copy uses a recommended ACL='bucket-owner-full-control', but
+    otherwise uses all default options for S3Client.copy_object().
+
+    :param src_s3_uri: a fully qualified S3 URI for the source
+    :param dst_s3_uri: a fully qualified S3 URI for the destination
+    :param s3_client: an optional botocore.client.BaseClient for s3
+    :return: boolean (True on success, False on failure)
+    """
+    # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/s3/client/copy_object.html
+    # CopySource={‘Bucket’: ‘bucket’, ‘Key’: ‘key’, ‘VersionId’: ‘id’}
+    # Note that the VersionId key is optional and may be omitted.
+    if s3_client is None:
+        s3_client = s3_io_client(*args, **kwargs)
+    try:
+        src_s3_uri = S3URI(src_s3_uri)
+        dst_s3_uri = S3URI(dst_s3_uri)
+        LOGGER.info(f"Copy: {src_s3_uri} to {dst_s3_uri}")
+        response = s3_client.copy_object(
+            ACL="bucket-owner-full-control",
+            Bucket=dst_s3_uri.bucket,
+            Key=dst_s3_uri.key,
+            CopySource={"Bucket": src_s3_uri.bucket, "Key": src_s3_uri.key},
+        )
+        if response:
+            return True
+    except botocore.exceptions.ClientError as err:
+        LOGGER.error(f"Failed S3 Copy: {src_s3_uri} to {dst_s3_uri}")
+        LOGGER.error(err)
+        return False
+
+
 def get_s3_content(s3_uri: str, *args, s3_client: BaseClient = None, **kwargs):
     """
     Read s3 URI
diff --git a/tests/test_s3_io.py b/tests/test_s3_io.py
index a7452b6..f63bd69 100644
--- a/tests/test_s3_io.py
+++ b/tests/test_s3_io.py
@@ -32,6 +32,7 @@
 from aio_aws.s3_io import get_s3_content
 from aio_aws.s3_io import json_s3_dump
 from aio_aws.s3_io import json_s3_load
+from aio_aws.s3_io import s3_file_copy
 from aio_aws.s3_io import s3_file_info
 from aio_aws.s3_io import s3_file_wait
 from aio_aws.s3_io import s3_files_info
@@ -41,6 +42,28 @@
 from aio_aws.s3_uri import S3Info
 
 
+def test_s3_file_copy(aws_s3_client, s3_uri_object, mocker):
+    assert_bucket_200(s3_uri_object.bucket, aws_s3_client)
+    assert_object_200(s3_uri_object.bucket, s3_uri_object.key, aws_s3_client)
+    src_s3_obj = s3_uri_object
+    dst_bucket = src_s3_obj.bucket
+    dst_file_path = "s3_file_test_copy"
+    dst_file_name = "s3_file_test_copy.txt"
+    dst_s3_uri = f"s3://{dst_bucket}/{dst_file_path}/{dst_file_name}"
+    spy_client = mocker.spy(boto3, "client")
+    spy_resource = mocker.spy(boto3, "resource")
+    success = s3_file_copy(src_s3_uri=s3_uri_object.s3_uri, dst_s3_uri=dst_s3_uri)
+    assert success
+    # the s3 client is used once to get the s3 object data
+    assert spy_client.call_count == 1
+    assert spy_resource.call_count == 0
+    # check the destination can be read OK
+    s3_info = s3_file_info(dst_s3_uri)
+    assert isinstance(s3_info, S3Info)
+    assert s3_info.s3_uri.key_file == dst_file_name
+    assert s3_info.s3_size > 0
+
+
 def test_s3_file_info(aws_s3_client, s3_uri_object, s3_object_text, mocker):
     assert_bucket_200(s3_uri_object.bucket, aws_s3_client)
     assert_object_200(s3_uri_object.bucket, s3_uri_object.key, aws_s3_client)