Add S3CopyPrefixOperator to copy all objects under a prefix (#68946)

S3CopyObjectOperator handles a single object at a time. Users who need
to copy all objects sharing a prefix must implement their own pagination,
error handling, and encryption support. This operator encapsulates that
pattern so it can be used directly in a Dag.

Author: bentorb

providers/amazon/docs/operators/s3/s3.rst

@@ -132,6 +132,24 @@ Ensure the role or user has the necessary permissions to use the key.
     :start-after: [START howto_operator_s3_copy_object]
     :end-before: [END howto_operator_s3_copy_object]
 
+.. _howto/operator:S3CopyPrefixOperator:
+
+Copy Amazon S3 objects by prefix
+================================
+
+To copy all Amazon S3 objects under a prefix from one bucket to another you can use
+:class:`~airflow.providers.amazon.aws.operators.s3.S3CopyPrefixOperator`.
+The Amazon S3 connection used here needs to have access to both source and destination bucket/prefix.
+You can also specify server-side encryption using AWS KMS if you do not want to use destination buckets default key.
+When using KMS, you must provide both the ``kms_key_id`` and ``kms_encryption_type`` parameters.
+Ensure the role or user has the necessary permissions to use the key.
+
+.. exampleinclude:: /../../amazon/tests/system/amazon/aws/example_s3.py
+    :language: python
+    :dedent: 4
+    :start-after: [START howto_operator_s3_copy_prefix]
+    :end-before: [END howto_operator_s3_copy_prefix]
+
 .. _howto/operator:S3DeleteObjectsOperator:
 
 Delete Amazon S3 objects

providers/amazon/src/airflow/providers/amazon/aws/operators/s3.py

@@ -385,6 +385,161 @@ def get_openlineage_facets_on_start(self):
         )
 
 
