Support serviceAccount attr for dataflow in the Apache beam

Author: Lukasz Wyszomirski

airflow/providers/apache/beam/operators/beam.py

@@ -89,6 +89,8 @@ def __get_dataflow_pipeline_options(
         pipeline_options = copy.deepcopy(pipeline_options)
         if job_name_key is not None:
             pipeline_options[job_name_key] = job_name
+        if self.dataflow_config.service_account:
+            pipeline_options["serviceAccount"] = self.dataflow_config.service_account
         pipeline_options["project"] = self.dataflow_config.project_id
         pipeline_options["region"] = self.dataflow_config.location
         pipeline_options.setdefault("labels", {}).update(
@@ -182,6 +184,7 @@ def _init_pipeline_options(
                 pipeline_options=pipeline_options,
                 job_name_variable_key=job_name_variable_key,
             )
+            self.log.info(pipeline_options)
 
         pipeline_options.update(self.pipeline_options)
 

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

@@ -124,6 +124,7 @@ class DataflowConfiguration:
         WaitForRun = wait until job finished and the run job.
         Supported only by:
         :py:class:`~airflow.providers.apache.beam.operators.beam.BeamRunJavaPipelineOperator`
+    :param service_account: Run the job as a specific service account, instead of the default GCE robot.
     """
 
     template_fields: Sequence[str] = ("job_name", "location")
@@ -144,6 +145,7 @@ def __init__(
         wait_until_finished: Optional[bool] = None,
         multiple_jobs: Optional[bool] = None,
         check_if_running: CheckJobRunning = CheckJobRunning.WaitForRun,
+        service_account: Optional[str] = None,
     ) -> None:
         self.job_name = job_name
         self.append_job_name = append_job_name
@@ -158,6 +160,7 @@ def __init__(
         self.wait_until_finished = wait_until_finished
         self.multiple_jobs = multiple_jobs
         self.check_if_running = check_if_running
+        self.service_account = service_account
 
 
 class DataflowCreateJavaJobOperator(BaseOperator):