[AIRFLOW-7099] Improve system test for cloud transfer service (#7794)

Author: turbaszek

airflow/providers/google/cloud/example_dags/example_cloud_storage_transfer_service.py renamed to airflow/providers/google/cloud/example_dags/example_cloud_storage_transfer_service_aws.py

@@ -28,8 +28,6 @@
   .. warning::
     You need to provide a large enough set of data so that operations do not execute too quickly.
     Otherwise, DAG will fail.
-* GCP_TRANSFER_FIRST_TARGET_BUCKET - Google Cloud Storage bucket to which files are copied from AWS.
-  It is also a source bucket in next step
 * GCP_TRANSFER_SECOND_TARGET_BUCKET - Google Cloud Storage bucket bucket to which files are copied
 * WAIT_FOR_OPERATION_POKE_INTERVAL - interval of what to check the status of the operation
   A smaller value than the default value accelerates the system test and ensures its correct execution with
@@ -40,20 +38,19 @@
 
 import os
 from datetime import datetime, timedelta
-from typing import Any, Dict
 
 from airflow import models
 from airflow.providers.google.cloud.hooks.cloud_storage_transfer_service import (
     ALREADY_EXISTING_IN_SINK, AWS_S3_DATA_SOURCE, BUCKET_NAME, DESCRIPTION, FILTER_JOB_NAMES,
-    FILTER_PROJECT_ID, GCS_DATA_SINK, GCS_DATA_SOURCE, PROJECT_ID, SCHEDULE, SCHEDULE_END_DATE,
-    SCHEDULE_START_DATE, START_TIME_OF_DAY, STATUS, TRANSFER_JOB, TRANSFER_JOB_FIELD_MASK, TRANSFER_OPTIONS,
-    TRANSFER_SPEC, GcpTransferJobsStatus, GcpTransferOperationStatus,
+    FILTER_PROJECT_ID, GCS_DATA_SINK, PROJECT_ID, SCHEDULE, SCHEDULE_END_DATE, SCHEDULE_START_DATE,
+    START_TIME_OF_DAY, STATUS, TRANSFER_OPTIONS, TRANSFER_SPEC, GcpTransferJobsStatus,
+    GcpTransferOperationStatus,
 )
 from airflow.providers.google.cloud.operators.cloud_storage_transfer_service import (
     CloudDataTransferServiceCancelOperationOperator, CloudDataTransferServiceCreateJobOperator,
     CloudDataTransferServiceDeleteJobOperator, CloudDataTransferServiceGetOperationOperator,
     CloudDataTransferServiceListOperationsOperator, CloudDataTransferServicePauseOperationOperator,
-    CloudDataTransferServiceResumeOperationOperator, CloudDataTransferServiceUpdateJobOperator,
+    CloudDataTransferServiceResumeOperationOperator,
 )
 from airflow.providers.google.cloud.sensors.cloud_storage_transfer_service import (
     CloudDataTransferServiceJobStatusSensor,
@@ -69,9 +66,6 @@
 GCP_TRANSFER_FIRST_TARGET_BUCKET = os.environ.get(
     'GCP_TRANSFER_FIRST_TARGET_BUCKET', 'gcp-transfer-first-target'
 )
-GCP_TRANSFER_SECOND_TARGET_BUCKET = os.environ.get(
-    'GCP_TRANSFER_SECOND_TARGET_BUCKET', 'gcp-transfer-second-target'
-)
 
 # [START howto_operator_gcp_transfer_create_job_body_aws]
 aws_to_gcs_transfer_body = {
@@ -91,40 +85,13 @@
 }
 # [END howto_operator_gcp_transfer_create_job_body_aws]
 
-# [START howto_operator_gcp_transfer_create_job_body_gcp]
-gcs_to_gcs_transfer_body = {
-    DESCRIPTION: GCP_DESCRIPTION,
-    STATUS: GcpTransferJobsStatus.ENABLED,
-    PROJECT_ID: GCP_PROJECT_ID,
-    SCHEDULE: {
-        SCHEDULE_START_DATE: datetime(2015, 1, 1).date(),
-        SCHEDULE_END_DATE: datetime(2030, 1, 1).date(),
-        START_TIME_OF_DAY: (datetime.utcnow() + timedelta(minutes=2)).time(),
-    },
-    TRANSFER_SPEC: {
-        GCS_DATA_SOURCE: {BUCKET_NAME: GCP_TRANSFER_FIRST_TARGET_BUCKET},
-        GCS_DATA_SINK: {BUCKET_NAME: GCP_TRANSFER_SECOND_TARGET_BUCKET},
-        TRANSFER_OPTIONS: {ALREADY_EXISTING_IN_SINK: True},
-    },
-}  # type: Dict[str, Any]
-# [END howto_operator_gcp_transfer_create_job_body_gcp]
-
-# [START howto_operator_gcp_transfer_update_job_body]
-update_body = {
-    PROJECT_ID: GCP_PROJECT_ID,
-    TRANSFER_JOB: {DESCRIPTION: "{}_updated".format(GCP_DESCRIPTION)},
-    TRANSFER_JOB_FIELD_MASK: "description",
-}
-# [END howto_operator_gcp_transfer_update_job_body]
-
-list_filter_dict = {FILTER_PROJECT_ID: GCP_PROJECT_ID, FILTER_JOB_NAMES: []}
 
 # [START howto_operator_gcp_transfer_default_args]
 default_args = {'start_date': days_ago(1)}
 # [END howto_operator_gcp_transfer_default_args]
 
 with models.DAG(
-    'example_gcp_transfer',
+    'example_gcp_transfer_aws',
     default_args=default_args,
     schedule_interval=None,  # Override to match your needs
     tags=['example'],
@@ -152,14 +119,6 @@
     )
     # [END howto_operator_gcp_transfer_pause_operation]
 
-    # [START howto_operator_gcp_transfer_update_job]
-    update_job = CloudDataTransferServiceUpdateJobOperator(
-        task_id="update_job",
-        job_name="{{task_instance.xcom_pull('create_transfer_job_from_aws')['name']}}",
-        body=update_body,
-    )
-    # [END howto_operator_gcp_transfer_update_job]
-
     # [START howto_operator_gcp_transfer_list_operations]
     list_operations = CloudDataTransferServiceListOperationsOperator(
         task_id="list_operations",
@@ -192,22 +151,6 @@
     )
     # [END howto_operator_gcp_transfer_wait_operation]
 
-    job_time = datetime.utcnow() + timedelta(minutes=2)
-
-    gcs_to_gcs_transfer_body['schedule']['startTimeOfDay'] = (datetime.utcnow() + timedelta(minutes=2)).time()
-
-    create_transfer_job_from_gcp = CloudDataTransferServiceCreateJobOperator(
-        task_id="create_transfer_job_from_gcp", body=gcs_to_gcs_transfer_body
-    )
-
-    wait_for_second_operation_to_start = CloudDataTransferServiceJobStatusSensor(
-        task_id="wait_for_second_operation_to_start",
-        job_name="{{ task_instance.xcom_pull('create_transfer_job_from_gcp')['name'] }}",
-        project_id=GCP_PROJECT_ID,
-        expected_statuses={GcpTransferOperationStatus.IN_PROGRESS},
-        poke_interval=WAIT_FOR_OPERATION_POKE_INTERVAL,
-    )
-
     # [START howto_operator_gcp_transfer_cancel_operation]
     cancel_operation = CloudDataTransferServiceCancelOperationOperator(
         task_id="cancel_operation",
@@ -224,13 +167,6 @@
     )
     # [END howto_operator_gcp_transfer_delete_job]
 
-    delete_transfer_from_gcp_job = CloudDataTransferServiceDeleteJobOperator(
-        task_id="delete_transfer_from_gcp_job",
-        job_name="{{task_instance.xcom_pull('create_transfer_job_from_gcp')['name']}}",
-        project_id=GCP_PROJECT_ID,
-    )
-
     create_transfer_job_from_aws >> wait_for_operation_to_start >> pause_operation >> \
         list_operations >> get_operation >> resume_operation >> wait_for_operation_to_end >> \
-        create_transfer_job_from_gcp >> wait_for_second_operation_to_start >> cancel_operation >> \
-        delete_transfer_from_aws_job >> delete_transfer_from_gcp_job
+        cancel_operation >> delete_transfer_from_aws_job

airflow/providers/google/cloud/example_dags/example_cloud_storage_transfer_service_gcp.py

@@ -0,0 +1,135 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Example Airflow DAG that demonstrates interactions with Google Cloud Transfer.
+
+
+This DAG relies on the following OS environment variables
+
+* GCP_PROJECT_ID - Google Cloud Project to use for the Google Cloud Transfer Service.
+* GCP_TRANSFER_FIRST_TARGET_BUCKET - Google Cloud Storage bucket to which files are copied from AWS.
+  It is also a source bucket in next step
+* GCP_TRANSFER_SECOND_TARGET_BUCKET - Google Cloud Storage bucket bucket to which files are copied
+"""
+
+import os
+from datetime import datetime, timedelta
+
+from airflow import models
+from airflow.providers.google.cloud.hooks.cloud_storage_transfer_service import (
+    ALREADY_EXISTING_IN_SINK, BUCKET_NAME, DESCRIPTION, FILTER_JOB_NAMES, FILTER_PROJECT_ID, GCS_DATA_SINK,
+    GCS_DATA_SOURCE, PROJECT_ID, SCHEDULE, SCHEDULE_END_DATE, SCHEDULE_START_DATE, START_TIME_OF_DAY, STATUS,
+    TRANSFER_JOB, TRANSFER_JOB_FIELD_MASK, TRANSFER_OPTIONS, TRANSFER_SPEC, GcpTransferJobsStatus,
+    GcpTransferOperationStatus,
+)
+from airflow.providers.google.cloud.operators.cloud_storage_transfer_service import (
+    CloudDataTransferServiceCreateJobOperator, CloudDataTransferServiceDeleteJobOperator,
+    CloudDataTransferServiceGetOperationOperator, CloudDataTransferServiceListOperationsOperator,
+    CloudDataTransferServiceUpdateJobOperator,
+)
+from airflow.providers.google.cloud.sensors.cloud_storage_transfer_service import (
+    CloudDataTransferServiceJobStatusSensor,
+)
+from airflow.utils.dates import days_ago
+
+GCP_PROJECT_ID = os.environ.get("GCP_PROJECT_ID", "example-project")
+GCP_TRANSFER_FIRST_TARGET_BUCKET = os.environ.get(
+    "GCP_TRANSFER_FIRST_TARGET_BUCKET", "gcp-transfer-first-target"
+)
+GCP_TRANSFER_SECOND_TARGET_BUCKET = os.environ.get(
+    "GCP_TRANSFER_SECOND_TARGET_BUCKET", "gcp-transfer-second-target"
+)
+
+# [START howto_operator_gcp_transfer_create_job_body_gcp]
+gcs_to_gcs_transfer_body = {
+    DESCRIPTION: "description",
+    STATUS: GcpTransferJobsStatus.ENABLED,
+    PROJECT_ID: GCP_PROJECT_ID,
+    SCHEDULE: {
+        SCHEDULE_START_DATE: datetime(2015, 1, 1).date(),
+        SCHEDULE_END_DATE: datetime(2030, 1, 1).date(),
+        START_TIME_OF_DAY: (datetime.utcnow() + timedelta(seconds=120)).time(),
+    },
+    TRANSFER_SPEC: {
+        GCS_DATA_SOURCE: {BUCKET_NAME: GCP_TRANSFER_FIRST_TARGET_BUCKET},
+        GCS_DATA_SINK: {BUCKET_NAME: GCP_TRANSFER_SECOND_TARGET_BUCKET},
+        TRANSFER_OPTIONS: {ALREADY_EXISTING_IN_SINK: True},
+    },
+}
+# [END howto_operator_gcp_transfer_create_job_body_gcp]
+
+# [START howto_operator_gcp_transfer_update_job_body]
+update_body = {
+    PROJECT_ID: GCP_PROJECT_ID,
+    TRANSFER_JOB: {DESCRIPTION: "description_updated"},
+    TRANSFER_JOB_FIELD_MASK: "description",
+}
+# [END howto_operator_gcp_transfer_update_job_body]
+
+default_args = {"start_date": days_ago(1)}
+
+with models.DAG(
+    "example_gcp_transfer",
+    default_args=default_args,
+    schedule_interval=None,  # Override to match your needs
+    tags=["example"],
+) as dag:
+
+    create_transfer = CloudDataTransferServiceCreateJobOperator(
+        task_id="create_transfer", body=gcs_to_gcs_transfer_body
+    )
+
+    # [START howto_operator_gcp_transfer_update_job]
+    update_transfer = CloudDataTransferServiceUpdateJobOperator(
+        task_id="update_transfer",
+        job_name="{{task_instance.xcom_pull('create_transfer')['name']}}",
+        body=update_body,
+    )
+    # [END howto_operator_gcp_transfer_update_job]
+
+    wait_for_transfer = CloudDataTransferServiceJobStatusSensor(
+        task_id="wait_for_transfer",
+        job_name="{{task_instance.xcom_pull('create_transfer')['name']}}",
+        project_id=GCP_PROJECT_ID,
+        expected_statuses={GcpTransferOperationStatus.SUCCESS},
+    )
+
+    list_operations = CloudDataTransferServiceListOperationsOperator(
+        task_id="list_operations",
+        request_filter={
+            FILTER_PROJECT_ID: GCP_PROJECT_ID,
+            FILTER_JOB_NAMES: [
+                "{{task_instance.xcom_pull('create_transfer')['name']}}"
+            ],
+        },
+    )
+
+    get_operation = CloudDataTransferServiceGetOperationOperator(
+        task_id="get_operation",
+        operation_name="{{task_instance.xcom_pull('list_operations')[0]['name']}}",
+    )
+
+    delete_transfer = CloudDataTransferServiceDeleteJobOperator(
+        task_id="delete_transfer_from_gcp_job",
+        job_name="{{task_instance.xcom_pull('create_transfer')['name']}}",
+        project_id=GCP_PROJECT_ID,
+    )
+
+    create_transfer >> wait_for_transfer >> update_transfer >> \
+        list_operations >> get_operation >> delete_transfer

docs/howto/operator/gcp/cloud_storage_transfer_service.rst

@@ -62,17 +62,17 @@ For parameter definition, take a look at
 Using the operator
 """"""""""""""""""
 
-.. exampleinclude:: ../../../../airflow/providers/google/cloud/example_dags/example_cloud_storage_transfer_service.py
+.. exampleinclude:: ../../../../airflow/providers/google/cloud/example_dags/example_cloud_storage_transfer_service_gcp.py
       :language: python
       :start-after: [START howto_operator_gcp_transfer_create_job_body_gcp]
       :end-before: [END howto_operator_gcp_transfer_create_job_body_gcp]
 
-.. exampleinclude:: ../../../../airflow/providers/google/cloud/example_dags/example_cloud_storage_transfer_service.py
+.. exampleinclude:: ../../../../airflow/providers/google/cloud/example_dags/example_cloud_storage_transfer_service_aws.py
       :language: python
       :start-after: [START howto_operator_gcp_transfer_create_job_body_aws]
       :end-before: [END howto_operator_gcp_transfer_create_job_body_aws]
 
-.. exampleinclude:: ../../../../airflow/providers/google/cloud/example_dags/example_cloud_storage_transfer_service.py
+.. exampleinclude:: ../../../../airflow/providers/google/cloud/example_dags/example_cloud_storage_transfer_service_aws.py
       :language: python
       :dedent: 4
       :start-after: [START howto_operator_gcp_transfer_create_job]
@@ -107,7 +107,7 @@ For parameter definition, take a look at
 Using the operator
 """"""""""""""""""
 
-.. exampleinclude:: ../../../../airflow/providers/google/cloud/example_dags/example_cloud_storage_transfer_service.py
+.. exampleinclude:: ../../../../airflow/providers/google/cloud/example_dags/example_cloud_storage_transfer_service_aws.py
       :language: python
       :dedent: 4
       :start-after: [START howto_operator_gcp_transfer_delete_job]
@@ -142,12 +142,12 @@ For parameter definition, take a look at
 Using the operator
 """"""""""""""""""
 
-.. exampleinclude:: ../../../../airflow/providers/google/cloud/example_dags/example_cloud_storage_transfer_service.py
+.. exampleinclude:: ../../../../airflow/providers/google/cloud/example_dags/example_cloud_storage_transfer_service_gcp.py
       :language: python
       :start-after: [START howto_operator_gcp_transfer_update_job_body]
       :end-before: [END howto_operator_gcp_transfer_update_job_body]
 
-.. exampleinclude:: ../../../../airflow/providers/google/cloud/example_dags/example_cloud_storage_transfer_service.py
+.. exampleinclude:: ../../../../airflow/providers/google/cloud/example_dags/example_cloud_storage_transfer_service_gcp.py
       :language: python
       :dedent: 4
       :start-after: [START howto_operator_gcp_transfer_update_job]
@@ -181,7 +181,7 @@ For parameter definition, take a look at
 Using the operator
 """"""""""""""""""
 
-.. exampleinclude:: ../../../../airflow/providers/google/cloud/example_dags/example_cloud_storage_transfer_service.py
+.. exampleinclude:: ../../../../airflow/providers/google/cloud/example_dags/example_cloud_storage_transfer_service_aws.py
       :language: python
       :dedent: 4
       :start-after: [START howto_operator_gcp_transfer_cancel_operation]
@@ -217,7 +217,7 @@ For parameter definition, take a look at
 Using the operator
 """"""""""""""""""
 
-.. exampleinclude:: ../../../../airflow/providers/google/cloud/example_dags/example_cloud_storage_transfer_service.py
+.. exampleinclude:: ../../../../airflow/providers/google/cloud/example_dags/example_cloud_storage_transfer_service_aws.py
       :language: python
       :dedent: 4
       :start-after: [START howto_operator_gcp_transfer_get_operation]
@@ -252,7 +252,7 @@ For parameter definition, take a look at
 Using the operator
 """"""""""""""""""
 
-.. exampleinclude:: ../../../../airflow/providers/google/cloud/example_dags/example_cloud_storage_transfer_service.py
+.. exampleinclude:: ../../../../airflow/providers/google/cloud/example_dags/example_cloud_storage_transfer_service_aws.py
       :language: python
       :dedent: 4
       :start-after: [START howto_operator_gcp_transfer_list_operations]
@@ -286,7 +286,7 @@ For parameter definition, take a look at
 Using the operator
 """"""""""""""""""
 
-.. exampleinclude:: ../../../../airflow/providers/google/cloud/example_dags/example_cloud_storage_transfer_service.py
+.. exampleinclude:: ../../../../airflow/providers/google/cloud/example_dags/example_cloud_storage_transfer_service_aws.py
       :language: python
       :dedent: 4
       :start-after: [START howto_operator_gcp_transfer_pause_operation]
@@ -320,7 +320,7 @@ For parameter definition, take a look at
 Using the operator
 """"""""""""""""""
 
-.. exampleinclude:: ../../../../airflow/providers/google/cloud/example_dags/example_cloud_storage_transfer_service.py
+.. exampleinclude:: ../../../../airflow/providers/google/cloud/example_dags/example_cloud_storage_transfer_service_aws.py
       :language: python
       :dedent: 4
       :start-after: [START howto_operator_gcp_transfer_resume_operation]
@@ -355,7 +355,7 @@ For parameter definition, take a look at
 Using the operator
 """"""""""""""""""
 
-.. exampleinclude:: ../../../../airflow/providers/google/cloud/example_dags/example_cloud_storage_transfer_service.py
+.. exampleinclude:: ../../../../airflow/providers/google/cloud/example_dags/example_cloud_storage_transfer_service_aws.py
       :language: python
       :dedent: 4
       :start-after: [START howto_operator_gcp_transfer_wait_operation]