Updating Google Cloud example DAGs to use XComArgs (#16875)

apache · Jul 26, 2021 · 448e50b · 448e50b
1 parent aaf44cc
commit 448e50b
Show file tree

Hide file tree

Showing 26 changed files with 271 additions and 177 deletions.
diff --git a/airflow/providers/google/cloud/example_dags/example_automl_nl_text_classification.py b/airflow/providers/google/cloud/example_dags/example_automl_nl_text_classification.py
@@ -67,7 +67,7 @@
         task_id="create_dataset_task", dataset=DATASET, location=GCP_AUTOML_LOCATION
     )
 
-    dataset_id = '{{ task_instance.xcom_pull("create_dataset_task", key="dataset_id") }}'
+    dataset_id = create_dataset_task.output['dataset_id']
 
     import_dataset_task = AutoMLImportDataOperator(
         task_id="import_dataset_task",
@@ -80,7 +80,7 @@
 
     create_model = AutoMLTrainModelOperator(task_id="create_model", model=MODEL, location=GCP_AUTOML_LOCATION)
 
-    model_id = "{{ task_instance.xcom_pull('create_model', key='model_id') }}"
+    model_id = create_model.output['model_id']
 
     delete_model_task = AutoMLDeleteModelOperator(
         task_id="delete_model_task",
@@ -96,4 +96,10 @@
         project_id=GCP_PROJECT_ID,
     )
 
-    create_dataset_task >> import_dataset_task >> create_model >> delete_model_task >> delete_datasets_task
+    import_dataset_task >> create_model
+    delete_model_task >> delete_datasets_task
+
+    # Task dependencies created via `XComArgs`:
+    #   create_dataset_task >> import_dataset_task
+    #   create_dataset_task >> create_model
+    #   create_dataset_task >> delete_datasets_task
diff --git a/airflow/providers/google/cloud/example_dags/example_automl_nl_text_extraction.py b/airflow/providers/google/cloud/example_dags/example_automl_nl_text_extraction.py
@@ -67,7 +67,7 @@
         task_id="create_dataset_task", dataset=DATASET, location=GCP_AUTOML_LOCATION
     )
 
-    dataset_id = '{{ task_instance.xcom_pull("create_dataset_task", key="dataset_id") }}'
+    dataset_id = create_dataset_task.output['dataset_id']
 
     import_dataset_task = AutoMLImportDataOperator(
         task_id="import_dataset_task",
@@ -80,7 +80,7 @@
 
     create_model = AutoMLTrainModelOperator(task_id="create_model", model=MODEL, location=GCP_AUTOML_LOCATION)
 
-    model_id = "{{ task_instance.xcom_pull('create_model', key='model_id') }}"
+    model_id = create_model.output['model_id']
 
     delete_model_task = AutoMLDeleteModelOperator(
         task_id="delete_model_task",
@@ -96,4 +96,11 @@
         project_id=GCP_PROJECT_ID,
     )
 
-    create_dataset_task >> import_dataset_task >> create_model >> delete_model_task >> delete_datasets_task
+    import_dataset_task >> create_model
+    delete_model_task >> delete_datasets_task
+
+    # Task dependencies created via `XComArgs`:
+    #   create_dataset_task >> import_dataset_task
+    #   create_dataset_task >> create_model
+    #   create_model >> delete_model_task
+    #   create_dataset_task >> delete_datasets_task
diff --git a/airflow/providers/google/cloud/example_dags/example_automl_nl_text_sentiment.py b/airflow/providers/google/cloud/example_dags/example_automl_nl_text_sentiment.py
@@ -68,7 +68,7 @@
         task_id="create_dataset_task", dataset=DATASET, location=GCP_AUTOML_LOCATION
     )
 
-    dataset_id = '{{ task_instance.xcom_pull("create_dataset_task", key="dataset_id") }}'
+    dataset_id = create_dataset_task.output['dataset_id']
 
     import_dataset_task = AutoMLImportDataOperator(
         task_id="import_dataset_task",
@@ -81,7 +81,7 @@
 
     create_model = AutoMLTrainModelOperator(task_id="create_model", model=MODEL, location=GCP_AUTOML_LOCATION)
 
-    model_id = "{{ task_instance.xcom_pull('create_model', key='model_id') }}"
+    model_id = create_model.output['model_id']
 
     delete_model_task = AutoMLDeleteModelOperator(
         task_id="delete_model_task",
@@ -97,4 +97,11 @@
         project_id=GCP_PROJECT_ID,
     )
 
-    create_dataset_task >> import_dataset_task >> create_model >> delete_model_task >> delete_datasets_task
+    import_dataset_task >> create_model
+    delete_model_task >> delete_datasets_task
+
+    # Task dependencies created via `XComArgs`:
+    #   create_dataset_task >> import_dataset_task
+    #   create_dataset_task >> create_model
+    #   create_model >> delete_model_task
+    #   create_dataset_task >> delete_datasets_task
diff --git a/airflow/providers/google/cloud/example_dags/example_automl_tables.py b/airflow/providers/google/cloud/example_dags/example_automl_tables.py
@@ -101,7 +101,7 @@ def get_target_column_spec(columns_specs: List[Dict], column_name: str) -> str:
         project_id=GCP_PROJECT_ID,
     )
 
