Skip to content

Commit 46f5e8d

Browse files
committed
Avoid deactivating configured Dag bundles on transient get bundle errors
Signed-off-by: PoAn Yang <payang@apache.org>
1 parent 9bb184a commit 46f5e8d

3 files changed

Lines changed: 120 additions & 1 deletion

File tree

‎airflow-core/src/airflow/dag_processing/bundles/manager.py‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -355,13 +355,18 @@ def _extract_and_sign_template(bundle_name: str) -> tuple[str | None, dict]:
355355
if not team:
356356
raise _bundle_item_exc(f"Team '{config.team_name}' does not exist")
357357

358+
# Pop configured bundles out of ``stored`` before ``_extract_and_sign_template``.
359+
# A transient connection error from that call must not deactivate the bundle,
360+
# which would make all Dags in it non-schedulable.
361+
bundle = stored.pop(name, None)
362+
358363
try:
359364
new_template, new_params = _extract_and_sign_template(name)
360365
except Exception as e:
361366
self.log.exception("Error creating bundle '%s': %s", name, e)
362367
continue
363368

364-
if bundle := stored.pop(name, None):
369+
if bundle:
365370
bundle.active = True
366371
if new_template != bundle.signed_url_template:
367372
bundle.signed_url_template = new_template

‎airflow-core/tests/unit/dag_processing/bundles/test_dag_bundle_manager.py‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -221,6 +221,36 @@ def _get_bundle_names_and_active():
221221
]
222222

223223

224+
@pytest.mark.db_test
225+
@conf_vars({("core", "LOAD_EXAMPLES"): "False"})
226+
def test_sync_bundles_to_db_keeps_configured_bundle_on_transient_creation_error(clear_db, session):
227+
"""A configured bundle failing to instantiate with transient error must not be deactivated."""
228+
with patch.dict(
229+
os.environ, {"AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST": json.dumps(BASIC_BUNDLE_CONFIG)}
230+
):
231+
manager = DagBundlesManager()
232+
manager.sync_bundles_to_db()
233+
234+
session.add(
235+
ParseImportError(
236+
bundle_name="my-test-bundle",
237+
filename="some_file.py",
238+
stacktrace="some error",
239+
)
240+
)
241+
session.flush()
242+
243+
manager = DagBundlesManager()
244+
with patch.object(
245+
DagBundlesManager, "get_bundle", autospec=True, side_effect=RuntimeError("transient error")
246+
):
247+
manager.sync_bundles_to_db()
248+
249+
bundle = session.scalar(select(DagBundleModel).where(DagBundleModel.name == "my-test-bundle"))
250+
assert bundle.active is True
251+
assert session.scalar(select(func.count(ParseImportError.id))) == 1
252+
253+
224254
@pytest.mark.db_test
225255
@conf_vars({("core", "LOAD_EXAMPLES"): "False"})
226256
def test_sync_bundles_to_db_does_not_log_removing_none_team(clear_db, caplog):

