diff --git a/task-sdk/src/airflow/sdk/importers/python_importer.py b/task-sdk/src/airflow/sdk/importers/python_importer.py index e6fae5b44529a..502745792f24e 100644 --- a/task-sdk/src/airflow/sdk/importers/python_importer.py +++ b/task-sdk/src/airflow/sdk/importers/python_importer.py @@ -250,10 +250,9 @@ def _handle_sigsegv(signum, frame): if mod_name in sys.modules: del sys.modules[mod_name] - DagContext.current_autoregister_module_name = mod_name - def parse() -> list[types.ModuleType]: try: + DagContext.current_autoregister_module_name = mod_name with definition.import_context(): loader: importlib.abc.Loader if get_file_suffix(definition) == ".pyc": @@ -285,6 +284,10 @@ def parse() -> list[types.ModuleType]: ) ) return [] + finally: + # Reset on every exit: DagContext.pop() looks this name up in sys.modules, which may not + # hold it if the load did not complete. + DagContext.current_autoregister_module_name = None dagbag_import_timeout: float try: @@ -323,7 +326,6 @@ def _process_modules( } top_level_dags.update(DagContext.autoregistered_dags) - DagContext.current_autoregister_module_name = None DagContext.autoregistered_dags.clear() for dag, _mod in top_level_dags: diff --git a/task-sdk/tests/task_sdk/importers/test_python_importer.py b/task-sdk/tests/task_sdk/importers/test_python_importer.py index 54cfdc94df664..eebbf86191672 100644 --- a/task-sdk/tests/task_sdk/importers/test_python_importer.py +++ b/task-sdk/tests/task_sdk/importers/test_python_importer.py @@ -33,6 +33,7 @@ import pytest +from airflow.sdk import DAG from airflow.sdk.exceptions import AirflowConfigException from airflow.sdk.importers import ( DagImportError, @@ -347,6 +348,42 @@ def test_invalid_dagbag_import_timeout_raises_custom_exception(self, mock_bundle bundle=mock_bundle, ) + @pytest.mark.parametrize( + ("timeout_policy", "expectation"), + [ + pytest.param( + {"return_value": "invalid_timeout_str"}, + pytest.raises(AirflowConfigException), + id="config-error-propagates", + ), + pytest.param( + {"side_effect": RuntimeError("policy failed")}, + contextlib.nullcontext(), + id="error-reported", + ), + pytest.param( + {"return_value": 0}, + pytest.raises(KeyboardInterrupt), + id="keyboard-interrupt-propagates", + ), + ], + ) + def test_failed_load_does_not_break_later_dags(self, mock_bundle, timeout_policy, expectation): + (mock_bundle.path / "dag.py").write_text("raise KeyboardInterrupt\n") + mock_settings = mock.MagicMock(spec=["get_dagbag_import_timeout"]) + mock_settings.get_dagbag_import_timeout.configure_mock(**timeout_policy) + mock_airflow = mock.MagicMock(spec=["settings"], settings=mock_settings) + + with mock.patch.dict("sys.modules", {"airflow": mock_airflow}), expectation: + PythonDagImporter().import_definition( + FilesystemDagDefinition(path=mock_bundle.path / "dag.py"), + bundle=mock_bundle, + ) + + with DAG("built_after_failed_load") as dag: + pass + assert dag.dag_id == "built_after_failed_load" + @mock.patch.object(PythonDagImporter, "_load_modules", side_effect=TypeError("unexpected None")) def test_unexpected_type_error_captured_in_result_errors(self, mock_load, mock_bundle): importer = PythonDagImporter()