Support impersonation service account parameter for Dataflow runner (#…

…23961)

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

@@ -51,6 +51,7 @@ class BeamDataflowMixin(metaclass=ABCMeta):
     dataflow_config: DataflowConfiguration
     gcp_conn_id: str
     delegate_to: Optional[str]
+    dataflow_support_impersonation: bool = True
 
     def _set_dataflow(
         self,
@@ -91,6 +92,13 @@ def __get_dataflow_pipeline_options(
             pipeline_options[job_name_key] = job_name
         if self.dataflow_config.service_account:
             pipeline_options["serviceAccount"] = self.dataflow_config.service_account
+        if self.dataflow_support_impersonation and self.dataflow_config.impersonation_chain:
+            if isinstance(self.dataflow_config.impersonation_chain, list):
+                pipeline_options["impersonateServiceAccount"] = ",".join(
+                    self.dataflow_config.impersonation_chain
+                )
+            else:
+                pipeline_options["impersonateServiceAccount"] = self.dataflow_config.impersonation_chain
         pipeline_options["project"] = self.dataflow_config.project_id
         pipeline_options["region"] = self.dataflow_config.location
         pipeline_options.setdefault("labels", {}).update(
@@ -550,6 +558,13 @@ def __init__(
             **kwargs,
         )
 
+        if self.dataflow_config.impersonation_chain:
+            self.log.info(
+                "Impersonation chain parameter is not supported for Apache Beam GO SDK and will be skipped "
+                "in the execution"
+            )
+        self.dataflow_support_impersonation = False
+
         self.go_file = go_file
         self.should_init_go_module = False
         self.pipeline_options.setdefault("labels", {}).update(

setup.py

@@ -202,7 +202,7 @@ def write_version(filename: str = os.path.join(*[my_dir, "airflow", "git_version
     'mypy-boto3-redshift-data>=1.21.0',
 ]
 apache_beam = [
-    'apache-beam>=2.33.0',
+    'apache-beam>=2.39.0',
 ]
 arangodb = ['python-arango>=7.3.2']
 asana = ['asana>=0.10']

tests/providers/apache/beam/operators/test_beam.py

@@ -47,6 +47,7 @@
     'output': 'gs://test/output',
     'labels': {'foo': 'bar', 'airflow-version': TEST_VERSION},
 }
+TEST_IMPERSONATION_ACCOUNT = "test@impersonation.com"
 
 
 class TestBeamRunPythonPipelineOperator(unittest.TestCase):
@@ -104,7 +105,7 @@ def test_exec_dataflow_runner(self, gcs_hook, dataflow_hook_mock, beam_hook_mock
         """Test DataflowHook is created and the right args are passed to
         start_python_dataflow.
         """
-        dataflow_config = DataflowConfiguration()
+        dataflow_config = DataflowConfiguration(impersonation_chain=TEST_IMPERSONATION_ACCOUNT)
         self.operator.runner = "DataflowRunner"
         self.operator.dataflow_config = dataflow_config
         gcs_provide_file = gcs_hook.return_value.provide_file
@@ -126,6 +127,7 @@ def test_exec_dataflow_runner(self, gcs_hook, dataflow_hook_mock, beam_hook_mock
             'output': 'gs://test/output',
             'labels': {'foo': 'bar', 'airflow-version': TEST_VERSION},
             'region': 'us-central1',
+            'impersonate_service_account': TEST_IMPERSONATION_ACCOUNT,
         }
         gcs_provide_file.assert_called_once_with(object_url=PY_FILE)
         persist_link_mock.assert_called_once_with(
@@ -223,7 +225,7 @@ def test_exec_dataflow_runner(self, gcs_hook, dataflow_hook_mock, beam_hook_mock
         """Test DataflowHook is created and the right args are passed to
         start_java_dataflow.
         """
-        dataflow_config = DataflowConfiguration()
+        dataflow_config = DataflowConfiguration(impersonation_chain="test@impersonation.com")
         self.operator.runner = "DataflowRunner"
         self.operator.dataflow_config = dataflow_config
         gcs_provide_file = gcs_hook.return_value.provide_file
@@ -248,6 +250,7 @@ def test_exec_dataflow_runner(self, gcs_hook, dataflow_hook_mock, beam_hook_mock
             'region': 'us-central1',
             'labels': {'foo': 'bar', 'airflow-version': TEST_VERSION},
             'output': 'gs://test/output',
+            'impersonateServiceAccount': TEST_IMPERSONATION_ACCOUNT,
         }
         persist_link_mock.assert_called_once_with(
             self.operator,
@@ -374,7 +377,7 @@ def test_exec_dataflow_runner(self, gcs_hook, dataflow_hook_mock, beam_hook_mock
         """Test DataflowHook is created and the right args are passed to
         start_go_dataflow.
         """
-        dataflow_config = DataflowConfiguration()
+        dataflow_config = DataflowConfiguration(impersonation_chain="test@impersonation.com")
         self.operator.runner = "DataflowRunner"
         self.operator.dataflow_config = dataflow_config
         gcs_provide_file = gcs_hook.return_value.provide_file