From 322d78a98b524fc3aa93e83ff34ee7927f3c95c4 Mon Sep 17 00:00:00 2001 From: Kacper Muda Date: Tue, 28 Jul 2026 16:52:09 +0200 Subject: [PATCH 1/2] Fix OpenLineage team_name access and other failing paths --- .../providers/openlineage/plugins/listener.py | 70 ++++++++--- .../providers/openlineage/utils/utils.py | 73 +++++++---- .../unit/openlineage/plugins/test_listener.py | 116 +++++++++++++++++- .../unit/openlineage/utils/test_utils.py | 59 +++++++++ 4 files changed, 277 insertions(+), 41 deletions(-) diff --git a/providers/openlineage/src/airflow/providers/openlineage/plugins/listener.py b/providers/openlineage/src/airflow/providers/openlineage/plugins/listener.py index a37c34ef54fd5..6772f8c2bd36f 100644 --- a/providers/openlineage/src/airflow/providers/openlineage/plugins/listener.py +++ b/providers/openlineage/src/airflow/providers/openlineage/plugins/listener.py @@ -21,7 +21,6 @@ import sys import threading from concurrent.futures import ProcessPoolExecutor -from concurrent.futures.process import BrokenProcessPool from datetime import datetime from functools import cache from typing import TYPE_CHECKING @@ -138,8 +137,12 @@ def _run_adapter_method(method_name: str, /, *args, **kwargs): background worker threads (e.g. the ``datadog`` transport, which always starts an async HTTP worker thread) are never closed, so this leaks one thread per event and steadily consumes scheduler CPU and memory until restart. + + Returns nothing: the emitted event is only of use inside the pool worker, and pickling it + back to the parent would turn a redacted event the pickler chokes on into a spurious + "failed to submit" warning for an emission that actually succeeded. """ - return getattr(_get_process_adapter(), method_name)(*args, **kwargs) + getattr(_get_process_adapter(), method_name)(*args, **kwargs) def _emit_manual_state_change_event(adapter_method_name: str, stats_key: str, **kwargs): @@ -149,11 +152,11 @@ def _emit_manual_state_change_event(adapter_method_name: str, stats_key: str, ** Module-level so it is picklable across the ProcessPoolExecutor boundary used by `_on_task_instance_manual_state_change` for scheduler-side "task state changed externally" emissions. The method is resolved on the per-process adapter so the - pool worker reuses one client across events (see ``_run_adapter_method``). + pool worker reuses one client across events (see ``_run_adapter_method``), and nothing is + returned so an unpicklable event cannot fail the future after a successful emission. """ event = getattr(_get_process_adapter(), adapter_method_name)(**kwargs) Stats.gauge(stats_key, len(Serde.to_json(event).encode("utf-8"))) - return event class OpenLineageListener: @@ -1058,14 +1061,21 @@ def _fork_execute(self, callable, callable_name: str): self._terminate_with_wait(process) self.log.debug("Process with pid %s finished - parent", pid) else: - setproctitle(getproctitle() + " - OpenLineage - " + callable_name) - if not AIRFLOW_V_3_0_PLUS: - configure_orm(disable_connection_pool=True) - self.log.debug("Executing OpenLineage process - %s - pid %s", callable_name, os.getpid()) + # Everything the child does lives in this try, and the exit lives in its finally: a child + # that returns instead of exiting becomes a second task runner. It would unwind into the + # hook's caller and keep executing task-runner code while sharing the supervisor socket + # with the real one, duplicating state writes and interleaving bytes on the channel. try: + setproctitle(getproctitle() + " - OpenLineage - " + callable_name) + if not AIRFLOW_V_3_0_PLUS: + configure_orm(disable_connection_pool=True) + self.log.debug("Executing OpenLineage process - %s - pid %s", callable_name, os.getpid()) callable() self.log.debug("Process with current pid finishes after %s", callable_name) - except Exception: + except BaseException: + # BaseException, not Exception: a SIGINT delivered to the process group reaches this + # child as KeyboardInterrupt, and Airflow's own AirflowTaskTimeout / TaskDeferred + # also derive from BaseException. self.log.warning( "OpenLineage %s process failed. This has no impact on actual task execution status.", callable_name, @@ -1076,8 +1086,13 @@ def _fork_execute(self, callable, callable_name: str): # logging so buffered records (including any warnings above) are flushed # before the process exits. Without this, the final log lines are silently # dropped, making failures invisible. - logging.shutdown() - os._exit(0) + # Nest os._exit in its own finally so that a raising logging.shutdown() + # (e.g. a remote handler whose close() throws) cannot skip the exit and + # unwind back into the task runner as a second process. + try: + logging.shutdown() + finally: + os._exit(0) @property def executor(self) -> ProcessPoolExecutor: @@ -1095,8 +1110,26 @@ def on_starting(self, component) -> None: @hookimpl def before_stopping(self, component) -> None: self.log.debug("before_stopping: %s", component.__class__.__name__) - with timeout(30): - self.executor.shutdown(wait=True) + if self._executor is None: + # Do not build a pool just to tear it down -- on the task runner this hook fires at the + # end of every task, where no pool was ever needed. + return + + # Detach before shutting down so a later event rebuilds a fresh pool. Left attached, every + # subsequent submission would raise "cannot schedule new futures after shutdown" and drop + # its event for the remaining lifetime of the process. + executor, self._executor = self._executor, None + try: + with timeout(30): + executor.shutdown(wait=True) + except BaseException: + # `timeout` is SIGALRM-based: it raises AirflowTaskTimeout when shutdown overruns, and + # ValueError when called off the main thread. Neither may escape a listener hook -- + # and BaseException is required, not Exception, because AirflowTaskTimeout derives from + # BaseException so that user code cannot swallow it. Every hook call site guards with + # `except Exception`, so letting it through would reach the caller. + self.log.warning("OpenLineage executor did not shut down cleanly.", exc_info=True) + executor.shutdown(wait=False) @hookimpl def on_dag_run_running(self, dag_run: DagRun, msg: str) -> None: @@ -1258,9 +1291,14 @@ def on_dag_run_failed(self, dag_run: DagRun, msg: str) -> None: def submit_callable(self, callable, *args, **kwargs): try: fut = self.executor.submit(callable, *args, **kwargs) - except BrokenProcessPool: - self.log.warning("ProcessPoolExecutor is broken; recreating and retrying submission.") - self._executor.shutdown(wait=False) + except RuntimeError: + # BrokenProcessPool subclasses RuntimeError, so this also covers a pool that was already + # shut down ("cannot schedule new futures after shutdown"). The retry is deliberately + # unguarded: every caller wraps this in `except BaseException`, and a second consecutive + # failure deserves to surface in their warning rather than be swallowed here. + self.log.warning("ProcessPoolExecutor is unusable; recreating and retrying submission.") + if self._executor is not None: + self._executor.shutdown(wait=False) self._executor = None fut = self.executor.submit(callable, *args, **kwargs) fut.add_done_callback(self.log_submit_error) diff --git a/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py b/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py index f40e9b12a0be7..5d58ecd9083e8 100644 --- a/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py +++ b/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py @@ -1045,15 +1045,18 @@ def dag_version_info(cls, dagrun: DagRun, key: str) -> str | int | None: dag_versions = safe_getattr(dagrun, "dag_versions", []) if not dag_versions: return None + # The DagVersion rows themselves are lazy-loaded, so reading their columns can still hit + # a detached session even though fetching the list above succeeded. current_version = dag_versions[-1] if key == "bundle_name": - return current_version.bundle_name + return safe_getattr(current_version, "bundle_name") if key == "bundle_version": - return current_version.bundle_version + return safe_getattr(current_version, "bundle_version") if key == "version_id": - return str(current_version.id) + version_id = safe_getattr(current_version, "id") + return str(version_id) if version_id is not None else None if key == "version_number": - return current_version.version_number + return safe_getattr(current_version, "version_number") raise ValueError(f"Unsupported key: {key}`") @classmethod @@ -1062,13 +1065,30 @@ def team_name(cls, dagrun: DagRun) -> str | None: if not AIRFLOW_V_3_3_PLUS or not airflow_conf.getboolean("core", "multi_team", fallback=False): return None - from airflow.models.dagbundle import DagBundleModel + # The Execution API delivers the team name on the DagRun payload it sends to the task + # runner, so task events resolve it from there rather than through the bundle lookup below, + # which needs a metadata DB session the task runner does not have. + # `hasattr` rather than a None check -- a team-less run legitimately carries None, while + # the scheduler's ORM DagRun has no such attribute at all. + if hasattr(dagrun, "team_name"): + return dagrun.team_name - bundle_name = cls.dag_version_info(dagrun, "bundle_name") - if not isinstance(bundle_name, str): - return None + try: + bundle_name = cls.dag_version_info(dagrun, "bundle_name") + if not isinstance(bundle_name, str): + return None - return DagBundleModel.get_team_name(bundle_name) + from airflow.models.dagbundle import DagBundleModel + + return DagBundleModel.get_team_name(bundle_name) + except Exception as e: + log.warning( + "OpenLineage failed to resolve the team name for dag `%s`: %s.", + safe_getattr(dagrun, "dag_id"), + e, + ) + log.debug("Exception details:", exc_info=True) + return None class TaskInstanceInfo(InfoJsonEncodable): @@ -1387,23 +1407,28 @@ def get_airflow_job_facet(dag_run: DagRun) -> dict[str, AirflowJobFacet]: def get_airflow_state_run_facet( dag_id: str, run_id: str, task_ids: list[str], dag_run_state: DagRunState ) -> dict[str, AirflowStateRunFacet]: - tis = DagRun.fetch_task_instances(dag_id=dag_id, run_id=run_id, task_ids=task_ids) + try: + tis = DagRun.fetch_task_instances(dag_id=dag_id, run_id=run_id, task_ids=task_ids) - def get_task_duration(ti): - if ti.duration is not None: - return ti.duration - if ti.end_date is not None and ti.start_date is not None: - return (ti.end_date - ti.start_date).total_seconds() - # Fallback to 0.0 for tasks with missing timestamps (e.g., skipped/terminated tasks) - return 0.0 + def get_task_duration(ti): + if ti.duration is not None: + return ti.duration + if ti.end_date is not None and ti.start_date is not None: + return (ti.end_date - ti.start_date).total_seconds() + # Fallback to 0.0 for tasks with missing timestamps (e.g., skipped/terminated tasks) + return 0.0 - return { - "airflowState": AirflowStateRunFacet( - dagRunState=dag_run_state, - tasksState={ti.task_id: ti.state for ti in tis}, - tasksDuration={ti.task_id: get_task_duration(ti) for ti in tis}, - ) - } + return { + "airflowState": AirflowStateRunFacet( + dagRunState=dag_run_state, + tasksState={ti.task_id: ti.state for ti in tis}, + tasksDuration={ti.task_id: get_task_duration(ti) for ti in tis}, + ) + } + except Exception as e: + log.warning("Failed to build AirflowStateRunFacet for DagRun %s/%s: %s.", dag_id, run_id, e) + log.debug("Exception details:", exc_info=True) + return {} def is_dag_run_asset_triggered( diff --git a/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py b/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py index 9168052311f21..e22a950cdc12b 100644 --- a/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py +++ b/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py @@ -36,7 +36,7 @@ from uuid6 import uuid7 from airflow.models import DAG, DagRun, TaskInstance -from airflow.providers.common.compat.sdk import BaseOperator +from airflow.providers.common.compat.sdk import AirflowTaskTimeout, BaseOperator from airflow.providers.openlineage.extractors.base import OperatorLineage from airflow.providers.openlineage.plugins.adapter import OpenLineageAdapter from airflow.providers.openlineage.plugins.listener import OpenLineageListener @@ -164,6 +164,26 @@ def test_process_adapter_reused_across_pool_submissions(self): assert pid_first == pid_second assert adapter_id_first == adapter_id_second + @patch("airflow.providers.openlineage.plugins.listener.Stats") + @patch("airflow.providers.openlineage.plugins.listener.Serde") + @patch("airflow.providers.openlineage.plugins.listener._get_process_adapter") + def test_pool_wrappers_return_nothing(self, mock_get_adapter, mock_serde, mock_stats): + """The emitted event must not be pickled back to the parent. + + An event the pickler chokes on would fail the future and be reported as a submission + failure, even though the emission itself succeeded. + """ + from airflow.providers.openlineage.plugins.listener import ( + _emit_manual_state_change_event, + _run_adapter_method, + ) + + assert _run_adapter_method("dag_started", dag_id="dag") is None + assert _emit_manual_state_change_event("fail_task", "ol.event.size.fail.op") is None + + mock_get_adapter.return_value.dag_started.assert_called_once_with(dag_id="dag") + mock_get_adapter.return_value.fail_task.assert_called_once_with() + class TestExecutorInitializer: """Tests for _executor_initializer function.""" @@ -2730,6 +2750,100 @@ def dummy_callable(): listener.log.warning.assert_called_once() assert "recreating" in listener.log.warning.call_args[0][0] + def test_submit_callable_recreates_executor_after_shutdown(self): + """A pool shut down by `before_stopping` must be replaced, not submitted to.""" + listener = OpenLineageListener() + dead_executor = MagicMock() + dead_executor.submit.side_effect = RuntimeError("cannot schedule new futures after shutdown") + new_executor = MagicMock() + listener._executor = dead_executor + listener.log = MagicMock() + + with mock.patch( + "airflow.providers.openlineage.plugins.listener.ProcessPoolExecutor", + return_value=new_executor, + ): + fut = listener.submit_callable(lambda: None) + + assert fut is new_executor.submit.return_value + assert listener._executor is new_executor + + def test_before_stopping_does_not_create_executor(self): + """On the task runner this hook fires per task; it must not spawn a pool to tear it down.""" + listener = OpenLineageListener() + + with mock.patch( + "airflow.providers.openlineage.plugins.listener.ProcessPoolExecutor" + ) as mock_pool_cls: + listener.before_stopping(MagicMock()) + + mock_pool_cls.assert_not_called() + assert listener._executor is None + + def test_before_stopping_detaches_executor(self): + executor = MagicMock() + listener = OpenLineageListener() + listener._executor = executor + + listener.before_stopping(MagicMock()) + + executor.shutdown.assert_called_once_with(wait=True) + assert listener._executor is None + + def test_before_stopping_swallows_shutdown_failure(self): + """A timed-out or off-main-thread shutdown must not escape the hook. + + `AirflowTaskTimeout` derives from `BaseException`, so the guard cannot be `except Exception`. + """ + executor = MagicMock() + executor.shutdown.side_effect = [AirflowTaskTimeout("timed out"), None] + listener = OpenLineageListener() + listener._executor = executor + listener.log = MagicMock() + + listener.before_stopping(MagicMock()) + + assert executor.shutdown.call_args_list == [mock.call(wait=True), mock.call(wait=False)] + assert listener._executor is None + listener.log.warning.assert_called_once() + + @mock.patch("airflow.providers.openlineage.plugins.listener.logging.shutdown") + @mock.patch("airflow.providers.openlineage.plugins.listener.os._exit") + @mock.patch("airflow.providers.openlineage.plugins.listener.os.fork", return_value=0) + def test_fork_execute_child_exits_on_base_exception(self, mock_fork, mock_exit, mock_log_shutdown): + """A SIGINT to the process group reaches the child as KeyboardInterrupt. + + The child must never return: doing so leaves a duplicate task runner behind, sharing the + supervisor connection with the real one. + """ + listener = OpenLineageListener() + listener.log = MagicMock() + + listener._fork_execute(mock.Mock(side_effect=KeyboardInterrupt("ctrl-c")), "on_running") + + mock_exit.assert_called_once_with(0) + mock_log_shutdown.assert_called_once() + listener.log.warning.assert_called_once() + + @mock.patch("airflow.providers.openlineage.plugins.listener.logging.shutdown") + @mock.patch("airflow.providers.openlineage.plugins.listener.os._exit") + @mock.patch("airflow.providers.openlineage.plugins.listener.os.fork", return_value=0) + def test_fork_execute_child_exits_when_setup_fails(self, mock_fork, mock_exit, mock_log_shutdown): + """A failure before the emission call must still exit, not unwind into the caller.""" + listener = OpenLineageListener() + listener.log = MagicMock() + callable_ = mock.Mock() + + with mock.patch( + "airflow.providers.openlineage.plugins.listener.getproctitle", + side_effect=OSError("cannot read proc title"), + ): + listener._fork_execute(callable_, "on_running") + + callable_.assert_not_called() + mock_exit.assert_called_once_with(0) + listener.log.warning.assert_called_once() + @pytest.mark.skipif(AIRFLOW_V_3_0_PLUS, reason="Airflow 2 tests") @pytest.mark.filterwarnings("ignore::airflow.exceptions.AirflowProviderDeprecationWarning") diff --git a/providers/openlineage/tests/unit/openlineage/utils/test_utils.py b/providers/openlineage/tests/unit/openlineage/utils/test_utils.py index e659e4f3de3da..2cd0f3dd3fad4 100644 --- a/providers/openlineage/tests/unit/openlineage/utils/test_utils.py +++ b/providers/openlineage/tests/unit/openlineage/utils/test_utils.py @@ -25,6 +25,7 @@ import pendulum import pytest from openlineage.client.facet_v2 import parent_run +from sqlalchemy.orm.exc import DetachedInstanceError from uuid6 import uuid7 from airflow import DAG @@ -317,6 +318,19 @@ def test_dag_run_version_no_versions(): assert result is None +@pytest.mark.parametrize("key", ["bundle_name", "bundle_version", "version_id", "version_number"]) +def test_dag_run_version_detached_version_row(key): + """The DagVersion rows are lazy-loaded, so reading their columns can hit a detached session.""" + version = MagicMock() + type(version).bundle_name = PropertyMock(side_effect=DetachedInstanceError) + type(version).bundle_version = PropertyMock(side_effect=DetachedInstanceError) + type(version).id = PropertyMock(side_effect=DetachedInstanceError) + type(version).version_number = PropertyMock(side_effect=DetachedInstanceError) + dag_run = MagicMock() + dag_run.dag_versions = [version] + assert DagRunInfo.dag_version_info(dag_run, key) is None + + @pytest.mark.parametrize("key", ["bundle_name", "bundle_version", "version_id", "version_number"]) @pytest.mark.db_test def test_dag_run_version(key): @@ -359,6 +373,37 @@ def test_dag_run_team_name( mock_get_team_name.assert_called_once_with("bundle_name") +@pytest.mark.db_test +@pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="multi-team requires Airflow 3.3+") +@pytest.mark.parametrize("team_name", ["team_a", None]) +@patch("airflow.models.dagbundle.DagBundleModel.get_team_name") +@patch("airflow.providers.openlineage.utils.utils.airflow_conf.getboolean", return_value=True) +def test_dag_run_team_name_from_execution_api_dag_run(mock_getboolean, mock_get_team_name, team_name): + """The task runner has no DB session, so a DagRun carrying `team_name` must be trusted as-is.""" + # A resolvable bundle plus a DB answer, so falling through to the lookup would be observable. + dagrun_mock = MagicMock(spec_set=["team_name", "dag_versions"]) + dagrun_mock.team_name = team_name + dagrun_mock.dag_versions = [MagicMock(bundle_name="bundle_name")] + mock_get_team_name.return_value = "from_db" + + assert DagRunInfo.team_name(dagrun_mock) == team_name + + mock_get_team_name.assert_not_called() + + +@pytest.mark.db_test +@pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="multi-team requires Airflow 3.3+") +@patch("airflow.models.dagbundle.DagBundleModel.get_team_name") +@patch("airflow.providers.openlineage.utils.utils.airflow_conf.getboolean", return_value=True) +def test_dag_run_team_name_lookup_failure_does_not_raise(mock_getboolean, mock_get_team_name): + """A failed lookup must degrade to None -- `_cast_fields` would otherwise lose the whole event.""" + dagrun_mock = MagicMock(DagRun) + dagrun_mock.dag_versions = [MagicMock(bundle_name="bundle_name")] + mock_get_team_name.side_effect = RuntimeError("Session must be set before!") + + assert DagRunInfo.team_name(dagrun_mock) is None + + @pytest.mark.db_test @pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="multi-team requires Airflow 3.3+") @patch("airflow.models.dagbundle.DagBundleModel.get_team_name") @@ -3588,6 +3633,20 @@ def test_task_with_none_timestamps_fallback_to_zero(self, dag_maker): assert result["airflowState"].tasksDuration["terminated_task"] == 0.0 + @patch( + "airflow.providers.openlineage.utils.utils.DagRun.fetch_task_instances", + side_effect=Exception("db hiccup"), + ) + def test_db_failure_returns_empty_facet(self, _mock_fetch): + """A DB error in the pool worker should drop only the facet, not the whole event.""" + result = get_airflow_state_run_facet( + dag_id="test_dag", + run_id="test_run", + task_ids=["test_task"], + dag_run_state=DagRunState.SUCCESS, + ) + assert result == {} + @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="Airflow 3 specific test") def test_is_dag_run_asset_triggered_af3(): From ae5613221a58d26b1b30a8ce350e9f6536e1f684 Mon Sep 17 00:00:00 2001 From: Kacper Muda Date: Wed, 29 Jul 2026 15:56:34 +0200 Subject: [PATCH 2/2] Add more tests --- .../unit/openlineage/plugins/test_listener.py | 51 +++++++++++++++---- 1 file changed, 40 insertions(+), 11 deletions(-) diff --git a/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py b/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py index e22a950cdc12b..a6b4fa4af6113 100644 --- a/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py +++ b/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py @@ -17,10 +17,11 @@ from __future__ import annotations import json +import logging import uuid from collections import defaultdict from collections.abc import Callable -from concurrent.futures import Future +from concurrent.futures import Future, ProcessPoolExecutor from contextlib import suppress from datetime import datetime from types import SimpleNamespace @@ -2753,11 +2754,11 @@ def dummy_callable(): def test_submit_callable_recreates_executor_after_shutdown(self): """A pool shut down by `before_stopping` must be replaced, not submitted to.""" listener = OpenLineageListener() - dead_executor = MagicMock() + dead_executor = MagicMock(spec=ProcessPoolExecutor) dead_executor.submit.side_effect = RuntimeError("cannot schedule new futures after shutdown") - new_executor = MagicMock() + new_executor = MagicMock(spec=ProcessPoolExecutor) listener._executor = dead_executor - listener.log = MagicMock() + listener.log = MagicMock(spec=logging.Logger) with mock.patch( "airflow.providers.openlineage.plugins.listener.ProcessPoolExecutor", @@ -2781,7 +2782,7 @@ def test_before_stopping_does_not_create_executor(self): assert listener._executor is None def test_before_stopping_detaches_executor(self): - executor = MagicMock() + executor = MagicMock(spec=ProcessPoolExecutor) listener = OpenLineageListener() listener._executor = executor @@ -2795,11 +2796,11 @@ def test_before_stopping_swallows_shutdown_failure(self): `AirflowTaskTimeout` derives from `BaseException`, so the guard cannot be `except Exception`. """ - executor = MagicMock() + executor = MagicMock(spec=ProcessPoolExecutor) executor.shutdown.side_effect = [AirflowTaskTimeout("timed out"), None] listener = OpenLineageListener() listener._executor = executor - listener.log = MagicMock() + listener.log = MagicMock(spec=logging.Logger) listener.before_stopping(MagicMock()) @@ -2817,9 +2818,11 @@ def test_fork_execute_child_exits_on_base_exception(self, mock_fork, mock_exit, supervisor connection with the real one. """ listener = OpenLineageListener() - listener.log = MagicMock() + listener.log = MagicMock(spec=logging.Logger) - listener._fork_execute(mock.Mock(side_effect=KeyboardInterrupt("ctrl-c")), "on_running") + listener._fork_execute( + mock.Mock(spec=Callable, side_effect=KeyboardInterrupt("ctrl-c")), "on_running" + ) mock_exit.assert_called_once_with(0) mock_log_shutdown.assert_called_once() @@ -2831,8 +2834,8 @@ def test_fork_execute_child_exits_on_base_exception(self, mock_fork, mock_exit, def test_fork_execute_child_exits_when_setup_fails(self, mock_fork, mock_exit, mock_log_shutdown): """A failure before the emission call must still exit, not unwind into the caller.""" listener = OpenLineageListener() - listener.log = MagicMock() - callable_ = mock.Mock() + listener.log = MagicMock(spec=logging.Logger) + callable_ = mock.Mock(spec=Callable) with mock.patch( "airflow.providers.openlineage.plugins.listener.getproctitle", @@ -2844,6 +2847,32 @@ def test_fork_execute_child_exits_when_setup_fails(self, mock_fork, mock_exit, m mock_exit.assert_called_once_with(0) listener.log.warning.assert_called_once() + @mock.patch( + "airflow.providers.openlineage.plugins.listener.logging.shutdown", + side_effect=RuntimeError("handler close failed"), + ) + @mock.patch("airflow.providers.openlineage.plugins.listener.os._exit") + @mock.patch("airflow.providers.openlineage.plugins.listener.os.fork", return_value=0) + def test_fork_execute_child_exits_when_logging_shutdown_raises( + self, mock_fork, mock_exit, mock_log_shutdown + ): + """logging.shutdown() raising must not bypass os._exit(0). + + A remote task-log handler (S3/GCS) whose close() raises, or a handler lock held + at fork time, can make logging.shutdown() propagate out of the finally block. + Without the nested finally the child unwinds into the task runner and runs as a + second process sharing the supervisor connection. + """ + listener = OpenLineageListener() + listener.log = MagicMock(spec=logging.Logger) + + # In the real process os._exit(0) terminates execution before the RuntimeError + # can propagate; with a mocked os._exit the exception surfaces in the test. + with pytest.raises(RuntimeError, match="handler close failed"): + listener._fork_execute(mock.Mock(spec=Callable), "on_running") + + mock_exit.assert_called_once_with(0) + @pytest.mark.skipif(AIRFLOW_V_3_0_PLUS, reason="Airflow 2 tests") @pytest.mark.filterwarnings("ignore::airflow.exceptions.AirflowProviderDeprecationWarning")