Skip to content

Reset Dag autoregistration state when a Python Dag file fails to load #74911

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
shahar1 merged 2 commits into from
Oct 24, 2026
Merged

Reset Dag autoregistration state when a Python Dag file fails to load #74911

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
eea37cf
Reset Dag autoregistration state when a Python Dag file fails to load
shahar1 Oct 17, 2026
e206681
Keep the autoregister module name set only while the Dag file executes
shahar1 Oct 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 5 additions & 3 deletions task-sdk/src/airflow/sdk/importers/python_importer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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":
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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:
Expand Down
37 changes: 37 additions & 0 deletions task-sdk/tests/task_sdk/importers/test_python_importer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@

import pytest

from airflow.sdk import DAG
from airflow.sdk.exceptions import AirflowConfigException
from airflow.sdk.importers import (
DagImportError,
Expand Down Expand Up @@ -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()
Expand Down