Skip to content

Commit 17776f4

Browse files
dilnazanlidjason810496
authored andcommitted
Handle Dags nested at any depth in Dag archives
(cherry picked from commit b4df58f)
1 parent 24eb53e commit 17776f4

4 files changed

Lines changed: 47 additions & 13 deletions

File tree

‎airflow-core/src/airflow/dag_processing/dagbag.py‎

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,7 @@
5050
from airflow.serialization.definitions.dag import SerializedDAG
5151
from airflow.serialization.definitions.notset import NOTSET, ArgNotSet, is_arg_set
5252
from airflow.serialization.serialized_objects import DagSerialization, LazyDeserializedDAG
53-
from airflow.utils.file import correct_maybe_zipped, find_enclosing_file
53+
from airflow.utils.file import find_enclosing_file
5454
from airflow.utils.log.logging_mixin import LoggingMixin
5555
from airflow.utils.session import NEW_SESSION, provide_session
5656

@@ -273,9 +273,7 @@ def get_dag(self, dag_id, *, session: Session = NEW_SESSION):
273273
self.dags.pop(dag_id, None)
274274
if dag is None or is_expired:
275275
# Reprocess source file.
276-
found_dags = self.process_file(
277-
filepath=correct_maybe_zipped(orm_dag.fileloc), only_if_updated=False
278-
)
276+
found_dags = self.process_file(filepath=orm_dag.fileloc, only_if_updated=False)
279277

280278
# If the source file no longer exports `dag_id`, delete it from self.dags
281279
if found_dags and dag_id in [found_dag.dag_id for found_dag in found_dags]:
@@ -284,8 +282,8 @@ def get_dag(self, dag_id, *, session: Session = NEW_SESSION):
284282
return self.dags.get(dag_id)
285283

286284
def process_file(self, filepath, only_if_updated=True, safe_mode=True):
287-
"""Process a DAG file and return found DAGs."""
288-
if filepath is None or not os.path.isfile(filepath):
285+
"""Process a Dag file, or a Dag definition nested in one, and return found Dags."""
286+
if filepath is None or os.path.isdir(filepath):
289287
return []
290288
return [
291289
dag

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

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -515,11 +515,17 @@ def deactivate_stale_dags(
515515
# Dag file's last_finish_time, the Dag is considered stale as has apparently been removed from the file,
516516
# This is especially relevant for Dag files that generate Dags in a dynamic manner.
517517
rel_path = Path(dag.relative_fileloc)
518-
file_info = DagFileInfo(rel_path=rel_path, bundle_name=dag.bundle_name)
519-
if file_info not in last_parsed:
520-
# Zip-packaged dags are keyed by the archive path, not the inner file, so try the parent as well
521-
file_info = DagFileInfo(rel_path=rel_path.parent, bundle_name=dag.bundle_name)
522-
if last_finish_time := last_parsed.get(file_info, None):
518+
# A Dag nested in a container (``archive.zip/sub/dag.py``) is parsed under the container.
519+
last_finish_time = next(
520+
(
521+
last_parsed[file_info]
522+
for candidate in (rel_path, *rel_path.parents)
523+
if (file_info := DagFileInfo(rel_path=candidate, bundle_name=dag.bundle_name))
524+
in last_parsed
525+
),
526+
None,
527+
)
528+
if last_finish_time:
523529
if dag.last_parsed_time + timedelta(seconds=self.stale_dag_threshold) < last_finish_time:
524530
self.log.info(
525531
"Deactivating stale DAG %s. Not parsed for %s seconds (last parsed: %s).",

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

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1137,6 +1137,13 @@ def basic_cycle():
11371137
self.validate_dags(test_dag, found_dags, dagbag, should_be_found=False)
11381138
assert file_path in dagbag.import_errors
11391139

1140+
def test_process_file_nested_definition(self, tmp_path, test_zip_path):
1141+
dagbag = DagBag(dag_folder=os.fspath(tmp_path), collect_dags=False)
1142+
1143+
found_dags = dagbag.process_file(os.path.join(test_zip_path, "test_zip.py"))
1144+
1145+
assert sorted(dag.dag_id for dag in found_dags) == ["test_zip_autoregister", "test_zip_dag"]
1146+
11401147
def test_process_file_with_none(self, tmp_path):
11411148
"""
11421149
test that process_file can handle Nones

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

Lines changed: 25 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1389,6 +1389,29 @@ def test_deactivate_stale_dags_marks_dags_in_inactive_bundles(self, session):
13891389
)
13901390
assert is_stale_by_dag == {"dag_in_inactive_bundle": True, "dag_in_active_bundle": False}
13911391

1392+
@pytest.mark.usefixtures("testing_dag_bundle")
1393+
def test_deactivate_stale_dags_matches_dags_nested_in_a_container(self, session):
1394+
session.add(
1395+
DagModel(
1396+
dag_id="nested_dag",
1397+
bundle_name="testing",
1398+
relative_fileloc="dags.zip/sub/nested_dag.py",
1399+
last_parsed_time=timezone.utcnow(),
1400+
is_stale=False,
1401+
)
1402+
)
1403+
session.flush()
1404+
1405+
manager = DagFileProcessorManager(max_runs=1, processor_timeout=10 * 60)
1406+
manager.deactivate_stale_dags(
1407+
last_parsed={
1408+
DagFileInfo(rel_path=Path("dags.zip"), bundle_name="testing"): timezone.utcnow()
1409+
+ timedelta(hours=1)
1410+
}
1411+
)
1412+
1413+
assert session.scalar(select(DagModel.is_stale).where(DagModel.dag_id == "nested_dag"))
1414+
13921415
@pytest.mark.usefixtures("testing_dag_bundle")
13931416
def test_deactivate_stale_dags_marks_dags_with_null_bundle_name(self, session):
13941417
"""Dags carried over from Airflow 2.x keep a NULL bundle_name and must still be deactivated.
@@ -4192,8 +4215,8 @@ def test_refresh_dag_bundles_discovery_failure_keeps_known_files_and_dags(self):
41924215
),
41934216
mock.patch.object(manager, "update_bundle_state"),
41944217
mock.patch.object(manager, "_find_files_in_bundle", side_effect=OSError("listing failed")),
4195-
mock.patch.object(manager, "deactivate_deleted_dags") as mock_deactivate,
4196-
mock.patch.object(manager, "clear_orphaned_import_errors") as mock_clear,
4218+
mock.patch.object(manager, "deactivate_deleted_dags", autospec=True) as mock_deactivate,
4219+
mock.patch.object(manager, "clear_orphaned_import_errors", autospec=True) as mock_clear,
41974220
mock.patch.object(manager, "handle_removed_files"),
41984221
mock.patch.object(manager, "_resort_file_queue"),
41994222
mock.patch.object(manager, "_add_new_files_to_queue"),

0 commit comments

Comments
 (0)