‎airflow-core/tests/unit/dag_processing/test_manager.py‎

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@
4545
from airflow._shared.timezones import timezone
4646
from airflow.callbacks.callback_requests import DagCallbackRequest
4747
from airflow.dag_processing.bundles.base import BaseDagBundle
48+
from airflow.dag_processing.bundles.local import LocalDagBundle
4849
from airflow.dag_processing.bundles.manager import DagBundlesManager
4950
from airflow.dag_processing.dagbag import DagBag
5051
from airflow.dag_processing.manager import (
@@ -83,6 +84,18 @@
8384

8485
logger = logging.getLogger(__name__)
8586
TEST_DAG_FOLDER = Path(__file__).parents[1].resolve() / "dags"
87+
88+
89+
class _HealthControlledDagBundle(LocalDagBundle):
90+
"""LocalDagBundle that fails to instantiate unless its health file contains "ok"."""
91+
92+
def __init__(self, *, health_file: str, **kwargs) -> None:
93+
state = Path(health_file).read_text().strip()
94+
if state != "ok":
95+
raise RuntimeError(f"simulated {state}")
96+
super().__init__(**kwargs)
97+
98+
8699
DEFAULT_DATE = timezone.datetime(2016, 1, 1)
87100

88101

@@ -1253,6 +1266,77 @@ def test_deactivate_stale_dags_logs_stuck_legacy_row_count(self, session):
12531266
assert len(legacy_log_calls) == 1
12541267
assert legacy_log_calls[0].args[1] == 3
12551268

1269+
def test_transient_bundle_creation_error_keeps_bundle_active_and_dags_fresh(self, tmp_path):
1270+
"""A bundle failing to instantiate with transient error must not be deactivated."""
1271+
dags_dir = tmp_path / "dags"
1272+
dags_dir.mkdir()
1273+
(dags_dir / "good_dag.py").write_text(
1274+
textwrap.dedent(
1275+
"""
1276+
from airflow.providers.standard.operators.empty import EmptyOperator
1277+
from airflow.sdk import DAG
1278+
1279+
with DAG(dag_id="dag_in_flaky_bundle", schedule=None):
1280+
EmptyOperator(task_id="noop")
1281+
"""
1282+
)
1283+
)
1284+
(dags_dir / "broken_dag.py").write_text("an invalid airflow DAG")
1285+
health_file = tmp_path / "health.txt"
1286+
1287+
bundle_config = [
1288+
{
1289+
"name": "flaky-bundle",
1290+
"classpath": f"{__name__}._HealthControlledDagBundle",
1291+
"kwargs": {
1292+
"path": str(dags_dir),
1293+
"health_file": str(health_file),
1294+
"refresh_interval": 0,
1295+
},
1296+
}
1297+
]
1298+
1299+
def get_db_state():
1300+
with create_session() as session:
1301+
active = session.scalar(
1302+
select(DagBundleModel.active).where(DagBundleModel.name == "flaky-bundle")
1303+
)
1304+
fresh_dags = session.scalar(
1305+
select(func.count())
1306+
.select_from(DagModel)
1307+
.where(DagModel.bundle_name == "flaky-bundle", ~DagModel.is_stale)
1308+
)
1309+
import_errors = session.scalar(select(func.count()).select_from(ParseImportError))
1310+
return active, fresh_dags, import_errors
1311+
1312+
with conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps(bundle_config)}):
1313+
# Phase 1: misconfigured from day one — sync never succeeds, so no bundle row exists at all
1314+
health_file.write_text("misconfiguration")
1315+
DagFileProcessorManager(max_runs=1, processor_timeout=365 * 86_400).run()
1316+
assert get_db_state() == (None, 0, 0)
1317+
1318+
# Phase 2: healthy restart — the bundle row appears and Dags get parsed
1319+
health_file.write_text("ok")
1320+
DagFileProcessorManager(max_runs=1, processor_timeout=365 * 86_400).run()
1321+
assert get_db_state() == (True, 1, 1)
1322+
1323+
# Phase 3: restart with transient connection error — must not deactivate the bundle
1324+
health_file.write_text("transient connection error")
1325+
DagFileProcessorManager(max_runs=1, processor_timeout=365 * 86_400).run()
1326+
assert get_db_state() == (True, 1, 1)
1327+
1328+
# Phase 4: healthy again
1329+
health_file.write_text("ok")
1330+
DagFileProcessorManager(max_runs=1, processor_timeout=365 * 86_400).run()
1331+
assert get_db_state() == (True, 1, 1)
1332+
1333+
# Phase 5: removing the (again failing) bundle from config must follow the removal
1334+
# path: deactivate it, mark its Dags stale, and drop its import errors.
1335+
health_file.write_text("transient connection error")
1336+
with conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps([])}):
1337+
DagFileProcessorManager(max_runs=1, processor_timeout=365 * 86_400).run()
1338+
assert get_db_state() == (False, 0, 0)
1339+
12561340
@mock.patch("airflow.dag_processing.manager.is_lock_not_available_error")
12571341
@pytest.mark.usefixtures("testing_dag_bundle")
12581342
def test_deactivate_stale_dags_handles_lock_timeout(self, mock_is_lock_not_available, session, caplog):

0 commit comments

Comments
 (0)