This is an automated email from the ASF dual-hosted git repository.
o-nikolas pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 5a8409cc527 Add S3CopyPrefixOperator to copy all objects under a
prefix (#68946)
5a8409cc527 is described below
commit 5a8409cc5275cc41978cd409a609f1fa9de04246
Author: bentorb <[email protected]>
AuthorDate: Sat Jun 27 01:33:37 2026 +0200
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.
---
providers/amazon/docs/operators/s3/s3.rst | 18 ++
.../airflow/providers/amazon/aws/operators/s3.py | 155 +++++++++++++
.../amazon/tests/system/amazon/aws/example_s3.py | 12 +
.../tests/unit/amazon/aws/operators/test_s3.py | 245 +++++++++++++++++++++
4 files changed, 430 insertions(+)
diff --git a/providers/amazon/docs/operators/s3/s3.rst
b/providers/amazon/docs/operators/s3/s3.rst
index d5bf4879606..5dec957b5ce 100644
--- a/providers/amazon/docs/operators/s3/s3.rst
+++ b/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
diff --git a/providers/amazon/src/airflow/providers/amazon/aws/operators/s3.py
b/providers/amazon/src/airflow/providers/amazon/aws/operators/s3.py
index 2a5b12052b6..b74fb60a1b7 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/operators/s3.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/s3.py
@@ -385,6 +385,161 @@ class S3CopyObjectOperator(AwsBaseOperator[S3Hook]):
)
+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.
diff --git a/providers/amazon/tests/system/amazon/aws/example_s3.py
b/providers/amazon/tests/system/amazon/aws/example_s3.py
index 184211e630d..19769a28696 100644
--- a/providers/amazon/tests/system/amazon/aws/example_s3.py
+++ b/providers/amazon/tests/system/amazon/aws/example_s3.py
@@ -20,6 +20,7 @@ from datetime import datetime
from airflow.providers.amazon.aws.operators.s3 import (
S3CopyObjectOperator,
+ S3CopyPrefixOperator,
S3CreateBucketOperator,
S3CreateObjectOperator,
S3DeleteBucketOperator,
@@ -254,6 +255,16 @@ with DAG(
)
# [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 @@ with DAG(
sensor_key_with_regex_deferrable,
],
copy_object,
+ copy_prefix,
file_transform,
sensor_keys_unchanged,
# TEST TEARDOWN
diff --git a/providers/amazon/tests/unit/amazon/aws/operators/test_s3.py
b/providers/amazon/tests/unit/amazon/aws/operators/test_s3.py
index f3a74a4dd3f..cd907171927 100644
--- a/providers/amazon/tests/unit/amazon/aws/operators/test_s3.py
+++ b/providers/amazon/tests/unit/amazon/aws/operators/test_s3.py
@@ -36,6 +36,7 @@ from airflow.models.taskinstance import TaskInstance
from airflow.providers.amazon.aws.hooks.s3 import S3Hook
from airflow.providers.amazon.aws.operators.s3 import (
S3CopyObjectOperator,
+ S3CopyPrefixOperator,
S3CreateBucketOperator,
S3CreateObjectOperator,
S3DeleteBucketOperator,
@@ -658,6 +659,250 @@ class TestS3CopyObjectOperator:
assert objects_in_dest_bucket["Contents"][0]["Key"] == self.dest_key
+class TestS3CopyPrefixOperator:
+ def setup_method(self):
+ self.source_bucket = "source-bucket"
+ self.source_prefix = "data/logs/"
+ self.dest_bucket = "dest-bucket"
+ self.dest_prefix = "backup/logs/"
+
+ self.source_s3_url = f"s3://{self.source_bucket}/{self.source_prefix}"
+ self.dest_s3_url = f"s3://{self.dest_bucket}/{self.dest_prefix}"
+
+ @staticmethod
+ def _create_s3_client():
+ return boto3.client("s3", region_name="us-east-1")
+
+ def _create_buckets(self, s3_client):
+ s3_client.create_bucket(Bucket=self.source_bucket)
+ s3_client.create_bucket(Bucket=self.dest_bucket)
+
+ def _upload_test_objects(self, s3_client, keys):
+ for key in keys:
+ s3_client.upload_fileobj(Bucket=self.source_bucket, Key=key,
Fileobj=BytesIO(b"test-content"))
+
+ @mock_aws
+ def test_s3_copy_prefix_basic(self):
+ s3_client = self._create_s3_client()
+ self._create_buckets(s3_client)
+ self._upload_test_objects(
+ s3_client,
+ [
+ f"{self.source_prefix}file1.txt",
+ f"{self.source_prefix}file2.txt",
+ f"{self.source_prefix}subdir/file3.txt",
+ ],
+ )
+
+ dest_objects = s3_client.list_objects_v2(Bucket=self.dest_bucket,
Prefix=self.dest_prefix)
+ assert "Contents" not in dest_objects
+
+ op = S3CopyPrefixOperator(
+ task_id="test_copy_prefix",
+ source_bucket_name=self.source_bucket,
+ source_bucket_prefix=self.source_prefix,
+ dest_bucket_name=self.dest_bucket,
+ dest_bucket_prefix=self.dest_prefix,
+ )
+
+ op.execute(None)
+
+ dest_objects = s3_client.list_objects_v2(Bucket=self.dest_bucket,
Prefix=self.dest_prefix)
+ assert len(dest_objects["Contents"]) == 3
+
+ copied_keys = [obj["Key"] for obj in dest_objects["Contents"]]
+ assert "backup/logs/file1.txt" in copied_keys
+ assert "backup/logs/file2.txt" in copied_keys
+ assert "backup/logs/subdir/file3.txt" in copied_keys
+
+ @mock_aws
+ def test_s3_copy_prefix_selective_copying(self):
+ s3_client = self._create_s3_client()
+ self._create_buckets(s3_client)
+ self._upload_test_objects(
+ s3_client,
+ [
+ f"{self.source_prefix}file1.txt",
+ f"{self.source_prefix}subdir/file2.txt",
+ "data/metrics/file3.txt",
+ "other/logs/file4.txt",
+ "archive/data/file5.txt",
+ ],
+ )
+
+ op = S3CopyPrefixOperator(
+ task_id="test_copy_prefix_selective",
+ source_bucket_name=self.source_bucket,
+ source_bucket_prefix=self.source_prefix,
+ dest_bucket_name=self.dest_bucket,
+ dest_bucket_prefix=self.dest_prefix,
+ )
+
+ op.execute(None)
+
+ dest_objects = s3_client.list_objects_v2(Bucket=self.dest_bucket,
Prefix=self.dest_prefix)
+ assert len(dest_objects["Contents"]) == 2
+
+ copied_keys = [obj["Key"] for obj in dest_objects["Contents"]]
+ assert "backup/logs/file1.txt" in copied_keys
+ assert "backup/logs/subdir/file2.txt" in copied_keys
+
+ @mock_aws
+ def test_s3_copy_prefix_s3_urls(self):
+ s3_client = self._create_s3_client()
+ self._create_buckets(s3_client)
+ self._upload_test_objects(
+ s3_client, [f"{self.source_prefix}file1.txt",
f"{self.source_prefix}file2.txt"]
+ )
+
+ op = S3CopyPrefixOperator(
+ task_id="test_copy_prefix_urls",
+ source_bucket_prefix=self.source_s3_url,
+ dest_bucket_prefix=self.dest_s3_url,
+ )
+
+ op.execute(None)
+
+ dest_objects = s3_client.list_objects_v2(Bucket=self.dest_bucket,
Prefix=self.dest_prefix)
+ assert len(dest_objects["Contents"]) == 2
+
+ def test_invalid_combination_bucket_with_s3_url(self):
+ op = S3CopyPrefixOperator(
+ task_id="test_invalid",
+ source_bucket_name=self.source_bucket,
+
source_bucket_prefix=f"s3://{self.source_bucket}/{self.source_prefix}",
+ dest_bucket_name=self.dest_bucket,
+ dest_bucket_prefix=self.dest_prefix,
+ )
+
+ with pytest.raises(TypeError, match="should be a relative path"):
+ op.execute(None)
+
+ @mock_aws
+ def test_s3_copy_prefix_same_bucket(self):
+ s3_client = self._create_s3_client()
+ s3_client.create_bucket(Bucket=self.source_bucket)
+ self._upload_test_objects(s3_client,
[f"{self.source_prefix}file1.txt"])
+
+ op = S3CopyPrefixOperator(
+ task_id="test_copy_prefix_same_bucket",
+ source_bucket_name=self.source_bucket,
+ source_bucket_prefix=self.source_prefix,
+ dest_bucket_name=self.source_bucket,
+ dest_bucket_prefix="archive/logs/",
+ )
+
+ op.execute(None)
+
+ all_objects = s3_client.list_objects_v2(Bucket=self.source_bucket)
+ keys = [obj["Key"] for obj in all_objects["Contents"]]
+ assert f"{self.source_prefix}file1.txt" in keys
+ assert "archive/logs/file1.txt" in keys
+
+ @mock_aws
+ def test_s3_copy_prefix_empty_result(self):
+ s3_client = self._create_s3_client()
+ self._create_buckets(s3_client)
+
+ op = S3CopyPrefixOperator(
+ task_id="test_copy_prefix_empty",
+ source_bucket_name=self.source_bucket,
+ source_bucket_prefix=self.source_prefix,
+ dest_bucket_name=self.dest_bucket,
+ dest_bucket_prefix=self.dest_prefix,
+ )
+
+ op.execute(None)
+
+ dest_objects = s3_client.list_objects_v2(Bucket=self.dest_bucket,
Prefix=self.dest_prefix)
+ assert "Contents" not in dest_objects
+
+ @mock_aws
+ def test_continue_on_failure_false(self):
+ s3_client = self._create_s3_client()
+ self._create_buckets(s3_client)
+ self._upload_test_objects(s3_client,
[f"{self.source_prefix}file1.txt"])
+
+ op = S3CopyPrefixOperator(
+ task_id="test_copy_prefix_fail_fast",
+ source_bucket_name=self.source_bucket,
+ source_bucket_prefix=self.source_prefix,
+ dest_bucket_name=self.dest_bucket,
+ dest_bucket_prefix=self.dest_prefix,
+ continue_on_failure=False,
+ )
+
+ with mock.patch.object(op.hook, "copy_object",
side_effect=Exception("Copy failed")):
+ with pytest.raises(
+ RuntimeError, match=f"Failed to copy
{self.source_prefix}file1.txt: Copy failed"
+ ):
+ op.execute(None)
+
+ @mock_aws
+ def test_continue_on_failure_true(self):
+ s3_client = self._create_s3_client()
+ self._create_buckets(s3_client)
+ self._upload_test_objects(
+ s3_client, [f"{self.source_prefix}file1.txt",
f"{self.source_prefix}file2.txt"]
+ )
+
+ op = S3CopyPrefixOperator(
+ task_id="test_copy_prefix_continue",
+ source_bucket_name=self.source_bucket,
+ source_bucket_prefix=self.source_prefix,
+ dest_bucket_name=self.dest_bucket,
+ dest_bucket_prefix=self.dest_prefix,
+ continue_on_failure=True,
+ )
+
+ def mock_copy_object(*args, **kwargs):
+ if "file1.txt" in kwargs.get("source_bucket_key", ""):
+ raise Exception("Copy failed for file1")
+ return None
+
+ with mock.patch.object(op.hook, "copy_object",
side_effect=mock_copy_object) as mock_copy:
+ with pytest.raises(RuntimeError, match=r"Failed to copy 1
object\(s\)"):
+ op.execute(None)
+
+ attempted_keys = [call.kwargs["source_bucket_key"] for call in
mock_copy.call_args_list]
+ assert f"{self.source_prefix}file2.txt" in attempted_keys
+
+ @pytest.mark.parametrize(
+ "op_kwargs",
+ [
+ {
+ "source_bucket_name": "source-bucket",
+ "source_bucket_prefix": "data/logs/",
+ "dest_bucket_name": "dest-bucket",
+ "dest_bucket_prefix": "backup/logs/",
+ },
+ {
+ "source_bucket_prefix": "s3://source-bucket/data/logs/",
+ "dest_bucket_prefix": "s3://dest-bucket/backup/logs/",
+ },
+ ],
+ ids=["bucket_and_prefix", "s3_urls"],
+ )
+ def test_get_openlineage_facets_on_start(self, op_kwargs):
+ from airflow.providers.common.compat.openlineage.facet import Dataset
+
+ op = S3CopyPrefixOperator(task_id="test", **op_kwargs)
+
+ lineage = op.get_openlineage_facets_on_start()
+ assert lineage.inputs == [Dataset(namespace="s3://source-bucket",
name="data/logs/")]
+ assert lineage.outputs == [Dataset(namespace="s3://dest-bucket",
name="backup/logs/")]
+
+ def test_template_fields(self):
+ operator = S3CopyPrefixOperator(
+ task_id="test_copy_prefix",
+ source_bucket_name=self.source_bucket,
+ source_bucket_prefix=self.source_prefix,
+ dest_bucket_name=self.dest_bucket,
+ dest_bucket_prefix=self.dest_prefix,
+ )
+ validate_template_fields(operator)
+
+
@mock_aws
class TestS3DeleteObjectsOperator:
def test_s3_delete_single_object(self):