-    dataset_id = "{{ task_instance.xcom_pull('create_dataset_task', key='dataset_id') }}"
+    dataset_id = create_dataset_task.output['dataset_id']
     # [END howto_operator_automl_create_dataset]
 
     MODEL["dataset_id"] = dataset_id
@@ -156,7 +156,7 @@ def get_target_column_spec(columns_specs: List[Dict], column_name: str) -> str:
         project_id=GCP_PROJECT_ID,
     )
 
-    model_id = "{{ task_instance.xcom_pull('create_model_task', key='model_id') }}"
+    model_id = create_model_task.output['model_id']
     # [END howto_operator_automl_create_model]
 
     # [START howto_operator_automl_delete_model]
@@ -176,15 +176,21 @@ def get_target_column_spec(columns_specs: List[Dict], column_name: str) -> str:
     )
 
     (
-        create_dataset_task
-        >> import_dataset_task
+        import_dataset_task
         >> list_tables_spec_task
         >> list_columns_spec_task
         >> update_dataset_task
         >> create_model_task
-        >> delete_model_task
-        >> delete_datasets_task
     )
+    delete_model_task >> delete_datasets_task
+
+    # Task dependencies created via `XComArgs`:
+    #   create_dataset_task >> import_dataset_task
+    #   create_dataset_task >> list_tables_spec_task
+    #   create_dataset_task >> list_columns_spec_task
+    #   create_dataset_task >> create_model_task
+    #   create_model_task >> delete_model_task
+    #   create_dataset_task >> delete_datasets_task
 
 
 # Example DAG for AutoML datasets operations
@@ -201,7 +207,7 @@ def get_target_column_spec(columns_specs: List[Dict], column_name: str) -> str:
         project_id=GCP_PROJECT_ID,
     )
 
