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 ae44d7f42bcaf..df3d80239faa7 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 @@ -125,6 +125,7 @@ def process_fd( log: Logger, process_line_callback: Callable[[str], None] | None = None, is_dataflow_job_id_exist_callback: Callable[[], bool] | None = None, + drain: bool = False, ): """ Print output to logs. @@ -134,6 +135,7 @@ def process_fd( :param process_line_callback: Optional callback which can be used to process stdout and stderr to detect job id. :param log: logger. + :param drain: Whether to read until EOF instead of returning after one line. """ if fd not in (proc.stdout, proc.stderr): raise AirflowException("No data in stderr or in stdout.") @@ -148,6 +150,8 @@ def process_fd( func_log(line.rstrip("\n")) if is_dataflow_job_id_exist_callback and is_dataflow_job_id_exist_callback(): return + if not drain: + return def run_beam_command( @@ -191,18 +195,31 @@ def run_beam_command( if is_dataflow_job_id_exist_callback and is_dataflow_job_id_exist_callback(): return + if is_dataflow_job_id_exist_callback and is_dataflow_job_id_exist_callback(): + return + if proc.poll() is not None: break # Corner case: check if more output was created between the last read and the process termination for readable_fd in reads: - process_fd(proc, readable_fd, log, process_line_callback, is_dataflow_job_id_exist_callback) + process_fd( + proc, + readable_fd, + log, + process_line_callback, + is_dataflow_job_id_exist_callback, + drain=True, + ) log.info("Process exited with return code: %s", proc.returncode) if proc.returncode != 0: raise AirflowException(f"Apache Beam process failed with return code {proc.returncode}") + if is_dataflow_job_id_exist_callback: + is_dataflow_job_id_exist_callback() + class BeamHook(BaseHook): """ 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 e42f1e1d60a39..119d6463b03a7 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 @@ -23,6 +23,7 @@ import os import stat import tempfile +import time from abc import ABC, ABCMeta, abstractmethod from collections.abc import Callable, Sequence from concurrent.futures import ThreadPoolExecutor, as_completed @@ -45,6 +46,7 @@ if TYPE_CHECKING: from airflow.providers.common.compat.sdk import Context GOOGLE_PROVIDER = ProvidersManager().providers.get("apache-airflow-providers-google") +_DATAFLOW_JOB_ID_LOOKUP_INTERVAL: float = 5.0 if GOOGLE_PROVIDER: @@ -102,7 +104,7 @@ def _set_dataflow( pipeline_options, dataflow_job_name, job_name_variable_key ) process_line_callback = self.__get_dataflow_process_callback() - is_dataflow_job_id_exist_callback = self.__is_dataflow_job_id_exist_callback() + is_dataflow_job_id_exist_callback = self.__get_dataflow_job_id_callback(dataflow_job_name) return dataflow_job_name, pipeline_options, process_line_callback, is_dataflow_job_id_exist_callback def __set_dataflow_hook(self) -> DataflowHook: @@ -152,11 +154,33 @@ def set_current_dataflow_job_id(job_id): on_new_job_id_callback=set_current_dataflow_job_id ) - def __is_dataflow_job_id_exist_callback(self) -> Callable[[], bool]: - def is_dataflow_job_id_exist() -> bool: - return True if self.dataflow_job_id else False + def __get_dataflow_job_id_callback(self, dataflow_job_name: str) -> Callable[[], bool]: + last_dataflow_job_id_lookup = -_DATAFLOW_JOB_ID_LOOKUP_INTERVAL - return is_dataflow_job_id_exist + def try_resolve_dataflow_job_id() -> bool: + nonlocal last_dataflow_job_id_lookup + + if self.dataflow_job_id: + return True + if not self.dataflow_hook: + return False + + now = time.monotonic() + if now - last_dataflow_job_id_lookup < _DATAFLOW_JOB_ID_LOOKUP_INTERVAL: + return False + last_dataflow_job_id_lookup = now + + job_id = self.dataflow_hook.fetch_job_id_by_name( + prefix_name=dataflow_job_name, + project_id=self.dataflow_config.project_id, + location=self.dataflow_config.location or DEFAULT_DATAFLOW_LOCATION, + ) + if job_id: + self.dataflow_job_id = job_id + return True + return False + + return try_resolve_dataflow_job_id class BeamBasePipelineOperator(BaseOperator, BeamDataflowMixin, ABC): @@ -467,7 +491,7 @@ def execute_on_dataflow(self, context: Context): 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, @@ -661,7 +685,7 @@ def execute_on_dataflow(self, context: Context): 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 e9750170280c0..0782933d4470b 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 @@ -32,6 +32,7 @@ BeamAsyncHook, BeamHook, beam_options_to_args, + process_fd, run_beam_command, ) from airflow.providers.common.compat.sdk import AirflowException @@ -408,12 +409,11 @@ def test_beam_wait_for_done_logging(self, mock_select, mock_popen, caplog): fake_stderr_fd.readline.side_effect = [ b"apache-beam-stderr-1", b"apache-beam-stderr-2", - StopIteration, b"apache-beam-stderr-3", - StopIteration, b"apache-beam-other-stderr", + b"", ] - fake_stdout_fd.readline.side_effect = [b"apache-beam-stdout", StopIteration] + fake_stdout_fd.readline.side_effect = [b"apache-beam-stdout", b""] mock_select.side_effect = [ ([fake_stderr_fd], None, None), (None, None, None), @@ -440,6 +440,36 @@ def test_beam_wait_for_done_logging(self, mock_select, mock_popen, caplog): assert "apache-beam-stderr-3" in warn_messages assert "apache-beam-other-stderr" in warn_messages + def test_process_fd_reads_one_line_without_draining(self): + fake_logger = logging.getLogger("fake-beam-process-fd-logger") + mock_proc = MagicMock(name="FakeProc") + mock_proc.stderr = MagicMock(name="FakeStderr") + mock_proc.stdout = MagicMock(name="FakeStdout") + mock_proc.stdout.readline.side_effect = [b"apache-beam-stdout-1", b"apache-beam-stdout-2"] + + process_fd(mock_proc, mock_proc.stdout, fake_logger) + + mock_proc.stdout.readline.assert_called_once_with() + + @mock.patch("subprocess.Popen") + @mock.patch("select.select") + def test_beam_wait_for_done_checks_dataflow_job_id_after_select_timeout(self, mock_select, mock_popen): + fake_logger = logging.getLogger("fake-beam-wait-for-done-logger") + mock_proc = MagicMock(name="FakeProc") + mock_proc.stderr = MagicMock(name="FakeStderr") + mock_proc.stdout = MagicMock(name="FakeStdout") + mock_proc.poll.side_effect = AssertionError("poll should not run after job id is resolved") + mock_popen.return_value = mock_proc + mock_select.return_value = ([], None, None) + is_dataflow_job_id_exist_callback = MagicMock(return_value=True) + + run_beam_command( + ["fake", "cmd"], fake_logger, is_dataflow_job_id_exist_callback=is_dataflow_job_id_exist_callback + ) + + is_dataflow_job_id_exist_callback.assert_called_once_with() + mock_proc.poll.assert_not_called() + 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 9c21fe8d2becb..a2ee102e7060c 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 @@ -1029,6 +1029,71 @@ def test_exec_dataflow_runner(self, gcs_hook_mock, dataflow_hook_mock, beam_hook ) beam_hook_mock.return_value.start_python_pipeline.assert_called_once() + @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_resolves_job_id_by_name_before_deferring( + self, gcs_hook_mock, dataflow_hook_mock, beam_hook_mock + ): + dataflow_config = DataflowConfiguration(impersonation_chain=TEST_IMPERSONATION_ACCOUNT) + op = BeamRunPythonPipelineOperator( + runner="DataflowRunner", + dataflow_config=dataflow_config, + **self.default_op_kwargs, + ) + dataflow_hook_mock.return_value.fetch_job_id_by_name.return_value = JOB_ID + + def start_python_pipeline(**kwargs): + assert kwargs["is_dataflow_job_id_exist_callback"]() is True + + beam_hook_mock.return_value.start_python_pipeline.side_effect = start_python_pipeline + + with pytest.raises(TaskDeferred) as exc: + op.execute(context=mock.MagicMock()) + + assert op.dataflow_job_id == JOB_ID + assert exc.value.trigger.job_id == JOB_ID + dataflow_hook_mock.return_value.fetch_job_id_by_name.assert_called_once_with( + prefix_name=dataflow_hook_mock.build_dataflow_job_name.return_value, + project_id=dataflow_hook_mock.return_value.project_id, + location="us-central1", + ) + + 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_without_job_id_falls_back_to_sync_wait( + self, gcs_hook_mock, dataflow_hook_mock, beam_hook_mock, persist_link_mock + ): + dataflow_config = DataflowConfiguration(impersonation_chain=TEST_IMPERSONATION_ACCOUNT) + op = BeamRunPythonPipelineOperator( + runner="DataflowRunner", + dataflow_config=dataflow_config, + **self.default_op_kwargs, + ) + dataflow_hook_mock.return_value.fetch_job_id_by_name.return_value = None + + def start_python_pipeline(**kwargs): + assert kwargs["is_dataflow_job_id_exist_callback"]() is False + + beam_hook_mock.return_value.start_python_pipeline.side_effect = start_python_pipeline + + result = op.execute(context=mock.MagicMock()) + + assert result == {"dataflow_job_id": None} + dataflow_hook_mock.return_value.wait_for_done.assert_called_once_with( + job_name=dataflow_hook_mock.build_dataflow_job_name.return_value, + location="us-central1", + job_id=None, + project_id=dataflow_hook_mock.return_value.project_id, + ) + persist_link_mock.assert_called_once() + + 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("GCSHook")) @@ -1157,6 +1222,70 @@ def test_exec_dataflow_runner(self, gcs_hook_mock, dataflow_hook_mock, beam_hook ) beam_hook_mock.return_value.start_python_pipeline.assert_not_called() + @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_resolves_job_id_by_name_before_deferring( + self, gcs_hook_mock, dataflow_hook_mock, beam_hook_mock + ): + dataflow_config = DataflowConfiguration(impersonation_chain=TEST_IMPERSONATION_ACCOUNT) + op = BeamRunJavaPipelineOperator( + runner="DataflowRunner", dataflow_config=dataflow_config, **self.default_op_kwargs + ) + dataflow_hook_mock.return_value.is_job_dataflow_running.return_value = False + dataflow_hook_mock.return_value.fetch_job_id_by_name.return_value = JOB_ID + + def start_java_pipeline(**kwargs): + assert kwargs["is_dataflow_job_id_exist_callback"]() is True + + beam_hook_mock.return_value.start_java_pipeline.side_effect = start_java_pipeline + + with pytest.raises(TaskDeferred) as exc: + op.execute(context=mock.MagicMock()) + + assert op.dataflow_job_id == JOB_ID + assert exc.value.trigger.job_id == JOB_ID + dataflow_hook_mock.return_value.fetch_job_id_by_name.assert_called_once_with( + prefix_name=dataflow_hook_mock.build_dataflow_job_name.return_value, + project_id=dataflow_hook_mock.return_value.project_id, + location="us-central1", + ) + + 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_without_job_id_falls_back_to_sync_wait( + self, gcs_hook_mock, dataflow_hook_mock, beam_hook_mock, persist_link_mock + ): + dataflow_config = DataflowConfiguration(impersonation_chain=TEST_IMPERSONATION_ACCOUNT) + op = BeamRunJavaPipelineOperator( + runner="DataflowRunner", dataflow_config=dataflow_config, **self.default_op_kwargs + ) + dataflow_hook_mock.return_value.is_job_dataflow_running.return_value = False + dataflow_hook_mock.return_value.fetch_job_id_by_name.return_value = None + + def start_java_pipeline(**kwargs): + assert kwargs["is_dataflow_job_id_exist_callback"]() is False + + beam_hook_mock.return_value.start_java_pipeline.side_effect = start_java_pipeline + + result = op.execute(context=mock.MagicMock()) + + assert result == {"dataflow_job_id": None} + dataflow_hook_mock.return_value.wait_for_done.assert_called_once_with( + job_name=dataflow_hook_mock.build_dataflow_job_name.return_value, + location="us-central1", + job_id=None, + multiple_jobs=False, + project_id=dataflow_hook_mock.return_value.project_id, + ) + persist_link_mock.assert_called_once() + + 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("GCSHook")) 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 ea1af01c1ad50..163f3fc244fa8 100644 --- a/providers/google/src/airflow/providers/google/cloud/hooks/dataflow.py +++ b/providers/google/src/airflow/providers/google/cloud/hooks/dataflow.py @@ -462,6 +462,12 @@ def _fetch_jobs_by_prefix_name(self, prefix_name: str) -> list[dict]: jobs = [job for job in jobs if job["name"].startswith(prefix_name)] return jobs + def fetch_job_id_by_name(self, prefix_name: str) -> str | None: + jobs = self._fetch_jobs_by_prefix_name(prefix_name) + if len(jobs) == 1: + return jobs[0]["id"] + return None + def _refresh_jobs(self) -> None: """ Get all jobs by name. @@ -1173,6 +1179,30 @@ def get_job( ) return jobs_controller.fetch_job_by_id(job_id) + @GoogleBaseHook.fallback_to_default_project_id + def fetch_job_id_by_name( + self, + prefix_name: str, + project_id: str = PROVIDE_PROJECT_ID, + location: str = DEFAULT_DATAFLOW_LOCATION, + ) -> str | None: + """ + Fetch the Dataflow job ID for a unique job name prefix. + + :param prefix_name: Job name prefix to search for. + :param project_id: Optional, the Google Cloud project ID in which to start a job. + If set to None or missing, the default project_id from the Google Cloud connection is used. + :param location: The location of the Dataflow job (for example europe-west1). See: + https://cloud.google.com/dataflow/docs/concepts/regional-endpoints + :return: the job ID if exactly one matching job exists, otherwise None. + """ + jobs_controller = _DataflowJobsController( + dataflow=self.get_conn(), + project_number=project_id, + location=location, + ) + return jobs_controller.fetch_job_id_by_name(prefix_name.lower()) + @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 28d06fde772ac..44599401856f7 100644 --- a/providers/google/tests/unit/google/cloud/hooks/test_dataflow.py +++ b/providers/google/tests/unit/google/cloud/hooks/test_dataflow.py @@ -249,6 +249,25 @@ def test_get_job(self, mock_conn, mock_dataflowjob): ) method_fetch_job_by_id.assert_called_once_with(TEST_JOB_ID) + @mock.patch(DATAFLOW_STRING.format("_DataflowJobsController")) + @mock.patch(DATAFLOW_STRING.format("DataflowHook.get_conn")) + def test_fetch_job_id_by_name(self, mock_conn, mock_dataflowjob): + method_fetch_job_id_by_name = mock_dataflowjob.return_value.fetch_job_id_by_name + method_fetch_job_id_by_name.return_value = TEST_JOB_ID + + result = self.dataflow_hook.fetch_job_id_by_name( + prefix_name=JOB_NAME.upper(), project_id=TEST_PROJECT_ID, location=TEST_LOCATION + ) + + assert result == TEST_JOB_ID + mock_conn.assert_called_once() + mock_dataflowjob.assert_called_once_with( + dataflow=mock_conn.return_value, + project_number=TEST_PROJECT_ID, + location=TEST_LOCATION, + ) + method_fetch_job_id_by_name.assert_called_once_with(JOB_NAME) + @mock.patch(DATAFLOW_STRING.format("_DataflowJobsController")) @mock.patch(DATAFLOW_STRING.format("DataflowHook.get_conn")) def test_fetch_job_metrics_by_id(self, mock_conn, mock_dataflowjob): @@ -703,6 +722,30 @@ def test_dataflow_job_init_without_job_id(self): mock_list.assert_called_once_with(projectId=TEST_PROJECT, location=TEST_LOCATION) + @pytest.mark.parametrize( + ("jobs", "expected_job_id"), + [ + pytest.param([], None, id="no-matches"), + pytest.param([{"id": TEST_JOB_ID, "name": UNIQUE_JOB_NAME}], TEST_JOB_ID, id="single-match"), + pytest.param( + [ + {"id": TEST_JOB_ID, "name": UNIQUE_JOB_NAME}, + {"id": "other-job-id", "name": f"{UNIQUE_JOB_NAME}-extra"}, + ], + None, + id="multiple-matches", + ), + ], + ) + def test_fetch_job_id_by_name_returns_unique_match(self, jobs, expected_job_id): + dataflow_job = _DataflowJobsController( + dataflow=self.mock_dataflow, + project_number=TEST_PROJECT, + location=TEST_LOCATION, + ) + with mock.patch.object(dataflow_job, "_fetch_all_jobs", return_value=jobs): + assert dataflow_job.fetch_job_id_by_name(UNIQUE_JOB_NAME) == expected_job_id + def test_dataflow_job_wait_for_multiple_jobs(self): job = { "id": TEST_JOB_ID,