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 eaeec367ead Add optional object-level retention support to
GcsToGcsOperator (#66204)
eaeec367ead is described below
commit eaeec367eadb14d3d004f84b4b3241163e09802b
Author: Chris Qiu <[email protected]>
AuthorDate: Wed May 6 01:08:31 2026 +0800
Add optional object-level retention support to GcsToGcsOperator (#66204)
---
.../airflow/providers/google/cloud/hooks/gcs.py | 35 +++++-
.../providers/google/cloud/transfers/gcs_to_gcs.py | 28 ++++-
.../tests/unit/google/cloud/hooks/test_gcs.py | 126 +++++++++++++++++++++
.../unit/google/cloud/transfers/test_gcs_to_gcs.py | 67 +++++++++++
4 files changed, 250 insertions(+), 6 deletions(-)
diff --git a/providers/google/src/airflow/providers/google/cloud/hooks/gcs.py
b/providers/google/src/airflow/providers/google/cloud/hooks/gcs.py
index 4dd7dd840e1..2b527365ec7 100644
--- a/providers/google/src/airflow/providers/google/cloud/hooks/gcs.py
+++ b/providers/google/src/airflow/providers/google/cloud/hooks/gcs.py
@@ -225,6 +225,8 @@ class GCSHook(GoogleBaseHook):
source_object: str,
destination_bucket: str,
destination_object: str | None = None,
+ retain_until_time: datetime | None = None,
+ retention_mode: str | None = None,
) -> None:
"""
Similar to copy; supports files over 5 TB, and copying between
locations and/or storage classes.
@@ -236,6 +238,13 @@ class GCSHook(GoogleBaseHook):
:param destination_bucket: The destination of the object to copied to.
:param destination_object: The (renamed) path of the object if given.
Can be omitted; then the same name is used.
+ :param retain_until_time: Optional datetime specifying until when the
destination
+ object should be retained. Requires the destination bucket to have
+ object retention enabled. The value is passed to the GCS client
as-is;
+ timezone handling follows the GCS client behavior.
+ :param retention_mode: Optional retention mode for the destination
object.
+ Must be ``"Locked"`` or ``"Unlocked"``. Defaults to ``"Unlocked"``
when
+ ``retain_until_time`` is set. Cannot be provided without
``retain_until_time``.
"""
destination_object = destination_object or source_object
if source_bucket == destination_bucket and source_object ==
destination_object:
@@ -245,24 +254,40 @@ class GCSHook(GoogleBaseHook):
)
if not source_bucket or not source_object:
raise ValueError("source_bucket and source_object cannot be
empty.")
+ if retention_mode is not None and retain_until_time is None:
+ raise ValueError("retention_mode cannot be set without
retain_until_time.")
+ if retention_mode is not None and retention_mode not in ("Locked",
"Unlocked"):
+ raise ValueError(f"retention_mode must be 'Locked' or 'Unlocked',
got {retention_mode!r}.")
client = self.get_conn()
source_bucket = client.bucket(source_bucket)
source_object = source_bucket.blob(blob_name=source_object) # type:
ignore[attr-defined]
destination_bucket = client.bucket(destination_bucket)
- token, bytes_rewritten, total_bytes = destination_bucket.blob( #
type: ignore[attr-defined]
+ destination_blob = destination_bucket.blob( # type:
ignore[attr-defined]
blob_name=destination_object
- ).rewrite(source=source_object)
+ )
+
+ token, bytes_rewritten, total_bytes =
destination_blob.rewrite(source=source_object)
self.log.info("Total Bytes: %s | Bytes Written: %s", total_bytes,
bytes_rewritten)
while token is not None:
- token, bytes_rewritten, total_bytes = destination_bucket.blob( #
type: ignore[attr-defined]
- blob_name=destination_object
- ).rewrite(source=source_object, token=token)
+ token, bytes_rewritten, total_bytes =
destination_blob.rewrite(source=source_object, token=token)
self.log.info("Total Bytes: %s | Bytes Written: %s", total_bytes,
bytes_rewritten)
+
+ if retain_until_time is not None:
+ destination_blob.retention.mode = retention_mode or "Unlocked"
+ destination_blob.retention.retain_until_time = retain_until_time
+ destination_blob.patch()
+ self.log.info(
+ "Applied retention (mode=%s, retain_until_time=%s) to object
%s in bucket %s",
+ destination_blob.retention.mode,
+ retain_until_time,
+ destination_object,
+ destination_bucket.name, # type: ignore[attr-defined]
+ )
get_hook_lineage_collector().add_input_asset(
context=self,
scheme="gs",
diff --git
a/providers/google/src/airflow/providers/google/cloud/transfers/gcs_to_gcs.py
b/providers/google/src/airflow/providers/google/cloud/transfers/gcs_to_gcs.py
index 12a75a5256c..aa027e71d85 100644
---
a/providers/google/src/airflow/providers/google/cloud/transfers/gcs_to_gcs.py
+++
b/providers/google/src/airflow/providers/google/cloud/transfers/gcs_to_gcs.py
@@ -31,6 +31,9 @@ from airflow.providers.google.version_compat import
BaseOperator
WILDCARD = "*"
if TYPE_CHECKING:
+ from datetime import datetime
+ from typing import Any
+
from airflow.providers.common.compat.sdk import Context
@@ -96,6 +99,14 @@ class GCSToGCSOperator(BaseOperator):
copied.
:param match_glob: (Optional) filters objects based on the glob pattern
given by the string (
e.g, ``'**/*/.json'``)
+ :param retain_until_time: (Optional) A datetime specifying until when the
destination
+ objects should be retained. Requires the destination bucket to have
object retention
+ enabled. The retention is applied after each object is copied/moved.
+ The value is passed to the GCS client as-is; timezone handling follows
the
+ GCS client behavior.
+ :param retention_mode: (Optional) The retention mode for destination
objects.
+ Must be ``"Locked"`` or ``"Unlocked"``. Defaults to ``"Unlocked"`` when
+ ``retain_until_time`` is set. Cannot be provided without
``retain_until_time``.
:Example:
@@ -199,6 +210,8 @@ class GCSToGCSOperator(BaseOperator):
source_object_required=False,
exact_match=False,
match_glob: str | None = None,
+ retain_until_time: datetime | None = None,
+ retention_mode: str | None = None,
**kwargs,
):
super().__init__(**kwargs)
@@ -237,6 +250,8 @@ class GCSToGCSOperator(BaseOperator):
self.source_object_required = source_object_required
self.exact_match = exact_match
self.match_glob = match_glob
+ self.retain_until_time = retain_until_time
+ self.retention_mode = retention_mode
def execute(self, context: Context) -> list[str]:
hook = GCSHook(
@@ -575,7 +590,18 @@ class GCSToGCSOperator(BaseOperator):
dest_bucket,
destination_object,
)
- hook.rewrite(self.source_bucket, source_object, dest_bucket,
destination_object)
+ rewrite_kwargs: dict[str, Any] = {}
+ if self.retain_until_time is not None:
+ rewrite_kwargs["retain_until_time"] = self.retain_until_time
+ if self.retention_mode is not None:
+ rewrite_kwargs["retention_mode"] = self.retention_mode
+ hook.rewrite(
+ self.source_bucket,
+ source_object,
+ dest_bucket,
+ destination_object,
+ **rewrite_kwargs,
+ )
if self.move_object:
hook.delete(self.source_bucket, source_object)
diff --git a/providers/google/tests/unit/google/cloud/hooks/test_gcs.py
b/providers/google/tests/unit/google/cloud/hooks/test_gcs.py
index d18aab10870..6d516889acd 100644
--- a/providers/google/tests/unit/google/cloud/hooks/test_gcs.py
+++ b/providers/google/tests/unit/google/cloud/hooks/test_gcs.py
@@ -501,6 +501,34 @@ class TestGCSHook:
destination_object=destination_object,
)
+ @pytest.mark.parametrize(
+ ("retain_until_time", "retention_mode", "expected_error"),
+ [
+ pytest.param(
+ None,
+ "Locked",
+ "retention_mode cannot be set without retain_until_time.",
+ id="mode_without_retain_until_time",
+ ),
+ pytest.param(
+ datetime(2027, 1, 1),
+ "Invalid",
+ "retention_mode must be 'Locked' or 'Unlocked'",
+ id="invalid_mode_value",
+ ),
+ ],
+ )
+ def test_rewrite_invalid_retention_params(self, retain_until_time,
retention_mode, expected_error):
+ with pytest.raises(ValueError, match=expected_error):
+ self.gcs_hook.rewrite(
+ source_bucket="test-source-bucket",
+ source_object="test-source-object",
+ destination_bucket="test-dest-bucket",
+ destination_object="test-dest-object",
+ retain_until_time=retain_until_time,
+ retention_mode=retention_mode,
+ )
+
@mock.patch(GCS_STRING.format("GCSHook.get_conn"))
def test_rewrite_exposes_lineage(self, mock_service,
hook_lineage_collector):
source_bucket_name = "test-source-bucket"
@@ -532,6 +560,104 @@ class TestGCSHook:
uri=f"gs://{destination_bucket_name}/{destination_object_name}"
)
+ @mock.patch("google.cloud.storage.Bucket")
+ @mock.patch(GCS_STRING.format("GCSHook.get_conn"))
+ def test_rewrite_with_retention(self, mock_service, mock_bucket):
+ source_bucket = "test-source-bucket"
+ source_object = "test-source-object"
+ destination_bucket = "test-dest-bucket"
+ destination_object = "test-dest-object"
+ retain_until = datetime(2027, 1, 1)
+
+ source_blob = mock_bucket.blob(source_object)
+
+ # Given
+ bucket_mock = mock_service.return_value.bucket
+ bucket_mock.return_value = mock_bucket
+ get_blob_method = bucket_mock.return_value.blob
+ destination_blob = get_blob_method.return_value
+ rewrite_method = destination_blob.rewrite
+ rewrite_method.side_effect = [(None, mock.ANY, mock.ANY)]
+
+ # When
+ self.gcs_hook.rewrite(
+ source_bucket=source_bucket,
+ source_object=source_object,
+ destination_bucket=destination_bucket,
+ destination_object=destination_object,
+ retain_until_time=retain_until,
+ retention_mode="Locked",
+ )
+
+ # Then
+ rewrite_method.assert_called_once_with(source=source_blob)
+ assert destination_blob.retention.mode == "Locked"
+ assert destination_blob.retention.retain_until_time == retain_until
+ destination_blob.patch.assert_called_once()
+
+ @mock.patch("google.cloud.storage.Bucket")
+ @mock.patch(GCS_STRING.format("GCSHook.get_conn"))
+ def test_rewrite_without_retention_does_not_patch(self, mock_service,
mock_bucket):
+ source_bucket = "test-source-bucket"
+ source_object = "test-source-object"
+ destination_bucket = "test-dest-bucket"
+ destination_object = "test-dest-object"
+
+ mock_bucket.blob(source_object)
+
+ # Given
+ bucket_mock = mock_service.return_value.bucket
+ bucket_mock.return_value = mock_bucket
+ get_blob_method = bucket_mock.return_value.blob
+ destination_blob = get_blob_method.return_value
+ rewrite_method = destination_blob.rewrite
+ rewrite_method.side_effect = [(None, mock.ANY, mock.ANY)]
+
+ # When
+ self.gcs_hook.rewrite(
+ source_bucket=source_bucket,
+ source_object=source_object,
+ destination_bucket=destination_bucket,
+ destination_object=destination_object,
+ )
+
+ # Then — no retention set, so patch should not be called
+ destination_blob.patch.assert_not_called()
+
+ @mock.patch("google.cloud.storage.Bucket")
+ @mock.patch(GCS_STRING.format("GCSHook.get_conn"))
+ def test_rewrite_with_retention_defaults_to_unlocked(self, mock_service,
mock_bucket):
+ source_bucket = "test-source-bucket"
+ source_object = "test-source-object"
+ destination_bucket = "test-dest-bucket"
+ destination_object = "test-dest-object"
+ retain_until = datetime(2027, 6, 1)
+
+ source_blob = mock_bucket.blob(source_object)
+
+ # Given
+ bucket_mock = mock_service.return_value.bucket
+ bucket_mock.return_value = mock_bucket
+ get_blob_method = bucket_mock.return_value.blob
+ destination_blob = get_blob_method.return_value
+ rewrite_method = destination_blob.rewrite
+ rewrite_method.side_effect = [(None, mock.ANY, mock.ANY)]
+
+ # When — retention_mode not specified, should default to "Unlocked"
+ self.gcs_hook.rewrite(
+ source_bucket=source_bucket,
+ source_object=source_object,
+ destination_bucket=destination_bucket,
+ destination_object=destination_object,
+ retain_until_time=retain_until,
+ )
+
+ # Then
+ rewrite_method.assert_called_once_with(source=source_blob)
+ assert destination_blob.retention.mode == "Unlocked"
+ assert destination_blob.retention.retain_until_time == retain_until
+ destination_blob.patch.assert_called_once()
+
@mock.patch("google.cloud.storage.Bucket")
@mock.patch(GCS_STRING.format("GCSHook.get_conn"))
def test_delete(self, mock_service, mock_bucket, caplog):
diff --git
a/providers/google/tests/unit/google/cloud/transfers/test_gcs_to_gcs.py
b/providers/google/tests/unit/google/cloud/transfers/test_gcs_to_gcs.py
index 1c1f4268a1b..7bbfa1c2285 100644
--- a/providers/google/tests/unit/google/cloud/transfers/test_gcs_to_gcs.py
+++ b/providers/google/tests/unit/google/cloud/transfers/test_gcs_to_gcs.py
@@ -1124,3 +1124,70 @@ class TestGoogleCloudStorageToCloudStorageOperator:
)
result = operator.execute(None)
assert result == [f"gs://{DESTINATION_BUCKET}/backup/file.txt"]
+
+ @mock.patch("airflow.providers.google.cloud.transfers.gcs_to_gcs.GCSHook")
+ def test_copy_single_file_with_retention(self, mock_hook):
+ retain_until = datetime(2027, 1, 1)
+ mock_hook.return_value.list.return_value = [SOURCE_OBJECT_NO_WILDCARD]
+ operator = GCSToGCSOperator(
+ task_id=TASK_ID,
+ source_bucket=TEST_BUCKET,
+ source_object=SOURCE_OBJECT_NO_WILDCARD,
+ destination_bucket=DESTINATION_BUCKET,
+ destination_object=SOURCE_OBJECT_NO_WILDCARD,
+ retain_until_time=retain_until,
+ retention_mode="Locked",
+ )
+
+ operator.execute(None)
+ mock_hook.return_value.rewrite.assert_called_once_with(
+ TEST_BUCKET,
+ SOURCE_OBJECT_NO_WILDCARD,
+ DESTINATION_BUCKET,
+ SOURCE_OBJECT_NO_WILDCARD,
+ retain_until_time=retain_until,
+ retention_mode="Locked",
+ )
+
+ @mock.patch("airflow.providers.google.cloud.transfers.gcs_to_gcs.GCSHook")
+ def test_copy_single_file_with_retention_without_mode(self, mock_hook):
+ retain_until = datetime(2027, 1, 1)
+ mock_hook.return_value.list.return_value = [SOURCE_OBJECT_NO_WILDCARD]
+ operator = GCSToGCSOperator(
+ task_id=TASK_ID,
+ source_bucket=TEST_BUCKET,
+ source_object=SOURCE_OBJECT_NO_WILDCARD,
+ destination_bucket=DESTINATION_BUCKET,
+ destination_object=SOURCE_OBJECT_NO_WILDCARD,
+ retain_until_time=retain_until,
+ )
+
+ operator.execute(None)
+ # When retention_mode is not set, only retain_until_time should be
passed
+ mock_hook.return_value.rewrite.assert_called_once_with(
+ TEST_BUCKET,
+ SOURCE_OBJECT_NO_WILDCARD,
+ DESTINATION_BUCKET,
+ SOURCE_OBJECT_NO_WILDCARD,
+ retain_until_time=retain_until,
+ )
+
+ @mock.patch("airflow.providers.google.cloud.transfers.gcs_to_gcs.GCSHook")
+ def test_copy_single_file_without_retention(self, mock_hook):
+ mock_hook.return_value.list.return_value = [SOURCE_OBJECT_NO_WILDCARD]
+ operator = GCSToGCSOperator(
+ task_id=TASK_ID,
+ source_bucket=TEST_BUCKET,
+ source_object=SOURCE_OBJECT_NO_WILDCARD,
+ destination_bucket=DESTINATION_BUCKET,
+ destination_object=SOURCE_OBJECT_NO_WILDCARD,
+ )
+
+ operator.execute(None)
+ # When no retention is set, rewrite should be called without retention
kwargs
+ mock_hook.return_value.rewrite.assert_called_once_with(
+ TEST_BUCKET,
+ SOURCE_OBJECT_NO_WILDCARD,
+ DESTINATION_BUCKET,
+ SOURCE_OBJECT_NO_WILDCARD,
+ )