-    dataset_id = '{{ task_instance.xcom_pull("create_dataset_task", key="dataset_id") }}'
+    dataset_id = create_dataset_task.output['dataset_id']
 
     import_dataset_task = AutoMLImportDataOperator(
         task_id="import_dataset_task",
@@ -243,14 +249,19 @@ def get_target_column_spec(columns_specs: List[Dict], column_name: str) -> str:
     # [END howto_operator_delete_dataset]
 
     (
-        create_dataset_task
-        >> import_dataset_task
+        import_dataset_task
         >> list_tables_spec_task
         >> list_columns_spec_task
         >> list_datasets_task
         >> delete_datasets_task
     )
 
+    # Task dependencies created via `XComArgs`:
+    # create_dataset_task >> import_dataset_task
+    # create_dataset_task >> list_tables_spec_task
+    # create_dataset_task >> list_columns_spec_task
+
+
 with models.DAG(
     "example_gcp_get_deploy",
     schedule_interval=None,  # Override to match your needs

diff --git a/airflow/providers/google/cloud/example_dags/example_automl_translation.py b/airflow/providers/google/cloud/example_dags/example_automl_translation.py
@@ -74,7 +74,7 @@
         task_id="create_dataset_task", dataset=DATASET, location=GCP_AUTOML_LOCATION
     )
 
-    dataset_id = '{{ task_instance.xcom_pull("create_dataset_task", key="dataset_id") }}'
+    dataset_id = create_dataset_task.output["dataset_id"]
 
     import_dataset_task = AutoMLImportDataOperator(
         task_id="import_dataset_task",
@@ -87,7 +87,7 @@
 
     create_model = AutoMLTrainModelOperator(task_id="create_model", model=MODEL, location=GCP_AUTOML_LOCATION)
 
-    model_id = "{{ task_instance.xcom_pull('create_model', key='model_id') }}"
+    model_id = create_model.output["model_id"]
 
     delete_model_task = AutoMLDeleteModelOperator(
         task_id="delete_model_task",
@@ -103,4 +103,11 @@
         project_id=GCP_PROJECT_ID,
     )
 
-    create_dataset_task >> import_dataset_task >> create_model >> delete_model_task >> delete_datasets_task
+    import_dataset_task >> create_model
+    delete_model_task >> delete_datasets_task
+
+    # Task dependencies created via `XComArgs`:
+    #   create_dataset_task >> import_dataset_task
+    #   create_dataset_task >> create_model
+    #   create_model >> delete_model_task
+    #   create_dataset_task >> delete_datasets_task
diff --git a/...w/providers/google/cloud/example_dags/example_automl_video_intelligence_classification.py b/...w/providers/google/cloud/example_dags/example_automl_video_intelligence_classification.py
@@ -71,7 +71,7 @@
         task_id="create_dataset_task", dataset=DATASET, location=GCP_AUTOML_LOCATION
     )
 
-    dataset_id = '{{ task_instance.xcom_pull("create_dataset_task", key="dataset_id") }}'
+    dataset_id = create_dataset_task.output["dataset_id"]
 
     import_dataset_task = AutoMLImportDataOperator(
         task_id="import_dataset_task",
@@ -84,7 +84,7 @@
 
     create_model = AutoMLTrainModelOperator(task_id="create_model", model=MODEL, location=GCP_AUTOML_LOCATION)
 
-    model_id = "{{ task_instance.xcom_pull('create_model', key='model_id') }}"
+    model_id = create_model.output["model_id"]
 
     delete_model_task = AutoMLDeleteModelOperator(
         task_id="delete_model_task",
@@ -100,4 +100,11 @@
         project_id=GCP_PROJECT_ID,
     )
 
-    create_dataset_task >> import_dataset_task >> create_model >> delete_model_task >> delete_datasets_task
+    import_dataset_task >> create_model
+    delete_model_task >> delete_datasets_task
+
+    # Task dependencies created via `XComArgs`:
+    #   create_dataset_task >> import_dataset_task
+    #   create_dataset_task >> create_model
+    #   create_model >> delete_model_task
+    #   create_dataset_task >> delete_datasets_task
diff --git a/airflow/providers/google/cloud/example_dags/example_automl_video_intelligence_tracking.py b/airflow/providers/google/cloud/example_dags/example_automl_video_intelligence_tracking.py
@@ -72,7 +72,7 @@
         task_id="create_dataset_task", dataset=DATASET, location=GCP_AUTOML_LOCATION
     )
 
-    dataset_id = '{{ task_instance.xcom_pull("create_dataset_task", key="dataset_id") }}'
+    dataset_id = create_dataset_task.output["dataset_id"]
 
     import_dataset_task = AutoMLImportDataOperator(
         task_id="import_dataset_task",
@@ -85,7 +85,7 @@
 
     create_model = AutoMLTrainModelOperator(task_id="create_model", model=MODEL, location=GCP_AUTOML_LOCATION)
 
-    model_id = "{{ task_instance.xcom_pull('create_model', key='model_id') }}"
+    model_id = create_model.output["model_id"]
 
     delete_model_task = AutoMLDeleteModelOperator(
         task_id="delete_model_task",
@@ -101,4 +101,11 @@
         project_id=GCP_PROJECT_ID,
     )
 
-    create_dataset_task >> import_dataset_task >> create_model >> delete_model_task >> delete_datasets_task
+    import_dataset_task >> create_model
+    delete_model_task >> delete_datasets_task
+
+    # Task dependencies created via `XComArgs`:
+    #   create_dataset_task >> import_dataset_task
+    #   create_dataset_task >> create_model
+    #   create_model >> delete_model_task
+    #   create_dataset_task >> delete_datasets_task
diff --git a/airflow/providers/google/cloud/example_dags/example_automl_vision_classification.py b/airflow/providers/google/cloud/example_dags/example_automl_vision_classification.py
@@ -69,7 +69,7 @@
         task_id="create_dataset_task", dataset=DATASET, location=GCP_AUTOML_LOCATION
     )
 
-    dataset_id = '{{ task_instance.xcom_pull("create_dataset_task", key="dataset_id") }}'
+    dataset_id = create_dataset_task.output["dataset_id"]
 
     import_dataset_task = AutoMLImportDataOperator(
         task_id="import_dataset_task",
@@ -82,7 +82,7 @@
 
     create_model = AutoMLTrainModelOperator(task_id="create_model", model=MODEL, location=GCP_AUTOML_LOCATION)
 
-    model_id = "{{ task_instance.xcom_pull('create_model', key='model_id') }}"
+    model_id = create_model.output["model_id"]
 
     delete_model_task = AutoMLDeleteModelOperator(
         task_id="delete_model_task",
@@ -98,4 +98,11 @@
         project_id=GCP_PROJECT_ID,
     )
 
-    create_dataset_task >> import_dataset_task >> create_model >> delete_model_task >> delete_datasets_task
+    import_dataset_task >> create_model
+    delete_model_task >> delete_datasets_task
+
+    # Task dependencies created via `XComArgs`:
+    #   create_dataset_task >> import_dataset_task
+    #   create_dataset_task >> create_model
+    #   create_model >> delete_model_task
+    #   create_dataset_task >> delete_datasets_task
diff --git a/airflow/providers/google/cloud/example_dags/example_automl_vision_object_detection.py b/airflow/providers/google/cloud/example_dags/example_automl_vision_object_detection.py
@@ -71,7 +71,7 @@
         task_id="create_dataset_task", dataset=DATASET, location=GCP_AUTOML_LOCATION
     )
 
-    dataset_id = '{{ task_instance.xcom_pull("create_dataset_task", key="dataset_id") }}'
+    dataset_id = create_dataset_task.output["dataset_id"]
 
     import_dataset_task = AutoMLImportDataOperator(
         task_id="import_dataset_task",
@@ -84,7 +84,7 @@
 
     create_model = AutoMLTrainModelOperator(task_id="create_model", model=MODEL, location=GCP_AUTOML_LOCATION)
 
-    model_id = "{{ task_instance.xcom_pull('create_model', key='model_id') }}"
+    model_id = create_model.output["model_id"]
 
     delete_model_task = AutoMLDeleteModelOperator(
         task_id="delete_model_task",
@@ -100,4 +100,11 @@
         project_id=GCP_PROJECT_ID,
     )
 
-    create_dataset_task >> import_dataset_task >> create_model >> delete_model_task >> delete_datasets_task
+    import_dataset_task >> create_model
+    delete_model_task >> delete_datasets_task
+
+    # Task dependencies created via `XComArgs`:
+    #   create_dataset_task >> import_dataset_task
+    #   create_dataset_task >> create_model
+    #   create_model >> delete_model_task
+    #   create_dataset_task >> delete_datasets_task
diff --git a/airflow/providers/google/cloud/example_dags/example_azure_fileshare_to_gcs.py b/airflow/providers/google/cloud/example_dags/example_azure_fileshare_to_gcs.py
@@ -24,19 +24,18 @@
 AZURE_SHARE_NAME = os.environ.get('AZURE_SHARE_NAME', 'test-azure-share')
 AZURE_DIRECTORY_NAME = "test-azure-dir"
 
-default_args = {
-    'owner': 'airflow',
-    'depends_on_past': False,
-    'email': ['airflow@example.com'],
-    'email_on_failure': False,
-    'email_on_retry': False,
-    'retries': 1,
-    'retry_delay': timedelta(minutes=5),
-}
 
 with DAG(
     dag_id='azure_fileshare_to_gcs_example',
-    default_args=default_args,
+    default_args={
+        'owner': 'airflow',
+        'depends_on_past': False,
+        'email': ['airflow@example.com'],
+        'email_on_failure': False,
+        'email_on_retry': False,
+        'retries': 1,
+        'retry_delay': timedelta(minutes=5),
+    },
     schedule_interval=None,
     start_date=datetime(2018, 11, 1),
     tags=['example'],

diff --git a/airflow/providers/google/cloud/example_dags/example_bigquery_dts.py b/airflow/providers/google/cloud/example_dags/example_bigquery_dts.py
@@ -75,9 +75,7 @@
         task_id="gcp_bigquery_create_transfer",
     )
 
-    transfer_config_id = (
-        "{{ task_instance.xcom_pull('gcp_bigquery_create_transfer', key='transfer_config_id') }}"
-    )
+    transfer_config_id = gcp_bigquery_create_transfer.output["transfer_config_id"]
     # [END howto_bigquery_create_data_transfer]
 
     # [START howto_bigquery_start_transfer]
@@ -86,14 +84,13 @@
         transfer_config_id=transfer_config_id,
         requested_run_time={"seconds": int(time.time() + 60)},
     )
-    run_id = "{{ task_instance.xcom_pull('gcp_bigquery_start_transfer', key='run_id') }}"
     # [END howto_bigquery_start_transfer]
 
     # [START howto_bigquery_dts_sensor]
     gcp_run_sensor = BigQueryDataTransferServiceTransferRunSensor(
         task_id="gcp_run_sensor",
         transfer_config_id=transfer_config_id,
-        run_id=run_id,
+        run_id=gcp_bigquery_start_transfer.output["run_id"],
         expected_statuses={"SUCCEEDED"},
     )
     # [END howto_bigquery_dts_sensor]
@@ -104,9 +101,10 @@
     )
     # [END howto_bigquery_delete_data_transfer]
 
-    (
-        gcp_bigquery_create_transfer
-        >> gcp_bigquery_start_transfer
-        >> gcp_run_sensor
-        >> gcp_bigquery_delete_transfer
-    )
+    gcp_run_sensor >> gcp_bigquery_delete_transfer
+
+    # Task dependencies created via `XComArgs`:
+    #   gcp_bigquery_create_transfer >> gcp_bigquery_start_transfer
+    #   gcp_bigquery_create_transfer >> gcp_run_sensor
+    #   gcp_bigquery_start_transfer >> gcp_run_sensor
+    #   gcp_bigquery_create_transfer >> gcp_bigquery_delete_transfer