Pass location using parmamter in Dataflow integration (#8382)

Author: mik-laj

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

@@ -30,7 +30,6 @@
 from airflow.providers.google.cloud.operators.gcs import GCSToLocalOperator
 from airflow.utils.dates import days_ago
 
-GCP_PROJECT_ID = os.environ.get('GCP_PROJECT_ID', 'example-project')
 GCS_TMP = os.environ.get('GCP_DATAFLOW_GCS_TMP', 'gs://test-dataflow-example/temp/')
 GCS_STAGING = os.environ.get('GCP_DATAFLOW_GCS_STAGING', 'gs://test-dataflow-example/staging/')
 GCS_OUTPUT = os.environ.get('GCP_DATAFLOW_GCS_OUTPUT', 'gs://test-dataflow-example/output')
@@ -44,7 +43,6 @@
 default_args = {
     "start_date": days_ago(1),
     'dataflow_default_options': {
-        'project': GCP_PROJECT_ID,
         'tempLocation': GCS_TMP,
         'stagingLocation': GCS_STAGING,
     }
@@ -68,6 +66,7 @@
         poll_sleep=10,
         job_class='org.apache.beam.examples.WordCount',
         check_if_running=CheckJobRunning.IgnoreJob,
+        location='europe-west3'
     )
     # [END howto_operator_start_java_job]
 
@@ -104,7 +103,8 @@
             'apache-beam[gcp]>=2.14.0'
         ],
         py_interpreter='python3',
-        py_system_site_packages=False
+        py_system_site_packages=False,
+        location='europe-west3'
     )
     # [END howto_operator_start_python_job]
 
@@ -130,4 +130,5 @@
             'inputFile': "gs://dataflow-samples/shakespeare/kinglear.txt",
             'output': GCS_OUTPUT
         },
+        location='europe-west3'
     )

airflow/providers/google/cloud/hooks/dataflow.py

@@ -25,6 +25,7 @@
 import subprocess
 import time
 import uuid
+import warnings
 from copy import deepcopy
 from tempfile import TemporaryDirectory
 from typing import Any, Callable, Dict, List, Optional, TypeVar
@@ -49,36 +50,44 @@
 RT = TypeVar('RT')  # pylint: disable=invalid-name
 
 
-def _fallback_to_project_id_from_variables(func: Callable[..., RT]) -> Callable[..., RT]:
-    """
-    Decorator that provides fallback for Google Cloud Platform project id.
+def _fallback_variable_parameter(parameter_name, variable_key_name):
 
-    :param func: function to wrap
-    :return: result of the function call
-    """
-    @functools.wraps(func)
-    def inner_wrapper(self: "DataflowHook", *args, **kwargs) -> RT:
-        if args:
-            raise AirflowException(
-                "You must use keyword arguments in this methods rather than positional")
-
-        parameter_project_id = kwargs.get('project_id')
-        variables_project_id = kwargs.get('variables', {}).get('project')
-
-        if parameter_project_id and variables_project_id:
-            raise AirflowException(
-                "The mutually exclusive parameter `project_id` and `project` key in `variables` parameters "
-                "are both present. Please remove one."
-            )
+    def _wrapper(func: Callable[..., RT]) -> Callable[..., RT]:
+        """
+        Decorator that provides fallback for location from `region` key in `variables` parameters.
+
+        :param func: function to wrap
+        :return: result of the function call
+        """
+        @functools.wraps(func)
+        def inner_wrapper(self: "DataflowHook", *args, **kwargs) -> RT:
+            if args:
+                raise AirflowException(
+                    "You must use keyword arguments in this methods rather than positional")
+
+            parameter_location = kwargs.get(parameter_name)
+            variables_location = kwargs.get('variables', {}).get(variable_key_name)
+
+            if parameter_location and variables_location:
+                raise AirflowException(
+                    f"The mutually exclusive parameter `{parameter_name}` and `{variable_key_name}` key "
+                    f"in `variables` parameter are both present. Please remove one."
+                )
+            if parameter_location or variables_location:
+                kwargs[parameter_name] = parameter_location or variables_location
+            if variables_location:
+                copy_variables = deepcopy(kwargs['variables'])
+                del copy_variables[variable_key_name]
+                kwargs['variables'] = copy_variables
+
+            return func(self, *args, **kwargs)
+        return inner_wrapper
 
-        kwargs['project_id'] = parameter_project_id or variables_project_id
-        if variables_project_id:
-            copy_variables = deepcopy(kwargs['variables'])
-            del copy_variables['project']
-            kwargs['variables'] = copy_variables
+    return _wrapper
 
