|
45 | 45 | from airflow._shared.timezones import timezone |
46 | 46 | from airflow.callbacks.callback_requests import DagCallbackRequest |
47 | 47 | from airflow.dag_processing.bundles.base import BaseDagBundle |
| 48 | +from airflow.dag_processing.bundles.local import LocalDagBundle |
48 | 49 | from airflow.dag_processing.bundles.manager import DagBundlesManager |
49 | 50 | from airflow.dag_processing.dagbag import DagBag |
50 | 51 | from airflow.dag_processing.manager import ( |
|
83 | 84 |
|
84 | 85 | logger = logging.getLogger(__name__) |
85 | 86 | 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 | + |
86 | 99 | DEFAULT_DATE = timezone.datetime(2016, 1, 1) |
87 | 100 |
|
88 | 101 |
|
@@ -1253,6 +1266,77 @@ def test_deactivate_stale_dags_logs_stuck_legacy_row_count(self, session): |
1253 | 1266 | assert len(legacy_log_calls) == 1 |
1254 | 1267 | assert legacy_log_calls[0].args[1] == 3 |
1255 | 1268 |
|
| 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 | + |
1256 | 1340 | @mock.patch("airflow.dag_processing.manager.is_lock_not_available_error") |
1257 | 1341 | @pytest.mark.usefixtures("testing_dag_bundle") |
1258 | 1342 | def test_deactivate_stale_dags_handles_lock_timeout(self, mock_is_lock_not_available, session, caplog): |
|
0 commit comments