This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 71a1f29213e Fix dataflow_config service_account for Beam Python and Go
pipelines (#73236)
71a1f29213e is described below
commit 71a1f29213e244f9f7b351d9f08df11c47cd259e
Author: PoAn Yang <[email protected]>
AuthorDate: Mon Sep 21 01:15:32 2026 +0800
Fix dataflow_config service_account for Beam Python and Go pipelines
(#73236)
Signed-off-by: PoAn Yang <[email protected]>
---
.../src/airflow/providers/apache/beam/operators/beam.py | 7 ++++++-
.../beam/tests/unit/apache/beam/operators/test_beam.py | 17 ++++++++++++++---
2 files changed, 20 insertions(+), 4 deletions(-)
diff --git
a/providers/apache/beam/src/airflow/providers/apache/beam/operators/beam.py
b/providers/apache/beam/src/airflow/providers/apache/beam/operators/beam.py
index e840940dad5..06de0c0e5b2 100644
--- a/providers/apache/beam/src/airflow/providers/apache/beam/operators/beam.py
+++ b/providers/apache/beam/src/airflow/providers/apache/beam/operators/beam.py
@@ -82,6 +82,9 @@ class BeamDataflowMixin(metaclass=ABCMeta):
dataflow_config: DataflowConfiguration
gcp_conn_id: str
dataflow_support_impersonation: bool = True
+ # Beam Java reads --serviceAccount, but the Python and Go SDKs read
--service_account_email.
+ # Java and Go reject any flag they do not define, so each operator sends
only the name its SDK reads.
+ dataflow_service_account_option: str = "serviceAccount"
def __init__(self):
if not GOOGLE_PROVIDER:
@@ -129,7 +132,7 @@ class BeamDataflowMixin(metaclass=ABCMeta):
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[self.dataflow_service_account_option] =
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(
@@ -351,6 +354,7 @@ class
BeamRunPythonPipelineOperator(BeamBasePipelineOperator):
)
template_fields_renderers = {"dataflow_config": "json",
"pipeline_options": "json"}
operator_extra_links = (DataflowJobLink(),) if GOOGLE_PROVIDER else ()
+ dataflow_service_account_option = "service_account_email"
def __init__(
self,
@@ -756,6 +760,7 @@ class BeamRunGoPipelineOperator(BeamBasePipelineOperator):
]
template_fields_renderers = {"dataflow_config": "json",
"pipeline_options": "json"}
operator_extra_links = (DataflowJobLink(),) if GOOGLE_PROVIDER else ()
+ dataflow_service_account_option = "service_account_email"
def __init__(
self,
diff --git
a/providers/apache/beam/tests/unit/apache/beam/operators/test_beam.py
b/providers/apache/beam/tests/unit/apache/beam/operators/test_beam.py
index 08176afab51..7e8fe2a1d51 100644
--- a/providers/apache/beam/tests/unit/apache/beam/operators/test_beam.py
+++ b/providers/apache/beam/tests/unit/apache/beam/operators/test_beam.py
@@ -54,6 +54,7 @@ OUTPUT_LOCATION = "gs://test/output"
TEST_VERSION = f"v{version.replace('.', '-').replace('+', '-')}"
TEST_IMPERSONATION_ACCOUNT = "[email protected]"
TEST_POLL_SLEEP = 30
+TEST_SERVICE_ACCOUNT = "[email protected]"
BEAM_OPERATOR_PATH = "airflow.providers.apache.beam.operators.beam.{}"
@@ -219,7 +220,10 @@ class TestBeamRunPythonPipelineOperator:
"""
gcs_provide_file = gcs_hook.return_value.provide_file
op = BeamRunPythonPipelineOperator(
- dataflow_config={"impersonation_chain":
TEST_IMPERSONATION_ACCOUNT},
+ dataflow_config={
+ "impersonation_chain": TEST_IMPERSONATION_ACCOUNT,
+ "service_account": TEST_SERVICE_ACCOUNT,
+ },
runner="DataflowRunner",
**self.default_op_kwargs,
)
@@ -244,6 +248,7 @@ class TestBeamRunPythonPipelineOperator:
"output": "gs://test/output",
"labels": {"foo": "bar", "airflow-version": TEST_VERSION},
"region": "us-central1",
+ "service_account_email": TEST_SERVICE_ACCOUNT,
"impersonate_service_account": TEST_IMPERSONATION_ACCOUNT,
"requirements_file":
gcs_provide_file.return_value.__enter__.return_value.name,
}
@@ -446,7 +451,9 @@ class TestBeamRunJavaPipelineOperator:
"""Test DataflowHook is created and the right args are passed to
start_java_dataflow.
"""
- dataflow_config =
DataflowConfiguration(impersonation_chain="[email protected]")
+ dataflow_config = DataflowConfiguration(
+ impersonation_chain="[email protected]",
service_account=TEST_SERVICE_ACCOUNT
+ )
op = BeamRunJavaPipelineOperator(
**self.default_op_kwargs, dataflow_config=dataflow_config,
runner="DataflowRunner"
)
@@ -473,6 +480,7 @@ class TestBeamRunJavaPipelineOperator:
"region": "us-central1",
"labels": {"foo": "bar"},
"output": "gs://test/output",
+ "serviceAccount": TEST_SERVICE_ACCOUNT,
"impersonateServiceAccount": TEST_IMPERSONATION_ACCOUNT,
}
persist_link_mock.assert_called_once_with(
@@ -777,7 +785,9 @@ class TestBeamRunGoPipelineOperator:
start_go_dataflow.
"""
gcs_download_method = gcs_hook.return_value.download
- dataflow_config =
DataflowConfiguration(impersonation_chain="[email protected]")
+ dataflow_config = DataflowConfiguration(
+ impersonation_chain="[email protected]",
service_account=TEST_SERVICE_ACCOUNT
+ )
op = BeamRunGoPipelineOperator(
runner="DataflowRunner",
dataflow_config=dataflow_config,
@@ -803,6 +813,7 @@ class TestBeamRunGoPipelineOperator:
"output": "gs://test/output",
"labels": {"foo": "bar", "airflow-version": TEST_VERSION},
"region": "us-central1",
+ "service_account_email": TEST_SERVICE_ACCOUNT,
}
persist_link_mock.assert_called_once_with(context={})
expected_go_file = "/tmp/apache-beam-go/main.go"
