This is an automated email from the ASF dual-hosted git repository.
shahar1 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 a6ae1b1f169 Fix impersonation_chain and poll_sleep ignored by
deferrable Beam Dataflow pipelines (#72146)
a6ae1b1f169 is described below
commit a6ae1b1f1690edc83d9b99f5b633558ef8884ab8
Author: rjgoyln <[email protected]>
AuthorDate: Sat Sep 19 02:30:07 2026 +0800
Fix impersonation_chain and poll_sleep ignored by deferrable Beam Dataflow
pipelines (#72146)
Once a Beam pipeline on the Dataflow runner defers, the trigger polls the
job through a hook of its own. It was built only from operator-level
fields, so the service account the user asked for in dataflow_config never
reached it and polling fell back to the connection's own identity —
pipelines that depend on impersonation failed with permission errors the
moment they deferred, even though the synchronous path worked. The
configured poll interval was dropped the same way.
---
.../providers/apache/beam/operators/beam.py | 4 +++
.../tests/unit/apache/beam/operators/test_beam.py | 32 ++++++++++++----------
2 files changed, 21 insertions(+), 15 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 e42f1e1d60a..e840940dad5 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
@@ -473,6 +473,8 @@ class
BeamRunPythonPipelineOperator(BeamBasePipelineOperator):
"project_id": self.dataflow_config.project_id,
"location": location,
"gcp_conn_id": self.gcp_conn_id,
+ "poll_sleep": self.dataflow_config.poll_sleep,
+ "impersonation_chain":
self.dataflow_config.impersonation_chain,
}
trigger: DataflowJobStatusTrigger | DataflowJobStateCompleteTrigger
@@ -667,6 +669,8 @@ class BeamRunJavaPipelineOperator(BeamBasePipelineOperator):
"project_id": self.dataflow_config.project_id,
"location": self.dataflow_config.location,
"gcp_conn_id": self.gcp_conn_id,
+ "poll_sleep": self.dataflow_config.poll_sleep,
+ "impersonation_chain":
self.dataflow_config.impersonation_chain,
}
trigger: DataflowJobStatusTrigger |
DataflowJobStateCompleteTrigger
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 9c21fe8d2be..08176afab51 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
@@ -53,6 +53,7 @@ STAGING_LOCATION = "gs://test/staging"
OUTPUT_LOCATION = "gs://test/output"
TEST_VERSION = f"v{version.replace('.', '-').replace('+', '-')}"
TEST_IMPERSONATION_ACCOUNT = "[email protected]"
+TEST_POLL_SLEEP = 30
BEAM_OPERATOR_PATH = "airflow.providers.apache.beam.operators.beam.{}"
@@ -992,12 +993,10 @@ class TestBeamRunPythonPipelineOperatorAsync:
@mock.patch(BEAM_OPERATOR_PATH.format("DataflowHook"))
@mock.patch(BEAM_OPERATOR_PATH.format("GCSHook"))
def test_exec_dataflow_runner(self, gcs_hook_mock, dataflow_hook_mock,
beam_hook_mock):
- """
- Test DataflowHook is created and the right args are passed to
- start_python_dataflow when executing Dataflow runner.
- """
-
- dataflow_config =
DataflowConfiguration(impersonation_chain=TEST_IMPERSONATION_ACCOUNT)
+ """Test the Dataflow hook and the deferral trigger both receive the
dataflow_config args."""
+ dataflow_config = DataflowConfiguration(
+ impersonation_chain=TEST_IMPERSONATION_ACCOUNT,
poll_sleep=TEST_POLL_SLEEP
+ )
op = BeamRunPythonPipelineOperator(
runner="DataflowRunner",
dataflow_config=dataflow_config,
@@ -1005,7 +1004,7 @@ class TestBeamRunPythonPipelineOperatorAsync:
)
magic_mock = mock.MagicMock()
if AIRFLOW_V_3_0_PLUS:
- with pytest.raises(TaskDeferred):
+ with pytest.raises(TaskDeferred) as exc:
op.execute(context=magic_mock)
else:
exception_msg = (
@@ -1014,11 +1013,13 @@ class TestBeamRunPythonPipelineOperatorAsync:
" completed!"
)
with (
- pytest.raises(TaskDeferred),
+ pytest.raises(TaskDeferred) as exc,
pytest.warns(AirflowProviderDeprecationWarning,
match=exception_msg),
):
op.execute(context=magic_mock)
+ assert exc.value.trigger.impersonation_chain ==
TEST_IMPERSONATION_ACCOUNT
+ assert exc.value.trigger.poll_sleep == TEST_POLL_SLEEP
dataflow_hook_mock.assert_called_once_with(
gcp_conn_id=dataflow_config.gcp_conn_id,
poll_sleep=dataflow_config.poll_sleep,
@@ -1122,18 +1123,17 @@ class TestBeamRunJavaPipelineOperatorAsync:
@mock.patch(BEAM_OPERATOR_PATH.format("DataflowHook"))
@mock.patch(BEAM_OPERATOR_PATH.format("GCSHook"))
def test_exec_dataflow_runner(self, gcs_hook_mock, dataflow_hook_mock,
beam_hook_mock):
- """
- Test DataflowHook is created and the right args are passed to
- start_java_pipeline when executing Dataflow runner.
- """
- dataflow_config =
DataflowConfiguration(impersonation_chain=TEST_IMPERSONATION_ACCOUNT)
+ """Test the Dataflow hook and the deferral trigger both receive the
dataflow_config args."""
+ dataflow_config = DataflowConfiguration(
+ impersonation_chain=TEST_IMPERSONATION_ACCOUNT,
poll_sleep=TEST_POLL_SLEEP
+ )
op = BeamRunJavaPipelineOperator(
runner="DataflowRunner", dataflow_config=dataflow_config,
**self.default_op_kwargs
)
dataflow_hook_mock.return_value.is_job_dataflow_running.return_value =
False
magic_mock = mock.MagicMock()
if AIRFLOW_V_3_0_PLUS:
- with pytest.raises(TaskDeferred):
+ with pytest.raises(TaskDeferred) as exc:
op.execute(context=magic_mock)
else:
exception_msg = (
@@ -1142,11 +1142,13 @@ class TestBeamRunJavaPipelineOperatorAsync:
" completed!"
)
with (
- pytest.raises(TaskDeferred),
+ pytest.raises(TaskDeferred) as exc,
pytest.warns(AirflowProviderDeprecationWarning,
match=exception_msg),
):
op.execute(context=magic_mock)
+ assert exc.value.trigger.impersonation_chain ==
TEST_IMPERSONATION_ACCOUNT
+ assert exc.value.trigger.poll_sleep == TEST_POLL_SLEEP
dataflow_hook_mock.assert_called_once_with(
gcp_conn_id=dataflow_config.gcp_conn_id,
poll_sleep=dataflow_config.poll_sleep,
