Skip to content

Commit 04f91fc

Browse files
SamWheatingpotiuk
authored andcommitted
[v3-3-test] Fix deactivation of stale zip-packaged DAGs (#70586)
(cherry picked from commit ef54cd6) Co-authored-by: Sam Wheating <samwheating@gmail.com>
1 parent f15c377 commit 04f91fc

2 files changed

Lines changed: 72 additions & 3 deletions

File tree

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

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -506,7 +506,11 @@ def deactivate_stale_dags(
506506
# When the Dag's last_parsed_time is more than the stale_dag_threshold older than the
507507
# Dag file's last_finish_time, the Dag is considered stale as has apparently been removed from the file,
508508
# This is especially relevant for Dag files that generate Dags in a dynamic manner.
509-
file_info = DagFileInfo(rel_path=Path(dag.relative_fileloc), bundle_name=dag.bundle_name)
509+
rel_path = Path(dag.relative_fileloc)
510+
file_info = DagFileInfo(rel_path=rel_path, bundle_name=dag.bundle_name)
511+
if file_info not in last_parsed:
512+
# Zip-packaged dags are keyed by the archive path, not the inner file, so try the parent as well
513+
file_info = DagFileInfo(rel_path=rel_path.parent, bundle_name=dag.bundle_name)
510514
if last_finish_time := last_parsed.get(file_info, None):
511515
if dag.last_parsed_time + timedelta(seconds=self.stale_dag_threshold) < last_finish_time:
512516
self.log.info(

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

Lines changed: 67 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -559,7 +559,6 @@ def test_files_sorted_random_seeded_by_host(self):
559559
assert manager._file_queue == expected
560560

561561
# Verify running it again produces same order
562-
manager._files = []
563562
manager.prepare_file_queue(known_files=known_files)
564563
assert manager._file_queue == expected
565564

@@ -999,7 +998,6 @@ def test_scan_stale_dags(self, session):
999998
run_count=1,
1000999
last_num_of_db_queries=1,
10011000
)
1002-
manager._files = [test_dag_path]
10031001
manager._file_stats[test_dag_path] = stat
10041002

10051003
active_dag_count = session.scalar(
@@ -1029,6 +1027,73 @@ def test_scan_stale_dags(self, session):
10291027
# SerializedDagModel gives history about Dags
10301028
assert serialized_dag_count == 1
10311029

1030+
@pytest.mark.usefixtures("testing_dag_bundle")
1031+
def test_scan_stale_dags_deactivates_zip_packaged_dags(self, session, test_zip_path):
1032+
"""
1033+
Ensure that zip-packaged DAGs are marked inactive when the file is parsed but the
1034+
DagModel.last_parsed_time is not updated, testing fallback to the parent path when
1035+
comparing DAG.relative_fileloc to last_parsed entries.
1036+
"""
1037+
manager = DagFileProcessorManager(
1038+
max_runs=1,
1039+
processor_timeout=10 * 60,
1040+
)
1041+
bundle = MagicMock()
1042+
bundle.name = "testing"
1043+
manager._dag_bundles = [bundle]
1044+
1045+
test_dag_path = DagFileInfo(
1046+
rel_path=Path("test_zip.zip"),
1047+
bundle_path=Path(test_zip_path).parent,
1048+
bundle_name="testing",
1049+
)
1050+
dagbag = DagBag(
1051+
test_dag_path.absolute_path,
1052+
bundle_path=test_dag_path.bundle_path,
1053+
)
1054+
1055+
# Add stale DAG to the DB
1056+
dag = dagbag.get_dag("test_zip_dag")
1057+
sync_dag_to_db(dag, session=session)
1058+
1059+
# Add DAG to the file_parsing_stats
1060+
stat = DagFileStat(
1061+
num_dags=1,
1062+
import_errors=0,
1063+
last_finish_time=timezone.utcnow() + timedelta(hours=1),
1064+
last_duration=1,
1065+
run_count=1,
1066+
last_num_of_db_queries=1,
1067+
)
1068+
manager._file_stats[test_dag_path] = stat
1069+
1070+
active_dag_count = session.scalar(
1071+
select(func.count(DagModel.dag_id)).where(
1072+
DagModel.dag_id == "test_zip_dag",
1073+
~DagModel.is_stale,
1074+
DagModel.relative_fileloc == str(test_dag_path.rel_path / "test_zip.py"),
1075+
)
1076+
)
1077+
assert active_dag_count == 1
1078+
1079+
manager._scan_stale_dags()
1080+
1081+
active_dag_count = session.scalar(
1082+
select(func.count(DagModel.dag_id)).where(
1083+
DagModel.dag_id == "test_zip_dag",
1084+
~DagModel.is_stale,
1085+
DagModel.relative_fileloc == str(test_dag_path.rel_path / "test_zip.py"),
1086+
)
1087+
)
1088+
assert active_dag_count == 0
1089+
1090+
serialized_dag_count = session.scalar(
1091+
select(func.count(SerializedDagModel.dag_id)).where(SerializedDagModel.dag_id == dag.dag_id)
1092+
)
1093+
# Deactivating the DagModel should not delete the SerializedDagModel
1094+
# SerializedDagModel gives history about Dags
1095+
assert serialized_dag_count == 1
1096+
10321097
@pytest.mark.usefixtures("testing_dag_bundle")
10331098
def test_deactivate_stale_dags_marks_dags_in_inactive_bundles(self, session):
10341099
"""Dags whose bundle is no longer active should be marked stale even without a parse signal."""

0 commit comments

Comments
 (0)