Fix location requirement in DataflowTemplatedJobStartOperator

apache · Feb 2, 2024 · 1a2f9f5 · 1a2f9f5
1 parent 3ec7819
commit 1a2f9f5
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 3 deletions.
diff --git a/airflow/providers/google/cloud/operators/dataflow.py b/airflow/providers/google/cloud/operators/dataflow.py
@@ -676,6 +676,9 @@ def set_current_job(current_job):
         options = self.dataflow_default_options
         options.update(self.options)
 
+        if not self.location:
+            self.location = DEFAULT_DATAFLOW_LOCATION
+
         self.job = self.hook.start_template_dataflow(
             job_name=self.job_name,
             variables=options,
@@ -703,7 +706,7 @@ def set_current_job(current_job):
             trigger=TemplateJobStartTrigger(
                 project_id=self.project_id,
                 job_id=job_id,
-                location=self.location if self.location else DEFAULT_DATAFLOW_LOCATION,
+                location=self.location,
                 gcp_conn_id=self.gcp_conn_id,
                 poll_sleep=self.poll_sleep,
                 impersonation_chain=self.impersonation_chain,

diff --git a/tests/providers/google/cloud/operators/test_dataflow.py b/tests/providers/google/cloud/operators/test_dataflow.py
@@ -24,7 +24,10 @@
 import pytest
 
 import airflow
-from airflow.providers.google.cloud.hooks.dataflow import DataflowJobStatus
+from airflow.providers.google.cloud.hooks.dataflow import (
+    DEFAULT_DATAFLOW_LOCATION,
+    DataflowJobStatus,
+)
 from airflow.providers.google.cloud.operators.dataflow import (
     CheckJobRunning,
     DataflowCreateJavaJobOperator,
@@ -554,7 +557,7 @@ def test_start_with_custom_region(self, dataflow_mock):
         assert dataflow_mock.called
         _, kwargs = dataflow_mock.call_args_list[0]
         assert kwargs["variables"]["region"] == TEST_REGION
-        assert kwargs["location"] is None
+        assert kwargs["location"] == DEFAULT_DATAFLOW_LOCATION
 
     @pytest.mark.db_test
     @mock.patch("airflow.providers.google.cloud.operators.dataflow.DataflowHook.start_template_dataflow")