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 ea7677ed172 Fix an issue with getting dataflow_job_id for Dataflow
Runner when logging is disabled in Pipeline (#72809)
ea7677ed172 is described below
commit ea7677ed17271db01fa5ce83d708b2c58c88984b
Author: Maksim <[email protected]>
AuthorDate: Fri Nov 13 08:11:26 2026 +0100
Fix an issue with getting dataflow_job_id for Dataflow Runner when logging
is disabled in Pipeline (#72809)
---
.../airflow/providers/apache/beam/hooks/beam.py | 30 ++++-
.../providers/apache/beam/operators/beam.py | 65 +++++++++-
.../beam/tests/unit/apache/beam/hooks/test_beam.py | 108 ++++++++++++++++-
.../tests/unit/apache/beam/operators/test_beam.py | 135 ++++++++++++++++++++-
.../providers/google/cloud/hooks/dataflow.py | 54 +++++++--
.../tests/unit/google/cloud/hooks/test_dataflow.py | 111 +++++++++++++++++
6 files changed, 488 insertions(+), 15 deletions(-)
diff --git
a/providers/apache/beam/src/airflow/providers/apache/beam/hooks/beam.py
b/providers/apache/beam/src/airflow/providers/apache/beam/hooks/beam.py
index 86ec5595503..647435778e1 100644
--- a/providers/apache/beam/src/airflow/providers/apache/beam/hooks/beam.py
+++ b/providers/apache/beam/src/airflow/providers/apache/beam/hooks/beam.py
@@ -142,8 +142,8 @@ def process_fd(
func_log = fd_to_log[fd]
for line_raw in iter(fd.readline, b""):
- line = line_raw.decode()
- if process_line_callback:
+ line = line_raw.decode(errors="replace")
+ if process_line_callback and line.endswith("\n"):
process_line_callback(line)
func_log(line.rstrip("\n"))
if is_dataflow_job_id_exist_callback and
is_dataflow_job_id_exist_callback():
@@ -156,6 +156,7 @@ def run_beam_command(
process_line_callback: Callable[[str], None] | None = None,
working_directory: str | None = None,
is_dataflow_job_id_exist_callback: Callable[[], bool] | None = None,
+ check_job_id_by_name_callback: Callable[[], bool] | None = None,
) -> None:
"""
Run pipeline command in subprocess.
@@ -176,13 +177,25 @@ def run_beam_command(
stderr=subprocess.PIPE,
close_fds=True,
)
+
+ if (
+ (is_dataflow_job_id_exist_callback or check_job_id_by_name_callback)
+ and hasattr(os, "set_blocking")
+ and proc.stdout
+ and proc.stderr
+ ):
+ os.set_blocking(proc.stdout.fileno(), False)
+ os.set_blocking(proc.stderr.fileno(), False)
+
# Waits for Apache Beam pipeline to complete.
log.info("Start waiting for Apache Beam process to complete.")
reads = [fd for fd in (proc.stderr, proc.stdout) if fd is not None]
while True:
# Wait for at least one available fd.
readable_fds, _, _ = select.select(reads, [], [], 5)
- if readable_fds is None:
+ if not readable_fds:
+ if check_job_id_by_name_callback and
check_job_id_by_name_callback():
+ return
log.info("Waiting for Apache Beam process to complete.")
continue
@@ -203,6 +216,11 @@ def run_beam_command(
if proc.returncode != 0:
raise AirflowException(f"Apache Beam process failed with return code
{proc.returncode}")
+ if check_job_id_by_name_callback and not (
+ is_dataflow_job_id_exist_callback and
is_dataflow_job_id_exist_callback()
+ ):
+ check_job_id_by_name_callback()
+
class BeamHook(BaseHook):
"""
@@ -228,6 +246,7 @@ class BeamHook(BaseHook):
process_line_callback: Callable[[str], None] | None = None,
working_directory: str | None = None,
is_dataflow_job_id_exist_callback: Callable[[], bool] | None = None,
+ check_job_id_by_name_callback: Callable[[], bool] | None = None,
) -> None:
cmd = [*command_prefix, f"--runner={self.runner}"]
if variables:
@@ -238,6 +257,7 @@ class BeamHook(BaseHook):
working_directory=working_directory,
log=self.log,
is_dataflow_job_id_exist_callback=is_dataflow_job_id_exist_callback,
+ check_job_id_by_name_callback=check_job_id_by_name_callback,
)
def start_python_pipeline(
@@ -250,6 +270,7 @@ class BeamHook(BaseHook):
py_system_site_packages: bool = False,
process_line_callback: Callable[[str], None] | None = None,
is_dataflow_job_id_exist_callback: Callable[[], bool] | None = None,
+ check_job_id_by_name_callback: Callable[[], bool] | None = None,
):
"""
Start Apache Beam python pipeline.
@@ -319,6 +340,7 @@ class BeamHook(BaseHook):
command_prefix=command_prefix,
process_line_callback=process_line_callback,
is_dataflow_job_id_exist_callback=is_dataflow_job_id_exist_callback,
+ check_job_id_by_name_callback=check_job_id_by_name_callback,
)
def start_java_pipeline(
@@ -328,6 +350,7 @@ class BeamHook(BaseHook):
job_class: str | None = None,
process_line_callback: Callable[[str], None] | None = None,
is_dataflow_job_id_exist_callback: Callable[[], bool] | None = None,
+ check_job_id_by_name_callback: Callable[[], bool] | None = None,
) -> None:
"""
Start Apache Beam Java pipeline.
@@ -347,6 +370,7 @@ class BeamHook(BaseHook):
command_prefix=command_prefix,
process_line_callback=process_line_callback,
is_dataflow_job_id_exist_callback=is_dataflow_job_id_exist_callback,
+ check_job_id_by_name_callback=check_job_id_by_name_callback,
)
def start_go_pipeline(
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 2e3f84cf508..68a49a66f29 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
@@ -161,6 +161,29 @@ class BeamDataflowMixin(metaclass=ABCMeta):
return is_dataflow_job_id_exist
+ def _check_dataflow_job_id_by_name_callback(self) -> Callable[[], bool]:
+ def check_dataflow_job_id_by_name() -> bool:
+ if (
+ not self.dataflow_job_id
+ and self.dataflow_hook
+ and hasattr(self.dataflow_hook, "get_job_id_of_running_job")
+ ):
+ try:
+ self.dataflow_job_id =
self.dataflow_hook.get_job_id_of_running_job(
+ job_name=self.dataflow_job_name, # type: ignore
+ project_id=self.dataflow_config.project_id,
+ location=self.dataflow_config.location,
+ )
+ except Exception:
+ self.log.warning( # type: ignore
+ "Failed to get Dataflow job ID by job name '%s'",
+ self.dataflow_job_name, # type: ignore
+ exc_info=True,
+ )
+ return True if self.dataflow_job_id else False
+
+ return check_dataflow_job_id_by_name
+
class BeamBasePipelineOperator(BaseOperator, BeamDataflowMixin, ABC):
"""
@@ -324,6 +347,24 @@ class
BeamRunPythonPipelineOperator(BeamBasePipelineOperator):
For more information on how to use this operator, take a look at the
guide:
:ref:`howto/operator:BeamRunPythonPipelineOperator`
+ .. warning::
+ When running with the ``DataflowRunner``, the operator stops reading
from and waiting
+ on the local Python subprocess as soon as the Dataflow Job ID is
detected, and monitors
+ the job via the Dataflow API instead (or defers when
``deferrable=True``).
+
+ This means any **additional code** in your ``py_file`` that is written
to execute
+ *after* the pipeline block (e.g., local post-processing, file uploads,
or local logging
+ on the Airflow worker) is not waited on and may not be executed,
especially in
+ deferrable mode where the worker task yields immediately upon deferral.
+
+ If you require local post-processing, split the post-processing into a
separate
+ Airflow task (e.g., using ``PythonOperator``).
+
+ Also note that when ``dataflow_config.append_job_name=False`` and the
launcher does not
+ log a Job ID within 5 seconds, the fallback lookup by job name may
attach to an already
+ running job with the same name from a concurrent run instead of
waiting for the launcher
+ to fail with a duplicate job name error.
+
.. seealso::
For more detail on Apache Beam have a look at the reference:
https://beam.apache.org/documentation/
@@ -464,6 +505,7 @@ class
BeamRunPythonPipelineOperator(BeamBasePipelineOperator):
py_system_site_packages=self.py_system_site_packages,
process_line_callback=self.process_line_callback,
is_dataflow_job_id_exist_callback=self.is_dataflow_job_id_exist_callback,
+
check_job_id_by_name_callback=self._check_dataflow_job_id_by_name_callback(),
)
location = self.dataflow_config.location or DEFAULT_DATAFLOW_LOCATION
@@ -474,7 +516,7 @@ class
BeamRunPythonPipelineOperator(BeamBasePipelineOperator):
project_id=self.dataflow_config.project_id,
)
- if self.deferrable:
+ if self.deferrable and self.dataflow_job_id:
trigger_args = {
"job_id": self.dataflow_job_id,
"project_id": self.dataflow_config.project_id,
@@ -532,6 +574,24 @@ class
BeamRunJavaPipelineOperator(BeamBasePipelineOperator):
For more information on how to use this operator, take a look at the
guide:
:ref:`howto/operator:BeamRunJavaPipelineOperator`
+ .. warning::
+ When running with the ``DataflowRunner``, the operator stops reading
from and waiting
+ on the local Java subprocess as soon as the Dataflow Job ID is
detected, and monitors
+ the job via the Dataflow API instead (or defers when
``deferrable=True``).
+
+ This means any **additional code** in your self-executing ``jar`` that
is written to
+ execute locally *after* the pipeline finishes on the worker is not
waited on and may
+ not be executed, especially in deferrable mode where the worker task
yields immediately
+ upon deferral.
+
+ If you require local post-processing, split the post-processing into a
separate
+ Airflow task.
+
+ Also note that when ``dataflow_config.append_job_name=False`` and the
launcher does not
+ log a Job ID within 5 seconds, the fallback lookup by job name may
attach to an already
+ running job with the same name from a concurrent run instead of
waiting for the launcher
+ to fail with a duplicate job name error.
+
.. seealso::
For more detail on Apache Beam have a look at the reference:
https://beam.apache.org/documentation/
@@ -663,6 +723,7 @@ class BeamRunJavaPipelineOperator(BeamBasePipelineOperator):
job_class=self.job_class,
process_line_callback=self.process_line_callback,
is_dataflow_job_id_exist_callback=self.is_dataflow_job_id_exist_callback,
+
check_job_id_by_name_callback=self._check_dataflow_job_id_by_name_callback(),
)
if self.dataflow_job_name and self.dataflow_config.location:
DataflowJobLink.persist(
@@ -671,7 +732,7 @@ class BeamRunJavaPipelineOperator(BeamBasePipelineOperator):
job_id=self.dataflow_job_id,
project_id=self.dataflow_config.project_id,
)
- if self.deferrable:
+ if self.deferrable and self.dataflow_job_id:
trigger_args = {
"job_id": self.dataflow_job_id,
"project_id": self.dataflow_config.project_id,
diff --git a/providers/apache/beam/tests/unit/apache/beam/hooks/test_beam.py
b/providers/apache/beam/tests/unit/apache/beam/hooks/test_beam.py
index 5b25e6c0929..8fffdcfd323 100644
--- a/providers/apache/beam/tests/unit/apache/beam/hooks/test_beam.py
+++ b/providers/apache/beam/tests/unit/apache/beam/hooks/test_beam.py
@@ -17,6 +17,7 @@
from __future__ import annotations
import copy
+import io
import logging
import os
import re
@@ -33,6 +34,7 @@ from airflow.providers.apache.beam.hooks.beam import (
BeamAsyncHook,
BeamHook,
beam_options_to_args,
+ process_fd,
run_beam_command,
)
from airflow.providers.common.compat.sdk import AirflowException
@@ -113,6 +115,7 @@ class TestBeamHook:
working_directory=None,
log=ANY,
is_dataflow_job_id_exist_callback=is_dataflow_job_id_exist_callback,
+ check_job_id_by_name_callback=None,
)
@mock.patch("airflow.providers.apache.beam.hooks.beam.subprocess.check_output",
return_value=b"2.35.0")
@@ -177,6 +180,7 @@ class TestBeamHook:
working_directory=None,
log=ANY,
is_dataflow_job_id_exist_callback=is_dataflow_job_id_exist_callback,
+ check_job_id_by_name_callback=None,
)
@pytest.mark.parametrize(
@@ -225,6 +229,7 @@ class TestBeamHook:
cmd=expected_cmd,
process_line_callback=process_line_callback,
is_dataflow_job_id_exist_callback=is_dataflow_job_id_exist_callback,
+ check_job_id_by_name_callback=None,
working_directory=None,
log=ANY,
)
@@ -283,6 +288,7 @@ class TestBeamHook:
working_directory=None,
log=ANY,
is_dataflow_job_id_exist_callback=None,
+ check_job_id_by_name_callback=None,
)
@mock.patch(BEAM_STRING.format("run_beam_command"))
@@ -312,6 +318,7 @@ class TestBeamHook:
working_directory=None,
log=ANY,
is_dataflow_job_id_exist_callback=None,
+ check_job_id_by_name_callback=None,
)
@mock.patch(BEAM_STRING.format("shutil.which"))
@@ -343,6 +350,7 @@ class TestBeamHook:
working_directory=go_workspace,
log=ANY,
is_dataflow_job_id_exist_callback=None,
+ check_job_id_by_name_callback=None,
)
@mock.patch(BEAM_STRING.format("shutil.which"))
@@ -389,13 +397,15 @@ class TestBeamHook:
working_directory=None,
log=ANY,
is_dataflow_job_id_exist_callback=None,
+ check_job_id_by_name_callback=None,
)
class TestBeamRunner:
+ @mock.patch("os.set_blocking")
@mock.patch("subprocess.Popen")
@mock.patch("select.select")
- def test_beam_wait_for_done_logging(self, mock_select, mock_popen, caplog):
+ def test_beam_wait_for_done_logging(self, mock_select, mock_popen,
mock_set_blocking, caplog):
logger_name = "fake-beam-wait-for-done-logger"
fake_logger = logging.getLogger(logger_name)
@@ -428,6 +438,7 @@ class TestBeamRunner:
with pytest.raises(AirflowException, match="Apache Beam process failed
with return code 1"):
run_beam_command(cmd, fake_logger)
+ mock_set_blocking.assert_not_called()
mock_popen.assert_called_once_with(
cmd, shell=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
close_fds=True, cwd=None
)
@@ -441,6 +452,101 @@ class TestBeamRunner:
assert "apache-beam-stderr-3" in warn_messages
assert "apache-beam-other-stderr" in warn_messages
+ @mock.patch("os.set_blocking")
+ @mock.patch("subprocess.Popen", autospec=True)
+ @mock.patch("select.select")
+ def test_run_beam_command_checks_job_id_by_name_on_timeout(
+ self, mock_select, mock_popen, mock_set_blocking
+ ):
+ fake_logger = MagicMock(spec=logging.Logger)
+ cmd = ["fake", "cmd"]
+ mock_proc = mock_popen.return_value
+ fake_stderr_fd = MagicMock(name="FakeStderr", spec=io.BufferedReader)
+ fake_stdout_fd = MagicMock(name="FakeStdout", spec=io.BufferedReader)
+ fake_stdout_fd.fileno.return_value = 10
+ fake_stderr_fd.fileno.return_value = 11
+ mock_proc.stderr = fake_stderr_fd
+ mock_proc.stdout = fake_stdout_fd
+
+ mock_select.side_effect = [
+ ([], [], []),
+ ([], [], []),
+ ]
+ check_job_id_by_name_callback = MagicMock(spec=callable,
side_effect=[False, True])
+
+ run_beam_command(
+ cmd,
+ fake_logger,
+ check_job_id_by_name_callback=check_job_id_by_name_callback,
+ )
+
+ mock_set_blocking.assert_has_calls([mock.call(10, False),
mock.call(11, False)])
+ assert check_job_id_by_name_callback.call_args_list == [
+ mock.call(),
+ mock.call(),
+ ]
+ mock_proc.poll.assert_not_called()
+
+ @mock.patch("os.set_blocking")
+ @mock.patch("subprocess.Popen", autospec=True)
+ @mock.patch("select.select")
+ def test_run_beam_command_checks_job_id_by_name_after_clean_exit(
+ self, mock_select, mock_popen, mock_set_blocking
+ ):
+ fake_logger = MagicMock(spec=logging.Logger)
+ cmd = ["fake", "cmd"]
+ mock_proc = mock_popen.return_value
+ fake_stderr_fd = MagicMock(name="FakeStderr", spec=io.BufferedReader)
+ fake_stdout_fd = MagicMock(name="FakeStdout", spec=io.BufferedReader)
+ fake_stdout_fd.fileno.return_value = 10
+ fake_stderr_fd.fileno.return_value = 11
+ fake_stdout_fd.readline.return_value = b""
+ fake_stderr_fd.readline.return_value = b""
+ mock_proc.stderr = fake_stderr_fd
+ mock_proc.stdout = fake_stdout_fd
+ mock_proc.poll.return_value = 0
+ mock_proc.returncode = 0
+
+ mock_select.return_value = ([fake_stdout_fd], [], [])
+ is_job_id_exist_callback = MagicMock(spec=callable, return_value=False)
+ check_job_id_by_name_callback = MagicMock(spec=callable,
return_value=True)
+
+ run_beam_command(
+ cmd,
+ fake_logger,
+ is_dataflow_job_id_exist_callback=is_job_id_exist_callback,
+ check_job_id_by_name_callback=check_job_id_by_name_callback,
+ )
+
+ check_job_id_by_name_callback.assert_called_once_with()
+
+ def
test_process_fd_decodes_invalid_utf8_and_skips_partial_lines_in_callback(self):
+ fake_logger = MagicMock(spec=logging.Logger)
+ mock_proc = MagicMock(name="FakeProc", spec=subprocess.Popen)
+ fake_stdout_fd = MagicMock(name="FakeStdout", spec=io.BufferedReader)
+ fake_stderr_fd = MagicMock(name="FakeStderr", spec=io.BufferedReader)
+ mock_proc.stdout = fake_stdout_fd
+ mock_proc.stderr = fake_stderr_fd
+ fake_stdout_fd.readline.side_effect = [
+ b"Submitted job: 2026-11-0",
+ b"bad \xff\xfe line\n",
+ b"",
+ ]
+ process_line_callback = MagicMock(spec=callable)
+
+ process_fd(
+ mock_proc,
+ fake_stdout_fd,
+ fake_logger,
+ process_line_callback=process_line_callback,
+ )
+
+ process_line_callback.assert_called_once_with("bad \ufffd\ufffd
line\n")
+ assert fake_logger.info.call_args_list == [
+ mock.call("Submitted job: 2026-11-0"),
+ mock.call("bad \ufffd\ufffd line"),
+ ]
+
class TestBeamOptionsToArgs:
@pytest.mark.parametrize(
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 7e8fe2a1d51..82a04917d45 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
@@ -138,6 +138,92 @@ class TestBeamBasePipelineOperator:
op.dataflow_job_id = "sample_df_job_same_value_id"
mock_ti.xcom_push.assert_not_called()
+ @mock.patch(BEAM_OPERATOR_PATH.format("BeamHook"))
+ @mock.patch(BEAM_OPERATOR_PATH.format("DataflowHook"))
+ def test_check_dataflow_job_id_by_name_callback(self, dataflow_hook_mock,
beam_hook_mock):
+ op = BeamBasePipelineOperator(
+ task_id=TASK_ID,
+ runner="DataflowRunner",
+ dataflow_config={"project_id": TEST_PROJECT, "location":
"us-central1"},
+ )
+ op._cast_dataflow_config()
+ (
+ _,
+ op.dataflow_job_name,
+ _,
+ _,
+ is_dataflow_job_id_exist_callback,
+ ) = op._init_pipeline_options(format_pipeline_options=True,
job_name_variable_key="job_name")
+ check_by_name_callback = op._check_dataflow_job_id_by_name_callback()
+
+ assert is_dataflow_job_id_exist_callback is not None
+ assert is_dataflow_job_id_exist_callback() is False
+
dataflow_hook_mock.return_value.get_job_id_of_running_job.assert_not_called()
+
+ dataflow_hook_mock.return_value.get_job_id_of_running_job.return_value
= None
+ assert check_by_name_callback() is False
+ assert op.dataflow_job_id is None
+
+ dataflow_hook_mock.return_value.get_job_id_of_running_job.return_value
= JOB_ID
+ assert check_by_name_callback() is True
+ assert is_dataflow_job_id_exist_callback() is True
+ assert op.dataflow_job_id == JOB_ID
+
dataflow_hook_mock.return_value.get_job_id_of_running_job.assert_called_with(
+ job_name=op.dataflow_job_name,
+ project_id=TEST_PROJECT,
+ location="us-central1",
+ )
+
+ @mock.patch(BEAM_OPERATOR_PATH.format("BeamHook"))
+ @mock.patch(BEAM_OPERATOR_PATH.format("DataflowHook"))
+ def test_check_dataflow_job_id_by_name_callback_handles_lookup_error(
+ self, dataflow_hook_mock, beam_hook_mock
+ ):
+ op = BeamBasePipelineOperator(
+ task_id=TASK_ID,
+ runner="DataflowRunner",
+ dataflow_config={"project_id": TEST_PROJECT, "location":
"us-central1"},
+ )
+ op._cast_dataflow_config()
+ (
+ _,
+ op.dataflow_job_name,
+ _,
+ _,
+ _,
+ ) = op._init_pipeline_options(format_pipeline_options=True,
job_name_variable_key="job_name")
+ check_by_name_callback = op._check_dataflow_job_id_by_name_callback()
+
+ dataflow_hook_mock.return_value.get_job_id_of_running_job.side_effect
= RuntimeError(
+ "missing dataflow.jobs.list permission"
+ )
+ assert check_by_name_callback() is False
+ assert op.dataflow_job_id is None
+
+ @mock.patch(BEAM_OPERATOR_PATH.format("BeamHook"))
+ @mock.patch(BEAM_OPERATOR_PATH.format("DataflowHook"))
+ def test_check_dataflow_job_id_by_name_callback_without_hook_method(
+ self, dataflow_hook_mock, beam_hook_mock
+ ):
+ op = BeamBasePipelineOperator(
+ task_id=TASK_ID,
+ runner="DataflowRunner",
+ dataflow_config={"project_id": TEST_PROJECT, "location":
"us-central1"},
+ )
+ op._cast_dataflow_config()
+ (
+ _,
+ op.dataflow_job_name,
+ _,
+ _,
+ _,
+ ) = op._init_pipeline_options(format_pipeline_options=True,
job_name_variable_key="job_name")
+ del dataflow_hook_mock.return_value.get_job_id_of_running_job
+ check_by_name_callback = op._check_dataflow_job_id_by_name_callback()
+
+ assert check_by_name_callback() is False
+ assert op.dataflow_job_id is None
+
class TestBeamRunPythonPipelineOperator:
@pytest.fixture(autouse=True)
@@ -269,6 +355,7 @@ class TestBeamRunPythonPipelineOperator:
py_system_site_packages=False,
process_line_callback=mock.ANY,
is_dataflow_job_id_exist_callback=mock.ANY,
+ check_job_id_by_name_callback=mock.ANY,
)
@mock.patch(BEAM_OPERATOR_PATH.format("DataflowJobLink.persist"))
@@ -495,6 +582,7 @@ class TestBeamRunJavaPipelineOperator:
job_class=JOB_CLASS,
process_line_callback=mock.ANY,
is_dataflow_job_id_exist_callback=mock.ANY,
+ check_job_id_by_name_callback=mock.ANY,
)
@mock.patch(BEAM_OPERATOR_PATH.format("DataflowJobLink.persist"))
@@ -1013,6 +1101,7 @@ class TestBeamRunPythonPipelineOperatorAsync:
dataflow_config=dataflow_config,
**self.default_op_kwargs,
)
+ op._dataflow_job_id = JOB_ID
magic_mock = mock.MagicMock()
if AIRFLOW_V_3_0_PLUS:
with pytest.raises(TaskDeferred) as exc:
@@ -1041,16 +1130,36 @@ class TestBeamRunPythonPipelineOperatorAsync:
)
beam_hook_mock.return_value.start_python_pipeline.assert_called_once()
+ @mock.patch(BEAM_OPERATOR_PATH.format("DataflowJobLink.persist"))
+ @mock.patch(BEAM_OPERATOR_PATH.format("BeamHook"))
+ @mock.patch(BEAM_OPERATOR_PATH.format("DataflowHook"))
+ @mock.patch(BEAM_OPERATOR_PATH.format("GCSHook"))
+ def
test_exec_dataflow_runner_falls_back_to_wait_for_done_when_job_id_missing(
+ self, gcs_hook_mock, dataflow_hook_mock, beam_hook_mock,
persist_link_mock
+ ):
+ op = BeamRunPythonPipelineOperator(
+ runner="DataflowRunner",
+ **self.default_op_kwargs,
+ )
+ res = op.execute(context=mock.MagicMock())
+ dataflow_hook_mock.return_value.wait_for_done.assert_called_once_with(
+ job_name=op.dataflow_job_name,
+ location=op.dataflow_config.location,
+ job_id=None,
+ project_id=op.dataflow_config.project_id,
+ )
+ assert res == {"dataflow_job_id": None}
+
@mock.patch(BEAM_OPERATOR_PATH.format("DataflowJobLink.persist"))
@mock.patch(BEAM_OPERATOR_PATH.format("BeamHook"))
@mock.patch(BEAM_OPERATOR_PATH.format("GCSHook"))
@mock.patch(BEAM_OPERATOR_PATH.format("DataflowHook"))
def test_on_kill_dataflow_runner(self, dataflow_hook_mock, _, __, ___):
op = BeamRunPythonPipelineOperator(runner="DataflowRunner",
**self.default_op_kwargs)
+ op._dataflow_job_id = JOB_ID
dataflow_cancel_job = dataflow_hook_mock.return_value.cancel_job
with pytest.raises(TaskDeferred):
op.execute(context=mock.MagicMock())
- op.dataflow_job_id = JOB_ID
op.on_kill()
dataflow_cancel_job.assert_called_once_with(
job_id=JOB_ID, project_id=op.dataflow_config.project_id,
location=op.dataflow_config.location
@@ -1062,6 +1171,7 @@ class TestBeamRunPythonPipelineOperatorAsync:
def test_on_kill_direct_runner(self, _, dataflow_mock, __):
dataflow_cancel_job = dataflow_mock.return_value.cancel_job
op = BeamRunPythonPipelineOperator(runner="DataflowRunner",
**self.default_op_kwargs)
+ op._dataflow_job_id = JOB_ID
if AIRFLOW_V_3_0_PLUS:
with pytest.raises(TaskDeferred):
op.execute(mock.MagicMock())
@@ -1076,6 +1186,7 @@ class TestBeamRunPythonPipelineOperatorAsync:
pytest.warns(AirflowProviderDeprecationWarning,
match=exception_msg),
):
op.execute(mock.MagicMock())
+ op._dataflow_job_id = None
op.on_kill()
dataflow_cancel_job.assert_not_called()
@@ -1141,6 +1252,7 @@ class TestBeamRunJavaPipelineOperatorAsync:
op = BeamRunJavaPipelineOperator(
runner="DataflowRunner", dataflow_config=dataflow_config,
**self.default_op_kwargs
)
+ op._dataflow_job_id = JOB_ID
dataflow_hook_mock.return_value.is_job_dataflow_running.return_value =
False
magic_mock = mock.MagicMock()
if AIRFLOW_V_3_0_PLUS:
@@ -1170,6 +1282,25 @@ class TestBeamRunJavaPipelineOperatorAsync:
)
beam_hook_mock.return_value.start_python_pipeline.assert_not_called()
+ @mock.patch(BEAM_OPERATOR_PATH.format("DataflowJobLink.persist"))
+ @mock.patch(BEAM_OPERATOR_PATH.format("BeamHook"))
+ @mock.patch(BEAM_OPERATOR_PATH.format("DataflowHook"))
+ @mock.patch(BEAM_OPERATOR_PATH.format("GCSHook"))
+ def
test_exec_dataflow_runner_falls_back_to_wait_for_done_when_job_id_missing(
+ self, gcs_hook_mock, dataflow_hook_mock, beam_hook_mock,
persist_link_mock
+ ):
+ dataflow_hook_mock.return_value.is_job_dataflow_running.return_value =
False
+ op = BeamRunJavaPipelineOperator(runner="DataflowRunner",
**self.default_op_kwargs)
+ res = op.execute(context=mock.MagicMock())
+ dataflow_hook_mock.return_value.wait_for_done.assert_called_once_with(
+ job_name=op.dataflow_job_name,
+ location=op.dataflow_config.location,
+ job_id=None,
+ multiple_jobs=False,
+ project_id=op.dataflow_config.project_id,
+ )
+ assert res == {"dataflow_job_id": None}
+
@mock.patch(BEAM_OPERATOR_PATH.format("DataflowJobLink.persist"))
@mock.patch(BEAM_OPERATOR_PATH.format("BeamHook"))
@mock.patch(BEAM_OPERATOR_PATH.format("GCSHook"))
@@ -1178,9 +1309,9 @@ class TestBeamRunJavaPipelineOperatorAsync:
dataflow_hook_mock.return_value.is_job_dataflow_running.return_value =
False
dataflow_cancel_job = dataflow_hook_mock.return_value.cancel_job
op = BeamRunJavaPipelineOperator(runner="DataflowRunner",
**self.default_op_kwargs)
+ op._dataflow_job_id = JOB_ID
with pytest.raises(TaskDeferred):
op.execute(context=mock.MagicMock())
- op.dataflow_job_id = JOB_ID
op.on_kill()
dataflow_cancel_job.assert_called_once_with(
job_id=JOB_ID, project_id=op.dataflow_config.project_id,
location=op.dataflow_config.location
diff --git
a/providers/google/src/airflow/providers/google/cloud/hooks/dataflow.py
b/providers/google/src/airflow/providers/google/cloud/hooks/dataflow.py
index e3fd6da93fc..c5d5a074ddb 100644
--- a/providers/google/src/airflow/providers/google/cloud/hooks/dataflow.py
+++ b/providers/google/src/airflow/providers/google/cloud/hooks/dataflow.py
@@ -434,13 +434,40 @@ class
_DataflowJobsController(DataflowJobTerminalStateHelper):
autoscaling_events.extend(response.get("autoscalingEvents", []))
return autoscaling_events
- def _fetch_all_jobs(self) -> list[dict]:
- request = (
- self._dataflow.projects()
- .locations()
- .jobs()
- .list(projectId=self._project_number, location=self._job_location)
- )
+ def fetch_job_id(self) -> str | None:
+ if self._job_id:
+ return self._job_id
+ if not self._job_name:
+ raise ValueError("Missing both dataflow job ID and name.")
+ jobs = self._fetch_all_jobs(jobs_filter="ACTIVE")
+ job_name = self._job_name.lower()
+ active_jobs = [
+ job
+ for job in jobs
+ if job.get("name") == job_name
+ and job.get("currentState") not in
DataflowJobStatus.TERMINAL_STATES
+ ]
+ if active_jobs:
+ for job in active_jobs:
+ self.log.info(
+ "Google Cloud DataFlow job %s is state: %s",
+ job["name"],
+ job["currentState"],
+ )
+ active_jobs.sort(key=lambda j: j.get("createTime", ""),
reverse=True)
+ self._job_id = active_jobs[0]["id"]
+ else:
+ self.log.info("Google Cloud DataFlow job not available yet..")
+ return self._job_id
+
+ def _fetch_all_jobs(self, jobs_filter: str | None = None) -> list[dict]:
+ list_kwargs: dict[str, Any] = {
+ "projectId": self._project_number,
+ "location": self._job_location,
+ }
+ if jobs_filter is not None:
+ list_kwargs["filter"] = jobs_filter
+ request =
self._dataflow.projects().locations().jobs().list(**list_kwargs)
all_jobs: list[dict] = []
while request is not None:
response = request.execute(num_retries=self._num_retries)
@@ -1124,6 +1151,19 @@ class DataflowHook(GoogleBaseHook):
)
return jobs_controller.fetch_job_by_id(job_id)
+ @GoogleBaseHook.fallback_to_default_project_id
+ def get_job_id_of_running_job(
+ self, job_name: str, project_id: str = PROVIDE_PROJECT_ID, location:
str = DEFAULT_DATAFLOW_LOCATION
+ ) -> str | None:
+ """Get the job_id for currently running Dataflow Job by job_name."""
+ jobs_controller = _DataflowJobsController(
+ dataflow=self.get_conn(),
+ project_number=project_id,
+ name=job_name,
+ location=location,
+ )
+ return jobs_controller.fetch_job_id()
+
@GoogleBaseHook.fallback_to_default_project_id
def fetch_job_metrics_by_id(
self,
diff --git a/providers/google/tests/unit/google/cloud/hooks/test_dataflow.py
b/providers/google/tests/unit/google/cloud/hooks/test_dataflow.py
index 0cbbddeabfd..3bae5c17407 100644
--- a/providers/google/tests/unit/google/cloud/hooks/test_dataflow.py
+++ b/providers/google/tests/unit/google/cloud/hooks/test_dataflow.py
@@ -672,6 +672,26 @@ class TestDataflowTemplateHook:
with pytest.raises(AirflowException):
self.dataflow_hook.extract_job_id({"not_id": True})
+ @mock.patch(DATAFLOW_STRING.format("_DataflowJobsController"))
+ @mock.patch(DATAFLOW_STRING.format("DataflowHook.get_conn"))
+ def test_get_job_id_of_running_job(self, mock_get_conn, jobs_controller):
+ jobs_controller.return_value.fetch_job_id.return_value = TEST_JOB_ID
+
+ result = self.dataflow_hook.get_job_id_of_running_job(
+ job_name=UNIQUE_JOB_NAME,
+ project_id=TEST_PROJECT,
+ location=TEST_LOCATION,
+ )
+
+ jobs_controller.assert_called_once_with(
+ dataflow=mock_get_conn.return_value,
+ name=UNIQUE_JOB_NAME,
+ project_number=TEST_PROJECT,
+ location=TEST_LOCATION,
+ )
+ jobs_controller.return_value.fetch_job_id.assert_called_once_with()
+ assert result == TEST_JOB_ID
+
class TestDataflowJob:
def setup_method(self):
@@ -1368,6 +1388,97 @@ class TestDataflowJob:
mock_fetch_responses.assert_called_once_with(job_id=TEST_JOB_ID)
assert result == ["event_1", "event_2"]
+ @pytest.mark.parametrize(
+ ("jobs", "expected_job_id"),
+ [
+ (
+ [
+ {
+ "id": "job-done-id",
+ "name": "test-job",
+ "currentState": DataflowJobStatus.JOB_STATE_DONE,
+ "createTime": "2026-06-25T16:00:00Z",
+ },
+ {
+ "id": "job-prefix-only-id",
+ "name": "test-job-other",
+ "currentState": DataflowJobStatus.JOB_STATE_RUNNING,
+ "createTime": "2026-06-25T16:10:00Z",
+ },
+ {
+ "id": "job-active-older-id",
+ "name": "test-job",
+ "currentState": DataflowJobStatus.JOB_STATE_PENDING,
+ "createTime": "2026-06-25T15:00:00Z",
+ },
+ {
+ "id": "job-active-newer-id",
+ "name": "test-job",
+ "currentState": DataflowJobStatus.JOB_STATE_RUNNING,
+ "createTime": "2026-06-25T15:30:00Z",
+ },
+ ],
+ "job-active-newer-id",
+ ),
+ (
+ [
+ {
+ "id": "job-done-id",
+ "name": "test-job",
+ "currentState": DataflowJobStatus.JOB_STATE_DONE,
+ "createTime": "2026-06-25T15:00:00Z",
+ }
+ ],
+ None,
+ ),
+ (
+ [
+ {
+ "id": "job-prefix-only-id",
+ "name": "test-job-suffix",
+ "currentState": DataflowJobStatus.JOB_STATE_RUNNING,
+ "createTime": "2026-06-25T15:00:00Z",
+ }
+ ],
+ None,
+ ),
+ ],
+ )
+ def test_fetch_job_id(self, jobs, expected_job_id):
+ mock_jobs =
self.mock_dataflow.projects.return_value.locations.return_value.jobs.return_value
+ mock_jobs.list.return_value.execute.return_value = {"jobs": jobs}
+ mock_jobs.list_next.return_value = None
+
+ jobs_controller = _DataflowJobsController(
+ dataflow=self.mock_dataflow,
+ project_number=TEST_PROJECT,
+ name="Test-Job",
+ location=TEST_LOCATION,
+ multiple_jobs=False,
+ )
+
+ result = jobs_controller.fetch_job_id()
+
+ mock_jobs.list.assert_called_once_with(
+ projectId=TEST_PROJECT,
+ location=TEST_LOCATION,
+ filter="ACTIVE",
+ )
+ assert result == expected_job_id
+ assert jobs_controller._job_id == expected_job_id
+
+ def test_fetch_job_id_without_job_name_returns_existing_job_id(self):
+ jobs_controller = _DataflowJobsController(
+ dataflow=self.mock_dataflow,
+ project_number=TEST_PROJECT,
+ name=None,
+ location=TEST_LOCATION,
+ job_id=TEST_JOB_ID,
+ )
+
+ assert jobs_controller.fetch_job_id() == TEST_JOB_ID
+ self.mock_dataflow.projects.assert_not_called()
+
@pytest.mark.db_test
class TestDataflowPipelineHook:
