This is an automated email from the ASF dual-hosted git repository.
shahar1 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 d3e7d369587 Return list of destination URIs from S3ToGCSOperator
(#64851)
d3e7d369587 is described below
commit d3e7d369587005959bc63c7f329176b19fdf37df
Author: Yuseok Jo <[email protected]>
AuthorDate: Mon Apr 20 23:29:51 2026 +0900
Return list of destination URIs from S3ToGCSOperator (#64851)
* Return list of destination URIs from S3ToGCSOperator
* Add return_gcs_uris to S3ToGCSOperator
---
.../providers/google/cloud/transfers/s3_to_gcs.py | 44 +++++++-
.../unit/google/cloud/transfers/test_s3_to_gcs.py | 115 +++++++++++++++++++--
2 files changed, 148 insertions(+), 11 deletions(-)
diff --git
a/providers/google/src/airflow/providers/google/cloud/transfers/s3_to_gcs.py
b/providers/google/src/airflow/providers/google/cloud/transfers/s3_to_gcs.py
index 9ac46273532..7b44556c044 100644
--- a/providers/google/src/airflow/providers/google/cloud/transfers/s3_to_gcs.py
+++ b/providers/google/src/airflow/providers/google/cloud/transfers/s3_to_gcs.py
@@ -17,6 +17,7 @@
# under the License.
from __future__ import annotations
+import warnings
from collections.abc import Sequence
from datetime import datetime, timezone
from tempfile import NamedTemporaryFile
@@ -58,6 +59,13 @@ except ImportError:
if TYPE_CHECKING:
from airflow.providers.common.compat.sdk import Context
+_RETURN_GCS_URIS_FALSE_DEPRECATION_MSG = (
+ "S3ToGCSOperator currently returns a list of S3 object keys to preserve
backward "
+ "compatibility. This default will change in a future release: the operator
will return "
+ "a list of destination 'gs://' URIs instead. Pass 'return_gcs_uris=True'
to opt into "
+ "the new behavior now and silence this warning."
+)
+
class S3ToGCSOperator(S3ListOperator):
"""
@@ -108,6 +116,11 @@ class S3ToGCSOperator(S3ListOperator):
:param deferrable: Run operator in the deferrable mode
:param poll_interval: time in seconds between polling for job completion.
The value is considered only when running in deferrable mode. Must be
greater than 0.
+ :param return_gcs_uris: If ``True``, ``execute`` and deferrable
``execute_complete`` return a list of
+ destination ``gs://`` URIs for copied objects. Defaults to ``False``
for backward
+ compatibility, which returns the raw S3 object keys and emits a
deprecation warning;
+ this default will change to ``True`` in a future release. Set
``return_gcs_uris=True``
+ to opt into the new behavior now and silence the warning.
**Example**:
@@ -154,6 +167,7 @@ class S3ToGCSOperator(S3ListOperator):
google_impersonation_chain: str | Sequence[str] | None = None,
deferrable=conf.getboolean("operators", "default_deferrable",
fallback=False),
poll_interval: int = 10,
+ return_gcs_uris: bool = False,
**kwargs,
):
super().__init__(bucket=bucket, prefix=prefix, delimiter=delimiter,
aws_conn_id=aws_conn_id, **kwargs)
@@ -168,6 +182,13 @@ class S3ToGCSOperator(S3ListOperator):
if poll_interval <= 0:
raise ValueError("Invalid value for poll_interval. Expected value
greater than 0")
self.poll_interval = poll_interval
+ self.return_gcs_uris = return_gcs_uris
+ if not return_gcs_uris:
+ warnings.warn(
+ _RETURN_GCS_URIS_FALSE_DEPRECATION_MSG,
+ FutureWarning,
+ stacklevel=2,
+ )
def _check_inputs(self) -> None:
if self.dest_gcs and not gcs_object_is_directory(self.dest_gcs):
@@ -180,6 +201,11 @@ class S3ToGCSOperator(S3ListOperator):
'The destination Google Cloud Storage path must end with a
slash "/" or be empty.'
)
+ def _destination_uris_for_s3_keys(self, s3_keys: list[str]) -> list[str]:
+ """Build ``gs://`` URIs for each transferred S3 object key."""
+ gcs_bucket, _ = _parse_gcs_url(self.dest_gcs)
+ return [f"gs://{gcs_bucket}/{self.s3_to_gcs_object(s3_object=k)}" for
k in s3_keys]
+
def _get_files(self, context: Context, gcs_hook: GCSHook) -> list[str]:
# use the super method to list all the files in an S3 bucket/key
s3_objects = super().execute(context)
@@ -189,7 +215,7 @@ class S3ToGCSOperator(S3ListOperator):
return s3_objects
- def execute(self, context: Context):
+ def execute(self, context: Context) -> list[str]:
self._check_inputs()
gcs_hook = GCSHook(
gcp_conn_id=self.gcp_conn_id,
@@ -206,6 +232,8 @@ class S3ToGCSOperator(S3ListOperator):
else:
self.transfer_files(s3_objects, gcs_hook, s3_hook)
+ if self.return_gcs_uris:
+ return self._destination_uris_for_s3_keys(s3_objects)
return s3_objects
def exclude_existing_objects(self, s3_objects: list[str], gcs_hook:
GCSHook) -> list[str]:
@@ -339,14 +367,20 @@ class S3ToGCSOperator(S3ListOperator):
"""
Handle the trigger callback when transfer jobs complete.
- Returns the list of copied file paths when available (deferrable mode
with
- files passed via trigger), so subsequent tasks can consume them via
XCom.
- Returns None when event does not contain files (e.g. legacy triggers).
+ Returns a list of copied files when available (deferrable mode with
files passed via trigger),
+ so subsequent tasks can consume them via XCom. When
``return_gcs_uris`` is ``True``, the list
+ contains destination ``gs://`` URIs; otherwise it contains the S3
object keys transferred.
+ Returns ``None`` when the event does not contain files (e.g. legacy
triggers).
"""
if event["status"] == "error":
raise AirflowException(event["message"])
self.log.info("%s completed with response %s ", self.task_id,
event["message"])
- return event.get("files")
+ files = event.get("files")
+ if files is None:
+ return None
+ if self.return_gcs_uris:
+ return self._destination_uris_for_s3_keys(files)
+ return files
def get_transfer_hook(self):
return CloudDataTransferServiceHook(
diff --git
a/providers/google/tests/unit/google/cloud/transfers/test_s3_to_gcs.py
b/providers/google/tests/unit/google/cloud/transfers/test_s3_to_gcs.py
index 1e17364b695..580123a69cf 100644
--- a/providers/google/tests/unit/google/cloud/transfers/test_s3_to_gcs.py
+++ b/providers/google/tests/unit/google/cloud/transfers/test_s3_to_gcs.py
@@ -25,7 +25,11 @@ import time_machine
from airflow.providers.common.compat.sdk import AirflowException, TaskDeferred
from airflow.providers.google.cloud.hooks.cloud_storage_transfer_service
import CloudDataTransferServiceHook
-from airflow.providers.google.cloud.transfers.s3_to_gcs import S3ToGCSOperator
+from airflow.providers.google.cloud.hooks.gcs import _parse_gcs_url
+from airflow.providers.google.cloud.transfers.s3_to_gcs import (
+ _RETURN_GCS_URIS_FALSE_DEPRECATION_MSG,
+ S3ToGCSOperator,
+)
from airflow.utils.timezone import utcnow
PROJECT_ID = "test-project-id"
@@ -84,6 +88,7 @@ class TestS3ToGoogleCloudStorageOperator:
apply_gcs_prefix=APPLY_GCS_PREFIX,
deferrable=DEFERRABLE,
poll_interval=POLL_INTERVAL,
+ return_gcs_uris=True,
)
assert operator.task_id == TASK_ID
@@ -96,6 +101,7 @@ class TestS3ToGoogleCloudStorageOperator:
assert operator.apply_gcs_prefix == APPLY_GCS_PREFIX
assert operator.deferrable == DEFERRABLE
assert operator.poll_interval == POLL_INTERVAL
+ assert operator.return_gcs_uris is True
@mock.patch("airflow.providers.google.cloud.transfers.s3_to_gcs.S3Hook")
@mock.patch("airflow.providers.google.cloud.transfers.s3_to_gcs.GCSHook")
@@ -110,6 +116,7 @@ class TestS3ToGoogleCloudStorageOperator:
gcp_conn_id=GCS_CONN_ID,
dest_gcs=GCS_PATH_PREFIX,
google_impersonation_chain=IMPERSONATION_CHAIN,
+ return_gcs_uris=True,
)
operator.hook = mock.MagicMock()
@@ -132,8 +139,71 @@ class TestS3ToGoogleCloudStorageOperator:
impersonation_chain=IMPERSONATION_CHAIN,
)
- # we expect MOCK_FILES to be uploaded
- assert sorted(MOCK_FILES) == sorted(uploaded_files)
+ expected_uris = [f"gs://{GCS_BUCKET}/{GCS_PREFIX}{f}" for f in
MOCK_FILES]
+ assert sorted(expected_uris) == sorted(uploaded_files)
+
+ @mock.patch("airflow.providers.google.cloud.transfers.s3_to_gcs.S3Hook")
+ @mock.patch("airflow.providers.google.cloud.transfers.s3_to_gcs.GCSHook")
+ def test_execute_legacy_returns_s3_keys(self, gcs_mock_hook, s3_mock_hook):
+ """When return_gcs_uris=False, execute returns raw S3 keys for
backward compatibility."""
+ with pytest.warns(FutureWarning, match="S3ToGCSOperator"):
+ operator = S3ToGCSOperator(
+ task_id=TASK_ID,
+ bucket=S3_BUCKET,
+ prefix=S3_PREFIX,
+ delimiter=S3_DELIMITER,
+ gcp_conn_id=GCS_CONN_ID,
+ dest_gcs=GCS_PATH_PREFIX,
+ google_impersonation_chain=IMPERSONATION_CHAIN,
+ return_gcs_uris=False,
+ )
+ operator.hook = mock.MagicMock()
+ operator.hook.list_keys.return_value = MOCK_FILES
+
+ uploaded_files = operator.execute(context={})
+ assert sorted(uploaded_files) == sorted(MOCK_FILES)
+
+ def
test_init_emits_deprecation_warning_when_return_gcs_uris_false_default(self):
+ """``return_gcs_uris`` defaults to ``False`` and emits a deprecation
warning at init."""
+ with pytest.warns(FutureWarning, match="S3ToGCSOperator"):
+ operator = S3ToGCSOperator(
+ task_id=TASK_ID,
+ bucket=S3_BUCKET,
+ prefix=S3_PREFIX,
+ delimiter=S3_DELIMITER,
+ gcp_conn_id=GCS_CONN_ID,
+ dest_gcs=GCS_PATH_PREFIX,
+ )
+ assert operator.return_gcs_uris is False
+
+ def
test_init_emits_deprecation_warning_when_return_gcs_uris_false_explicit(self):
+ """Explicit ``return_gcs_uris=False`` still emits the deprecation
warning."""
+ with pytest.warns(FutureWarning, match="S3ToGCSOperator"):
+ S3ToGCSOperator(
+ task_id=TASK_ID,
+ bucket=S3_BUCKET,
+ prefix=S3_PREFIX,
+ delimiter=S3_DELIMITER,
+ gcp_conn_id=GCS_CONN_ID,
+ dest_gcs=GCS_PATH_PREFIX,
+ return_gcs_uris=False,
+ )
+
+ def test_init_no_warning_when_return_gcs_uris_true(self, recwarn):
+ """``return_gcs_uris=True`` opts into the new behavior and silences
the warning."""
+ S3ToGCSOperator(
+ task_id=TASK_ID,
+ bucket=S3_BUCKET,
+ prefix=S3_PREFIX,
+ delimiter=S3_DELIMITER,
+ gcp_conn_id=GCS_CONN_ID,
+ dest_gcs=GCS_PATH_PREFIX,
+ return_gcs_uris=True,
+ )
+ assert not any(
+ issubclass(w.category, FutureWarning) and
_RETURN_GCS_URIS_FALSE_DEPRECATION_MSG in str(w.message)
+ for w in recwarn
+ )
@mock.patch("airflow.providers.google.cloud.transfers.s3_to_gcs.S3Hook")
@mock.patch("airflow.providers.google.cloud.transfers.s3_to_gcs.GCSHook")
@@ -245,6 +315,7 @@ class TestS3ToGoogleCloudStorageOperator:
dest_gcs=gcs_destination,
google_impersonation_chain=IMPERSONATION_CHAIN,
apply_gcs_prefix=apply_gcs_prefix,
+ return_gcs_uris=True,
)
operator.hook = mock.MagicMock()
operator.hook.list_keys.return_value = [s3_prefix + s3_object]
@@ -264,7 +335,9 @@ class TestS3ToGoogleCloudStorageOperator:
impersonation_chain=IMPERSONATION_CHAIN,
)
- assert sorted([s3_prefix + s3_object]) == sorted(uploaded_files)
+ dest_bucket, _ = _parse_gcs_url(gcs_destination)
+ expected_uri = f"gs://{dest_bucket}/{gcs_object}"
+ assert uploaded_files == [expected_uri]
@pytest.mark.parametrize(
("s3_prefix", "gcs_destination", "apply_gcs_prefix", "expected_input",
"expected_output"),
@@ -499,14 +572,44 @@ class TestS3ToGoogleCloudStorageOperatorDeferrable:
"airflow.providers.google.cloud.transfers.s3_to_gcs.S3ToGCSOperator.log",
new_callable=PropertyMock
)
def test_execute_complete_success_returns_copied_files(self, mock_log):
- """Deferrable mode returns list of copied files for use in subsequent
tasks via XCom."""
+ """Deferrable mode returns destination GCS URIs for XCom when
return_gcs_uris=True."""
expected_files = [MOCK_FILE_1, MOCK_FILE_2]
event = {
"status": "success",
"message": "Transfer completed",
"files": expected_files,
}
- operator = S3ToGCSOperator(task_id=TASK_ID, bucket=S3_BUCKET)
+ operator = S3ToGCSOperator(
+ task_id=TASK_ID,
+ bucket=S3_BUCKET,
+ prefix=S3_PREFIX,
+ dest_gcs=GCS_PATH_PREFIX,
+ return_gcs_uris=True,
+ )
+ result = operator.execute_complete(context={}, event=event)
+
+ expected_uris = [f"gs://{GCS_BUCKET}/{GCS_PREFIX}{f}" for f in
expected_files]
+ assert result == expected_uris
+
+ @mock.patch(
+
"airflow.providers.google.cloud.transfers.s3_to_gcs.S3ToGCSOperator.log",
new_callable=PropertyMock
+ )
+ def test_execute_complete_legacy_returns_s3_keys(self, mock_log):
+ """When return_gcs_uris is False, execute_complete returns raw S3 keys
for backward compatibility."""
+ expected_files = [MOCK_FILE_1, MOCK_FILE_2]
+ event = {
+ "status": "success",
+ "message": "Transfer completed",
+ "files": expected_files,
+ }
+ with pytest.warns(FutureWarning, match="S3ToGCSOperator"):
+ operator = S3ToGCSOperator(
+ task_id=TASK_ID,
+ bucket=S3_BUCKET,
+ prefix=S3_PREFIX,
+ dest_gcs=GCS_PATH_PREFIX,
+ return_gcs_uris=False,
+ )
result = operator.execute_complete(context={}, event=event)
assert result == expected_files
