use the proper key to retrieve the dataflow job_id (#27336)

apache · Oct 31, 2022 · bcb026b · bcb026b
1 parent f9b1ae9
commit bcb026b
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/airflow/providers/google/cloud/example_dags/example_dataflow.py b/airflow/providers/google/cloud/example_dags/example_dataflow.py
@@ -175,7 +175,7 @@
     # [START howto_sensor_wait_for_job_status]
     wait_for_python_job_async_done = DataflowJobStatusSensor(
         task_id="wait-for-python-job-async-done",
-        job_id="{{task_instance.xcom_pull('start-python-job-async')['dataflow_job_id']}}",
+        job_id="{{task_instance.xcom_pull('start-python-job-async')['id']}}",
         expected_statuses={DataflowJobStatus.JOB_STATE_DONE},
         location="europe-west3",
     )
@@ -199,7 +199,7 @@ def callback(metrics: list[dict]) -> bool:
 
     wait_for_python_job_async_metric = DataflowJobMetricsSensor(
         task_id="wait-for-python-job-async-metric",
-        job_id="{{task_instance.xcom_pull('start-python-job-async')['dataflow_job_id']}}",
+        job_id="{{task_instance.xcom_pull('start-python-job-async')['id']}}",
         location="europe-west3",
         callback=check_metric_scalar_gte(metric_name="Service-cpu_num_seconds", value=100),
         fail_on_terminal_state=False,
@@ -216,7 +216,7 @@ def check_message(messages: list[dict]) -> bool:
 
     wait_for_python_job_async_message = DataflowJobMessagesSensor(
         task_id="wait-for-python-job-async-message",
-        job_id="{{task_instance.xcom_pull('start-python-job-async')['dataflow_job_id']}}",
+        job_id="{{task_instance.xcom_pull('start-python-job-async')['id']}}",
         location="europe-west3",
         callback=check_message,
         fail_on_terminal_state=False,
@@ -233,7 +233,7 @@ def check_autoscaling_event(autoscaling_events: list[dict]) -> bool:
 
     wait_for_python_job_async_autoscaling_event = DataflowJobAutoScalingEventsSensor(
         task_id="wait-for-python-job-async-autoscaling-event",
-        job_id="{{task_instance.xcom_pull('start-python-job-async')['dataflow_job_id']}}",
+        job_id="{{task_instance.xcom_pull('start-python-job-async')['id']}}",
         location="europe-west3",
         callback=check_autoscaling_event,
         fail_on_terminal_state=False,