Add param to CloudDataTransferServiceOperator (#14118)

When a one-time job is created with `CloudDataTransferServiceS3(GCS)ToGCSOperator`, the job remains on the GCP console even after the job is completed. This is a specification of the data transfer service, but I would like to add this parameter because there are normally cases where don't want to leave a one-time job.
apache · Feb 10, 2021 · 02288cf · 02288cf
1 parent ebd39ca
commit 02288cf
Show file tree

Hide file tree

Showing 2 changed files with 135 additions and 2 deletions.
diff --git a/airflow/providers/google/cloud/operators/cloud_storage_transfer_service.py b/airflow/providers/google/cloud/operators/cloud_storage_transfer_service.py
@@ -37,6 +37,7 @@
     HTTP_DATA_SOURCE,
     MINUTES,
     MONTH,
+    NAME,
     OBJECT_CONDITIONS,
     PROJECT_ID,
     SCHEDULE,
@@ -798,7 +799,8 @@ class CloudDataTransferServiceS3ToGCSOperator(BaseOperator):
     :param transfer_options: Optional transfer service transfer options; see
         https://cloud.google.com/storage-transfer/docs/reference/rest/v1/TransferSpec
     :type transfer_options: dict
-    :param wait: Wait for transfer to finish
+    :param wait: Wait for transfer to finish. It must be set to True, if
+        'delete_job_after_completion' is set to True.
     :type wait: bool
     :param timeout: Time to wait for the operation to end in seconds. Defaults to 60 seconds if not specified.
     :type timeout: Optional[Union[float, timedelta]]
@@ -811,6 +813,9 @@ class CloudDataTransferServiceS3ToGCSOperator(BaseOperator):
         Service Account Token Creator IAM role to the directly preceding identity, with first
         account from the list granting this role to the originating account (templated).
     :type google_impersonation_chain: Union[str, Sequence[str]]
+    :param delete_job_after_completion: If True, delete the job after complete.
+        If set to True, 'wait' must be set to True.
+    :type delete_job_after_completion: bool
     """
 
     template_fields = (
@@ -840,6 +845,7 @@ def __init__(  # pylint: disable=too-many-arguments
         wait: bool = True,
         timeout: Optional[float] = None,
         google_impersonation_chain: Optional[Union[str, Sequence[str]]] = None,
+        delete_job_after_completion: bool = False,
         **kwargs,
     ) -> None:
 
@@ -857,6 +863,12 @@ def __init__(  # pylint: disable=too-many-arguments
         self.wait = wait
         self.timeout = timeout
         self.google_impersonation_chain = google_impersonation_chain
+        self.delete_job_after_completion = delete_job_after_completion
+        self._validate_inputs()
+
+    def _validate_inputs(self) -> None:
+        if self.delete_job_after_completion and not self.wait:
+            raise AirflowException("If 'delete_job_after_completion' is True, then 'wait' must also be True.")
 
     def execute(self, context) -> None:
         hook = CloudDataTransferServiceHook(
@@ -872,6 +884,8 @@ def execute(self, context) -> None:
 
         if self.wait:
             hook.wait_for_transfer_job(job, timeout=self.timeout)
+            if self.delete_job_after_completion:
+                hook.delete_transfer_job(job_name=job[NAME], project_id=self.project_id)
 
     def _create_body(self) -> dict:
         body = {
@@ -955,7 +969,8 @@ class CloudDataTransferServiceGCSToGCSOperator(BaseOperator):
     :param transfer_options: Optional transfer service transfer options; see
         https://cloud.google.com/storage-transfer/docs/reference/rest/v1/TransferSpec#TransferOptions
     :type transfer_options: dict
-    :param wait: Wait for transfer to finish; defaults to `True`
+    :param wait: Wait for transfer to finish. It must be set to True, if
+        'delete_job_after_completion' is set to True.
     :type wait: bool
     :param timeout: Time to wait for the operation to end in seconds. Defaults to 60 seconds if not specified.
     :type timeout: Optional[Union[float, timedelta]]
@@ -968,6 +983,9 @@ class CloudDataTransferServiceGCSToGCSOperator(BaseOperator):
         Service Account Token Creator IAM role to the directly preceding identity, with first
         account from the list granting this role to the originating account (templated).
     :type google_impersonation_chain: Union[str, Sequence[str]]
+    :param delete_job_after_completion: If True, delete the job after complete.
+        If set to True, 'wait' must be set to True.
+    :type delete_job_after_completion: bool
     """
 
     template_fields = (
@@ -996,6 +1014,7 @@ def __init__(  # pylint: disable=too-many-arguments
         wait: bool = True,
         timeout: Optional[float] = None,
         google_impersonation_chain: Optional[Union[str, Sequence[str]]] = None,
+        delete_job_after_completion: bool = False,
         **kwargs,
     ) -> None:
 
@@ -1012,6 +1031,12 @@ def __init__(  # pylint: disable=too-many-arguments
         self.wait = wait
         self.timeout = timeout
         self.google_impersonation_chain = google_impersonation_chain
+        self.delete_job_after_completion = delete_job_after_completion
+        self._validate_inputs()
+
+    def _validate_inputs(self) -> None:
+        if self.delete_job_after_completion and not self.wait:
+            raise AirflowException("If 'delete_job_after_completion' is True, then 'wait' must also be True.")
 
     def execute(self, context) -> None:
         hook = CloudDataTransferServiceHook(
@@ -1028,6 +1053,8 @@ def execute(self, context) -> None:
 
         if self.wait:
             hook.wait_for_transfer_job(job, timeout=self.timeout)
+            if self.delete_job_after_completion:
+                hook.delete_transfer_job(job_name=job[NAME], project_id=self.project_id)
 
     def _create_body(self) -> dict:
         body = {

diff --git a/tests/providers/google/cloud/operators/test_cloud_storage_transfer_service.py b/tests/providers/google/cloud/operators/test_cloud_storage_transfer_service.py
@@ -778,6 +778,7 @@ def test_execute(self, mock_aws_hook, mock_transfer_hook):
         )
 
         assert mock_transfer_hook.return_value.wait_for_transfer_job.called
+        assert not mock_transfer_hook.return_value.delete_transfer_job.called
 
     @mock.patch(
         'airflow.providers.google.cloud.operators.cloud_storage_transfer_service.CloudDataTransferServiceHook'
@@ -804,6 +805,63 @@ def test_execute_skip_wait(self, mock_aws_hook, mock_transfer_hook):
         )
 
         assert not mock_transfer_hook.return_value.wait_for_transfer_job.called
+        assert not mock_transfer_hook.return_value.delete_transfer_job.called
+
+    @mock.patch(
+        'airflow.providers.google.cloud.operators.cloud_storage_transfer_service.CloudDataTransferServiceHook'
+    )
+    @mock.patch('airflow.providers.google.cloud.operators.cloud_storage_transfer_service.AwsBaseHook')
+    def test_execute_delete_job_after_completion(self, mock_aws_hook, mock_transfer_hook):
+        mock_aws_hook.return_value.get_credentials.return_value = Credentials(
+            TEST_AWS_ACCESS_KEY_ID, TEST_AWS_ACCESS_SECRET, None
+        )
+
+        operator = CloudDataTransferServiceS3ToGCSOperator(
+            task_id=TASK_ID,
+            s3_bucket=AWS_BUCKET_NAME,
+            gcs_bucket=GCS_BUCKET_NAME,
+            description=DESCRIPTION,
+            schedule=SCHEDULE_DICT,
+            wait=True,
+            delete_job_after_completion=True,
+        )
+
+        operator.execute(None)
+
+        mock_transfer_hook.return_value.create_transfer_job.assert_called_once_with(
+            body=VALID_TRANSFER_JOB_AWS_RAW
+        )
+
+        assert mock_transfer_hook.return_value.wait_for_transfer_job.called
+        assert mock_transfer_hook.return_value.delete_transfer_job.called
+
+    @mock.patch(
+        'airflow.providers.google.cloud.operators.cloud_storage_transfer_service.CloudDataTransferServiceHook'
+    )
+    @mock.patch('airflow.providers.google.cloud.operators.cloud_storage_transfer_service.AwsBaseHook')
+    def test_execute_should_throw_ex_when_delete_job_without_wait(self, mock_aws_hook, mock_transfer_hook):
+        mock_aws_hook.return_value.get_credentials.return_value = Credentials(
+            TEST_AWS_ACCESS_KEY_ID, TEST_AWS_ACCESS_SECRET, None
+        )
+
+        with pytest.raises(AirflowException) as ctx:
+
+            operator = CloudDataTransferServiceS3ToGCSOperator(
+                task_id=TASK_ID,
+                s3_bucket=AWS_BUCKET_NAME,
+                gcs_bucket=GCS_BUCKET_NAME,
+                description=DESCRIPTION,
+                schedule=SCHEDULE_DICT,
+                wait=False,
+                delete_job_after_completion=True,
+            )
+
+            operator.execute(None)
+
+        err = ctx.value
+        assert "If 'delete_job_after_completion' is True, then 'wait' must also be True." in str(err)
+        mock_aws_hook.assert_not_called()
+        mock_transfer_hook.assert_not_called()
 
 
 class TestGoogleCloudStorageToGoogleCloudStorageTransferOperator(unittest.TestCase):
@@ -873,6 +931,7 @@ def test_execute(self, mock_transfer_hook):
             body=VALID_TRANSFER_JOB_GCS_RAW
         )
         assert mock_transfer_hook.return_value.wait_for_transfer_job.called
+        assert not mock_transfer_hook.return_value.delete_transfer_job.called
 
     @mock.patch(
         'airflow.providers.google.cloud.operators.cloud_storage_transfer_service.CloudDataTransferServiceHook'
@@ -893,3 +952,50 @@ def test_execute_skip_wait(self, mock_transfer_hook):
             body=VALID_TRANSFER_JOB_GCS_RAW
         )
         assert not mock_transfer_hook.return_value.wait_for_transfer_job.called
+        assert not mock_transfer_hook.return_value.delete_transfer_job.called
+
+    @mock.patch(
+        'airflow.providers.google.cloud.operators.cloud_storage_transfer_service.CloudDataTransferServiceHook'
+    )
+    def test_execute_delete_job_after_completion(self, mock_transfer_hook):
+
+        operator = CloudDataTransferServiceGCSToGCSOperator(
+            task_id=TASK_ID,
+            source_bucket=GCS_BUCKET_NAME,
+            destination_bucket=GCS_BUCKET_NAME,
+            description=DESCRIPTION,
+            schedule=SCHEDULE_DICT,
+            wait=True,
+            delete_job_after_completion=True,
+        )
+
+        operator.execute(None)
+
+        mock_transfer_hook.return_value.create_transfer_job.assert_called_once_with(
+            body=VALID_TRANSFER_JOB_GCS_RAW
+        )
+        assert mock_transfer_hook.return_value.wait_for_transfer_job.called
+        assert mock_transfer_hook.return_value.delete_transfer_job.called
+
+    @mock.patch(
+        'airflow.providers.google.cloud.operators.cloud_storage_transfer_service.CloudDataTransferServiceHook'
+    )
+    def test_execute_should_throw_ex_when_delete_job_without_wait(self, mock_transfer_hook):
+
+        with pytest.raises(AirflowException) as ctx:
+
+            operator = CloudDataTransferServiceS3ToGCSOperator(
+                task_id=TASK_ID,
+                s3_bucket=AWS_BUCKET_NAME,
+                gcs_bucket=GCS_BUCKET_NAME,
+                description=DESCRIPTION,
+                schedule=SCHEDULE_DICT,
+                wait=False,
+                delete_job_after_completion=True,
+            )
+
+            operator.execute(None)
+
+        err = ctx.value
+        assert "If 'delete_job_after_completion' is True, then 'wait' must also be True." in str(err)
+        mock_transfer_hook.assert_not_called()