-        return func(self, *args, **kwargs)
-    return inner_wrapper
+
+_fallback_to_location_from_variables = _fallback_variable_parameter('location', 'region')
+_fallback_to_project_id_from_variables = _fallback_variable_parameter('project_id', 'project')
 
 
 class DataflowJobStatus:
@@ -425,9 +434,9 @@ def _start_dataflow(
         label_formatter: Callable[[Dict], List[str]],
         project_id: str,
         multiple_jobs: bool = False,
-        on_new_job_id_callback: Optional[Callable[[str], None]] = None
+        on_new_job_id_callback: Optional[Callable[[str], None]] = None,
+        location: str = DEFAULT_DATAFLOW_LOCATION
     ) -> None:
-        variables = self._set_variables(variables)
         cmd = command_prefix + self._build_cmd(variables, label_formatter, project_id)
         runner = _DataflowRunner(
             cmd=cmd,
@@ -438,20 +447,15 @@ def _start_dataflow(
             dataflow=self.get_conn(),
             project_number=project_id,
             name=name,
-            location=variables['region'],
+            location=location,
             poll_sleep=self.poll_sleep,
             job_id=job_id,
             num_retries=self.num_retries,
             multiple_jobs=multiple_jobs
         )
         job_controller.wait_for_done()
 
-    @staticmethod
-    def _set_variables(variables: Dict) -> Dict:
-        if 'region' not in variables.keys():
-            variables['region'] = DEFAULT_DATAFLOW_LOCATION
-        return variables
-
+    @_fallback_to_location_from_variables
     @_fallback_to_project_id_from_variables
     @GoogleBaseHook.fallback_to_default_project_id
     def start_java_dataflow(
@@ -463,7 +467,8 @@ def start_java_dataflow(
         job_class: Optional[str] = None,
         append_job_name: bool = True,
         multiple_jobs: bool = False,
-        on_new_job_id_callback: Optional[Callable[[str], None]] = None
+        on_new_job_id_callback: Optional[Callable[[str], None]] = None,
+        location: str = DEFAULT_DATAFLOW_LOCATION
     ) -> None:
         """
         Starts Dataflow java job.
@@ -484,9 +489,12 @@ def start_java_dataflow(
         :type multiple_jobs: bool
         :param on_new_job_id_callback: Callback called when the job ID is known.
         :type on_new_job_id_callback: callable
+        :param location: Job location.
+        :type location: str
         """
         name = self._build_dataflow_job_name(job_name, append_job_name)
         variables['jobName'] = name
+        variables['region'] = location
 
         def label_formatter(labels_dict):
             return ['--labels={}'.format(
@@ -501,9 +509,11 @@ def label_formatter(labels_dict):
             label_formatter=label_formatter,
             project_id=project_id,
             multiple_jobs=multiple_jobs,
-            on_new_job_id_callback=on_new_job_id_callback
+            on_new_job_id_callback=on_new_job_id_callback,
+            location=location
         )
 
+    @_fallback_to_location_from_variables
     @_fallback_to_project_id_from_variables
     @GoogleBaseHook.fallback_to_default_project_id
     def start_template_dataflow(
@@ -514,7 +524,8 @@ def start_template_dataflow(
         dataflow_template: str,
         project_id: str,
         append_job_name: bool = True,
-        on_new_job_id_callback: Optional[Callable[[str], None]] = None
+        on_new_job_id_callback: Optional[Callable[[str], None]] = None,
+        location: str = DEFAULT_DATAFLOW_LOCATION
     ) -> Dict:
         """
         Starts Dataflow template job.
@@ -533,8 +544,9 @@ def start_template_dataflow(
         :type append_job_name: bool
         :param on_new_job_id_callback: Callback called when the job ID is known.
         :type on_new_job_id_callback: callable
+        :param location: Job location.
+        :type location: str
         """
-        variables = self._set_variables(variables)
         name = self._build_dataflow_job_name(job_name, append_job_name)
         # Builds RuntimeEnvironment from variables dictionary
         # https://cloud.google.com/dataflow/docs/reference/rest/v1b3/RuntimeEnvironment
@@ -550,7 +562,7 @@ def start_template_dataflow(
         service = self.get_conn()
         request = service.projects().locations().templates().launch(  # pylint: disable=no-member
             projectId=project_id,
-            location=variables['region'],
+            location=location,
             gcsPath=dataflow_template,
             body=body
         )
@@ -560,18 +572,18 @@ def start_template_dataflow(
         if on_new_job_id_callback:
             on_new_job_id_callback(job_id)
 
-        variables = self._set_variables(variables)
         jobs_controller = _DataflowJobsController(
             dataflow=self.get_conn(),
             project_number=project_id,
             name=name,
             job_id=job_id,
-            location=variables['region'],
+            location=location,
             poll_sleep=self.poll_sleep,
             num_retries=self.num_retries)
         jobs_controller.wait_for_done()
         return response["job"]
 
+    @_fallback_to_location_from_variables
     @_fallback_to_project_id_from_variables
     @GoogleBaseHook.fallback_to_default_project_id
     def start_python_dataflow(  # pylint: disable=too-many-arguments
@@ -585,7 +597,8 @@ def start_python_dataflow(  # pylint: disable=too-many-arguments
         py_requirements: Optional[List[str]] = None,
         py_system_site_packages: bool = False,
         append_job_name: bool = True,
-        on_new_job_id_callback: Optional[Callable[[str], None]] = None
+        on_new_job_id_callback: Optional[Callable[[str], None]] = None,
+        location: str = DEFAULT_DATAFLOW_LOCATION
     ):
         """
         Starts Dataflow job.
@@ -620,9 +633,12 @@ def start_python_dataflow(  # pylint: disable=too-many-arguments
             If set to None or missing, the default project_id from the GCP connection is used.
         :param on_new_job_id_callback: Callback called when the job ID is known.
         :type on_new_job_id_callback: callable
+        :param location: Job location.
+        :type location: str
         """
         name = self._build_dataflow_job_name(job_name, append_job_name)
         variables['job_name'] = name
+        variables['region'] = location
 
         def label_formatter(labels_dict):
             return ['--labels={}={}'.format(key, value)
@@ -644,7 +660,8 @@ def label_formatter(labels_dict):
                     command_prefix=command_prefix,
                     label_formatter=label_formatter,
                     project_id=project_id,
-                    on_new_job_id_callback=on_new_job_id_callback
+                    on_new_job_id_callback=on_new_job_id_callback,
+                    location=location
                 )
         else:
             command_prefix = [py_interpreter] + py_options + [dataflow]
@@ -655,7 +672,8 @@ def label_formatter(labels_dict):
                 command_prefix=command_prefix,
                 label_formatter=label_formatter,
                 project_id=project_id,
-                on_new_job_id_callback=on_new_job_id_callback
+                on_new_job_id_callback=on_new_job_id_callback,
+                location=location
             )
 
     @staticmethod
@@ -700,27 +718,38 @@ def _build_cmd(variables: Dict, label_formatter: Callable, project_id: str) -> L
                 command.append(f"--{attr}={value}")
         return command
 
+    @_fallback_to_location_from_variables
     @_fallback_to_project_id_from_variables
     @GoogleBaseHook.fallback_to_default_project_id
-    def is_job_dataflow_running(self, name: str, variables: Dict, project_id: str) -> bool:
+    def is_job_dataflow_running(
+        self,
+        name: str,
+        project_id: str,
+        location: str = DEFAULT_DATAFLOW_LOCATION,
+        variables: Optional[Dict] = None
+    ) -> bool:
         """
         Helper method to check if jos is still running in dataflow
 
         :param name: The name of the job.
         :type name: str
-        :param variables: Variables passed to the job.
-        :type variables: dict
         :param project_id: Optional, the GCP project ID in which to start a job.
             If set to None or missing, the default project_id from the GCP connection is used.
+        :type project_id: str
+        :param location: Job location.
+        :type location: str
         :return: True if job is running.
         :rtype: bool
         """
-        variables = self._set_variables(variables)
+        if variables:
+            warnings.warn(
+                "The variables parameter has been deprecated. You should pass location using "
+                "the location parameter.", DeprecationWarning, stacklevel=4)
         jobs_controller = _DataflowJobsController(
             dataflow=self.get_conn(),
             project_number=project_id,
             name=name,
-            location=variables['region'],
+            location=location,
             poll_sleep=self.poll_sleep
         )
         return jobs_controller.is_job_running()
@@ -731,7 +760,7 @@ def cancel_job(
         project_id: str,
         job_name: Optional[str] = None,
         job_id: Optional[str] = None,
-        location: Optional[str] = None,
+        location: str = DEFAULT_DATAFLOW_LOCATION,
     ) -> None:
         """
         Cancels the job with the specified name prefix or Job ID.
@@ -753,7 +782,7 @@ def cancel_job(
             project_number=project_id,
             name=job_name,
             job_id=job_id,
-            location=location or DEFAULT_DATAFLOW_LOCATION,
+            location=location,
             poll_sleep=self.poll_sleep
         )
         jobs_controller.cancel()