Add optional object-level retention support to GcsToGcsOperator

providers/google/src/airflow/providers/google/cloud/hooks/gcs.py

@@ -222,6 +222,8 @@ def rewrite(
         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.
@@ -233,6 +235,13 @@ def rewrite(
         :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:
@@ -242,24 +251,40 @@ def rewrite(
             )
         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",

providers/google/src/airflow/providers/google/cloud/transfers/gcs_to_gcs.py

@@ -31,6 +31,9 @@
 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 @@ def __init__(
         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 @@ def __init__(
         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(
@@ -566,7 +581,18 @@ def _copy_single_object(self, hook, source_object, destination_object) -> str |
             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)

providers/google/tests/unit/google/cloud/hooks/test_gcs.py

@@ -495,6 +495,34 @@ def test_rewrite_empty_source_object(self):
                 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"
@@ -526,6 +554,104 @@ def test_rewrite_exposes_lineage(self, mock_service, hook_lineage_collector):
             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):

providers/google/tests/unit/google/cloud/transfers/test_gcs_to_gcs.py

@@ -1062,3 +1062,70 @@ def test_execute_excludes_directory_uri_from_return(self, mock_hook):
         )
         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,
+        )