+class S3CopyPrefixOperator(AwsBaseOperator[S3Hook]):
+    """
+    Creates a copy of all objects under a prefix already stored in S3.
+
+    Note: the S3 connection used here needs to have access to both
+    source and destination bucket/prefix.
+
+    .. seealso::
+        For more information on how to use this operator, take a look at the guide:
+        :ref:`howto/operator:S3CopyPrefixOperator`
+
+    :param source_bucket_prefix: The prefix in the source bucket. (templated)
+        It can be either full s3:// style url or relative path from root level.
+        When it's specified as a full s3:// url, please omit source_bucket_name.
+    :param dest_bucket_prefix: The prefix in the destination to copy to. (templated)
+        The convention to specify `dest_bucket_prefix` is the same as `source_bucket_prefix`.
+    :param source_bucket_name: Name of the S3 bucket where the source objects are in. (templated)
+        It should be omitted when `source_bucket_prefix` is provided as a full s3:// url.
+    :param dest_bucket_name: Name of the S3 bucket to where the objects are copied. (templated)
+        It should be omitted when `dest_bucket_prefix` is provided as a full s3:// url.
+    :param kms_key_id: The ARN, id or alias of the AWS KMS key to use for encrypting the destination object.
+        Required if using KMS-based server-side encryption with a non-default key. (templated)
+    :param kms_encryption_type: Type of KMS encryption to use for the object.
+        Can be either "aws:kms" (standard KMS) or "aws:kms:dsse" (double-shielded KMS).
+    :param continue_on_failure: If False, stop and fail the task on the first copy error.
+        If True, try to copy every object in the prefix and then fail the task on any error.
+        Default is False.
+    :param acl_policy: String specifying the canned ACL policy for the file being
+        uploaded to the S3 bucket.
+    :param meta_data_directive: Whether to `COPY` the metadata from the source object or `REPLACE` it with
+        metadata that's provided in the request.
+    :param aws_conn_id: The Airflow connection used for AWS credentials.
+        If this is ``None`` or empty then the default boto3 behaviour is used. If
+        running Airflow in a distributed manner and aws_conn_id is None or
+        empty, then default boto3 configuration would be used (and must be
+        maintained on each worker node).
+    :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+    :param verify: Whether or not to verify SSL certificates. See:
+        https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html
+    :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+        https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+    """
+
+    template_fields: Sequence[str] = aws_template_fields(
+        "source_bucket_prefix",
+        "dest_bucket_prefix",
+        "source_bucket_name",
+        "dest_bucket_name",
+        "kms_key_id",
+    )
+    aws_hook_class = S3Hook
+
+    def __init__(
+        self,
+        *,
+        source_bucket_prefix: str,
+        dest_bucket_prefix: str,
+        source_bucket_name: str | None = None,
+        dest_bucket_name: str | None = None,
+        kms_key_id: str | None = None,
+        kms_encryption_type: str | None = None,
+        continue_on_failure: bool = False,
+        acl_policy: str | None = None,
+        meta_data_directive: str | None = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.source_bucket_prefix = source_bucket_prefix
+        self.dest_bucket_prefix = dest_bucket_prefix
+        self.source_bucket_name = source_bucket_name
+        self.dest_bucket_name = dest_bucket_name
+        self.kms_key_id = kms_key_id
+        self.kms_encryption_type = kms_encryption_type
+        self.continue_on_failure = continue_on_failure
+        self.acl_policy = acl_policy
+        self.meta_data_directive = meta_data_directive
+
+    def execute(self, context: Context):
+        source_bucket_name, source_bucket_prefix = self.hook.get_s3_bucket_key(
+            self.source_bucket_name, self.source_bucket_prefix, "source_bucket_name", "source_bucket_prefix"
+        )
+
+        dest_bucket_name, dest_bucket_prefix = self.hook.get_s3_bucket_key(
+            self.dest_bucket_name, self.dest_bucket_prefix, "dest_bucket_name", "dest_bucket_prefix"
+        )
+
+        s3_client = self.hook.get_conn()
+
+        paginator = s3_client.get_paginator("list_objects_v2")
+        pages = paginator.paginate(
+            Bucket=source_bucket_name,
+            Prefix=source_bucket_prefix,
+        )
+
+        copied_object_count = 0
+        failed_object_count = 0
+        for page in pages:
+            if "Contents" in page:
+                for obj in page["Contents"]:
+                    source_key = obj["Key"]
+                    dest_key = dest_bucket_prefix + source_key[len(source_bucket_prefix) :]
+
+                    try:
+                        self.hook.copy_object(
+                            source_bucket_key=source_key,
+                            dest_bucket_key=dest_key,
+                            source_bucket_name=source_bucket_name,
+                            dest_bucket_name=dest_bucket_name,
+                            kms_key_id=self.kms_key_id,
+                            kms_encryption_type=self.kms_encryption_type,
+                            acl_policy=self.acl_policy,
+                            meta_data_directive=self.meta_data_directive,
+                        )
+
+                        copied_object_count += 1
+                    except Exception as e:
+                        if self.continue_on_failure:
+                            self.log.error("Failed to copy %s: %s", source_key, e)
+                            failed_object_count += 1
+                        else:
+                            raise RuntimeError(f"Failed to copy {source_key}: {e}") from e
+
+        self.log.info("Successfully copied %s object(s)", copied_object_count)
+
+        if failed_object_count > 0:
+            raise RuntimeError(f"Failed to copy {failed_object_count} object(s)")
+
+    def get_openlineage_facets_on_start(self):
+        from airflow.providers.common.compat.openlineage.facet import Dataset
+        from airflow.providers.openlineage.extractors import OperatorLineage
+
+        source_bucket_name, source_bucket_prefix = self.hook.get_s3_bucket_key(
+            self.source_bucket_name, self.source_bucket_prefix, "source_bucket_name", "source_bucket_prefix"
+        )
+
+        dest_bucket_name, dest_bucket_prefix = self.hook.get_s3_bucket_key(
+            self.dest_bucket_name, self.dest_bucket_prefix, "dest_bucket_name", "dest_bucket_prefix"
+        )
+
+        input_dataset = Dataset(
+            namespace=f"s3://{source_bucket_name}",
+            name=source_bucket_prefix,
+        )
+        output_dataset = Dataset(
+            namespace=f"s3://{dest_bucket_name}",
+            name=dest_bucket_prefix,
+        )
+
+        return OperatorLineage(
+            inputs=[input_dataset],
+            outputs=[output_dataset],
+        )
+
+
 class S3CreateObjectOperator(AwsBaseOperator[S3Hook]):
     """
     Creates a new object from `data` as string or bytes.

providers/amazon/tests/system/amazon/aws/example_s3.py

@@ -20,6 +20,7 @@
 
 from airflow.providers.amazon.aws.operators.s3 import (
     S3CopyObjectOperator,
+    S3CopyPrefixOperator,
     S3CreateBucketOperator,
     S3CreateObjectOperator,
     S3DeleteBucketOperator,
@@ -254,6 +255,16 @@ def check_fn(files: list, **kwargs) -> bool:
     )
     # [END howto_operator_s3_copy_object]
 
+    # [START howto_operator_s3_copy_prefix]
+    copy_prefix = S3CopyPrefixOperator(
+        task_id="copy_prefix",
+        source_bucket_name=bucket_name,
+        source_bucket_prefix=f"{env_id}-",
+        dest_bucket_name=bucket_name_2,
+        dest_bucket_prefix=f"{env_id}-copied-",
+    )
+    # [END howto_operator_s3_copy_prefix]
+
     # [START howto_operator_s3_file_transform]
     file_transform = S3FileTransformOperator(
         task_id="file_transform",
@@ -321,6 +332,7 @@ def check_fn(files: list, **kwargs) -> bool:
             sensor_key_with_regex_deferrable,
         ],
         copy_object,
+        copy_prefix,
         file_transform,
         sensor_keys_unchanged,
         # TEST TEARDOWN