From 116461e55b9f7c202215a3afef3d72477b6c5a21 Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Fri, 14 Aug 2026 16:12:43 +0100 Subject: [PATCH 01/11] Pin the statement budget for persisting Dag parse results The Dag processor's manager issues these statements once per parsed file, and how many was written down nowhere, so a change could add round trips to that path unnoticed. Measure it: a call costs 10 statements per file plus 3 per Dag the file defines, and a sweep pays that fixed price once per persistence call -- currently one call per file. Those counts are calibrated against Postgres, and the tests carrying them are marked for it, since statement counts differ by dialect. The sweep cost is measured through _collect_results and derives both prices from its own measurements, so it carries no dialect-specific number, runs on any backend, and tracks whatever strategy the manager uses. Timing this proved unreliable: on a throttled single-core container, repeat runs of identical code vary by more than a third, far larger than the savings under consideration. Counting statements is exact. --- .../test_parse_result_query_budget.py | 255 ++++++++++++++++++ 1 file changed, 255 insertions(+) create mode 100644 airflow-core/tests/unit/dag_processing/test_parse_result_query_budget.py diff --git a/airflow-core/tests/unit/dag_processing/test_parse_result_query_budget.py b/airflow-core/tests/unit/dag_processing/test_parse_result_query_budget.py new file mode 100644 index 0000000000000..a5443beff26fa --- /dev/null +++ b/airflow-core/tests/unit/dag_processing/test_parse_result_query_budget.py @@ -0,0 +1,255 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +""" +Statement budget for persisting Dag parse results. + +The Dag processor's manager issues these once per parsed file, so a change adding round trips has +to move a number here and account for it in review. + +Counts are calibrated against Postgres and the tests carrying them are marked for it, because +statement counts differ by dialect: ``_update_import_errors`` deletes with a predicate the ORM +cannot evaluate in Python, which costs one statement where ``DELETE ... RETURNING`` exists and two +on MySQL. The sweep test hard-codes only the call count and measures the rest, so it runs anywhere. +""" + +from __future__ import annotations + +import re +import time +from collections import Counter +from contextlib import contextmanager, suppress +from pathlib import Path +from socket import socket, socketpair +from unittest.mock import MagicMock + +import pytest +from sqlalchemy import event +from uuid6 import uuid7 + +from airflow.dag_processing.collection import update_dag_parsing_results_in_db +from airflow.dag_processing.manager import DagFileInfo, DagFileProcessorManager, DagFileStat +from airflow.dag_processing.processor import DagFileParsingResult, DagFileProcessorProcess +from airflow.providers.standard.operators.empty import EmptyOperator +from airflow.sdk import DAG +from airflow.serialization.serialized_objects import LazyDeserializedDAG + +from tests_common.test_utils.db import clear_db_dags, clear_db_serialized_dags + +pytestmark = pytest.mark.db_test + +BUNDLE = "testing" +DAG_FILE = "budget_dags.py" + +# Per persistence call, and per Dag in the file. +FIXED_PER_FILE = 10 +PER_DAG = 3 + +SWEEP_FILES = 4 +# Calls the manager takes for that sweep: one per file today, 1 if a sweep is ever batched. +SWEEP_CALLS = 4 + + +def _classify(statement: str) -> tuple[str, str]: + """Reduce a statement to (operation, table) so a failure says what changed, not just by how much.""" + collapsed = " ".join(statement.split()).lower() + operation = collapsed.split(" ", 1)[0] + patterns = { + "select": r"\bfrom\s+([a-z_][a-z0-9_]*)", + "delete": r"\bfrom\s+([a-z_][a-z0-9_]*)", + "insert": r"\binto\s+([a-z_][a-z0-9_]*)", + "update": r"\bupdate\s+([a-z_][a-z0-9_]*)", + } + match = re.search(patterns[operation], collapsed) if operation in patterns else None + return operation, (match.group(1) if match else "?") + + +@contextmanager +def _count_statements(session): + counts: Counter[tuple[str, str]] = Counter() + + def _capture(conn, cursor, statement, parameters, context, executemany): + counts[_classify(statement)] += 1 + + bind = session.get_bind() + event.listen(bind, "before_cursor_execute", _capture) + try: + yield counts + finally: + event.remove(bind, "before_cursor_execute", _capture) + + +def _breakdown(counts: Counter[tuple[str, str]]) -> str: + return "\n".join(f" {n:>3} {op.upper():<6} {table}" for (op, table), n in sorted(counts.items())) + + +@pytest.fixture(autouse=True) +def clean_db(): + yield + clear_db_serialized_dags() + clear_db_dags() + + +@pytest.fixture +def sockets(): + """Socket ends to close on teardown; ``DagFileProcessorProcess.close()`` leaves stdin open.""" + open_sockets: list[socket] = [] + yield open_sockets + for sock in open_sockets: + with suppress(OSError): + sock.close() + + +def _make_dags(dag_file: Path, dag_ids: list[str], rel_path: str) -> list[LazyDeserializedDAG]: + # DagCode reads the source off disk; without a real file the Dags fail to serialize and the + # measured statements stop resembling a real parse. + dag_file.write_text("# statement budget fixture\n") + dags = [] + for dag_id in dag_ids: + dag = DAG(dag_id=dag_id, schedule="@daily") + EmptyOperator(task_id="task1", dag=dag) + dag.fileloc = str(dag_file) + dag.relative_fileloc = rel_path + dags.append(LazyDeserializedDAG.from_dag(dag)) + return dags + + +def _measure_call(session, dags: list[LazyDeserializedDAG]) -> Counter[tuple[str, str]]: + """Count one steady-state call: the warm pass inserts the rows, the counted pass re-reads them.""" + files_parsed = {(BUNDLE, DAG_FILE)} + errors: dict = {} + + update_dag_parsing_results_in_db( + BUNDLE, None, dags, errors, 0.1, set(), session, files_parsed=files_parsed + ) + session.commit() + assert not errors, f"fixture Dags must serialize cleanly: {errors}" + + with _count_statements(session) as counts: + update_dag_parsing_results_in_db( + BUNDLE, None, dags, errors, 0.1, set(), session, files_parsed=files_parsed + ) + session.flush() + return counts + + +@pytest.mark.backend("postgres") +@pytest.mark.parametrize( + ("n_dags", "expected"), + [ + pytest.param(1, FIXED_PER_FILE + PER_DAG, id="one-dag"), + pytest.param(5, FIXED_PER_FILE + 5 * PER_DAG, id="five-dags"), + ], +) +def test_call_statement_budget(n_dags, expected, session, testing_dag_bundle, tmp_path): + """What one file's parse result costs to persist.""" + dags = _make_dags(tmp_path / DAG_FILE, [f"budget_dag_{i}" for i in range(n_dags)], DAG_FILE) + + counts = _measure_call(session, dags) + + total = sum(counts.values()) + assert total == expected, ( + f"a {n_dags}-Dag file costs {total} statements, expected {expected} " + f"({FIXED_PER_FILE} per file + {PER_DAG} per Dag).\n{_breakdown(counts)}" + ) + + +@pytest.mark.backend("postgres") +def test_per_dag_cost_matches_budget(session, testing_dag_bundle, tmp_path): + """A change trading fixed cost for per-Dag cost leaves the one-Dag total intact; slope catches it.""" + totals = {} + for n_dags in (1, 5): + dags = _make_dags(tmp_path / DAG_FILE, [f"budget_dag_{i}" for i in range(n_dags)], DAG_FILE) + totals[n_dags] = sum(_measure_call(session, dags).values()) + clear_db_serialized_dags() + clear_db_dags() + + slope = (totals[5] - totals[1]) / 4 + assert slope == PER_DAG, f"per-Dag cost is now {slope}, expected {PER_DAG}: {totals}" + + +def _ready_processor(rel_path: str, dag_file: Path, dag_ids: list[str], sockets: list[socket]): + """A finished parser subprocess, as the manager sees it. Mirrors ``mock_processor`` in test_manager.""" + read_end, write_end = socketpair() + sockets += [read_end, write_end] + processor = DagFileProcessorProcess( + process_log=MagicMock(), + id=uuid7(), + pid=1234, + process=MagicMock(wait=MagicMock(return_value=0)), + stdin=write_end, + logger_filehandle=MagicMock(), + client=MagicMock(), + bundle_name=BUNDLE, + dag_file_rel_path=rel_path, + ) + processor._open_sockets.clear() + processor.start_time = time.monotonic() - 1 + processor.had_callbacks = False + processor.parsing_result = DagFileParsingResult( + fileloc=str(dag_file), serialized_dags=_make_dags(dag_file, dag_ids, rel_path) + ) + return processor + + +def _register(manager, sweep_dir: Path, n_files: int, sockets: list[socket], dags_per_file: int) -> None: + for i in range(n_files): + rel_path = f"file_{i}.py" + file = DagFileInfo(bundle_name=BUNDLE, rel_path=Path(rel_path), bundle_path=sweep_dir) + manager._file_stats.setdefault(file, DagFileStat()) + manager._processors[file] = _ready_processor( + rel_path, sweep_dir / rel_path, [f"dag_{i}_{d}" for d in range(dags_per_file)], sockets + ) + + +def _measure_sweep(session, tmp_path: Path, n_files: int, sockets, name: str, dags_per_file=1) -> int: + """Count a steady-state sweep through ``_collect_results``.""" + manager = DagFileProcessorManager(max_runs=1) + manager._bundle_versions[BUNDLE] = None + sweep_dir = tmp_path / name + sweep_dir.mkdir() + + # The manager persists on sessions of its own, so release ours rather than contend with them. + _register(manager, sweep_dir, n_files, sockets, dags_per_file) + session.commit() + manager._collect_results() + + # Collecting consumed the processors, so register a second set for the counted sweep. + _register(manager, sweep_dir, n_files, sockets, dags_per_file) + with _count_statements(session) as counts: + manager._collect_results() + return sum(counts.values()) + + +def test_sweep_pays_fixed_cost_once_per_call(session, testing_dag_bundle, tmp_path, sockets): + """ + How a sweep scales with the number of persistence calls it takes. + + Batching a sweep into one call moves ``SWEEP_CALLS`` to 1. Both prices are measured here, so the + assertion holds on any backend. + """ + one_dag = _measure_sweep(session, tmp_path, 1, sockets, "one") + two_dags = _measure_sweep(session, tmp_path, 1, sockets, "two", dags_per_file=2) + per_dag = two_dags - one_dag + fixed = one_dag - per_dag + + sweep = _measure_sweep(session, tmp_path, SWEEP_FILES, sockets, "sweep") + + expected = SWEEP_CALLS * fixed + SWEEP_FILES * per_dag + assert sweep == expected, ( + f"a {SWEEP_FILES}-file sweep costs {sweep} statements, expected {expected} " + f"({SWEEP_CALLS} x {fixed} fixed + {SWEEP_FILES} x {per_dag} per Dag)." + ) From bd35a7e1cbf00080177c99067421e0f96debe404 Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Tue, 18 Aug 2026 19:15:02 +0100 Subject: [PATCH 02/11] Apply suggestions from code review --- .../test_parse_result_query_budget.py | 78 +++++++++---------- 1 file changed, 37 insertions(+), 41 deletions(-) diff --git a/airflow-core/tests/unit/dag_processing/test_parse_result_query_budget.py b/airflow-core/tests/unit/dag_processing/test_parse_result_query_budget.py index a5443beff26fa..2fa44226cd80e 100644 --- a/airflow-core/tests/unit/dag_processing/test_parse_result_query_budget.py +++ b/airflow-core/tests/unit/dag_processing/test_parse_result_query_budget.py @@ -15,15 +15,11 @@ # specific language governing permissions and limitations # under the License. """ -Statement budget for persisting Dag parse results. - -The Dag processor's manager issues these once per parsed file, so a change adding round trips has +Statement budget for persisting Dag parse results: a change that adds round trips to this path has to move a number here and account for it in review. -Counts are calibrated against Postgres and the tests carrying them are marked for it, because -statement counts differ by dialect: ``_update_import_errors`` deletes with a predicate the ORM -cannot evaluate in Python, which costs one statement where ``DELETE ... RETURNING`` exists and two -on MySQL. The sweep test hard-codes only the call count and measures the rest, so it runs anywhere. +Counts are calibrated against Postgres and the tests carrying them are marked for it, since +statement counts differ by dialect. The sweep test measures its prices, so it runs anywhere. """ from __future__ import annotations @@ -47,6 +43,7 @@ from airflow.sdk import DAG from airflow.serialization.serialized_objects import LazyDeserializedDAG +from tests_common.test_utils.config import conf_vars from tests_common.test_utils.db import clear_db_dags, clear_db_serialized_dags pytestmark = pytest.mark.db_test @@ -54,9 +51,12 @@ BUNDLE = "testing" DAG_FILE = "budget_dags.py" -# Per persistence call, and per Dag in the file. -FIXED_PER_FILE = 10 -PER_DAG = 3 +# Per persistence call, and per Dag in the file. A call leaves the serialized Dag alone while the +# content is unchanged; once the hash has moved and [core] min_serialized_dag_update_interval has +# lapsed it rewrites it, which costs two more statements per Dag and nothing extra per call. +FIXED_PER_CALL = 10 +UNCHANGED_PER_DAG = 3 +REWRITE_PER_DAG = 5 SWEEP_FILES = 4 # Calls the manager takes for that sweep: one per file today, 1 if a sweep is ever batched. @@ -113,22 +113,27 @@ def sockets(): sock.close() -def _make_dags(dag_file: Path, dag_ids: list[str], rel_path: str) -> list[LazyDeserializedDAG]: +def _make_dags( + dag_file: Path, dag_ids: list[str], rel_path: str, n_tasks: int = 1 +) -> list[LazyDeserializedDAG]: # DagCode reads the source off disk; without a real file the Dags fail to serialize and the # measured statements stop resembling a real parse. dag_file.write_text("# statement budget fixture\n") dags = [] for dag_id in dag_ids: dag = DAG(dag_id=dag_id, schedule="@daily") - EmptyOperator(task_id="task1", dag=dag) + for t in range(n_tasks): + EmptyOperator(task_id=f"task{t}", dag=dag) dag.fileloc = str(dag_file) dag.relative_fileloc = rel_path dags.append(LazyDeserializedDAG.from_dag(dag)) return dags -def _measure_call(session, dags: list[LazyDeserializedDAG]) -> Counter[tuple[str, str]]: - """Count one steady-state call: the warm pass inserts the rows, the counted pass re-reads them.""" +def _measure_call( + session, dags: list[LazyDeserializedDAG], counted: list[LazyDeserializedDAG] +) -> Counter[tuple[str, str]]: + """Count one steady-state call: the first pass inserts the rows, the counted pass re-persists.""" files_parsed = {(BUNDLE, DAG_FILE)} errors: dict = {} @@ -140,47 +145,40 @@ def _measure_call(session, dags: list[LazyDeserializedDAG]) -> Counter[tuple[str with _count_statements(session) as counts: update_dag_parsing_results_in_db( - BUNDLE, None, dags, errors, 0.1, set(), session, files_parsed=files_parsed + BUNDLE, None, counted, errors, 0.1, set(), session, files_parsed=files_parsed ) session.flush() return counts @pytest.mark.backend("postgres") +@pytest.mark.parametrize("n_dags", [1, 5]) @pytest.mark.parametrize( - ("n_dags", "expected"), + ("rewrite", "per_dag"), [ - pytest.param(1, FIXED_PER_FILE + PER_DAG, id="one-dag"), - pytest.param(5, FIXED_PER_FILE + 5 * PER_DAG, id="five-dags"), + pytest.param(False, UNCHANGED_PER_DAG, id="unchanged"), + pytest.param(True, REWRITE_PER_DAG, id="rewrite"), ], ) -def test_call_statement_budget(n_dags, expected, session, testing_dag_bundle, tmp_path): - """What one file's parse result costs to persist.""" - dags = _make_dags(tmp_path / DAG_FILE, [f"budget_dag_{i}" for i in range(n_dags)], DAG_FILE) - - counts = _measure_call(session, dags) - +def test_call_statement_budget(rewrite, per_dag, n_dags, session, testing_dag_bundle, tmp_path): + """What one file's parse result costs to persist, unchanged and rewritten.""" + dag_ids = [f"budget_dag_{i}" for i in range(n_dags)] + dags = _make_dags(tmp_path / DAG_FILE, dag_ids, DAG_FILE) + # A moved hash is what sends the call down the write path; the update interval only gates how + # soon it can get there. + counted = _make_dags(tmp_path / DAG_FILE, dag_ids, DAG_FILE, n_tasks=2) if rewrite else dags + + with conf_vars({("core", "min_serialized_dag_update_interval"): "0" if rewrite else "30"}): + counts = _measure_call(session, dags, counted) + + expected = FIXED_PER_CALL + n_dags * per_dag total = sum(counts.values()) assert total == expected, ( f"a {n_dags}-Dag file costs {total} statements, expected {expected} " - f"({FIXED_PER_FILE} per file + {PER_DAG} per Dag).\n{_breakdown(counts)}" + f"({FIXED_PER_CALL} per call + {per_dag} per Dag).\n{_breakdown(counts)}" ) -@pytest.mark.backend("postgres") -def test_per_dag_cost_matches_budget(session, testing_dag_bundle, tmp_path): - """A change trading fixed cost for per-Dag cost leaves the one-Dag total intact; slope catches it.""" - totals = {} - for n_dags in (1, 5): - dags = _make_dags(tmp_path / DAG_FILE, [f"budget_dag_{i}" for i in range(n_dags)], DAG_FILE) - totals[n_dags] = sum(_measure_call(session, dags).values()) - clear_db_serialized_dags() - clear_db_dags() - - slope = (totals[5] - totals[1]) / 4 - assert slope == PER_DAG, f"per-Dag cost is now {slope}, expected {PER_DAG}: {totals}" - - def _ready_processor(rel_path: str, dag_file: Path, dag_ids: list[str], sockets: list[socket]): """A finished parser subprocess, as the manager sees it. Mirrors ``mock_processor`` in test_manager.""" read_end, write_end = socketpair() @@ -222,9 +220,7 @@ def _measure_sweep(session, tmp_path: Path, n_files: int, sockets, name: str, da sweep_dir = tmp_path / name sweep_dir.mkdir() - # The manager persists on sessions of its own, so release ours rather than contend with them. _register(manager, sweep_dir, n_files, sockets, dags_per_file) - session.commit() manager._collect_results() # Collecting consumed the processors, so register a second set for the counted sweep. From ddd8f55a5d4736ad730d884bf3061a34123ecfb9 Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Tue, 18 Aug 2026 14:56:39 +0100 Subject: [PATCH 03/11] Persist a sweep of Dag parse results in one pass The Dag processor's manager wrote parse results one file at a time. Most of what that costs is fixed per call rather than per Dag, so a sweep of single-Dag files paid the fixed cost once per file. Persisting a sweep together pays it once. Measured over 60 files at 32 parsing processes, where sweeps held six files on average: 120 persistence calls become 21 or 22, and statements fall from 15.0 to between 5.1 and 5.2 per file. The baseline reproduces to the hundredth; what the change costs varies with how many sweeps a run happens to produce. What that saves in wall clock depends on how much of a file's cost is fixed. The Dags measured carry 200 tasks each, so the per-Dag work of hashing and serializing them dominates, and the saving sits inside the run-to-run variance of the machine it was measured on. Files defining smaller Dags pay proportionally more of the fixed cost and gain more of it back. The saving is also worth nothing at the default two parsing processes, where a sweep only ever holds one file: deployments running a wide Dag processor are the ones this is for. Parse duration is per file, so batching means it can no longer be a single value for the call; it is carried per Dag instead. Callers persisting one file still pass one value. A sweep is split into groups: one bundle each, since the bundle determines the version and version data the write needs, and each dag_id claimed at most once, since writing two files that define the same dag_id together would merge them into one Dag and lose the duplicate warning. Groups are written in the order they are returned, and a file is only held back to a group after the one holding the file it duplicates, so the same file wins a duplicated dag_id as when each was written on its own. A group is also capped, because the saving flattens out quickly while the rows one transaction holds locked do not. Each group is persisted by its own call, and so in its own transaction: update_dag_parsing_results_in_db rolls the session back before retrying a database error, which would otherwise discard a group already written alongside it while its files were still recorded as persisted. A group that fails is retried one file at a time, so one unwritable file does not discard the results of the others, and finished processors are closed whether or not the sweep succeeded. Persistence moves out of handle_parsing_result, which keeps the stat and metric handling. That method and persist_parsing_result were released in 3.3 as the seams for deployments forwarding results somewhere other than the metadata DB. Both keep their released contract and are still called once per file, and both are now deprecated in favour of persist_parsing_results, which is handed a whole group at once. Overriding a deprecated seam costs that deployment its batching and nothing else. From @seanmuth's investigation of the Airflow 2-to-3 dag-processor CPU regression: github.com/seanmuth/af2-af3-scheduler-cpu-repro. This is the per-sweep result batching that investigation identified as the largest lever. --- airflow-core/newsfragments/71618.misc.rst | 1 + .../src/airflow/dag_processing/collection.py | 21 +- .../src/airflow/dag_processing/manager.py | 377 ++++++++++++++-- .../airflow/serialization/definitions/dag.py | 4 +- .../tests/unit/dag_processing/test_manager.py | 404 ++++++++++++++++-- .../test_parse_result_query_budget.py | 231 +++++++++- 6 files changed, 952 insertions(+), 86 deletions(-) create mode 100644 airflow-core/newsfragments/71618.misc.rst diff --git a/airflow-core/newsfragments/71618.misc.rst b/airflow-core/newsfragments/71618.misc.rst new file mode 100644 index 0000000000000..6956495512533 --- /dev/null +++ b/airflow-core/newsfragments/71618.misc.rst @@ -0,0 +1 @@ +Deprecate ``DagFileProcessorManager.handle_parsing_result`` and ``persist_parsing_result``; override ``persist_parsing_results`` instead, which is handed every file that finished parsing together. Overriding either deprecated method still works but stops parse results being persisted a sweep at a time. diff --git a/airflow-core/src/airflow/dag_processing/collection.py b/airflow-core/src/airflow/dag_processing/collection.py index 8361cc43b0c46..b971271e5ccb1 100644 --- a/airflow-core/src/airflow/dag_processing/collection.py +++ b/airflow-core/src/airflow/dag_processing/collection.py @@ -28,6 +28,7 @@ from __future__ import annotations import traceback +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, NamedTuple, TypeVar import structlog @@ -203,6 +204,20 @@ def calculate(cls, dag: LazyDeserializedDAG, *, session: Session) -> Self: return cls(latest_run, active_run_counts.get(dag.dag_id, 0)) +def _resolve_parse_duration( + parse_duration: float | Mapping[str, float | None] | None, dag_id: str +) -> float | None: + """ + Pick the parse duration for one Dag. + + Persisting several files in one call means the duration differs per Dag, so callers doing that + pass a mapping. A single file still passes one value that applies to every Dag it defines. + """ + if isinstance(parse_duration, Mapping): + return parse_duration.get(dag_id) + return parse_duration + + def _update_dag_tags(tag_names: set[str], dm: DagModel, *, session: Session) -> None: orm_tags = {t.name: t for t in dm.tags} tags_to_delete = [] @@ -471,7 +486,7 @@ def update_dag_parsing_results_in_db( bundle_version: str | None, dags: Collection[LazyDeserializedDAG], import_errors: dict[tuple[str, str], str], - parse_duration: float | None, + parse_duration: float | Mapping[str, float | None] | None, warnings: set[DagWarning], session: Session, *, @@ -612,7 +627,7 @@ def add_dags(self, *, session: Session) -> dict[str, DagModel]: def update_dags( self, orm_dags: dict[str, DagModel], - parse_duration: float | None, + parse_duration: float | Mapping[str, float | None] | None, *, session: Session, ) -> None: @@ -626,7 +641,7 @@ def update_dags( dm.is_stale = False dm.has_import_errors = False dm.last_parsed_time = utcnow() - dm.last_parse_duration = parse_duration + dm.last_parse_duration = _resolve_parse_duration(parse_duration, dag_id) if hasattr(dag, "_dag_display_property_value"): dm._dag_display_property_value = dag._dag_display_property_value elif dag.dag_display_name != dag.dag_id: diff --git a/airflow-core/src/airflow/dag_processing/manager.py b/airflow-core/src/airflow/dag_processing/manager.py index 395ec25b37155..9944405384c1b 100644 --- a/airflow-core/src/airflow/dag_processing/manager.py +++ b/airflow-core/src/airflow/dag_processing/manager.py @@ -29,6 +29,7 @@ import signal import sys import time +import warnings import zipfile from collections import OrderedDict, defaultdict from dataclasses import dataclass, field @@ -106,6 +107,15 @@ def _make_execution_api() -> InProcessExecutionAPI: return InProcessExecutionAPI() +MAX_FILES_PER_PERSISTENCE_GROUP = 16 +""" +How many files at most are persisted together in one call, and so in one transaction. + +The saving from grouping is the fixed per-call cost spread across the group, so it flattens out +well before this bound; the ``dag`` rows the transaction holds locked do not. +""" + + class DagParsingStat(NamedTuple): """Information on processing progress.""" @@ -158,6 +168,16 @@ def normalized_file_path_for_stats(self) -> str: return normalize_name_for_stats(str(self.rel_path)) +class FileParseResult(NamedTuple): + """One file's finished parse, carried from stat handling to persistence.""" + + file: DagFileInfo + parsing_result: DagFileParsingResult + run_duration: float + stat: DagFileStat + """The stat to record once this file's results are persisted.""" + + def _config_int_factory(section: str, key: str): return functools.partial(conf.getint, section, key) @@ -372,6 +392,7 @@ def before_run(self) -> None: self.register_exit_signals() self.log.info("Processing files using up to %s processes at a time ", self._parallelism) self.log.info("Process each file at most once every %s seconds", self._file_process_interval) + self._warn_if_batching_is_disabled() self.prepare_bundles() self._symlink_latest_log_directory() # To prevent COW in forked process parsing dag file @@ -1252,10 +1273,13 @@ def handle_parsing_result( """ Post-process a single finished parse result. + .. deprecated:: 3.4.0 + Override :meth:`persist_parsing_results` instead. A subclass overriding this method + still receives every file, one at a time, exactly as before -- but a sweep can then no + longer be persisted in one pass, so batching is skipped for the whole Dag processor. + Detects callback-only processing, updates file stats, emits metrics, and persists DAGs/import-errors via :meth:`persist_parsing_result`. - Extracted from ``_collect_results`` to keep result handling and - persistence separate. Owns its own DB session via ``@provide_session`` so subclasses that forward results without touching the metadata DB (e.g. AIP-92 API-backed @@ -1267,6 +1291,38 @@ def handle_parsing_result( throttles immediate retries, so other files in the same ``_collect_results`` cycle still run. """ + result = self._build_parse_result(file, proc) + if result is None: + return + + try: + self.persist_parsing_result( + bundle_name=file.bundle_name, + bundle_version=self._bundle_versions[file.bundle_name], + version_data=self._bundle_version_data.get(file.bundle_name), + parsing_result=result.parsing_result, + run_duration=result.run_duration, + relative_fileloc=str(file.rel_path), + session=session, + ) + except Exception: + self._throttle_after_failed_persist(result) + return + + self._file_stats[file] = result.stat + + def _build_parse_result( + self, + file: DagFileInfo, + proc: DagFileProcessorProcess, + ) -> FileParseResult | None: + """ + Work out what a finished parse leaves to persist. + + Detects callback-only processing, emits metrics, and works out the stat this file should + record. Files with nothing to persist -- callback-only runs and failed parses -- have their + stat recorded here and return ``None``; the rest carry their Dags to whoever writes them. + """ is_callback_only = proc.had_callbacks and proc.parsing_result is None if is_callback_only: self.log.debug("Detected callback-only processing for %s", file) @@ -1285,37 +1341,16 @@ def handle_parsing_result( team_name=team_name, ) - if proc.parsing_result is not None: - try: - self.persist_parsing_result( - bundle_name=file.bundle_name, - bundle_version=self._bundle_versions[file.bundle_name], - version_data=self._bundle_version_data.get(file.bundle_name), - parsing_result=proc.parsing_result, - run_duration=run_duration, - relative_fileloc=str(file.rel_path), - session=session, - ) - except Exception: - self.log.exception( - "Failed to persist parsing result for %s in bundle %s; " - "keeping previous persisted stats while throttling retries. " - "Other files in this cycle are still processed.", - str(file.rel_path), - file.bundle_name, - ) - current_stat = self._file_stats[file] - self._file_stats[file] = DagFileStat( - num_dags=current_stat.num_dags, - import_errors=current_stat.import_errors, - last_finish_time=finish_time, - last_duration=run_duration, - run_count=current_stat.run_count + 1, - last_num_of_db_queries=current_stat.last_num_of_db_queries, - ) - return + if proc.parsing_result is None: + self._file_stats[file] = next_stat + return None - self._file_stats[file] = next_stat + return FileParseResult( + file=file, + parsing_result=proc.parsing_result, + run_duration=run_duration, + stat=next_stat, + ) def persist_parsing_result( self, @@ -1328,7 +1363,14 @@ def persist_parsing_result( relative_fileloc: str | None, session: Session, ) -> None: - """Persist parsed DAG data to the metadata database.""" + """ + Persist parsed DAG data to the metadata database. + + .. deprecated:: 3.4.0 + Override :meth:`persist_parsing_results` instead. A subclass overriding this method + still receives every file, one at a time, exactly as before -- but a sweep can then no + longer be persisted in one pass, so batching is skipped for the whole Dag processor. + """ import_errors: dict[tuple[str, str], str] = {} if parsing_result.import_errors: import_errors = { @@ -1358,18 +1400,269 @@ def persist_parsing_result( files_parsed=files_parsed, ) + @classmethod + def _overrides_per_file_persist(cls) -> bool: + """ + Report whether a subclass has replaced the per-file persistence hook. + + Batching would otherwise write straight past such an override, silently sending to the + metadata DB the results a deployment had arranged to send elsewhere. A subclass that also + replaced the batch hook has said where a whole sweep should go, so that one is used and + this reports ``False``. + """ + if cls.persist_parsing_results is not DagFileProcessorManager.persist_parsing_results: + return False + return cls.persist_parsing_result is not DagFileProcessorManager.persist_parsing_result + + def _warn_if_batching_is_disabled(self) -> None: + """Say once, at startup, that an override is costing this Dag processor its batched writes.""" + if self._overrides_handle_parsing_result(): + replaced = "handle_parsing_result" + elif self._overrides_per_file_persist(): + replaced = "persist_parsing_result" + else: + return + warnings.warn( + f"{type(self).__name__} overrides {replaced}, which is deprecated and prevents parse " + "results being persisted a sweep at a time. Override persist_parsing_results instead, " + "which is handed every file that finished together.", + DeprecationWarning, + stacklevel=2, + ) + + @classmethod + def _overrides_handle_parsing_result(cls) -> bool: + """ + Report whether a subclass has replaced the per-file result handler. + + Released in 3.3 as the seam for deployments that forward results instead of writing them to + the metadata DB. Such an override persists a file itself, so it is still called once per + file and a sweep is left with nothing to batch. + """ + return cls.handle_parsing_result is not DagFileProcessorManager.handle_parsing_result + + @provide_session + def persist_parsing_results( + self, + results: Sequence[FileParseResult], + *, + session: Session = NEW_SESSION, + ) -> None: + """ + Persist several files' parse results in one pass. + + Everything :func:`update_dag_parsing_results_in_db` does besides writing the Dags themselves + is paid once per call, so persisting a whole sweep together costs far fewer statements than + persisting file by file. Results are grouped by bundle, since the bundle is what determines + the version and version data the call needs. + + Deployments that send parse results somewhere other than the metadata DB (e.g. AIP-92) can + override either this method or the per-file :meth:`persist_parsing_result`. A subclass that + overrides only the per-file hook keeps its existing behaviour: it is called once per file, + batching is skipped, and it is handed this method's session rather than being quietly + bypassed by a batched write it never sees. + + Note that raising from here discards the whole group -- the caller falls back to persisting + each file on its own so one bad file cannot take the others down with it. + """ + if self._overrides_per_file_persist(): + for item in results: + self.persist_parsing_result( + bundle_name=item.file.bundle_name, + bundle_version=self._bundle_versions[item.file.bundle_name], + version_data=self._bundle_version_data.get(item.file.bundle_name), + parsing_result=item.parsing_result, + run_duration=item.run_duration, + relative_fileloc=str(item.file.rel_path), + session=session, + ) + return + + for group in self.build_persistence_groups(results): + self._persist_bundle_group(group[0].file.bundle_name, group, session=session) + + def build_persistence_groups(self, results: Sequence[FileParseResult]) -> list[list[FileParseResult]]: + """ + Split a sweep into units that can each be written in a single call. + + A group holds one bundle's files, since the bundle determines the version and version data + the write needs, and claims a dag_id at most once: writing two files that define the same + dag_id together would merge them into one Dag and lose the duplicate warning, which is + raised by comparing an incoming Dag against the file already recorded in the DB. + + Groups are returned in the order they must be written, and a file is only ever held back to + a group after the one holding the file it duplicates -- so the duplicate still sees what it + duplicates, and the same file wins as when each was written on its own. Only the files + actually in conflict are held back; one repeated dag_id must not cost every other file its + place in a group. + + Groups are capped at ``MAX_FILES_PER_PERSISTENCE_GROUP`` files. + """ + groups: list[list[FileParseResult]] = [] + bundles: list[str] = [] + # The last group to claim each dag_id, which is the earliest a file repeating it may go. + claimed_by: dict[str, int] = {} + + for item in results: + dag_ids = {dag.dag_id for dag in item.parsing_result.serialized_dags} + bundle_name = item.file.bundle_name + earliest = max((claimed_by[dag_id] + 1 for dag_id in dag_ids if dag_id in claimed_by), default=0) + + for index in range(earliest, len(groups)): + if bundles[index] == bundle_name and len(groups[index]) < MAX_FILES_PER_PERSISTENCE_GROUP: + break + else: + index = len(groups) + groups.append([]) + bundles.append(bundle_name) + + groups[index].append(item) + claimed_by.update(dict.fromkeys(dag_ids, index)) + return groups + + def _persist_bundle_group( + self, + bundle_name: str, + items: Sequence[FileParseResult], + *, + session: Session, + ) -> None: + """Merge one bundle's parse results into a single write.""" + dags: list = [] + import_errors: dict[tuple[str, str], str] = {} + files_parsed: set[tuple[str, str]] = set() + # Duration is per file, but the Dags of several files are written together, so it has to be + # carried per Dag rather than as one value for the call. + parse_durations: dict[str, float] = {} + warnings: set[DagWarning] = set() + + for item in items: + parsing_result = item.parsing_result + relative_fileloc = str(item.file.rel_path) + + file_errors = { + (bundle_name, rel_path): error + for rel_path, error in (parsing_result.import_errors or {}).items() + } + import_errors.update(file_errors) + # Include the parsed file even when it defines no Dags, so its stale import errors + # still get cleared. + files_parsed.add((bundle_name, relative_fileloc)) + files_parsed.update(file_errors) + + dags.extend(parsing_result.serialized_dags) + for dag in parsing_result.serialized_dags: + parse_durations[dag.dag_id] = item.run_duration + + file_warnings = parsing_result.warnings or [] + if file_warnings and isinstance(file_warnings[0], dict): + file_warnings = [DagWarning(**warn) for warn in file_warnings] + warnings.update(file_warnings) + + update_dag_parsing_results_in_db( + bundle_name=bundle_name, + bundle_version=self._bundle_versions[bundle_name], + version_data=self._bundle_version_data.get(bundle_name), + dags=dags, + import_errors=import_errors, + parse_duration=parse_durations, + warnings=warnings, + session=session, + files_parsed=files_parsed, + ) + def _collect_results(self): finished = [] - for file, proc in self._processors.items(): - if not proc.is_ready: - # This processor hasn't finished yet, or we haven't read all the output from it yet + to_persist: list[FileParseResult] = [] + # An override owns the whole of a file's handling, persistence included, so there is + # nothing left for a sweep to write together. + handles_each_file = self._overrides_handle_parsing_result() + try: + for file, proc in self._processors.items(): + if not proc.is_ready: + # This processor hasn't finished yet, or we haven't read all the output from it yet + continue + finished.append(file) + if handles_each_file: + self.handle_parsing_result(file, proc) + elif (result := self._build_parse_result(file, proc)) is not None: + to_persist.append(result) + + if to_persist: + self._persist_sweep(to_persist) + finally: + # Whatever went wrong, these processes are done with; leaving them open leaks their + # sockets and keeps them queued as though they were still running. + for file in finished: + processor = self._processors.pop(file) + processor.close() + + def _persist_sweep(self, to_persist: list[FileParseResult]) -> None: + """ + Persist a sweep one group at a time, falling back to single files when a group fails. + + Each group gets its own call, and so its own transaction: ``update_dag_parsing_results_in_db`` + rolls the session back before retrying an ``OperationalError``, which would otherwise discard + groups that had already succeeded in the same transaction while their files were still + recorded as persisted. + + A subclass handling files one at a time is given single-file groups, so it never receives a + file twice: a retry after a grouped attempt would hand it results it had already accepted. + """ + if self._overrides_per_file_persist(): + groups = [[item] for item in to_persist] + else: + groups = self.build_persistence_groups(to_persist) + + for group in groups: + if len(group) == 1: + self._persist_single(group[0]) continue - finished.append(file) - self.handle_parsing_result(file, proc) + try: + self.persist_parsing_results(group) + except Exception: + self.log.exception( + "Failed to persist %d parse results as a group; retrying them individually.", + len(group), + ) + for item in group: + self._persist_single(item) + else: + for item in group: + self._file_stats[item.file] = item.stat - for file in finished: - processor = self._processors.pop(file) - processor.close() + def _persist_single(self, item: FileParseResult) -> None: + """Persist one file, throttling its retries if that fails rather than claiming success.""" + try: + self.persist_parsing_results([item]) + except Exception: + self._throttle_after_failed_persist(item) + else: + self._file_stats[item.file] = item.stat + + def _throttle_after_failed_persist(self, item: FileParseResult) -> None: + """ + Record a failed write without claiming its results. + + Keeps the counts of whatever was last persisted and only moves the timestamps, so the file + is not parsed again immediately while the rest of the cycle carries on. + """ + self.log.exception( + "Failed to persist parsing result for %s in bundle %s; " + "keeping previous persisted stats while throttling retries. " + "Other files in this cycle are still processed.", + str(item.file.rel_path), + item.file.bundle_name, + ) + current_stat = self._file_stats[item.file] + self._file_stats[item.file] = DagFileStat( + num_dags=current_stat.num_dags, + import_errors=current_stat.import_errors, + last_finish_time=item.stat.last_finish_time, + last_duration=item.stat.last_duration, + run_count=current_stat.run_count + 1, + last_num_of_db_queries=current_stat.last_num_of_db_queries, + ) def _get_log_dir(self) -> str: return os.path.join(self.base_log_dir, timezone.utcnow().strftime("%Y-%m-%d")) @@ -1768,7 +2061,7 @@ def process_parse_results( Create a DagFileStat from parsing results and emit metrics. This function handles stat creation and metrics only — database persistence - is handled separately by ``DagFileProcessorManager.persist_parsing_result``. + is handled separately by ``DagFileProcessorManager.persist_parsing_results``. """ if is_callback_only: # Callback-only processing - don't update timestamps to avoid stale DAG detection issues diff --git a/airflow-core/src/airflow/serialization/definitions/dag.py b/airflow-core/src/airflow/serialization/definitions/dag.py index 8ed0fee2ccabd..94acd282225fa 100644 --- a/airflow-core/src/airflow/serialization/definitions/dag.py +++ b/airflow-core/src/airflow/serialization/definitions/dag.py @@ -63,7 +63,7 @@ if TYPE_CHECKING: import datetime - from collections.abc import Collection, Iterable, Sequence + from collections.abc import Collection, Iterable, Mapping, Sequence from typing import Any, Literal from pendulum.tz.timezone import FixedTimezone, Timezone @@ -194,7 +194,7 @@ def bulk_write_to_db( bundle_name: str, bundle_version: str | None, dags: Collection[DAG | LazyDeserializedDAG], - parse_duration: float | None = None, + parse_duration: float | Mapping[str, float | None] | None = None, *, session: Session = NEW_SESSION, ) -> None: diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index 6d83d9c69abbb..bf6d43d2f330f 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py @@ -48,10 +48,12 @@ from airflow.dag_processing.bundles.manager import DagBundlesManager from airflow.dag_processing.dagbag import DagBag from airflow.dag_processing.manager import ( + MAX_FILES_PER_PERSISTENCE_GROUP, BundleState, DagFileInfo, DagFileProcessorManager, DagFileStat, + FileParseResult, ) from airflow.dag_processing.processor import DagFileParsingResult, DagFileProcessorProcess from airflow.models import DagModel, DbCallbackRequest @@ -61,6 +63,8 @@ from airflow.models.dagcode import DagCode from airflow.models.serialized_dag import SerializedDagModel from airflow.models.team import Team +from airflow.sdk import DAG as SdkDAG +from airflow.serialization.serialized_objects import LazyDeserializedDAG from airflow.utils.net import get_hostname from airflow.utils.session import create_session @@ -1520,24 +1524,27 @@ def test_terminate_normalizes_file_path_stats_tag(self): ) processor.kill.assert_called_once_with(signal.SIGTERM, escalation_delay=5.0) - def test_handle_parsing_result_provides_its_own_session_when_caller_omits(self): - """``handle_parsing_result`` is wrapped in ``@provide_session`` so subclasses overriding it can run without a caller-supplied session.""" + def test_persist_parsing_results_provides_its_own_session_when_caller_omits(self): + """``persist_parsing_results`` is wrapped in ``@provide_session`` so subclasses overriding it can run without a caller-supplied session.""" manager = DagFileProcessorManager(max_runs=1) file = DagFileInfo(bundle_name="testing", rel_path=Path("abc.txt"), bundle_path=TEST_DAGS_FOLDER) manager._file_stats[file] = DagFileStat() manager._bundle_versions["testing"] = "v1" - processor, _ = self.mock_processor(start_time=time.monotonic() - 1) - processor.had_callbacks = False - processor.parsing_result = DagFileParsingResult(fileloc="abc.txt", serialized_dags=[]) + item = FileParseResult( + file=file, + parsing_result=DagFileParsingResult(fileloc="abc.txt", serialized_dags=[]), + run_duration=1.0, + stat=DagFileStat(), + ) - with mock.patch.object(manager, "persist_parsing_result") as mock_persist: - manager.handle_parsing_result(file, processor) + with mock.patch("airflow.dag_processing.manager.update_dag_parsing_results_in_db") as mock_update: + manager.persist_parsing_results([item]) - mock_persist.assert_called_once() - assert mock_persist.call_args.kwargs["session"] is not None + mock_update.assert_called_once() + assert mock_update.call_args.kwargs["session"] is not None - def test_handle_parsing_result_throttles_retry_when_first_persist_fails(self, session): + def test_a_failed_write_throttles_the_retry_without_claiming_success(self, session): """Persist errors should throttle retries without claiming persistence succeeded.""" manager = DagFileProcessorManager(max_runs=1) file = DagFileInfo(bundle_name="testing", rel_path=Path("abc.txt"), bundle_path=TEST_DAGS_FOLDER) @@ -1550,8 +1557,9 @@ def test_handle_parsing_result_throttles_retry_when_first_persist_fails(self, se processor.had_callbacks = False processor.parsing_result = DagFileParsingResult(fileloc="abc.txt", serialized_dags=[]) - with mock.patch.object(manager, "persist_parsing_result", side_effect=RuntimeError("boom")): - manager.handle_parsing_result(file, processor, session=session) + with mock.patch.object(manager, "persist_parsing_results", side_effect=RuntimeError("boom")): + result = manager._build_parse_result(file, processor) + manager._persist_sweep([result]) assert manager._file_stats[file] is not original_stat assert manager._file_stats[file].num_dags == 0 @@ -1561,7 +1569,7 @@ def test_handle_parsing_result_throttles_retry_when_first_persist_fails(self, se assert manager._file_stats[file].last_duration is not None assert manager.processed_recently(timezone.utcnow(), file) is True - def test_handle_parsing_result_updates_stats_after_successful_persist(self, session): + def test_a_written_file_records_the_stat_its_parse_produced(self, session): manager = DagFileProcessorManager(max_runs=1) file = DagFileInfo(bundle_name="testing", rel_path=Path("abc.txt"), bundle_path=TEST_DAGS_FOLDER) original_stat = DagFileStat( @@ -1579,24 +1587,357 @@ def test_handle_parsing_result_updates_stats_after_successful_persist(self, sess processor.had_callbacks = False processor.parsing_result = DagFileParsingResult(fileloc="abc.txt", serialized_dags=[]) - with mock.patch.object(manager, "persist_parsing_result") as mock_persist: - manager.handle_parsing_result(file, processor, session=session) + with mock.patch.object(manager, "persist_parsing_results") as mock_persist: + result = manager._build_parse_result(file, processor) + manager._persist_sweep([result]) - mock_persist.assert_called_once_with( - bundle_name="testing", - bundle_version="v1", - version_data=None, - parsing_result=processor.parsing_result, - run_duration=mock.ANY, - relative_fileloc="abc.txt", - session=session, - ) + mock_persist.assert_called_once_with([result]) + assert result.file is file + assert result.parsing_result is processor.parsing_result assert manager._file_stats[file] is not original_stat assert manager._file_stats[file].run_count == 4 assert manager._file_stats[file].last_finish_time is not None assert manager._file_stats[file].last_finish_time > original_stat.last_finish_time assert manager._file_stats[file].num_dags == 0 + def _ready_processor(self, manager, rel_path: str, num_dags: int = 0): + """Register a finished processor for ``rel_path`` and return its file.""" + file = DagFileInfo(bundle_name="testing", rel_path=Path(rel_path), bundle_path=TEST_DAGS_FOLDER) + manager._file_stats.setdefault(file, DagFileStat()) + processor, _ = self.mock_processor(start_time=time.monotonic() - 1) + processor.had_callbacks = False + processor.parsing_result = DagFileParsingResult( + fileloc=rel_path, + serialized_dags=[self._lazy_dag(f"{Path(rel_path).stem}_dag_{i}") for i in range(num_dags)], + ) + manager._processors[file] = processor + return file + + def test_collect_results_persists_the_whole_sweep_in_one_call(self): + """Every file ready in a sweep is handed to persistence together, not one call each.""" + manager = DagFileProcessorManager(max_runs=1) + manager._bundle_versions["testing"] = "v1" + + files = [self._ready_processor(manager, name) for name in ("a.py", "b.py", "c.py")] + + with mock.patch.object(manager, "persist_parsing_results") as mock_persist: + manager._collect_results() + + mock_persist.assert_called_once() + assert [item.file for item in mock_persist.call_args.args[0]] == files + assert manager._processors == {}, "finished processors should be closed and dropped" + for file in files: + assert manager._file_stats[file].run_count == 1 + + @staticmethod + def _lazy_dag(dag_id: str): + """A real LazyDeserializedDAG; DagFileParsingResult validates this field.""" + return LazyDeserializedDAG.from_dag(SdkDAG(dag_id=dag_id, schedule=None)) + + def test_files_sharing_a_dag_id_are_written_one_at_a_time(self): + """ + Two files defining the same dag_id must not be merged into a single write. + + The duplicate warning is raised by comparing an incoming Dag against the file already + recorded in the DB, so batching the pair would collapse them, skip the warning, and let one + file silently win. + """ + manager = DagFileProcessorManager(max_runs=1) + manager._bundle_versions["testing"] = "v1" + + shared = self._lazy_dag("shared_dag") + items = [] + for name in ("first.py", "second.py"): + file = DagFileInfo(bundle_name="testing", rel_path=Path(name), bundle_path=TEST_DAGS_FOLDER) + items.append( + FileParseResult( + file=file, + parsing_result=DagFileParsingResult(fileloc=name, serialized_dags=[shared]), + run_duration=1.0, + stat=DagFileStat(), + ) + ) + + with mock.patch.object(manager, "_persist_bundle_group") as mock_group: + manager.persist_parsing_results(items, session=mock.MagicMock()) + + assert [len(call.args[1]) for call in mock_group.call_args_list] == [1, 1], ( + "conflicting files must be written separately so the second sees the first" + ) + + def test_files_with_distinct_dag_ids_are_written_together(self): + """The split only applies to conflicts; everything else still batches.""" + manager = DagFileProcessorManager(max_runs=1) + manager._bundle_versions["testing"] = "v1" + + items = [] + for i, name in enumerate(("first.py", "second.py")): + file = DagFileInfo(bundle_name="testing", rel_path=Path(name), bundle_path=TEST_DAGS_FOLDER) + items.append( + FileParseResult( + file=file, + parsing_result=DagFileParsingResult( + fileloc=name, serialized_dags=[self._lazy_dag(f"dag_{i}")] + ), + run_duration=1.0, + stat=DagFileStat(), + ) + ) + + with mock.patch.object(manager, "_persist_bundle_group") as mock_group: + manager.persist_parsing_results(items, session=mock.MagicMock()) + + assert [len(call.args[1]) for call in mock_group.call_args_list] == [2] + + def _item(self, name: str, dag_ids: list[str], bundle_name: str = "testing") -> FileParseResult: + file = DagFileInfo(bundle_name=bundle_name, rel_path=Path(name), bundle_path=TEST_DAGS_FOLDER) + return FileParseResult( + file=file, + parsing_result=DagFileParsingResult( + fileloc=name, serialized_dags=[self._lazy_dag(dag_id) for dag_id in dag_ids] + ), + run_duration=1.0, + stat=DagFileStat(), + ) + + def test_one_duplicate_dag_id_does_not_split_the_rest_of_the_sweep(self): + """Only the conflicting file is held back; the others keep their place in the group.""" + manager = DagFileProcessorManager(max_runs=1) + + items = [ + self._item("first.py", ["shared_dag"]), + self._item("second.py", ["shared_dag"]), + *(self._item(f"other_{i}.py", [f"dag_{i}"]) for i in range(3)), + ] + + groups = manager.build_persistence_groups(items) + + assert [len(group) for group in groups] == [4, 1] + assert [str(item.file.rel_path) for item in groups[1]] == ["second.py"], ( + "the later of the two files must be the one held back, so it sees the earlier one" + ) + + def test_a_group_is_capped_so_one_transaction_cannot_lock_a_whole_sweep(self): + """Grouping saves a fixed per-call cost that flattens out; locked rows do not.""" + manager = DagFileProcessorManager(max_runs=1) + + items = [self._item(f"file_{i}.py", [f"dag_{i}"]) for i in range(MAX_FILES_PER_PERSISTENCE_GROUP + 3)] + + groups = manager.build_persistence_groups(items) + + assert [len(group) for group in groups] == [MAX_FILES_PER_PERSISTENCE_GROUP, 3] + assert [item for group in groups for item in group] == items, "no file may be lost or reordered" + + @pytest.mark.parametrize( + ("items", "expected"), + [ + pytest.param( + [("a.py", ["x"]), ("b.py", ["y"]), ("c.py", ["y", "z"]), ("d.py", ["z"])], + [["a.py", "b.py"], ["c.py"], ["d.py"]], + id="a-file-may-not-overtake-the-file-it-duplicates", + ), + pytest.param( + [("a.py", ["x"]), ("b.py", ["x"]), ("c.py", ["x"])], + [["a.py"], ["b.py"], ["c.py"]], + id="three-files-sharing-one-dag-id", + ), + ], + ) + def test_a_duplicate_is_never_written_before_what_it_duplicates(self, items, expected): + """ + The file written last wins the dag_id, so sweep order has to survive grouping. + + Packing each file into the first group with room reorders them: a file bumped out by one + conflict lets a later file it conflicts with take the place it lost. + """ + manager = DagFileProcessorManager(max_runs=1) + + groups = manager.build_persistence_groups([self._item(name, dag_ids) for name, dag_ids in items]) + + assert [[str(item.file.rel_path) for item in group] for group in groups] == expected + + @pytest.mark.parametrize("persist_fails", [False, True], ids=["written", "unwritable"]) + def test_the_released_handler_still_persists_a_file_on_its_own(self, persist_fails): + """A subclass may call this rather than replace it, so it has to keep writing what it is given.""" + manager = DagFileProcessorManager(max_runs=1) + manager._bundle_versions["testing"] = "v1" + manager._file_stats[ + file := DagFileInfo(bundle_name="testing", rel_path=Path("a.py"), bundle_path=TEST_DAGS_FOLDER) + ] = DagFileStat(num_dags=7, run_count=3) + processor, _ = self.mock_processor(start_time=time.monotonic() - 1) + processor.had_callbacks = False + processor.parsing_result = DagFileParsingResult(fileloc="a.py", serialized_dags=[]) + + side_effect = RuntimeError("boom") if persist_fails else None + with mock.patch.object(manager, "persist_parsing_result", side_effect=side_effect) as persist: + manager.handle_parsing_result(file, processor) + + assert persist.call_args.kwargs["relative_fileloc"] == "a.py" + assert persist.call_args.kwargs["session"] is not None, "the override seam owns its own session" + assert manager._file_stats[file].run_count == 4 + # A failed write keeps the counts of whatever was persisted last rather than claiming these. + assert manager._file_stats[file].num_dags == (7 if persist_fails else 0) + + def test_an_override_of_the_released_handler_still_handles_every_file(self): + """ + 3.3 shipped handle_parsing_result as the seam for forwarding results outside the DB. + + Such an override persists the file itself and returns nothing, so it has to keep being + called once per file rather than being asked for something to batch. + """ + handled: list[str] = [] + + class ApiBackedManager(DagFileProcessorManager): + def handle_parsing_result(self, file, proc, *, session=None): + handled.append(str(file.rel_path)) + + manager = ApiBackedManager(max_runs=1) + manager._bundle_versions["testing"] = "v1" + for name in ("a.py", "b.py"): + self._ready_processor(manager, name, num_dags=1) + + with mock.patch.object(manager, "_persist_sweep", autospec=True) as sweep: + manager._collect_results() + + assert handled == ["a.py", "b.py"] + sweep.assert_not_called() + assert manager._processors == {}, "finished processors should still be closed and dropped" + + @mock.patch("airflow.dag_processing.manager.update_dag_parsing_results_in_db", autospec=True) + def test_an_override_delegating_to_super_still_persists(self, mock_write): + """The released implementation stays usable, so an override can add to it rather than replace it.""" + + class WrappingManager(DagFileProcessorManager): + def handle_parsing_result(self, file, proc, *, session=None): + super().handle_parsing_result(file, proc) + + manager = WrappingManager(max_runs=1) + manager._bundle_versions["testing"] = "v1" + file = self._ready_processor(manager, "wrapped.py", num_dags=1) + + manager._collect_results() + + mock_write.assert_called_once() + assert manager._file_stats[file].run_count == 1 + + @pytest.mark.parametrize( + ("overrides", "expected"), + [ + pytest.param((), False, id="no-override-batches"), + pytest.param(("handle_parsing_result",), True, id="released-handler"), + pytest.param(("persist_parsing_result",), True, id="released-per-file-persist"), + pytest.param(("persist_parsing_results",), False, id="batch-seam-still-batches"), + pytest.param(("persist_parsing_result", "persist_parsing_results"), False, id="both-persist"), + ], + ) + def test_only_a_deprecated_override_gives_up_batching(self, overrides, expected): + """A subclass that adopted the batch seam keeps batching, even while it still carries the old one.""" + subclass = type( + "Subclass", (DagFileProcessorManager,), {name: lambda *a, **kw: None for name in overrides} + ) + + manager = subclass(max_runs=1) + + if expected: + with pytest.warns(DeprecationWarning, match="persist_parsing_results"): + manager._warn_if_batching_is_disabled() + else: + manager._warn_if_batching_is_disabled() + + def test_finished_processors_are_closed_even_when_persistence_raises(self): + """A processor left open leaks its sockets and stays queued as though it were still running.""" + manager = DagFileProcessorManager(max_runs=1) + manager._bundle_versions["testing"] = "v1" + file = self._ready_processor(manager, "boom.py", num_dags=1) + processor = manager._processors[file] + + with mock.patch.object(processor, "close") as close: + with mock.patch.object(manager, "_persist_sweep", side_effect=RuntimeError("boom")): + with pytest.raises(RuntimeError, match="boom"): + manager._collect_results() + + assert manager._processors == {} + close.assert_called_once() + + def test_groups_of_different_bundles_keep_their_place_in_the_sweep(self): + """A dag_id is unique across bundles, so duplicates of one usually arrive in different bundles.""" + manager = DagFileProcessorManager(max_runs=1) + + items = [ + self._item("x.py", ["z"], bundle_name="b"), + self._item("y.py", ["shared"], bundle_name="a"), + self._item("z.py", ["shared"], bundle_name="b"), + ] + + groups = manager.build_persistence_groups(items) + + assert [[str(item.file.rel_path) for item in group] for group in groups] == [ + ["x.py"], + ["y.py"], + ["z.py"], + ], "collecting a bundle's files together must not let them overtake another bundle's" + + def test_collect_results_leaves_unfinished_processors_alone(self): + """A processor that has not finished must not be persisted or dropped.""" + manager = DagFileProcessorManager(max_runs=1) + manager._bundle_versions["testing"] = "v1" + + done = self._ready_processor(manager, "done.py") + pending = self._ready_processor(manager, "pending.py") + # An open socket means output is still being read, so the processor is not ready. Keep a + # strong reference: _open_sockets is a WeakKeyDictionary. + open_socket = MagicMock() + manager._processors[pending]._open_sockets[open_socket] = MagicMock() + + with mock.patch.object(manager, "persist_parsing_results") as mock_persist: + manager._collect_results() + + assert [item.file for item in mock_persist.call_args.args[0]] == [done] + assert list(manager._processors) == [pending] + + def test_one_unwritable_file_does_not_discard_its_neighbours(self): + """ + The batch is all-or-nothing, so a failure inside it must be retried file by file. + + Without that retry a single unwritable file would throw away the parse results of every + other file that happened to finish in the same sweep. + """ + manager = DagFileProcessorManager(max_runs=1) + manager._bundle_versions["testing"] = "v1" + manager._file_process_interval = 60 + + good = DagFileInfo(bundle_name="testing", rel_path=Path("good.py"), bundle_path=TEST_DAGS_FOLDER) + bad = DagFileInfo(bundle_name="testing", rel_path=Path("bad.py"), bundle_path=TEST_DAGS_FOLDER) + manager._file_stats[good] = DagFileStat() + # Counts from the last successful parse, which the throttle path must preserve. + manager._file_stats[bad] = DagFileStat(num_dags=7, import_errors=2, run_count=3) + + items = [] + for file in (good, bad): + processor, _ = self.mock_processor(start_time=time.monotonic() - 1) + processor.had_callbacks = False + processor.parsing_result = DagFileParsingResult(fileloc=str(file.rel_path), serialized_dags=[]) + items.append(manager._build_parse_result(file, processor)) + + persisted: list[DagFileInfo] = [] + + def persist_unless_bad_is_present(results, **kwargs): + if any(item.file == bad for item in results): + raise RuntimeError("simulated write failure") + persisted.extend(item.file for item in results) + + with mock.patch.object(manager, "persist_parsing_results", side_effect=persist_unless_bad_is_present): + manager._persist_sweep(items) + + assert persisted == [good], "the healthy file must still be written by the per-file retry" + assert manager._file_stats[good] is items[0].stat + + throttled = manager._file_stats[bad] + assert throttled is not items[1].stat, "a failed write must not record a successful parse" + assert (throttled.num_dags, throttled.import_errors) == (7, 2), "previous counts preserved" + assert throttled.run_count == 4 + assert manager.processed_recently(timezone.utcnow(), bad) is True + def test_collect_results_processes_remaining_files_when_one_persist_fails(self, session): manager = DagFileProcessorManager(max_runs=1) file_a = DagFileInfo(bundle_name="testing", rel_path=Path("a.py"), bundle_path=TEST_DAGS_FOLDER) @@ -1624,13 +1965,18 @@ def test_collect_results_processes_remaining_files_when_one_persist_fails(self, stat_a_before = manager._file_stats[file_a] stat_b_before = manager._file_stats[file_b] - with mock.patch.object( - manager, - "persist_parsing_result", - side_effect=[RuntimeError("boom"), None], - ): + # Patch the batched seam: _collect_results routes through it, and a per-file patch on the + # instance would not be seen by the class-level override check. + def fail_for_a(results, **kwargs): + if any(item.file == file_a for item in results): + raise RuntimeError("boom") + + with mock.patch.object(manager, "persist_parsing_results", side_effect=fail_for_a) as mock_persist: manager._collect_results() + # The grouped attempt, then one retry per file. + assert [len(call.args[0]) for call in mock_persist.call_args_list] == [2, 1, 1] + assert manager._file_stats[file_a] is not stat_a_before assert manager._file_stats[file_a].num_dags == stat_a_before.num_dags assert manager._file_stats[file_a].import_errors == stat_a_before.import_errors diff --git a/airflow-core/tests/unit/dag_processing/test_parse_result_query_budget.py b/airflow-core/tests/unit/dag_processing/test_parse_result_query_budget.py index 2fa44226cd80e..426f124dbba22 100644 --- a/airflow-core/tests/unit/dag_processing/test_parse_result_query_budget.py +++ b/airflow-core/tests/unit/dag_processing/test_parse_result_query_budget.py @@ -28,23 +28,40 @@ import time from collections import Counter from contextlib import contextmanager, suppress +from io import BytesIO from pathlib import Path from socket import socket, socketpair +from unittest import mock from unittest.mock import MagicMock import pytest -from sqlalchemy import event +from sqlalchemy import event, select +from sqlalchemy.exc import OperationalError from uuid6 import uuid7 from airflow.dag_processing.collection import update_dag_parsing_results_in_db -from airflow.dag_processing.manager import DagFileInfo, DagFileProcessorManager, DagFileStat +from airflow.dag_processing.manager import ( + DagFileInfo, + DagFileProcessorManager, + DagFileStat, + FileParseResult, +) from airflow.dag_processing.processor import DagFileParsingResult, DagFileProcessorProcess +from airflow.models.dag import DagModel +from airflow.models.dagwarning import DagWarning, DagWarningType +from airflow.models.errors import ParseImportError from airflow.providers.standard.operators.empty import EmptyOperator from airflow.sdk import DAG +from airflow.sdk.api.client import Client +from airflow.sdk.execution_time.supervisor import ProcessTracker from airflow.serialization.serialized_objects import LazyDeserializedDAG from tests_common.test_utils.config import conf_vars -from tests_common.test_utils.db import clear_db_dags, clear_db_serialized_dags +from tests_common.test_utils.db import ( + clear_db_dags, + clear_db_import_errors, + clear_db_serialized_dags, +) pytestmark = pytest.mark.db_test @@ -59,8 +76,8 @@ REWRITE_PER_DAG = 5 SWEEP_FILES = 4 -# Calls the manager takes for that sweep: one per file today, 1 if a sweep is ever batched. -SWEEP_CALLS = 4 +# Calls the manager takes for that sweep: the whole sweep is persisted together. +SWEEP_CALLS = 1 def _classify(statement: str) -> tuple[str, str]: @@ -100,6 +117,7 @@ def _breakdown(counts: Counter[tuple[str, str]]) -> str: def clean_db(): yield clear_db_serialized_dags() + clear_db_import_errors() clear_db_dags() @@ -183,14 +201,16 @@ def _ready_processor(rel_path: str, dag_file: Path, dag_ids: list[str], sockets: """A finished parser subprocess, as the manager sees it. Mirrors ``mock_processor`` in test_manager.""" read_end, write_end = socketpair() sockets += [read_end, write_end] + process = MagicMock(spec=ProcessTracker) + process.wait.return_value = 0 processor = DagFileProcessorProcess( process_log=MagicMock(), id=uuid7(), pid=1234, - process=MagicMock(wait=MagicMock(return_value=0)), + process=process, stdin=write_end, - logger_filehandle=MagicMock(), - client=MagicMock(), + logger_filehandle=BytesIO(), + client=MagicMock(spec=Client), bundle_name=BUNDLE, dag_file_rel_path=rel_path, ) @@ -234,8 +254,8 @@ def test_sweep_pays_fixed_cost_once_per_call(session, testing_dag_bundle, tmp_pa """ How a sweep scales with the number of persistence calls it takes. - Batching a sweep into one call moves ``SWEEP_CALLS`` to 1. Both prices are measured here, so the - assertion holds on any backend. + The sweep is persisted in one call, so the fixed price is paid once rather than per file. Both + prices are measured here, so the assertion holds on any backend. """ one_dag = _measure_sweep(session, tmp_path, 1, sockets, "one") two_dags = _measure_sweep(session, tmp_path, 1, sockets, "two", dags_per_file=2) @@ -249,3 +269,194 @@ def test_sweep_pays_fixed_cost_once_per_call(session, testing_dag_bundle, tmp_pa f"a {SWEEP_FILES}-file sweep costs {sweep} statements, expected {expected} " f"({SWEEP_CALLS} x {fixed} fixed + {SWEEP_FILES} x {per_dag} per Dag)." ) + + +def _parse_result(tmp_path: Path, dag_id: str, run_duration: float = 0.5) -> FileParseResult: + rel_path = f"{dag_id}.py" + dag_file = tmp_path / rel_path + return FileParseResult( + file=DagFileInfo(bundle_name=BUNDLE, rel_path=Path(rel_path), bundle_path=tmp_path), + parsing_result=DagFileParsingResult( + fileloc=str(dag_file), serialized_dags=_make_dags(dag_file, [dag_id], rel_path) + ), + run_duration=run_duration, + stat=DagFileStat(), + ) + + +def test_batched_sweep_keeps_each_files_own_parse_duration(session, testing_dag_bundle, tmp_path): + """Duration is per file, so writing several files together must not level them out.""" + durations = {"sweep_a": 1.5, "sweep_b": 4.25} + manager = DagFileProcessorManager(max_runs=1) + manager._bundle_versions[BUNDLE] = None + + manager.persist_parsing_results( + [_parse_result(tmp_path, dag_id, d) for dag_id, d in durations.items()], session=session + ) + session.commit() + + for dag_id, duration in durations.items(): + assert session.get(DagModel, dag_id).last_parse_duration == duration + + +def test_per_file_override_still_replaces_the_database(tmp_path): + """Batching past an existing per-file override would send its results to the DB, silently.""" + calls: list[str] = [] + + class ApiBackedManager(DagFileProcessorManager): + def persist_parsing_result(self, *, relative_fileloc, session, **kwargs): + assert session is not None, "override must be handed a session it did not create" + calls.append(relative_fileloc) + + manager = ApiBackedManager(max_runs=1) + manager._bundle_versions[BUNDLE] = None + batch = [_parse_result(tmp_path, "override_a"), _parse_result(tmp_path, "override_b")] + + with mock.patch("airflow.dag_processing.manager.update_dag_parsing_results_in_db") as db_write: + manager.persist_parsing_results(batch) + + assert calls == ["override_a.py", "override_b.py"] + db_write.assert_not_called() + + +def test_batch_override_replaces_the_database(tmp_path): + """Overriding the batch seam keeps the metadata DB out of the path entirely.""" + seen: list[int] = [] + + class BatchApiManager(DagFileProcessorManager): + def persist_parsing_results(self, results, *, session=None): + seen.append(len(results)) + + manager = BatchApiManager(max_runs=1) + manager._bundle_versions[BUNDLE] = None + batch = [_parse_result(tmp_path, "batch_a"), _parse_result(tmp_path, "batch_b")] + + with mock.patch("airflow.dag_processing.manager.update_dag_parsing_results_in_db") as db_write: + manager.persist_parsing_results(batch) + + assert seen == [2], "the batch override should see the whole sweep at once" + db_write.assert_not_called() + + +def test_default_manager_takes_the_batched_path(tmp_path): + assert DagFileProcessorManager._overrides_per_file_persist() is False + + +def _parse_result_with( + tmp_path: Path, + dag_id: str, + *, + import_errors: dict[str, str] | None = None, + warnings: list | None = None, +) -> FileParseResult: + rel_path = f"{dag_id}.py" + dag_file = tmp_path / rel_path + return FileParseResult( + file=DagFileInfo(bundle_name=BUNDLE, rel_path=Path(rel_path), bundle_path=tmp_path), + parsing_result=DagFileParsingResult( + fileloc=str(dag_file), + serialized_dags=_make_dags(dag_file, [dag_id], rel_path), + import_errors=import_errors, + warnings=warnings, + ), + run_duration=0.5, + stat=DagFileStat(), + ) + + +def test_batched_sweep_keeps_each_files_import_errors(session, testing_dag_bundle, tmp_path): + """Merging a sweep must not lose one file's import errors, nor attribute them to another.""" + manager = DagFileProcessorManager(max_runs=1) + manager._bundle_versions[BUNDLE] = None + + manager.persist_parsing_results( + [ + _parse_result_with(tmp_path, "err_a", import_errors={"err_a.py": "boom a"}), + _parse_result_with(tmp_path, "err_b"), + ], + session=session, + ) + session.commit() + + recorded = {e.filename: e.stacktrace for e in session.scalars(select(ParseImportError))} + assert recorded == {"err_a.py": "boom a"} + + +def test_batched_sweep_clears_import_errors_for_files_that_now_parse(session, testing_dag_bundle, tmp_path): + """A file in the sweep with no errors must have its stale error cleared, not left behind.""" + manager = DagFileProcessorManager(max_runs=1) + manager._bundle_versions[BUNDLE] = None + + manager.persist_parsing_results( + [_parse_result_with(tmp_path, "fixed", import_errors={"fixed.py": "was broken"})], + session=session, + ) + session.commit() + assert session.scalars(select(ParseImportError)).all() + + manager.persist_parsing_results([_parse_result_with(tmp_path, "fixed")], session=session) + session.commit() + + assert not session.scalars(select(ParseImportError)).all(), "stale error should have been cleared" + + +def test_batched_sweep_records_warnings_from_every_file(session, testing_dag_bundle, tmp_path): + """Warnings are merged across the sweep, so each file's must survive the merge.""" + manager = DagFileProcessorManager(max_runs=1) + manager._bundle_versions[BUNDLE] = None + + def warning_for(dag_id: str) -> dict: + return { + "dag_id": dag_id, + "warning_type": DagWarningType.NONEXISTENT_POOL, + "message": f"{dag_id} wants a missing pool", + } + + manager.persist_parsing_results( + [ + _parse_result_with(tmp_path, "warn_a", warnings=[warning_for("warn_a")]), + _parse_result_with(tmp_path, "warn_b", warnings=[warning_for("warn_b")]), + ], + session=session, + ) + session.commit() + + assert {w.dag_id for w in session.scalars(select(DagWarning))} == {"warn_a", "warn_b"} + + +def test_a_failing_group_does_not_discard_one_that_already_succeeded(tmp_path): + """ + Each group is persisted in its own transaction. + + ``update_dag_parsing_results_in_db`` rolls the session back before retrying an OperationalError. + Sharing one transaction across groups would let that rollback discard a group already written, + while its files were still recorded as persisted. + """ + manager = DagFileProcessorManager(max_runs=1) + manager._bundle_versions.update({"bundle_a": None, "bundle_b": None}) + + def result_in(bundle: str, dag_id: str) -> FileParseResult: + item = _parse_result_with(tmp_path, dag_id) + return item._replace( + file=DagFileInfo(bundle_name=bundle, rel_path=item.file.rel_path, bundle_path=tmp_path) + ) + + good = result_in("bundle_a", "group_good") + bad = result_in("bundle_b", "group_bad") + manager._file_stats.update({good.file: DagFileStat(), bad.file: DagFileStat(run_count=2)}) + + calls: list[str] = [] + + def persist(results, **kwargs): + bundle = results[0].file.bundle_name + calls.append(bundle) + if bundle == "bundle_b": + raise OperationalError("simulated contention", None, Exception()) + + with mock.patch.object(manager, "persist_parsing_results", side_effect=persist): + manager._persist_sweep([good, bad]) + + assert calls == ["bundle_a", "bundle_b"], "each bundle must be persisted by its own call" + assert manager._file_stats[good.file] is good.stat, "the written group keeps its parse stat" + assert manager._file_stats[bad.file] is not bad.stat, "the failed group must not claim success" + assert manager._file_stats[bad.file].run_count == 3 From 364e0694135fbab415eae09b9c561c1e085041a3 Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Tue, 18 Aug 2026 15:59:05 +0100 Subject: [PATCH 04/11] Detect a replaced parse-result hook on the instance as well as the class Replacing one of the persistence hooks on a single manager left the detection reading the class, which still reported the default. The manager went on writing to the metadata DB while looking, to whoever replaced the hook, as though it no longer did. Typing the Dags a group merges lets the checker see what reaches the write, which is the one place merging several files could go wrong. The sweep budget derives the number of calls it expects from the group cap rather than assuming a sweep fits in one, and measures a sweep past that cap, so splitting a sweep into groups is exercised rather than only asserted. Naming a local set of Dag warnings `warnings` shadowed the module of the same name, which the deprecation warning now needs. --- .../{71618.misc.rst => 71771.misc.rst} | 0 .../src/airflow/dag_processing/manager.py | 47 ++++++++++++------- .../tests/unit/dag_processing/test_manager.py | 7 ++- .../test_parse_result_query_budget.py | 32 ++++++++----- 4 files changed, 55 insertions(+), 31 deletions(-) rename airflow-core/newsfragments/{71618.misc.rst => 71771.misc.rst} (100%) diff --git a/airflow-core/newsfragments/71618.misc.rst b/airflow-core/newsfragments/71771.misc.rst similarity index 100% rename from airflow-core/newsfragments/71618.misc.rst rename to airflow-core/newsfragments/71771.misc.rst diff --git a/airflow-core/src/airflow/dag_processing/manager.py b/airflow-core/src/airflow/dag_processing/manager.py index 9944405384c1b..b62697930d853 100644 --- a/airflow-core/src/airflow/dag_processing/manager.py +++ b/airflow-core/src/airflow/dag_processing/manager.py @@ -95,6 +95,7 @@ from airflow.callbacks.callback_requests import CallbackRequest from airflow.dag_processing.bundles.base import BaseDagBundle from airflow.sdk.api.client import Client + from airflow.serialization.serialized_objects import LazyDeserializedDAG def _make_execution_api() -> InProcessExecutionAPI: @@ -1384,9 +1385,9 @@ def persist_parsing_result( files_parsed = {(bundle_name, relative_fileloc)} files_parsed.update(import_errors.keys()) - warnings = parsing_result.warnings or [] - if warnings and isinstance(warnings[0], dict): - warnings = [DagWarning(**warn) for warn in warnings] + dag_warnings = parsing_result.warnings or [] + if dag_warnings and isinstance(dag_warnings[0], dict): + dag_warnings = [DagWarning(**warn) for warn in dag_warnings] update_dag_parsing_results_in_db( bundle_name=bundle_name, @@ -1395,24 +1396,35 @@ def persist_parsing_result( dags=parsing_result.serialized_dags, import_errors=import_errors, parse_duration=run_duration, - warnings=set(warnings), + warnings=set(dag_warnings), session=session, files_parsed=files_parsed, ) - @classmethod - def _overrides_per_file_persist(cls) -> bool: + def _overrides(self, name: str) -> bool: """ - Report whether a subclass has replaced the per-file persistence hook. + Report whether this manager has its own version of one of the hooks. + + Looks at the instance as well as the class, so that replacing a hook on a single manager + counts. Reading only the class would leave such a manager writing to the metadata DB while + looking, to whoever replaced the hook, as though it no longer did. + """ + if name in getattr(self, "__dict__", ()): + return True + return getattr(type(self), name) is not getattr(DagFileProcessorManager, name) + + def _overrides_per_file_persist(self) -> bool: + """ + Report whether the per-file persistence hook has been replaced. Batching would otherwise write straight past such an override, silently sending to the - metadata DB the results a deployment had arranged to send elsewhere. A subclass that also + metadata DB the results a deployment had arranged to send elsewhere. A manager that also replaced the batch hook has said where a whole sweep should go, so that one is used and this reports ``False``. """ - if cls.persist_parsing_results is not DagFileProcessorManager.persist_parsing_results: + if self._overrides("persist_parsing_results"): return False - return cls.persist_parsing_result is not DagFileProcessorManager.persist_parsing_result + return self._overrides("persist_parsing_result") def _warn_if_batching_is_disabled(self) -> None: """Say once, at startup, that an override is costing this Dag processor its batched writes.""" @@ -1430,16 +1442,15 @@ def _warn_if_batching_is_disabled(self) -> None: stacklevel=2, ) - @classmethod - def _overrides_handle_parsing_result(cls) -> bool: + def _overrides_handle_parsing_result(self) -> bool: """ - Report whether a subclass has replaced the per-file result handler. + Report whether the per-file result handler has been replaced. Released in 3.3 as the seam for deployments that forward results instead of writing them to the metadata DB. Such an override persists a file itself, so it is still called once per file and a sweep is left with nothing to batch. """ - return cls.handle_parsing_result is not DagFileProcessorManager.handle_parsing_result + return self._overrides("handle_parsing_result") @provide_session def persist_parsing_results( @@ -1528,13 +1539,13 @@ def _persist_bundle_group( session: Session, ) -> None: """Merge one bundle's parse results into a single write.""" - dags: list = [] + dags: list[LazyDeserializedDAG] = [] import_errors: dict[tuple[str, str], str] = {} files_parsed: set[tuple[str, str]] = set() # Duration is per file, but the Dags of several files are written together, so it has to be # carried per Dag rather than as one value for the call. parse_durations: dict[str, float] = {} - warnings: set[DagWarning] = set() + dag_warnings: set[DagWarning] = set() for item in items: parsing_result = item.parsing_result @@ -1557,7 +1568,7 @@ def _persist_bundle_group( file_warnings = parsing_result.warnings or [] if file_warnings and isinstance(file_warnings[0], dict): file_warnings = [DagWarning(**warn) for warn in file_warnings] - warnings.update(file_warnings) + dag_warnings.update(file_warnings) update_dag_parsing_results_in_db( bundle_name=bundle_name, @@ -1566,7 +1577,7 @@ def _persist_bundle_group( dags=dags, import_errors=import_errors, parse_duration=parse_durations, - warnings=warnings, + warnings=dag_warnings, session=session, files_parsed=files_parsed, ) diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index bf6d43d2f330f..3fb7df345e671 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py @@ -1999,10 +1999,13 @@ def test_collect_results_tolerates_stale_file_handle_on_close(self): proc.logger_filehandle.close.side_effect = OSError(116, "Stale file handle") manager._processors = {file: proc} - with mock.patch.object(manager, "persist_parsing_result"): - manager._collect_results() + with mock.patch("airflow.dag_processing.manager.update_dag_parsing_results_in_db") as db_write: + with mock.patch.object(manager, "persist_parsing_result") as persist: + manager._collect_results() assert len(manager._processors) == 0 + persist.assert_called_once() + assert db_write.call_count == 0, "replacing the hook on one manager has to keep the DB out" @pytest.mark.usefixtures("testing_dag_bundle") @pytest.mark.parametrize( diff --git a/airflow-core/tests/unit/dag_processing/test_parse_result_query_budget.py b/airflow-core/tests/unit/dag_processing/test_parse_result_query_budget.py index 426f124dbba22..d28bf158783aa 100644 --- a/airflow-core/tests/unit/dag_processing/test_parse_result_query_budget.py +++ b/airflow-core/tests/unit/dag_processing/test_parse_result_query_budget.py @@ -19,11 +19,13 @@ to move a number here and account for it in review. Counts are calibrated against Postgres and the tests carrying them are marked for it, since -statement counts differ by dialect. The sweep test measures its prices, so it runs anywhere. +statement counts differ by dialect. The sweep test derives its call count from the group cap and +measures its prices, so it runs anywhere. """ from __future__ import annotations +import math import re import time from collections import Counter @@ -41,6 +43,7 @@ from airflow.dag_processing.collection import update_dag_parsing_results_in_db from airflow.dag_processing.manager import ( + MAX_FILES_PER_PERSISTENCE_GROUP, DagFileInfo, DagFileProcessorManager, DagFileStat, @@ -76,8 +79,6 @@ REWRITE_PER_DAG = 5 SWEEP_FILES = 4 -# Calls the manager takes for that sweep: the whole sweep is persisted together. -SWEEP_CALLS = 1 def _classify(statement: str) -> tuple[str, str]: @@ -250,24 +251,33 @@ def _measure_sweep(session, tmp_path: Path, n_files: int, sockets, name: str, da return sum(counts.values()) -def test_sweep_pays_fixed_cost_once_per_call(session, testing_dag_bundle, tmp_path, sockets): +@pytest.mark.parametrize( + "n_files", + [ + pytest.param(SWEEP_FILES, id="one-group"), + pytest.param(MAX_FILES_PER_PERSISTENCE_GROUP + 1, id="over-the-group-cap"), + ], +) +def test_sweep_pays_fixed_cost_once_per_call(n_files, session, testing_dag_bundle, tmp_path, sockets): """ How a sweep scales with the number of persistence calls it takes. - The sweep is persisted in one call, so the fixed price is paid once rather than per file. Both - prices are measured here, so the assertion holds on any backend. + A sweep is persisted a group at a time, so the fixed price is paid once per group rather than + once per file -- and a sweep past the cap really does split, rather than only being asserted to. + Both prices are measured here, so the assertion holds on any backend. """ one_dag = _measure_sweep(session, tmp_path, 1, sockets, "one") two_dags = _measure_sweep(session, tmp_path, 1, sockets, "two", dags_per_file=2) per_dag = two_dags - one_dag fixed = one_dag - per_dag - sweep = _measure_sweep(session, tmp_path, SWEEP_FILES, sockets, "sweep") + sweep = _measure_sweep(session, tmp_path, n_files, sockets, f"sweep_{n_files}") - expected = SWEEP_CALLS * fixed + SWEEP_FILES * per_dag + calls = math.ceil(n_files / MAX_FILES_PER_PERSISTENCE_GROUP) + expected = calls * fixed + n_files * per_dag assert sweep == expected, ( - f"a {SWEEP_FILES}-file sweep costs {sweep} statements, expected {expected} " - f"({SWEEP_CALLS} x {fixed} fixed + {SWEEP_FILES} x {per_dag} per Dag)." + f"a {n_files}-file sweep costs {sweep} statements, expected {expected} " + f"({calls} x {fixed} fixed + {n_files} x {per_dag} per Dag)." ) @@ -339,7 +349,7 @@ def persist_parsing_results(self, results, *, session=None): def test_default_manager_takes_the_batched_path(tmp_path): - assert DagFileProcessorManager._overrides_per_file_persist() is False + assert DagFileProcessorManager(max_runs=1)._overrides_per_file_persist() is False def _parse_result_with( From 641fa7b4605fedf5c230e10423b2e468ec1854a2 Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Tue, 18 Aug 2026 16:30:38 +0100 Subject: [PATCH 05/11] Cover what a batched sweep does that no test reached Overriding the batch seam was checked by calling it directly, which only shows that Python dispatches to an override; the manager was never involved, so nothing said it routes a sweep through the seam at all. It is now driven through a sweep, with a default manager doing the same as its negative. A file can parse to no Dags and is still a file the sweep parsed, so the error it left behind last time has to be cleared with the rest. Removing the line that records such a file left every test passing. A group carries one bundle's version and version data, and nothing wrote a sweep spanning two bundles to check neither was crossed over. The stale-error test asserts about the file it is checking rather than about the whole table, which other tests leave rows in. --- .../test_parse_result_query_budget.py | 103 +++++++++++++++--- 1 file changed, 90 insertions(+), 13 deletions(-) diff --git a/airflow-core/tests/unit/dag_processing/test_parse_result_query_budget.py b/airflow-core/tests/unit/dag_processing/test_parse_result_query_budget.py index d28bf158783aa..e6ce7b19c06a9 100644 --- a/airflow-core/tests/unit/dag_processing/test_parse_result_query_budget.py +++ b/airflow-core/tests/unit/dag_processing/test_parse_result_query_budget.py @@ -51,6 +51,8 @@ ) from airflow.dag_processing.processor import DagFileParsingResult, DagFileProcessorProcess from airflow.models.dag import DagModel +from airflow.models.dag_version import DagVersion +from airflow.models.dagbundle import DagBundleModel from airflow.models.dagwarning import DagWarning, DagWarningType from airflow.models.errors import ParseImportError from airflow.providers.standard.operators.empty import EmptyOperator @@ -69,6 +71,7 @@ pytestmark = pytest.mark.db_test BUNDLE = "testing" +OTHER_BUNDLE = "testing-other" DAG_FILE = "budget_dags.py" # Per persistence call, and per Dag in the file. A call leaves the serialized Dag alone while the @@ -281,11 +284,13 @@ def test_sweep_pays_fixed_cost_once_per_call(n_files, session, testing_dag_bundl ) -def _parse_result(tmp_path: Path, dag_id: str, run_duration: float = 0.5) -> FileParseResult: +def _parse_result( + tmp_path: Path, dag_id: str, run_duration: float = 0.5, bundle_name: str = BUNDLE +) -> FileParseResult: rel_path = f"{dag_id}.py" dag_file = tmp_path / rel_path return FileParseResult( - file=DagFileInfo(bundle_name=BUNDLE, rel_path=Path(rel_path), bundle_path=tmp_path), + file=DagFileInfo(bundle_name=bundle_name, rel_path=Path(rel_path), bundle_path=tmp_path), parsing_result=DagFileParsingResult( fileloc=str(dag_file), serialized_dags=_make_dags(dag_file, [dag_id], rel_path) ), @@ -329,27 +334,44 @@ def persist_parsing_result(self, *, relative_fileloc, session, **kwargs): db_write.assert_not_called() -def test_batch_override_replaces_the_database(tmp_path): - """Overriding the batch seam keeps the metadata DB out of the path entirely.""" - seen: list[int] = [] +def _collect_a_two_file_sweep(manager, tmp_path: Path, sockets, name: str, session) -> mock.MagicMock: + """Run a sweep of two files all the way through the manager, and report what reached the DB.""" + sweep_dir = tmp_path / name + sweep_dir.mkdir() + _register(manager, sweep_dir, 2, sockets, dags_per_file=1) + # The manager persists on sessions of its own, so release ours rather than contend with them. + session.commit() + + with mock.patch("airflow.dag_processing.manager.update_dag_parsing_results_in_db") as db_write: + manager._collect_results() + return db_write + + +def test_a_batch_override_is_handed_the_sweep_by_the_manager(session, testing_dag_bundle, tmp_path, sockets): + """Dispatching to an override proves nothing unless the manager is the one routing through it.""" + seen: list[list[str]] = [] class BatchApiManager(DagFileProcessorManager): def persist_parsing_results(self, results, *, session=None): - seen.append(len(results)) + seen.append([str(item.file.rel_path) for item in results]) manager = BatchApiManager(max_runs=1) manager._bundle_versions[BUNDLE] = None - batch = [_parse_result(tmp_path, "batch_a"), _parse_result(tmp_path, "batch_b")] - with mock.patch("airflow.dag_processing.manager.update_dag_parsing_results_in_db") as db_write: - manager.persist_parsing_results(batch) + db_write = _collect_a_two_file_sweep(manager, tmp_path, sockets, "batch_override", session) - assert seen == [2], "the batch override should see the whole sweep at once" + assert seen == [["file_0.py", "file_1.py"]], "the whole sweep should arrive in one call" db_write.assert_not_called() -def test_default_manager_takes_the_batched_path(tmp_path): - assert DagFileProcessorManager(max_runs=1)._overrides_per_file_persist() is False +def test_the_default_manager_writes_the_sweep_to_the_database(session, testing_dag_bundle, tmp_path, sockets): + """The negative of the override case: with nothing replaced, a sweep still reaches the DB once.""" + manager = DagFileProcessorManager(max_runs=1) + manager._bundle_versions[BUNDLE] = None + + db_write = _collect_a_two_file_sweep(manager, tmp_path, sockets, "default_path", session) + + db_write.assert_called_once() def _parse_result_with( @@ -358,6 +380,7 @@ def _parse_result_with( *, import_errors: dict[str, str] | None = None, warnings: list | None = None, + dag_ids: list[str] | None = None, ) -> FileParseResult: rel_path = f"{dag_id}.py" dag_file = tmp_path / rel_path @@ -365,7 +388,7 @@ def _parse_result_with( file=DagFileInfo(bundle_name=BUNDLE, rel_path=Path(rel_path), bundle_path=tmp_path), parsing_result=DagFileParsingResult( fileloc=str(dag_file), - serialized_dags=_make_dags(dag_file, [dag_id], rel_path), + serialized_dags=_make_dags(dag_file, [dag_id] if dag_ids is None else dag_ids, rel_path), import_errors=import_errors, warnings=warnings, ), @@ -410,6 +433,60 @@ def test_batched_sweep_clears_import_errors_for_files_that_now_parse(session, te assert not session.scalars(select(ParseImportError)).all(), "stale error should have been cleared" +def test_batched_sweep_clears_a_stale_error_for_a_file_that_now_defines_no_dags( + session, testing_dag_bundle, tmp_path +): + """A file can stop defining Dags altogether, and is still a file the sweep parsed.""" + manager = DagFileProcessorManager(max_runs=1) + manager._bundle_versions[BUNDLE] = None + + manager.persist_parsing_results( + [_parse_result_with(tmp_path, "emptied", import_errors={"emptied.py": "was broken"})], + session=session, + ) + session.commit() + assert session.scalars(select(ParseImportError)).all() + + manager.persist_parsing_results( + [_parse_result_with(tmp_path, "healthy"), _parse_result_with(tmp_path, "emptied", dag_ids=[])], + session=session, + ) + session.commit() + + remaining = {(error.bundle_name, error.filename) for error in session.scalars(select(ParseImportError))} + assert (BUNDLE, "emptied.py") not in remaining, ( + f"the emptied file was parsed, so its stale error should have been cleared: {remaining}" + ) + + +def test_a_sweep_writes_each_bundles_files_under_its_own_version(session, testing_dag_bundle, tmp_path): + """A group carries one bundle's version, so a sweep spanning two must not cross them over.""" + session.add(DagBundleModel(name=OTHER_BUNDLE)) + session.commit() + + manager = DagFileProcessorManager(max_runs=1) + manager._bundle_versions.update({BUNDLE: "v-testing", OTHER_BUNDLE: "v-other"}) + manager._bundle_version_data.update({BUNDLE: {"sha": "aaa"}, OTHER_BUNDLE: {"sha": "bbb"}}) + + manager.persist_parsing_results( + [ + _parse_result(tmp_path, "in_testing"), + _parse_result(tmp_path, "in_other", bundle_name=OTHER_BUNDLE), + ], + session=session, + ) + session.commit() + + assert session.get(DagModel, "in_testing").bundle_name == BUNDLE + assert session.get(DagModel, "in_other").bundle_name == OTHER_BUNDLE + + versions = {version.dag_id: version for version in session.scalars(select(DagVersion))} + assert versions["in_testing"].bundle_version == "v-testing" + assert versions["in_other"].bundle_version == "v-other" + assert versions["in_testing"].version_data == {"sha": "aaa"} + assert versions["in_other"].version_data == {"sha": "bbb"} + + def test_batched_sweep_records_warnings_from_every_file(session, testing_dag_bundle, tmp_path): """Warnings are merged across the sweep, so each file's must survive the merge.""" manager = DagFileProcessorManager(max_runs=1) From 6b3cadc01263cc2f91634a1d720cd4d64a15bb32 Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Tue, 18 Aug 2026 18:24:06 +0100 Subject: [PATCH 06/11] Bound a persistence group by the Dags it carries, not the files Grouping saves a fixed cost once per file, but what one transaction holds locked grows with the Dags in it. Counting files bounds neither: sixteen files that generate their Dags dynamically can be thousands of rows locked for the length of the write, which is the case the bound exists for. Counting Dags does not bind on files defining one or two of them, which is the shape a sweep usually has, so this costs nothing in the common case. A file is still never split, since the record of it having been parsed and the import errors keyed to it belong to it as a whole. --- .../src/airflow/dag_processing/manager.py | 26 +++++++++++---- .../tests/unit/dag_processing/test_manager.py | 26 ++++++++++++--- .../test_parse_result_query_budget.py | 32 +++++++++++-------- 3 files changed, 60 insertions(+), 24 deletions(-) diff --git a/airflow-core/src/airflow/dag_processing/manager.py b/airflow-core/src/airflow/dag_processing/manager.py index b62697930d853..e42a51ac23787 100644 --- a/airflow-core/src/airflow/dag_processing/manager.py +++ b/airflow-core/src/airflow/dag_processing/manager.py @@ -108,12 +108,18 @@ def _make_execution_api() -> InProcessExecutionAPI: return InProcessExecutionAPI() -MAX_FILES_PER_PERSISTENCE_GROUP = 16 +MAX_DAGS_PER_PERSISTENCE_GROUP = 32 """ -How many files at most are persisted together in one call, and so in one transaction. +How many Dags at most are persisted together in one call, and so in one transaction. -The saving from grouping is the fixed per-call cost spread across the group, so it flattens out -well before this bound; the ``dag`` rows the transaction holds locked do not. +Grouping saves a fixed cost once per file, but what one transaction holds locked grows with the +Dags in it, not the files -- a few files generating Dags dynamically can be thousands. Counting +Dags bounds the thing that actually grows, and at this size it does not bind on files defining +one or two Dags, which is the shape a sweep usually has. + +A file is never split across groups, since the record of it having been parsed and the import +errors keyed to it belong to it as a whole. One file defining more Dags than this is still +written on its own. """ @@ -1507,27 +1513,35 @@ def build_persistence_groups(self, results: Sequence[FileParseResult]) -> list[l actually in conflict are held back; one repeated dag_id must not cost every other file its place in a group. - Groups are capped at ``MAX_FILES_PER_PERSISTENCE_GROUP`` files. + A group carries at most ``MAX_DAGS_PER_PERSISTENCE_GROUP`` Dags. """ groups: list[list[FileParseResult]] = [] bundles: list[str] = [] + dag_counts: list[int] = [] # The last group to claim each dag_id, which is the earliest a file repeating it may go. claimed_by: dict[str, int] = {} for item in results: dag_ids = {dag.dag_id for dag in item.parsing_result.serialized_dags} + n_dags = len(item.parsing_result.serialized_dags) bundle_name = item.file.bundle_name earliest = max((claimed_by[dag_id] + 1 for dag_id in dag_ids if dag_id in claimed_by), default=0) for index in range(earliest, len(groups)): - if bundles[index] == bundle_name and len(groups[index]) < MAX_FILES_PER_PERSISTENCE_GROUP: + if ( + bundles[index] == bundle_name + and dag_counts[index] + n_dags <= MAX_DAGS_PER_PERSISTENCE_GROUP + ): break else: + # A new group takes the file whatever its size: a file cannot be split. index = len(groups) groups.append([]) bundles.append(bundle_name) + dag_counts.append(0) groups[index].append(item) + dag_counts[index] += n_dags claimed_by.update(dict.fromkeys(dag_ids, index)) return groups diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index 3fb7df345e671..3c999de1a0fd0 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py @@ -48,7 +48,7 @@ from airflow.dag_processing.bundles.manager import DagBundlesManager from airflow.dag_processing.dagbag import DagBag from airflow.dag_processing.manager import ( - MAX_FILES_PER_PERSISTENCE_GROUP, + MAX_DAGS_PER_PERSISTENCE_GROUP, BundleState, DagFileInfo, DagFileProcessorManager, @@ -1717,17 +1717,33 @@ def test_one_duplicate_dag_id_does_not_split_the_rest_of_the_sweep(self): "the later of the two files must be the one held back, so it sees the earlier one" ) - def test_a_group_is_capped_so_one_transaction_cannot_lock_a_whole_sweep(self): - """Grouping saves a fixed per-call cost that flattens out; locked rows do not.""" + def test_a_group_is_capped_by_dags_so_one_transaction_cannot_lock_a_whole_sweep(self): + """What one transaction holds locked grows with the Dags in it, not with the files.""" manager = DagFileProcessorManager(max_runs=1) + per_file = 10 + fits = MAX_DAGS_PER_PERSISTENCE_GROUP // per_file - items = [self._item(f"file_{i}.py", [f"dag_{i}"]) for i in range(MAX_FILES_PER_PERSISTENCE_GROUP + 3)] + items = [ + self._item(f"file_{i}.py", [f"dag_{i}_{d}" for d in range(per_file)]) for i in range(fits + 1) + ] groups = manager.build_persistence_groups(items) - assert [len(group) for group in groups] == [MAX_FILES_PER_PERSISTENCE_GROUP, 3] + assert [len(group) for group in groups] == [fits, 1] assert [item for group in groups for item in group] == items, "no file may be lost or reordered" + def test_a_file_defining_more_dags_than_the_cap_is_still_written(self): + """A file is never split: the record of it being parsed and its import errors are its own.""" + manager = DagFileProcessorManager(max_runs=1) + oversized = [f"huge_dag_{i}" for i in range(MAX_DAGS_PER_PERSISTENCE_GROUP + 5)] + + groups = manager.build_persistence_groups( + [self._item("small.py", ["small_dag"]), self._item("huge.py", oversized)] + ) + + assert [len(group) for group in groups] == [1, 1] + assert len(groups[1][0].parsing_result.serialized_dags) == len(oversized) + @pytest.mark.parametrize( ("items", "expected"), [ diff --git a/airflow-core/tests/unit/dag_processing/test_parse_result_query_budget.py b/airflow-core/tests/unit/dag_processing/test_parse_result_query_budget.py index e6ce7b19c06a9..b7c5e72aa0daf 100644 --- a/airflow-core/tests/unit/dag_processing/test_parse_result_query_budget.py +++ b/airflow-core/tests/unit/dag_processing/test_parse_result_query_budget.py @@ -19,8 +19,8 @@ to move a number here and account for it in review. Counts are calibrated against Postgres and the tests carrying them are marked for it, since -statement counts differ by dialect. The sweep test derives its call count from the group cap and -measures its prices, so it runs anywhere. +statement counts differ by dialect. The sweep test derives its call count from the Dag budget a +group carries and measures its prices, so it runs anywhere. """ from __future__ import annotations @@ -43,7 +43,7 @@ from airflow.dag_processing.collection import update_dag_parsing_results_in_db from airflow.dag_processing.manager import ( - MAX_FILES_PER_PERSISTENCE_GROUP, + MAX_DAGS_PER_PERSISTENCE_GROUP, DagFileInfo, DagFileProcessorManager, DagFileStat, @@ -255,18 +255,21 @@ def _measure_sweep(session, tmp_path: Path, n_files: int, sockets, name: str, da @pytest.mark.parametrize( - "n_files", + ("n_files", "dags_per_file"), [ - pytest.param(SWEEP_FILES, id="one-group"), - pytest.param(MAX_FILES_PER_PERSISTENCE_GROUP + 1, id="over-the-group-cap"), + pytest.param(SWEEP_FILES, 1, id="one-group"), + pytest.param(5, 8, id="over-the-dag-cap"), ], ) -def test_sweep_pays_fixed_cost_once_per_call(n_files, session, testing_dag_bundle, tmp_path, sockets): +def test_sweep_pays_fixed_cost_once_per_call( + n_files, dags_per_file, session, testing_dag_bundle, tmp_path, sockets +): """ How a sweep scales with the number of persistence calls it takes. A sweep is persisted a group at a time, so the fixed price is paid once per group rather than - once per file -- and a sweep past the cap really does split, rather than only being asserted to. + once per file -- and a sweep carrying more Dags than a group takes really does split, rather + than only being asserted to. Both prices are measured here, so the assertion holds on any backend. """ one_dag = _measure_sweep(session, tmp_path, 1, sockets, "one") @@ -274,13 +277,16 @@ def test_sweep_pays_fixed_cost_once_per_call(n_files, session, testing_dag_bundl per_dag = two_dags - one_dag fixed = one_dag - per_dag - sweep = _measure_sweep(session, tmp_path, n_files, sockets, f"sweep_{n_files}") + sweep = _measure_sweep( + session, tmp_path, n_files, sockets, f"sweep_{n_files}x{dags_per_file}", dags_per_file=dags_per_file + ) - calls = math.ceil(n_files / MAX_FILES_PER_PERSISTENCE_GROUP) - expected = calls * fixed + n_files * per_dag + total_dags = n_files * dags_per_file + calls = math.ceil(total_dags / MAX_DAGS_PER_PERSISTENCE_GROUP) + expected = calls * fixed + total_dags * per_dag assert sweep == expected, ( - f"a {n_files}-file sweep costs {sweep} statements, expected {expected} " - f"({calls} x {fixed} fixed + {n_files} x {per_dag} per Dag)." + f"a {n_files}-file sweep of {total_dags} Dags costs {sweep} statements, expected {expected} " + f"({calls} x {fixed} fixed + {total_dags} x {per_dag} per Dag)." ) From 34304f457f369d3020c24aeda65bf8909a3fc579 Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Tue, 18 Aug 2026 18:37:56 +0100 Subject: [PATCH 07/11] Say less in the comments around persisting a sweep The reasoning was spread over paragraphs where a sentence carries it, repeated between the two deprecated seams, and in places narrated what the next line already says. What stays is what the code cannot say for itself: why each group needs its own transaction, why files sharing a dag_id are ordered rather than only separated, and why a hook is looked for on the instance as well as the class. --- .../src/airflow/dag_processing/manager.py | 104 +++++++----------- .../tests/unit/dag_processing/test_manager.py | 3 +- 2 files changed, 39 insertions(+), 68 deletions(-) diff --git a/airflow-core/src/airflow/dag_processing/manager.py b/airflow-core/src/airflow/dag_processing/manager.py index e42a51ac23787..da01e6553612b 100644 --- a/airflow-core/src/airflow/dag_processing/manager.py +++ b/airflow-core/src/airflow/dag_processing/manager.py @@ -110,16 +110,10 @@ def _make_execution_api() -> InProcessExecutionAPI: MAX_DAGS_PER_PERSISTENCE_GROUP = 32 """ -How many Dags at most are persisted together in one call, and so in one transaction. +Dags at most per persistence call, bounding what one transaction holds locked. -Grouping saves a fixed cost once per file, but what one transaction holds locked grows with the -Dags in it, not the files -- a few files generating Dags dynamically can be thousands. Counting -Dags bounds the thing that actually grows, and at this size it does not bind on files defining -one or two Dags, which is the shape a sweep usually has. - -A file is never split across groups, since the record of it having been parsed and the import -errors keyed to it belong to it as a whole. One file defining more Dags than this is still -written on its own. +Counted in Dags rather than files because that is what grows: a few files generating Dags +dynamically can be thousands. """ @@ -1281,9 +1275,8 @@ def handle_parsing_result( Post-process a single finished parse result. .. deprecated:: 3.4.0 - Override :meth:`persist_parsing_results` instead. A subclass overriding this method - still receives every file, one at a time, exactly as before -- but a sweep can then no - longer be persisted in one pass, so batching is skipped for the whole Dag processor. + Override :meth:`persist_parsing_results` instead. This still receives every file, one + at a time, but overriding it skips batching for the whole Dag processor. Detects callback-only processing, updates file stats, emits metrics, and persists DAGs/import-errors via :meth:`persist_parsing_result`. @@ -1326,9 +1319,8 @@ def _build_parse_result( """ Work out what a finished parse leaves to persist. - Detects callback-only processing, emits metrics, and works out the stat this file should - record. Files with nothing to persist -- callback-only runs and failed parses -- have their - stat recorded here and return ``None``; the rest carry their Dags to whoever writes them. + Files with nothing to write -- callback-only runs and failed parses -- have their stat + recorded here and return ``None``; the rest carry their Dags to whoever writes them. """ is_callback_only = proc.had_callbacks and proc.parsing_result is None if is_callback_only: @@ -1371,12 +1363,11 @@ def persist_parsing_result( session: Session, ) -> None: """ - Persist parsed DAG data to the metadata database. + Persist parsed Dag data to the metadata database. .. deprecated:: 3.4.0 - Override :meth:`persist_parsing_results` instead. A subclass overriding this method - still receives every file, one at a time, exactly as before -- but a sweep can then no - longer be persisted in one pass, so batching is skipped for the whole Dag processor. + Override :meth:`persist_parsing_results` instead. This still receives every file, one + at a time, but overriding it skips batching for the whole Dag processor. """ import_errors: dict[tuple[str, str], str] = {} if parsing_result.import_errors: @@ -1411,9 +1402,7 @@ def _overrides(self, name: str) -> bool: """ Report whether this manager has its own version of one of the hooks. - Looks at the instance as well as the class, so that replacing a hook on a single manager - counts. Reading only the class would leave such a manager writing to the metadata DB while - looking, to whoever replaced the hook, as though it no longer did. + Reads the instance as well as the class, so replacing a hook on one manager counts. """ if name in getattr(self, "__dict__", ()): return True @@ -1423,10 +1412,9 @@ def _overrides_per_file_persist(self) -> bool: """ Report whether the per-file persistence hook has been replaced. - Batching would otherwise write straight past such an override, silently sending to the - metadata DB the results a deployment had arranged to send elsewhere. A manager that also - replaced the batch hook has said where a whole sweep should go, so that one is used and - this reports ``False``. + Batching would otherwise write past such an override, sending to the metadata DB results a + deployment had arranged to send elsewhere. Replacing the batch hook too says where a whole + sweep should go, so that one wins and this reports ``False``. """ if self._overrides("persist_parsing_results"): return False @@ -1452,9 +1440,8 @@ def _overrides_handle_parsing_result(self) -> bool: """ Report whether the per-file result handler has been replaced. - Released in 3.3 as the seam for deployments that forward results instead of writing them to - the metadata DB. Such an override persists a file itself, so it is still called once per - file and a sweep is left with nothing to batch. + Released in 3.3 for deployments that forward results rather than write them. Such an + override persists a file itself, leaving a sweep nothing to batch. """ return self._overrides("handle_parsing_result") @@ -1468,19 +1455,14 @@ def persist_parsing_results( """ Persist several files' parse results in one pass. - Everything :func:`update_dag_parsing_results_in_db` does besides writing the Dags themselves - is paid once per call, so persisting a whole sweep together costs far fewer statements than - persisting file by file. Results are grouped by bundle, since the bundle is what determines - the version and version data the call needs. + Everything :func:`update_dag_parsing_results_in_db` does besides writing the Dags is paid + once per call, so a whole sweep together costs far fewer statements than file by file. - Deployments that send parse results somewhere other than the metadata DB (e.g. AIP-92) can - override either this method or the per-file :meth:`persist_parsing_result`. A subclass that - overrides only the per-file hook keeps its existing behaviour: it is called once per file, - batching is skipped, and it is handed this method's session rather than being quietly - bypassed by a batched write it never sees. + This is the seam to override for deployments sending results somewhere other than the + metadata DB (e.g. AIP-92); one that overrides only the per-file + :meth:`persist_parsing_result` is still called once per file instead, with batching skipped. - Note that raising from here discards the whole group -- the caller falls back to persisting - each file on its own so one bad file cannot take the others down with it. + Raising from here discards the whole group; the caller then persists each file on its own. """ if self._overrides_per_file_persist(): for item in results: @@ -1502,23 +1484,16 @@ def build_persistence_groups(self, results: Sequence[FileParseResult]) -> list[l """ Split a sweep into units that can each be written in a single call. - A group holds one bundle's files, since the bundle determines the version and version data - the write needs, and claims a dag_id at most once: writing two files that define the same - dag_id together would merge them into one Dag and lose the duplicate warning, which is - raised by comparing an incoming Dag against the file already recorded in the DB. - - Groups are returned in the order they must be written, and a file is only ever held back to - a group after the one holding the file it duplicates -- so the duplicate still sees what it - duplicates, and the same file wins as when each was written on its own. Only the files - actually in conflict are held back; one repeated dag_id must not cost every other file its - place in a group. - - A group carries at most ``MAX_DAGS_PER_PERSISTENCE_GROUP`` Dags. + A group holds one bundle's files, which fixes the version the write needs, and claims a + dag_id at most once: writing two files defining the same dag_id together would merge them + and lose the duplicate warning, which comes from comparing an incoming Dag against the file + already recorded. Groups are returned in the order they must be written, so a file held back + still sees the one it duplicates and the same file wins. """ groups: list[list[FileParseResult]] = [] bundles: list[str] = [] dag_counts: list[int] = [] - # The last group to claim each dag_id, which is the earliest a file repeating it may go. + # Last group to claim each dag_id: the earliest a file repeating it may go. claimed_by: dict[str, int] = {} for item in results: @@ -1534,7 +1509,7 @@ def build_persistence_groups(self, results: Sequence[FileParseResult]) -> list[l ): break else: - # A new group takes the file whatever its size: a file cannot be split. + # A new group takes the file whatever its size; a file cannot be split. index = len(groups) groups.append([]) bundles.append(bundle_name) @@ -1599,8 +1574,7 @@ def _persist_bundle_group( def _collect_results(self): finished = [] to_persist: list[FileParseResult] = [] - # An override owns the whole of a file's handling, persistence included, so there is - # nothing left for a sweep to write together. + # Such an override owns persistence too, leaving a sweep nothing to write together. handles_each_file = self._overrides_handle_parsing_result() try: for file, proc in self._processors.items(): @@ -1616,8 +1590,7 @@ def _collect_results(self): if to_persist: self._persist_sweep(to_persist) finally: - # Whatever went wrong, these processes are done with; leaving them open leaks their - # sockets and keeps them queued as though they were still running. + # Leaving these open leaks their sockets and keeps them queued as if still running. for file in finished: processor = self._processors.pop(file) processor.close() @@ -1626,13 +1599,12 @@ def _persist_sweep(self, to_persist: list[FileParseResult]) -> None: """ Persist a sweep one group at a time, falling back to single files when a group fails. - Each group gets its own call, and so its own transaction: ``update_dag_parsing_results_in_db`` - rolls the session back before retrying an ``OperationalError``, which would otherwise discard - groups that had already succeeded in the same transaction while their files were still - recorded as persisted. + Each group gets its own transaction: ``update_dag_parsing_results_in_db`` rolls the session + back before retrying an ``OperationalError``, which would otherwise discard groups already + written alongside it while their files stayed recorded as persisted. - A subclass handling files one at a time is given single-file groups, so it never receives a - file twice: a retry after a grouped attempt would hand it results it had already accepted. + A subclass handling files one at a time gets single-file groups, so the fallback cannot hand + it a file it has already accepted. """ if self._overrides_per_file_persist(): groups = [[item] for item in to_persist] @@ -1669,8 +1641,8 @@ def _throttle_after_failed_persist(self, item: FileParseResult) -> None: """ Record a failed write without claiming its results. - Keeps the counts of whatever was last persisted and only moves the timestamps, so the file - is not parsed again immediately while the rest of the cycle carries on. + Keeps the last persisted counts and only moves the timestamps, so the file is not parsed + again immediately while the rest of the cycle carries on. """ self.log.exception( "Failed to persist parsing result for %s in bundle %s; " diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index 3c999de1a0fd0..c7ac8e7146e32 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py @@ -1525,7 +1525,7 @@ def test_terminate_normalizes_file_path_stats_tag(self): processor.kill.assert_called_once_with(signal.SIGTERM, escalation_delay=5.0) def test_persist_parsing_results_provides_its_own_session_when_caller_omits(self): - """``persist_parsing_results`` is wrapped in ``@provide_session`` so subclasses overriding it can run without a caller-supplied session.""" + """An override must be able to run without a session the caller never had.""" manager = DagFileProcessorManager(max_runs=1) file = DagFileInfo(bundle_name="testing", rel_path=Path("abc.txt"), bundle_path=TEST_DAGS_FOLDER) manager._file_stats[file] = DagFileStat() @@ -1821,7 +1821,6 @@ def handle_parsing_result(self, file, proc, *, session=None): @mock.patch("airflow.dag_processing.manager.update_dag_parsing_results_in_db", autospec=True) def test_an_override_delegating_to_super_still_persists(self, mock_write): - """The released implementation stays usable, so an override can add to it rather than replace it.""" class WrappingManager(DagFileProcessorManager): def handle_parsing_result(self, file, proc, *, session=None): From bfcea9dbc7c62250410d4663c1e2d678910892d4 Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Wed, 19 Aug 2026 08:23:20 +0100 Subject: [PATCH 08/11] Contain a failure to build a parse result to the file it came from Working out what a file leaves to persist can reach the database, for the team its bundle belongs to. Persisting a sweep together moves that work ahead of every write in it, so one such failure discarded the results of files that had already been handled, while their processors were closed as though they had been written. Writing a file at a time, as before, meant only the failing file was lost. The replacement seam is handed one group of a sweep, not the whole of it: a sweep spanning bundles, or carrying more Dags than a group takes, arrives as more than one call. The deprecation warning and the newsfragment said otherwise, and no test had a sweep that split, so nothing caught the claim. --- airflow-core/newsfragments/71771.misc.rst | 2 +- .../src/airflow/dag_processing/manager.py | 40 ++++++++++++++----- .../tests/unit/dag_processing/test_manager.py | 31 ++++++++++++++ .../test_parse_result_query_budget.py | 26 +++++++++++- 4 files changed, 88 insertions(+), 11 deletions(-) diff --git a/airflow-core/newsfragments/71771.misc.rst b/airflow-core/newsfragments/71771.misc.rst index 6956495512533..89b0c81c64f09 100644 --- a/airflow-core/newsfragments/71771.misc.rst +++ b/airflow-core/newsfragments/71771.misc.rst @@ -1 +1 @@ -Deprecate ``DagFileProcessorManager.handle_parsing_result`` and ``persist_parsing_result``; override ``persist_parsing_results`` instead, which is handed every file that finished parsing together. Overriding either deprecated method still works but stops parse results being persisted a sweep at a time. +Deprecate ``DagFileProcessorManager.handle_parsing_result`` and ``persist_parsing_result``; override ``persist_parsing_results`` instead, which is handed a group of the files that finished parsing together. Overriding either deprecated method still works but stops parse results being persisted a sweep at a time. diff --git a/airflow-core/src/airflow/dag_processing/manager.py b/airflow-core/src/airflow/dag_processing/manager.py index da01e6553612b..f978e0c4aa0f0 100644 --- a/airflow-core/src/airflow/dag_processing/manager.py +++ b/airflow-core/src/airflow/dag_processing/manager.py @@ -1430,8 +1430,8 @@ def _warn_if_batching_is_disabled(self) -> None: return warnings.warn( f"{type(self).__name__} overrides {replaced}, which is deprecated and prevents parse " - "results being persisted a sweep at a time. Override persist_parsing_results instead, " - "which is handed every file that finished together.", + "results being persisted several files at a time. Override persist_parsing_results " + "instead, which is handed a group of the files that finished together.", DeprecationWarning, stacklevel=2, ) @@ -1453,10 +1453,12 @@ def persist_parsing_results( session: Session = NEW_SESSION, ) -> None: """ - Persist several files' parse results in one pass. + Persist one group of a sweep's parse results in a single pass. Everything :func:`update_dag_parsing_results_in_db` does besides writing the Dags is paid - once per call, so a whole sweep together costs far fewer statements than file by file. + once per call, so several files together cost far fewer statements than file by file. A + sweep is split by :meth:`build_persistence_groups` first, so a sweep spanning bundles, or + carrying more Dags than a group takes, arrives as more than one call. This is the seam to override for deployments sending results somewhere other than the metadata DB (e.g. AIP-92); one that overrides only the per-file @@ -1584,7 +1586,21 @@ def _collect_results(self): finished.append(file) if handles_each_file: self.handle_parsing_result(file, proc) - elif (result := self._build_parse_result(file, proc)) is not None: + continue + try: + result = self._build_parse_result(file, proc) + except Exception: + # Working out what a file leaves to persist can reach the DB, for the team a + # bundle belongs to. Losing that file must not lose the sweep it arrived in. + self.log.exception( + "Failed to handle the parse result for %s in bundle %s; " + "the rest of the sweep is still persisted.", + str(file.rel_path), + file.bundle_name, + ) + self._throttle_retry(file, timezone.utcnow(), time.monotonic() - proc.start_time) + continue + if result is not None: to_persist.append(result) if to_persist: @@ -1651,12 +1667,18 @@ def _throttle_after_failed_persist(self, item: FileParseResult) -> None: str(item.file.rel_path), item.file.bundle_name, ) - current_stat = self._file_stats[item.file] - self._file_stats[item.file] = DagFileStat( + self._throttle_retry(item.file, item.stat.last_finish_time, item.stat.last_duration) + + def _throttle_retry( + self, file: DagFileInfo, finish_time: datetime | None, run_duration: float | None + ) -> None: + """Record that a file was handled without claiming results it did not produce.""" + current_stat = self._file_stats[file] + self._file_stats[file] = DagFileStat( num_dags=current_stat.num_dags, import_errors=current_stat.import_errors, - last_finish_time=item.stat.last_finish_time, - last_duration=item.stat.last_duration, + last_finish_time=finish_time, + last_duration=run_duration, run_count=current_stat.run_count + 1, last_num_of_db_queries=current_stat.last_num_of_db_queries, ) diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index c7ac8e7146e32..b6ecd067c3b1e 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py @@ -1859,6 +1859,37 @@ def test_only_a_deprecated_override_gives_up_batching(self, overrides, expected) else: manager._warn_if_batching_is_disabled() + def test_a_file_that_cannot_be_handled_does_not_discard_its_neighbours(self): + """ + Working out what a file leaves to persist can reach the DB, for the team its bundle is in. + + Persisting a sweep together means one such failure arrives before any of it is written, so + it has to be contained to its own file rather than take the sweep down with it. + """ + manager = DagFileProcessorManager(max_runs=1) + manager._bundle_versions["testing"] = "v1" + good = self._ready_processor(manager, "good.py", num_dags=1) + bad = self._ready_processor(manager, "bad.py", num_dags=1) + # Counts from the last successful parse, which the throttled retry must preserve. + manager._file_stats[bad] = DagFileStat(num_dags=7, import_errors=2, run_count=3) + build = manager._build_parse_result + + def fail_for_bad(file, proc): + if file == bad: + raise RuntimeError("team lookup failed") + return build(file, proc) + + with mock.patch.object(manager, "_build_parse_result", side_effect=fail_for_bad): + with mock.patch.object(manager, "_persist_sweep", autospec=True) as sweep: + manager._collect_results() + + assert [item.file for item in sweep.call_args.args[0]] == [good], ( + "the file that could be handled must still be persisted" + ) + assert manager._processors == {} + assert manager._file_stats[bad].run_count == 4, "the failed file must be counted as run" + assert manager._file_stats[bad].num_dags == 7, "and must not claim results it never produced" + def test_finished_processors_are_closed_even_when_persistence_raises(self): """A processor left open leaks its sockets and stays queued as though it were still running.""" manager = DagFileProcessorManager(max_runs=1) diff --git a/airflow-core/tests/unit/dag_processing/test_parse_result_query_budget.py b/airflow-core/tests/unit/dag_processing/test_parse_result_query_budget.py index b7c5e72aa0daf..97aba73117e82 100644 --- a/airflow-core/tests/unit/dag_processing/test_parse_result_query_budget.py +++ b/airflow-core/tests/unit/dag_processing/test_parse_result_query_budget.py @@ -366,10 +366,34 @@ def persist_parsing_results(self, results, *, session=None): db_write = _collect_a_two_file_sweep(manager, tmp_path, sockets, "batch_override", session) - assert seen == [["file_0.py", "file_1.py"]], "the whole sweep should arrive in one call" + assert seen == [["file_0.py", "file_1.py"]], "files that can share a group arrive in one call" db_write.assert_not_called() +def test_a_batch_override_is_called_once_per_group_not_once_per_sweep(tmp_path): + """A sweep spanning bundles cannot be one write, so the seam has to be documented per group.""" + seen: list[list[str]] = [] + + class BatchApiManager(DagFileProcessorManager): + def persist_parsing_results(self, results, *, session=None): + seen.append([str(item.file.rel_path) for item in results]) + + manager = BatchApiManager(max_runs=1) + manager._bundle_versions.update({BUNDLE: None, OTHER_BUNDLE: None}) + + manager._persist_sweep( + [ + _parse_result(tmp_path, "in_testing_a"), + _parse_result(tmp_path, "in_other", bundle_name=OTHER_BUNDLE), + _parse_result(tmp_path, "in_testing_b"), + ] + ) + + assert seen == [["in_testing_a.py", "in_testing_b.py"], ["in_other.py"]], ( + "each bundle is its own group, so the override sees one call per group" + ) + + def test_the_default_manager_writes_the_sweep_to_the_database(session, testing_dag_bundle, tmp_path, sockets): """The negative of the override case: with nothing replaced, a sweep still reaches the DB once.""" manager = DagFileProcessorManager(max_runs=1) From 02ec79a0e6440f94d18452b74fd1c9876a210f47 Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Wed, 19 Aug 2026 12:56:38 +0100 Subject: [PATCH 09/11] Stop a batched sweep leaving stale a Dag it had just written An import error is recorded against a file after the Dags filed under that file are written, so a file reporting an error and a file whose Dag is filed under it could not share a write: the error marked the Dag stale, hiding it from the scheduler and the UI until something parsed it again. Grouping already held the two apart in one direction and not the other. Working out how to group a sweep reads the Dags a parser sent, so a malformed one took the whole cycle down with it rather than the file that carried it. That is the failure the step before it was already hardened against. A file handled on its own was written through the seam being deprecated rather than the one deployments are told to move to, so a deployment that had moved found its results in the metadata DB anyway. Two tests added a bundle of the same name without either removing it, so whichever ran second failed on a suite that runs them in a random order. A group is one transaction only while the metadata database is the only thing written; the FAB auth manager commits Dag permissions partway through one. Said where the retry that assumes otherwise is documented. Each of these was a separate way merging a sweep could differ from writing it a file at a time, found one at a time by reading. What the sweep leaves behind is now compared against what one file at a time leaves behind, over the shapes where merging has something to get wrong, so the next one does not need to be thought of first. --- airflow-core/newsfragments/71771.misc.rst | 2 +- .../src/airflow/dag_processing/manager.py | 184 +++++++---- .../test_batched_persistence_equivalence.py | 229 ++++++++++++++ .../tests/unit/dag_processing/test_manager.py | 292 +++++++++++++++--- .../test_parse_result_query_budget.py | 62 +++- 5 files changed, 669 insertions(+), 100 deletions(-) create mode 100644 airflow-core/tests/unit/dag_processing/test_batched_persistence_equivalence.py diff --git a/airflow-core/newsfragments/71771.misc.rst b/airflow-core/newsfragments/71771.misc.rst index 89b0c81c64f09..82bc13e0d143b 100644 --- a/airflow-core/newsfragments/71771.misc.rst +++ b/airflow-core/newsfragments/71771.misc.rst @@ -1 +1 @@ -Deprecate ``DagFileProcessorManager.handle_parsing_result`` and ``persist_parsing_result``; override ``persist_parsing_results`` instead, which is handed a group of the files that finished parsing together. Overriding either deprecated method still works but stops parse results being persisted a sweep at a time. +Deprecate ``DagFileProcessorManager.persist_parsing_result``; override ``persist_parsing_results`` instead, which is handed a group of the files that finished parsing together. Overriding ``handle_parsing_result`` is still supported, but stops parse results being persisted a group at a time. diff --git a/airflow-core/src/airflow/dag_processing/manager.py b/airflow-core/src/airflow/dag_processing/manager.py index f978e0c4aa0f0..b7a80720be8cd 100644 --- a/airflow-core/src/airflow/dag_processing/manager.py +++ b/airflow-core/src/airflow/dag_processing/manager.py @@ -177,6 +177,9 @@ class FileParseResult(NamedTuple): run_duration: float stat: DagFileStat """The stat to record once this file's results are persisted.""" + bundle_version: str | None = None + version_data: dict | None = None + """The bundle's version when the file was collected, which a write needs and DagFileInfo lacks.""" def _config_int_factory(section: str, key: str): @@ -1274,12 +1277,14 @@ def handle_parsing_result( """ Post-process a single finished parse result. - .. deprecated:: 3.4.0 - Override :meth:`persist_parsing_results` instead. This still receives every file, one - at a time, but overriding it skips batching for the whole Dag processor. + Overriding this still receives every file, one at a time, and is still supported -- but a + sweep can then no longer be persisted a group at a time, so batching is skipped for the + whole Dag processor. :meth:`persist_parsing_results` is the seam to prefer where it is + enough, since it keeps batching. - Detects callback-only processing, updates file stats, emits metrics, - and persists DAGs/import-errors via :meth:`persist_parsing_result`. + Detects callback-only processing, updates file stats, emits metrics, and persists the + file's Dags and import errors through :meth:`persist_parsing_results`, so a subclass that + adopted that seam is not written past. Owns its own DB session via ``@provide_session`` so subclasses that forward results without touching the metadata DB (e.g. AIP-92 API-backed @@ -1296,15 +1301,7 @@ def handle_parsing_result( return try: - self.persist_parsing_result( - bundle_name=file.bundle_name, - bundle_version=self._bundle_versions[file.bundle_name], - version_data=self._bundle_version_data.get(file.bundle_name), - parsing_result=result.parsing_result, - run_duration=result.run_duration, - relative_fileloc=str(file.rel_path), - session=session, - ) + self.persist_parsing_results([result], session=session) except Exception: self._throttle_after_failed_persist(result) return @@ -1349,6 +1346,8 @@ def _build_parse_result( parsing_result=proc.parsing_result, run_duration=run_duration, stat=next_stat, + bundle_version=self._bundle_versions[file.bundle_name], + version_data=self._bundle_version_data.get(file.bundle_name), ) def persist_parsing_result( @@ -1421,20 +1420,23 @@ def _overrides_per_file_persist(self) -> bool: return self._overrides("persist_parsing_result") def _warn_if_batching_is_disabled(self) -> None: - """Say once, at startup, that an override is costing this Dag processor its batched writes.""" - if self._overrides_handle_parsing_result(): - replaced = "handle_parsing_result" - elif self._overrides_per_file_persist(): - replaced = "persist_parsing_result" - else: - return - warnings.warn( - f"{type(self).__name__} overrides {replaced}, which is deprecated and prevents parse " - "results being persisted several files at a time. Override persist_parsing_results " - "instead, which is handed a group of the files that finished together.", - DeprecationWarning, - stacklevel=2, - ) + """Say once, at startup, what an override is costing this Dag processor.""" + if self._overrides_per_file_persist(): + warnings.warn( + f"{type(self).__name__} overrides persist_parsing_result, which is deprecated. " + "Override persist_parsing_results instead, which is handed a group of the files " + "that finished together.", + DeprecationWarning, + stacklevel=2, + ) + elif self._overrides_handle_parsing_result(): + # Not deprecated: nothing replaces handling a file in full, so say what it costs + # somewhere an operator will actually see it. + self.log.warning( + "%s overrides handle_parsing_result, so parse results are persisted one file at a " + "time rather than a group at a time.", + type(self).__name__, + ) def _overrides_handle_parsing_result(self) -> bool: """ @@ -1464,14 +1466,29 @@ def persist_parsing_results( metadata DB (e.g. AIP-92); one that overrides only the per-file :meth:`persist_parsing_result` is still called once per file instead, with batching skipped. - Raising from here discards the whole group; the caller then persists each file on its own. + Raising from the built-in write discards the whole group, and the caller then persists each + file on its own, since the group was one transaction and nothing was kept. A replaced write + gets no such retry: it may have accepted the group before failing, and its files are left to + be parsed again rather than sent twice. + + The one transaction holds while the metadata DB is the only thing being written. Under the + FAB auth manager it does not: syncing a Dag's permissions commits on this session partway + through the group, so a later failure can leave the Dags written before it committed and the + rest rolled back. The per-file retry then rewrites what was already kept, which the writes + themselves tolerate, but a listener watching import errors sees the same file announced + twice. + + Note that a replaced write is not reached without touching the metadata DB in multi-team + deployments: the team a bundle belongs to is looked up while working out what a file leaves + to persist, and a file whose lookup fails never arrives here. Callback-only completions do + not arrive here either; :meth:`handle_parsing_result` sees every file. """ if self._overrides_per_file_persist(): for item in results: self.persist_parsing_result( bundle_name=item.file.bundle_name, - bundle_version=self._bundle_versions[item.file.bundle_name], - version_data=self._bundle_version_data.get(item.file.bundle_name), + bundle_version=item.bundle_version, + version_data=item.version_data, parsing_result=item.parsing_result, run_duration=item.run_duration, relative_fileloc=str(item.file.rel_path), @@ -1486,40 +1503,57 @@ def build_persistence_groups(self, results: Sequence[FileParseResult]) -> list[l """ Split a sweep into units that can each be written in a single call. - A group holds one bundle's files, which fixes the version the write needs, and claims a - dag_id at most once: writing two files defining the same dag_id together would merge them - and lose the duplicate warning, which comes from comparing an incoming Dag against the file - already recorded. Groups are returned in the order they must be written, so a file held back - still sees the one it duplicates and the same file wins. + Groups are contiguous runs of the sweep, so writing them in turn writes every file in the + order it arrived. Files share far more than a dag_id -- assets, aliases and triggers are + shared across bundles by design -- and for all of it the last write wins, so a file allowed + to overtake another would quietly change which definition survives. + + A run ends where the next file cannot join it: a different bundle, since the bundle fixes + the version the write needs; a dag_id the run already claims, since merging two files + defining one dag_id would lose the duplicate warning that comes from comparing an incoming + Dag against the file already recorded; a file the run and the incoming file speak + for in opposite directions, since an error recorded against a file is applied after the + Dags filed under it are written and would leave them stale; or + ``MAX_DAGS_PER_PERSISTENCE_GROUP``. """ groups: list[list[FileParseResult]] = [] bundles: list[str] = [] dag_counts: list[int] = [] - # Last group to claim each dag_id: the earliest a file repeating it may go. - claimed_by: dict[str, int] = {} + claimed_dag_ids: list[set[str]] = [] + # Where a run files its Dags, and which files it reports on: a collision either way round + # has to end the run. + dag_locs_claimed: list[set[str]] = [] + file_locs_claimed: list[set[str]] = [] for item in results: - dag_ids = {dag.dag_id for dag in item.parsing_result.serialized_dags} - n_dags = len(item.parsing_result.serialized_dags) + dags = item.parsing_result.serialized_dags + dag_ids = {dag.dag_id for dag in dags} + dag_locs = {dag.relative_fileloc for dag in dags if dag.relative_fileloc} + file_locs = {str(item.file.rel_path), *(item.parsing_result.import_errors or ())} bundle_name = item.file.bundle_name - earliest = max((claimed_by[dag_id] + 1 for dag_id in dag_ids if dag_id in claimed_by), default=0) - - for index in range(earliest, len(groups)): - if ( - bundles[index] == bundle_name - and dag_counts[index] + n_dags <= MAX_DAGS_PER_PERSISTENCE_GROUP - ): - break - else: + + joins_run = ( + bool(groups) + and bundles[-1] == bundle_name + and dag_counts[-1] + len(dags) <= MAX_DAGS_PER_PERSISTENCE_GROUP + and claimed_dag_ids[-1].isdisjoint(dag_ids) + and dag_locs_claimed[-1].isdisjoint(file_locs) + and file_locs_claimed[-1].isdisjoint(dag_locs) + ) + if not joins_run: # A new group takes the file whatever its size; a file cannot be split. - index = len(groups) groups.append([]) bundles.append(bundle_name) dag_counts.append(0) - - groups[index].append(item) - dag_counts[index] += n_dags - claimed_by.update(dict.fromkeys(dag_ids, index)) + claimed_dag_ids.append(set()) + dag_locs_claimed.append(set()) + file_locs_claimed.append(set()) + + groups[-1].append(item) + dag_counts[-1] += len(dags) + claimed_dag_ids[-1].update(dag_ids) + dag_locs_claimed[-1].update(dag_locs) + file_locs_claimed[-1].update(file_locs) return groups def _persist_bundle_group( @@ -1546,6 +1580,11 @@ def _persist_bundle_group( (bundle_name, rel_path): error for rel_path, error in (parsing_result.import_errors or {}).items() } + # A file's own parse is the last word on it, so it drops an error another file in the + # sweep reported against it. Merging without that keeps an error the file has fixed. + # A group therefore records where its files ended up, not how they got there: an error + # raised and resolved inside one sweep is never written, so no listener hears of it. + import_errors.pop((bundle_name, relative_fileloc), None) import_errors.update(file_errors) # Include the parsed file even when it defines no Dags, so its stale import errors # still get cleared. @@ -1559,12 +1598,16 @@ def _persist_bundle_group( file_warnings = parsing_result.warnings or [] if file_warnings and isinstance(file_warnings[0], dict): file_warnings = [DagWarning(**warn) for warn in file_warnings] + # Likewise for the Dags this file defines: what it says now replaces what was said + # about them earlier in the sweep, including saying nothing. + defined = {dag.dag_id for dag in parsing_result.serialized_dags} + dag_warnings = {warning for warning in dag_warnings if warning.dag_id not in defined} dag_warnings.update(file_warnings) update_dag_parsing_results_in_db( bundle_name=bundle_name, - bundle_version=self._bundle_versions[bundle_name], - version_data=self._bundle_version_data.get(bundle_name), + bundle_version=items[0].bundle_version, + version_data=items[0].version_data, dags=dags, import_errors=import_errors, parse_duration=parse_durations, @@ -1617,15 +1660,26 @@ def _persist_sweep(self, to_persist: list[FileParseResult]) -> None: Each group gets its own transaction: ``update_dag_parsing_results_in_db`` rolls the session back before retrying an ``OperationalError``, which would otherwise discard groups already - written alongside it while their files stayed recorded as persisted. + written alongside it while their files stayed recorded as persisted. Under the FAB auth + manager a group is not quite one transaction -- see :meth:`persist_parsing_results`. A subclass handling files one at a time gets single-file groups, so the fallback cannot hand - it a file it has already accepted. + it a file it has already accepted. One that replaced the batch write gets no fallback at + all, since only the built-in write is known to keep nothing when it raises. """ if self._overrides_per_file_persist(): groups = [[item] for item in to_persist] else: - groups = self.build_persistence_groups(to_persist) + try: + groups = self.build_persistence_groups(to_persist) + except Exception: + # Grouping reads the Dags the parser sent, so a malformed one lands here. A file at + # a time hits the same failure inside _persist_single, which contains it. + self.log.exception( + "Failed to group %d parse results; persisting them a file at a time.", + len(to_persist), + ) + groups = [[item] for item in to_persist] for group in groups: if len(group) == 1: @@ -1634,6 +1688,18 @@ def _persist_sweep(self, to_persist: list[FileParseResult]) -> None: try: self.persist_parsing_results(group) except Exception: + if self._overrides("persist_parsing_results"): + # Only the built-in write is known to have kept nothing when it raises. A + # replaced one may have accepted the group and failed afterwards, so sending + # its files again would write them twice. + self.log.exception( + "Failed to persist %d parse results as a group; the hook that took them " + "may have kept some, so they are not sent again.", + len(group), + ) + for item in group: + self._throttle_retry(item.file, item.stat.last_finish_time, item.stat.last_duration) + continue self.log.exception( "Failed to persist %d parse results as a group; retrying them individually.", len(group), diff --git a/airflow-core/tests/unit/dag_processing/test_batched_persistence_equivalence.py b/airflow-core/tests/unit/dag_processing/test_batched_persistence_equivalence.py new file mode 100644 index 0000000000000..07a2cb87105f2 --- /dev/null +++ b/airflow-core/tests/unit/dag_processing/test_batched_persistence_equivalence.py @@ -0,0 +1,229 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +""" +Persisting a sweep together must leave the database where writing it a file at a time would. + +Everything a sweep carries used to be scoped to one file per write: import errors, warnings, +Dag rows, versions. Merging them makes that state shared, and each kind of it can be merged +wrongly on its own. Rather than argue each kind through, write the same sweep both ways against +a real database and compare what is in it afterwards. + +Both halves go through ``_persist_sweep`` so the only difference is how much is handed over at +once, and the shapes below are the ones where merging has something to get wrong. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +from sqlalchemy import delete, select + +from airflow.dag_processing.manager import DagFileInfo, DagFileProcessorManager, DagFileStat, FileParseResult +from airflow.dag_processing.processor import DagFileParsingResult +from airflow.models.dag import DagModel +from airflow.models.dag_version import DagVersion +from airflow.models.dagbundle import DagBundleModel +from airflow.models.dagwarning import DagWarning +from airflow.models.errors import ParseImportError +from airflow.providers.standard.operators.empty import EmptyOperator +from airflow.sdk import DAG +from airflow.serialization.serialized_objects import LazyDeserializedDAG + +from tests_common.test_utils.db import ( + clear_db_dag_bundles, + clear_db_dags, + clear_db_import_errors, + clear_db_serialized_dags, +) + +pytestmark = pytest.mark.db_test + +BUNDLE = "equiv" +OTHER_BUNDLE = "equiv-other" + + +def _dag(tmp_path: Path, dag_id: str, rel_path: str) -> LazyDeserializedDAG: + # DagCode reads the source off disk; without a real file the Dag fails to serialize. + (tmp_path / rel_path).parent.mkdir(parents=True, exist_ok=True) + (tmp_path / rel_path).write_text("# equivalence fixture\n") + dag = DAG(dag_id=dag_id, schedule="@daily") + EmptyOperator(task_id="task", dag=dag) + dag.fileloc = str(tmp_path / rel_path) + dag.relative_fileloc = rel_path + return LazyDeserializedDAG.from_dag(dag) + + +def _file( + tmp_path: Path, + rel_path: str, + dags: list[tuple[str, str]] | None = None, + errors: dict[str, str] | None = None, + warnings: list[dict] | None = None, + bundle: str = BUNDLE, + version: str | None = "v1", +) -> FileParseResult: + """One finished parse. ``dags`` is (dag_id, the file the Dag is filed under).""" + (tmp_path / rel_path).parent.mkdir(parents=True, exist_ok=True) + (tmp_path / rel_path).write_text("# equivalence fixture\n") + return FileParseResult( + file=DagFileInfo(bundle_name=bundle, rel_path=Path(rel_path), bundle_path=tmp_path), + parsing_result=DagFileParsingResult( + fileloc=str(tmp_path / rel_path), + serialized_dags=[_dag(tmp_path, dag_id, under) for dag_id, under in (dags or [])], + import_errors=errors, + warnings=warnings, + ), + run_duration=0.5, + stat=DagFileStat(), + bundle_version=version, + version_data={"sha": version} if version else None, + ) + + +def _snapshot(session) -> dict[str, list[tuple[str, ...]]]: + """Everything a sweep writes that a reader can tell apart.""" + + def rows(stmt): + return sorted(tuple(str(column) for column in row) for row in session.execute(stmt).all()) + + session.expire_all() + return { + "dag": rows( + select( + DagModel.dag_id, + DagModel.bundle_name, + DagModel.bundle_version, + DagModel.relative_fileloc, + DagModel.is_stale, + DagModel.has_import_errors, + ) + ), + "dag_version": rows(select(DagVersion.dag_id, DagVersion.bundle_version, DagVersion.version_data)), + "import_error": rows( + select(ParseImportError.bundle_name, ParseImportError.filename, ParseImportError.stacktrace) + ), + "warning": rows(select(DagWarning.dag_id, DagWarning.warning_type, DagWarning.message)), + } + + +def _reset(session) -> None: + """Clear what a sweep writes, leaving the bundles it is written against.""" + session.execute(delete(DagWarning)) + session.commit() + clear_db_serialized_dags() + clear_db_import_errors() + clear_db_dags() + for name in (BUNDLE, OTHER_BUNDLE): + session.merge(DagBundleModel(name=name)) + session.commit() + + +def _persist(sweep: list[FileParseResult], *, batched: bool) -> None: + manager = DagFileProcessorManager(max_runs=1) + manager._bundle_versions.update({BUNDLE: "v1", OTHER_BUNDLE: "v1"}) + if batched: + manager._persist_sweep(sweep) + else: + for item in sweep: + manager._persist_sweep([item]) + + +def _plain(tmp_path): + return [_file(tmp_path, f"plain_{i}.py", dags=[(f"plain_{i}", f"plain_{i}.py")]) for i in range(3)] + + +def _duplicate_dag_id(tmp_path): + return [ + _file(tmp_path, "first.py", dags=[("shared", "first.py")]), + _file(tmp_path, "second.py", dags=[("shared", "second.py")]), + ] + + +def _error_against_a_file_that_then_parses(tmp_path): + return [ + _file(tmp_path, "blamer.py", errors={"blamed.py": "reported by another file"}), + _file(tmp_path, "blamed.py", dags=[("blamed_dag", "blamed.py")]), + ] + + +def _dag_filed_under_a_file_that_errored(tmp_path): + return [ + _file(tmp_path, "broken.py", errors={"broken.py": "it broke"}), + _file(tmp_path, "healthy.py", dags=[("filed_elsewhere", "broken.py")]), + ] + + +def _file_that_now_defines_nothing(tmp_path): + return [ + _file(tmp_path, "emptied.py", errors={"emptied.py": "was broken"}), + _file(tmp_path, "healthy.py", dags=[("healthy_dag", "healthy.py")]), + _file(tmp_path, "emptied.py"), + ] + + +def _two_bundles_interleaved(tmp_path): + return [ + _file(tmp_path, "a.py", dags=[("a_dag", "a.py")]), + _file(tmp_path, "b.py", dags=[("b_dag", "b.py")], bundle=OTHER_BUNDLE), + _file(tmp_path, "c.py", dags=[("c_dag", "c.py")]), + ] + + +def _warning_then_a_clean_parse(tmp_path): + return [ + _file( + tmp_path, + "warner.py", + warnings=[{"dag_id": "warned_dag", "warning_type": "non-existent pool", "message": "gone"}], + ), + _file(tmp_path, "owner.py", dags=[("warned_dag", "owner.py")]), + ] + + +@pytest.mark.parametrize( + "build_sweep", + [ + _plain, + _duplicate_dag_id, + _error_against_a_file_that_then_parses, + _dag_filed_under_a_file_that_errored, + _file_that_now_defines_nothing, + _two_bundles_interleaved, + _warning_then_a_clean_parse, + ], + ids=lambda fn: fn.__name__.strip("_"), +) +def test_a_sweep_persisted_together_lands_where_one_file_at_a_time_would(build_sweep, session, tmp_path): + _reset(session) + _persist(build_sweep(tmp_path / "sequential"), batched=False) + sequential = _snapshot(session) + + _reset(session) + _persist(build_sweep(tmp_path / "batched"), batched=True) + batched = _snapshot(session) + + assert batched == sequential + + +@pytest.fixture(autouse=True) +def clean_db(): + yield + clear_db_serialized_dags() + clear_db_import_errors() + clear_db_dags() + clear_db_dag_bundles() diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index b6ecd067c3b1e..c046e894092c5 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py @@ -27,6 +27,7 @@ import signal import textwrap import time +import warnings import zipfile from collections import OrderedDict, defaultdict, namedtuple from datetime import datetime, timedelta @@ -1587,7 +1588,7 @@ def test_a_written_file_records_the_stat_its_parse_produced(self, session): processor.had_callbacks = False processor.parsing_result = DagFileParsingResult(fileloc="abc.txt", serialized_dags=[]) - with mock.patch.object(manager, "persist_parsing_results") as mock_persist: + with mock.patch.object(manager, "persist_parsing_results", autospec=True) as mock_persist: result = manager._build_parse_result(file, processor) manager._persist_sweep([result]) @@ -1620,7 +1621,7 @@ def test_collect_results_persists_the_whole_sweep_in_one_call(self): files = [self._ready_processor(manager, name) for name in ("a.py", "b.py", "c.py")] - with mock.patch.object(manager, "persist_parsing_results") as mock_persist: + with mock.patch.object(manager, "persist_parsing_results", autospec=True) as mock_persist: manager._collect_results() mock_persist.assert_called_once() @@ -1658,7 +1659,7 @@ def test_files_sharing_a_dag_id_are_written_one_at_a_time(self): ) ) - with mock.patch.object(manager, "_persist_bundle_group") as mock_group: + with mock.patch.object(manager, "_persist_bundle_group", autospec=True) as mock_group: manager.persist_parsing_results(items, session=mock.MagicMock()) assert [len(call.args[1]) for call in mock_group.call_args_list] == [1, 1], ( @@ -1684,18 +1685,25 @@ def test_files_with_distinct_dag_ids_are_written_together(self): ) ) - with mock.patch.object(manager, "_persist_bundle_group") as mock_group: + with mock.patch.object(manager, "_persist_bundle_group", autospec=True) as mock_group: manager.persist_parsing_results(items, session=mock.MagicMock()) assert [len(call.args[1]) for call in mock_group.call_args_list] == [2] - def _item(self, name: str, dag_ids: list[str], bundle_name: str = "testing") -> FileParseResult: + def _item( + self, + name: str, + dag_ids: list[str], + bundle_name: str = "testing", + filed_under: str | None = None, + ) -> FileParseResult: file = DagFileInfo(bundle_name=bundle_name, rel_path=Path(name), bundle_path=TEST_DAGS_FOLDER) + dags = [self._lazy_dag(dag_id) for dag_id in dag_ids] + for dag in dags: + dag.data["dag"]["relative_fileloc"] = filed_under or name return FileParseResult( file=file, - parsing_result=DagFileParsingResult( - fileloc=name, serialized_dags=[self._lazy_dag(dag_id) for dag_id in dag_ids] - ), + parsing_result=DagFileParsingResult(fileloc=name, serialized_dags=dags), run_duration=1.0, stat=DagFileStat(), ) @@ -1712,10 +1720,9 @@ def test_one_duplicate_dag_id_does_not_split_the_rest_of_the_sweep(self): groups = manager.build_persistence_groups(items) - assert [len(group) for group in groups] == [4, 1] - assert [str(item.file.rel_path) for item in groups[1]] == ["second.py"], ( - "the later of the two files must be the one held back, so it sees the earlier one" - ) + assert [len(group) for group in groups] == [1, 4] + assert [str(item.file.rel_path) for item in groups[0]] == ["first.py"] + assert str(groups[1][0].file.rel_path) == "second.py", "the duplicate is written after it" def test_a_group_is_capped_by_dags_so_one_transaction_cannot_lock_a_whole_sweep(self): """What one transaction holds locked grows with the Dags in it, not with the files.""" @@ -1835,29 +1842,244 @@ def handle_parsing_result(self, file, proc, *, session=None): mock_write.assert_called_once() assert manager._file_stats[file].run_count == 1 + def test_the_startup_check_runs_before_the_parsing_loop(self): + """Nothing else calls it, so without this the notice would silently never fire.""" + manager = DagFileProcessorManager(max_runs=1) + + with mock.patch.object(manager, "_warn_if_batching_is_disabled", autospec=True) as warn: + with mock.patch.object(manager, "prepare_bundles", autospec=True): + with mock.patch.object(manager, "_symlink_latest_log_directory", autospec=True): + manager.before_run() + + warn.assert_called_once() + @pytest.mark.parametrize( - ("overrides", "expected"), + ("overrides", "warns"), [ - pytest.param((), False, id="no-override-batches"), - pytest.param(("handle_parsing_result",), True, id="released-handler"), - pytest.param(("persist_parsing_result",), True, id="released-per-file-persist"), - pytest.param(("persist_parsing_results",), False, id="batch-seam-still-batches"), - pytest.param(("persist_parsing_result", "persist_parsing_results"), False, id="both-persist"), + pytest.param((), False, id="nothing-overridden"), + pytest.param(("handle_parsing_result",), False, id="handling-a-file-has-no-replacement"), + pytest.param(("persist_parsing_result",), True, id="the-per-file-write-is-replaced"), + pytest.param(("persist_parsing_results",), False, id="the-batch-write-is-the-replacement"), + pytest.param(("persist_parsing_result", "persist_parsing_results"), False, id="both-writes"), ], ) - def test_only_a_deprecated_override_gives_up_batching(self, overrides, expected): - """A subclass that adopted the batch seam keeps batching, even while it still carries the old one.""" + def test_only_the_write_with_a_replacement_is_deprecated(self, overrides, warns): + """Handling a file in full has nothing to move to, so overriding it is supported, not deprecated.""" subclass = type( "Subclass", (DagFileProcessorManager,), {name: lambda *a, **kw: None for name in overrides} ) manager = subclass(max_runs=1) - if expected: + if warns: with pytest.warns(DeprecationWarning, match="persist_parsing_results"): manager._warn_if_batching_is_disabled() else: - manager._warn_if_batching_is_disabled() + with warnings.catch_warnings(): + warnings.simplefilter("error", DeprecationWarning) + manager._warn_if_batching_is_disabled() + + def test_a_sweep_is_written_in_the_order_it_arrived(self): + """ + Files share assets, aliases and triggers across bundles, and for all of it the last write + wins. Letting a file overtake another to fill an earlier group would change which + definition survives. + """ + manager = DagFileProcessorManager(max_runs=1) + + groups = manager.build_persistence_groups( + [ + self._item("a.py", ["dag_a"]), + self._item("b.py", ["dag_b"], bundle_name="other"), + self._item("c.py", ["dag_c"]), + ] + ) + + assert [[str(item.file.rel_path) for item in group] for group in groups] == [ + ["a.py"], + ["b.py"], + ["c.py"], + ], "c.py may not join a.py's group, which would write it before b.py" + + def test_a_files_own_parse_clears_an_error_another_file_reported_against_it(self): + """Merging a sweep must apply what a later file says, including that it is now clean.""" + manager = DagFileProcessorManager(max_runs=1) + blamed = self._item("blamed.py", ["blamed_dag"]) + blamer = FileParseResult( + file=DagFileInfo(bundle_name="testing", rel_path=Path("blamer.py"), bundle_path=TEST_DAGS_FOLDER), + parsing_result=DagFileParsingResult( + fileloc="blamer.py", serialized_dags=[], import_errors={"blamed.py": "stale"} + ), + run_duration=1.0, + stat=DagFileStat(), + ) + + with mock.patch( + "airflow.dag_processing.manager.update_dag_parsing_results_in_db", autospec=True + ) as write: + manager._persist_bundle_group("testing", [blamer, blamed], session=mock.MagicMock()) + + assert ("testing", "blamed.py") not in write.call_args.kwargs["import_errors"], ( + "the file's own parse is the last word on it" + ) + + def test_a_files_own_parse_clears_a_warning_carried_for_its_dag(self): + """A later result saying nothing about a Dag it defines has to retract what was said.""" + manager = DagFileProcessorManager(max_runs=1) + warned = FileParseResult( + file=DagFileInfo(bundle_name="testing", rel_path=Path("warner.py"), bundle_path=TEST_DAGS_FOLDER), + parsing_result=DagFileParsingResult( + fileloc="warner.py", + serialized_dags=[], + warnings=[{"dag_id": "owned_dag", "warning_type": "non-existent pool", "message": "gone"}], + ), + run_duration=1.0, + stat=DagFileStat(), + ) + owner = self._item("owner.py", ["owned_dag"]) + + with mock.patch( + "airflow.dag_processing.manager.update_dag_parsing_results_in_db", autospec=True + ) as write: + manager._persist_bundle_group("testing", [warned, owner], session=mock.MagicMock()) + + assert not write.call_args.kwargs["warnings"], "the Dag's own file said nothing about it" + + def test_a_replaced_batch_write_is_not_sent_the_same_files_twice(self): + """A hook that failed may still have kept the group; only the built-in write is known not to.""" + manager = DagFileProcessorManager(max_runs=1) + manager._bundle_versions["testing"] = "v1" + items = [self._item(f"file_{i}.py", [f"dag_{i}"]) for i in range(2)] + calls: list[int] = [] + + def take_then_fail(results, **kwargs): + calls.append(len(results)) + raise RuntimeError("accepted, then the connection dropped") + + with mock.patch.object(manager, "persist_parsing_results", side_effect=take_then_fail): + manager._persist_sweep(items) + + assert calls == [2], "the failed group must not be resent a file at a time" + for item in items: + assert manager._file_stats[item.file].run_count == 1, "each file is still counted as run" + + def test_a_file_a_dag_is_filed_under_is_not_written_alongside_it(self): + """ + Serializing a Dag can fail and record an error against the file it is filed under. + + That happens inside the write, after a group has merged, so a file reported as parsed + cleanly in the same group would have the error re-added on top of its own clear. + """ + manager = DagFileProcessorManager(max_runs=1) + + groups = manager.build_persistence_groups( + [self._item("a.py", ["dag_a"], filed_under="b.py"), self._item("b.py", ["dag_b"])] + ) + + assert [[str(item.file.rel_path) for item in group] for group in groups] == [ + ["a.py"], + ["b.py"], + ], "b.py must be written after the Dag that could blame it" + + def test_a_dag_is_not_written_alongside_an_error_against_the_file_it_is_filed_under(self): + """ + An error is applied after the Dags filed under that file are written. + + Grouping the two would set is_stale on a Dag the same sweep had just written as healthy, + hiding it from the scheduler until the next parse. + """ + manager = DagFileProcessorManager(max_runs=1) + blamer = FileParseResult( + file=DagFileInfo(bundle_name="testing", rel_path=Path("a.py"), bundle_path=TEST_DAGS_FOLDER), + parsing_result=DagFileParsingResult( + fileloc="a.py", serialized_dags=[], import_errors={"a.py": "boom"} + ), + run_duration=1.0, + stat=DagFileStat(), + ) + + groups = manager.build_persistence_groups([blamer, self._item("b.py", ["dag_b"], filed_under="a.py")]) + + assert [[str(item.file.rel_path) for item in group] for group in groups] == [ + ["a.py"], + ["b.py"], + ], "the Dag filed under a.py must be written after the error against it" + + def test_an_error_raised_and_resolved_inside_one_sweep_is_never_written(self): + """A group records where its files ended up, not how they got there. Deliberate.""" + manager = DagFileProcessorManager(max_runs=1) + blamer = FileParseResult( + file=DagFileInfo(bundle_name="testing", rel_path=Path("blamer.py"), bundle_path=TEST_DAGS_FOLDER), + parsing_result=DagFileParsingResult( + fileloc="blamer.py", serialized_dags=[], import_errors={"blamed.py": "transient"} + ), + run_duration=1.0, + stat=DagFileStat(), + ) + + with mock.patch( + "airflow.dag_processing.manager.update_dag_parsing_results_in_db", autospec=True + ) as write: + manager._persist_bundle_group( + "testing", [blamer, self._item("blamed.py", ["blamed_dag"])], session=mock.MagicMock() + ) + + assert write.call_args.kwargs["import_errors"] == {}, ( + "nothing is written for it, so no listener hears of an error the sweep resolved" + ) + + def test_a_batch_override_is_given_the_bundle_version_the_file_was_collected_under(self): + """The payload has to carry the version context, since DagFileInfo does not.""" + seen: list[tuple[str | None, dict | None]] = [] + + class BatchApiManager(DagFileProcessorManager): + def persist_parsing_results(self, results, *, session=None): + seen.extend((item.bundle_version, item.version_data) for item in results) + + manager = BatchApiManager(max_runs=1) + manager._bundle_versions["testing"] = "v-collected" + manager._bundle_version_data["testing"] = {"sha": "abc"} + self._ready_processor(manager, "a.py", num_dags=1) + + manager._collect_results() + + assert seen == [("v-collected", {"sha": "abc"})] + + def test_a_bundle_refreshing_mid_sweep_does_not_relabel_what_was_already_collected(self): + """The version travels on the result precisely so a refresh cannot rewrite history.""" + manager = DagFileProcessorManager(max_runs=1) + manager._bundle_versions["testing"] = "v-old" + manager._bundle_version_data["testing"] = {"sha": "old"} + file = self._ready_processor(manager, "a.py", num_dags=1) + collected = manager._build_parse_result(file, manager._processors[file]) + + # The bundle moves on before the sweep is written. + manager._bundle_versions["testing"] = "v-new" + manager._bundle_version_data["testing"] = {"sha": "new"} + + with mock.patch( + "airflow.dag_processing.manager.update_dag_parsing_results_in_db", autospec=True + ) as write: + manager._persist_bundle_group("testing", [collected], session=mock.MagicMock()) + + assert write.call_args.kwargs["bundle_version"] == "v-old" + assert write.call_args.kwargs["version_data"] == {"sha": "old"} + + def test_a_per_file_override_is_not_sent_a_file_twice_when_a_later_one_fails(self): + """It gets single-file groups so the retry cannot hand it something it already accepted.""" + manager = DagFileProcessorManager(max_runs=1) + manager._bundle_versions["testing"] = "v1" + handed: list[str] = [] + + def take_then_fail_on_the_second(*, relative_fileloc, **kwargs): + handed.append(relative_fileloc) + if relative_fileloc == "file_1.py": + raise RuntimeError("accepted, then failed") + + manager.persist_parsing_result = take_then_fail_on_the_second # type: ignore[method-assign] + manager._persist_sweep([self._item(f"file_{i}.py", [f"dag_{i}"]) for i in range(2)]) + + assert handed == ["file_0.py", "file_1.py"], "no file may be handed over twice" def test_a_file_that_cannot_be_handled_does_not_discard_its_neighbours(self): """ @@ -1879,7 +2101,7 @@ def fail_for_bad(file, proc): raise RuntimeError("team lookup failed") return build(file, proc) - with mock.patch.object(manager, "_build_parse_result", side_effect=fail_for_bad): + with mock.patch.object(manager, "_build_parse_result", autospec=True, side_effect=fail_for_bad): with mock.patch.object(manager, "_persist_sweep", autospec=True) as sweep: manager._collect_results() @@ -1935,7 +2157,7 @@ def test_collect_results_leaves_unfinished_processors_alone(self): open_socket = MagicMock() manager._processors[pending]._open_sockets[open_socket] = MagicMock() - with mock.patch.object(manager, "persist_parsing_results") as mock_persist: + with mock.patch.object(manager, "persist_parsing_results", autospec=True) as mock_persist: manager._collect_results() assert [item.file for item in mock_persist.call_args.args[0]] == [done] @@ -1967,12 +2189,12 @@ def test_one_unwritable_file_does_not_discard_its_neighbours(self): persisted: list[DagFileInfo] = [] - def persist_unless_bad_is_present(results, **kwargs): - if any(item.file == bad for item in results): + def persist_unless_bad_is_present(bundle_name, group, *, session): + if any(item.file == bad for item in group): raise RuntimeError("simulated write failure") - persisted.extend(item.file for item in results) + persisted.extend(item.file for item in group) - with mock.patch.object(manager, "persist_parsing_results", side_effect=persist_unless_bad_is_present): + with mock.patch.object(manager, "_persist_bundle_group", side_effect=persist_unless_bad_is_present): manager._persist_sweep(items) assert persisted == [good], "the healthy file must still be written by the per-file retry" @@ -2011,17 +2233,17 @@ def test_collect_results_processes_remaining_files_when_one_persist_fails(self, stat_a_before = manager._file_stats[file_a] stat_b_before = manager._file_stats[file_b] - # Patch the batched seam: _collect_results routes through it, and a per-file patch on the - # instance would not be seen by the class-level override check. - def fail_for_a(results, **kwargs): - if any(item.file == file_a for item in results): + # Fail inside the built-in write. Replacing the seam itself would say the write is not + # known to be transactional, and the retry is deliberately not offered in that case. + def fail_for_a(bundle_name, group, *, session): + if any(item.file == file_a for item in group): raise RuntimeError("boom") - with mock.patch.object(manager, "persist_parsing_results", side_effect=fail_for_a) as mock_persist: + with mock.patch.object(manager, "_persist_bundle_group", side_effect=fail_for_a) as mock_persist: manager._collect_results() # The grouped attempt, then one retry per file. - assert [len(call.args[0]) for call in mock_persist.call_args_list] == [2, 1, 1] + assert [len(call.args[1]) for call in mock_persist.call_args_list] == [2, 1, 1] assert manager._file_stats[file_a] is not stat_a_before assert manager._file_stats[file_a].num_dags == stat_a_before.num_dags @@ -2046,7 +2268,7 @@ def test_collect_results_tolerates_stale_file_handle_on_close(self): manager._processors = {file: proc} with mock.patch("airflow.dag_processing.manager.update_dag_parsing_results_in_db") as db_write: - with mock.patch.object(manager, "persist_parsing_result") as persist: + with mock.patch.object(manager, "persist_parsing_result", autospec=True) as persist: manager._collect_results() assert len(manager._processors) == 0 diff --git a/airflow-core/tests/unit/dag_processing/test_parse_result_query_budget.py b/airflow-core/tests/unit/dag_processing/test_parse_result_query_budget.py index 97aba73117e82..3b60355ed183b 100644 --- a/airflow-core/tests/unit/dag_processing/test_parse_result_query_budget.py +++ b/airflow-core/tests/unit/dag_processing/test_parse_result_query_budget.py @@ -63,6 +63,7 @@ from tests_common.test_utils.config import conf_vars from tests_common.test_utils.db import ( + clear_db_dag_bundles, clear_db_dags, clear_db_import_errors, clear_db_serialized_dags, @@ -123,6 +124,8 @@ def clean_db(): clear_db_serialized_dags() clear_db_import_errors() clear_db_dags() + # Tests here add bundles of their own, and two adding the same one collide without this. + clear_db_dag_bundles() @pytest.fixture @@ -291,7 +294,12 @@ def test_sweep_pays_fixed_cost_once_per_call( def _parse_result( - tmp_path: Path, dag_id: str, run_duration: float = 0.5, bundle_name: str = BUNDLE + tmp_path: Path, + dag_id: str, + run_duration: float = 0.5, + bundle_name: str = BUNDLE, + bundle_version: str | None = None, + version_data: dict | None = None, ) -> FileParseResult: rel_path = f"{dag_id}.py" dag_file = tmp_path / rel_path @@ -302,6 +310,8 @@ def _parse_result( ), run_duration=run_duration, stat=DagFileStat(), + bundle_version=bundle_version, + version_data=version_data, ) @@ -370,6 +380,42 @@ def persist_parsing_results(self, results, *, session=None): db_write.assert_not_called() +def test_a_group_that_rolls_back_keeps_what_an_earlier_group_committed(session, testing_dag_bundle, tmp_path): + """ + Each group is its own transaction, which is the whole reason a sweep is split into them. + + Mocking the write proves the split; only a real one proves the earlier group survived the + later rollback rather than sharing its fate. + """ + session.add(DagBundleModel(name=OTHER_BUNDLE)) + session.commit() + + manager = DagFileProcessorManager(max_runs=1) + manager._bundle_versions.update({BUNDLE: None, OTHER_BUNDLE: None}) + real_write = update_dag_parsing_results_in_db + + def fail_for_the_second_bundle(*args, **kwargs): + if kwargs["bundle_name"] == OTHER_BUNDLE: + raise OperationalError("simulated", {}, Exception("write failed")) + return real_write(*args, **kwargs) + + with mock.patch( + "airflow.dag_processing.manager.update_dag_parsing_results_in_db", + side_effect=fail_for_the_second_bundle, + ): + manager._persist_sweep( + [ + _parse_result(tmp_path, "committed_dag"), + _parse_result(tmp_path, "rolled_back_dag", bundle_name=OTHER_BUNDLE), + ] + ) + + assert session.get(DagModel, "committed_dag") is not None, ( + "the group that succeeded must not be undone by the one that followed it" + ) + assert session.get(DagModel, "rolled_back_dag") is None + + def test_a_batch_override_is_called_once_per_group_not_once_per_sweep(tmp_path): """A sweep spanning bundles cannot be one write, so the seam has to be documented per group.""" seen: list[list[str]] = [] @@ -389,8 +435,8 @@ def persist_parsing_results(self, results, *, session=None): ] ) - assert seen == [["in_testing_a.py", "in_testing_b.py"], ["in_other.py"]], ( - "each bundle is its own group, so the override sees one call per group" + assert seen == [["in_testing_a.py"], ["in_other.py"], ["in_testing_b.py"]], ( + "groups are contiguous runs, so an interleaved sweep is written in the order it arrived" ) @@ -500,8 +546,14 @@ def test_a_sweep_writes_each_bundles_files_under_its_own_version(session, testin manager.persist_parsing_results( [ - _parse_result(tmp_path, "in_testing"), - _parse_result(tmp_path, "in_other", bundle_name=OTHER_BUNDLE), + _parse_result(tmp_path, "in_testing", bundle_version="v-testing", version_data={"sha": "aaa"}), + _parse_result( + tmp_path, + "in_other", + bundle_name=OTHER_BUNDLE, + bundle_version="v-other", + version_data={"sha": "bbb"}, + ), ], session=session, ) From dfa12f265c9316daf6cc51617b0aa774bb00e74d Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Wed, 19 Aug 2026 13:36:27 +0100 Subject: [PATCH 10/11] Compare a sweep's assets against writing it a file at a time Assets, aliases and the references to them are shared across files and bundles by design, and the last write decides what they look like. That made them the one kind of shared state the comparison did not cover, and the ordering a sweep is written in exists mostly for their sake. Closing a finished processor could raise where a hook had already dropped it, so the cleanup became the failure that hid whatever had brought it there. A callback-only run leaves its timestamps alone, because the sweep that deactivates unparsed Dags reads them. Failing to handle one stamped a parse that never happened. The cap does not bound what a transaction holds: a file is never split, so one defining more Dags than the cap is written in one go as it always was. What it bounds is how much grouping adds to that. Said so where the number is, and said once, at the branch that decides it, why a write someone replaced is not sent its files a second time. --- .../src/airflow/dag_processing/manager.py | 30 ++++-- .../test_batched_persistence_equivalence.py | 96 ++++++++++++++++++- .../tests/unit/dag_processing/test_manager.py | 74 ++++++++++++-- .../test_parse_result_query_budget.py | 11 ++- 4 files changed, 182 insertions(+), 29 deletions(-) diff --git a/airflow-core/src/airflow/dag_processing/manager.py b/airflow-core/src/airflow/dag_processing/manager.py index b7a80720be8cd..5f7564c500d79 100644 --- a/airflow-core/src/airflow/dag_processing/manager.py +++ b/airflow-core/src/airflow/dag_processing/manager.py @@ -110,10 +110,12 @@ def _make_execution_api() -> InProcessExecutionAPI: MAX_DAGS_PER_PERSISTENCE_GROUP = 32 """ -Dags at most per persistence call, bounding what one transaction holds locked. +How many Dags a group takes on before the next file starts a new one. -Counted in Dags rather than files because that is what grows: a few files generating Dags -dynamically can be thousands. +Counted in Dags rather than files because that is what a write costs and what it locks: a few +files generating Dags dynamically can be thousands. It bounds what grouping adds, not the write +itself -- a file is never split, so one file defining more than this is still written in one go, +exactly as it was before sweeps were grouped at all. """ @@ -1466,10 +1468,14 @@ def persist_parsing_results( metadata DB (e.g. AIP-92); one that overrides only the per-file :meth:`persist_parsing_result` is still called once per file instead, with batching skipped. + Whatever is handed over is written on the one session, so a caller passing results that do + not form a single group gets them in a single transaction. The manager never does: it splits + a sweep with :meth:`build_persistence_groups` and calls this once per group, which is what + gives each group a transaction of its own. + Raising from the built-in write discards the whole group, and the caller then persists each file on its own, since the group was one transaction and nothing was kept. A replaced write - gets no such retry: it may have accepted the group before failing, and its files are left to - be parsed again rather than sent twice. + gets no such retry -- see :meth:`_persist_sweep`. The one transaction holds while the metadata DB is the only thing being written. Under the FAB auth manager it does not: syncing a Dag's permissions commits on this session partway @@ -1641,7 +1647,10 @@ def _collect_results(self): str(file.rel_path), file.bundle_name, ) - self._throttle_retry(file, timezone.utcnow(), time.monotonic() - proc.start_time) + if not (proc.had_callbacks and proc.parsing_result is None): + self._throttle_retry(file, timezone.utcnow(), time.monotonic() - proc.start_time) + # A callback-only run leaves the timestamps alone; failing to handle one must + # too, or it advertises a parse that never happened and its Dags go stale. continue if result is not None: to_persist.append(result) @@ -1650,9 +1659,11 @@ def _collect_results(self): self._persist_sweep(to_persist) finally: # Leaving these open leaks their sockets and keeps them queued as if still running. + # A hook that ran above may have dropped one already, and this must not become the + # failure that hides whatever brought us here. for file in finished: - processor = self._processors.pop(file) - processor.close() + if (processor := self._processors.pop(file, None)) is not None: + processor.close() def _persist_sweep(self, to_persist: list[FileParseResult]) -> None: """ @@ -1664,8 +1675,7 @@ def _persist_sweep(self, to_persist: list[FileParseResult]) -> None: manager a group is not quite one transaction -- see :meth:`persist_parsing_results`. A subclass handling files one at a time gets single-file groups, so the fallback cannot hand - it a file it has already accepted. One that replaced the batch write gets no fallback at - all, since only the built-in write is known to keep nothing when it raises. + it a file it has already accepted. """ if self._overrides_per_file_persist(): groups = [[item] for item in to_persist] diff --git a/airflow-core/tests/unit/dag_processing/test_batched_persistence_equivalence.py b/airflow-core/tests/unit/dag_processing/test_batched_persistence_equivalence.py index 07a2cb87105f2..4841969740090 100644 --- a/airflow-core/tests/unit/dag_processing/test_batched_persistence_equivalence.py +++ b/airflow-core/tests/unit/dag_processing/test_batched_persistence_equivalence.py @@ -35,13 +35,20 @@ from airflow.dag_processing.manager import DagFileInfo, DagFileProcessorManager, DagFileStat, FileParseResult from airflow.dag_processing.processor import DagFileParsingResult +from airflow.models.asset import ( + AssetActive, + AssetAliasModel, + AssetModel, + DagScheduleAssetReference, + TaskOutletAssetReference, +) from airflow.models.dag import DagModel from airflow.models.dag_version import DagVersion from airflow.models.dagbundle import DagBundleModel from airflow.models.dagwarning import DagWarning from airflow.models.errors import ParseImportError from airflow.providers.standard.operators.empty import EmptyOperator -from airflow.sdk import DAG +from airflow.sdk import DAG, Asset from airflow.serialization.serialized_objects import LazyDeserializedDAG from tests_common.test_utils.db import ( @@ -57,12 +64,18 @@ OTHER_BUNDLE = "equiv-other" -def _dag(tmp_path: Path, dag_id: str, rel_path: str) -> LazyDeserializedDAG: +def _dag( + tmp_path: Path, + dag_id: str, + rel_path: str, + schedule_on: Asset | None = None, + outlet: Asset | None = None, +) -> LazyDeserializedDAG: # DagCode reads the source off disk; without a real file the Dag fails to serialize. (tmp_path / rel_path).parent.mkdir(parents=True, exist_ok=True) (tmp_path / rel_path).write_text("# equivalence fixture\n") - dag = DAG(dag_id=dag_id, schedule="@daily") - EmptyOperator(task_id="task", dag=dag) + dag = DAG(dag_id=dag_id, schedule=[schedule_on] if schedule_on else "@daily") + EmptyOperator(task_id="task", dag=dag, outlets=[outlet] if outlet else []) dag.fileloc = str(tmp_path / rel_path) dag.relative_fileloc = rel_path return LazyDeserializedDAG.from_dag(dag) @@ -72,6 +85,7 @@ def _file( tmp_path: Path, rel_path: str, dags: list[tuple[str, str]] | None = None, + assets: list[tuple[Asset | None, Asset | None]] | None = None, errors: dict[str, str] | None = None, warnings: list[dict] | None = None, bundle: str = BUNDLE, @@ -84,7 +98,10 @@ def _file( file=DagFileInfo(bundle_name=bundle, rel_path=Path(rel_path), bundle_path=tmp_path), parsing_result=DagFileParsingResult( fileloc=str(tmp_path / rel_path), - serialized_dags=[_dag(tmp_path, dag_id, under) for dag_id, under in (dags or [])], + serialized_dags=[ + _dag(tmp_path, dag_id, under, *(assets or [(None, None)] * len(dags or []))[index]) + for index, (dag_id, under) in enumerate(dags or []) + ], import_errors=errors, warnings=warnings, ), @@ -118,6 +135,19 @@ def rows(stmt): select(ParseImportError.bundle_name, ParseImportError.filename, ParseImportError.stacktrace) ), "warning": rows(select(DagWarning.dag_id, DagWarning.warning_type, DagWarning.message)), + # Assets are shared across files and bundles, so whichever file is written last decides + # what they look like. + "asset": rows(select(AssetModel.name, AssetModel.uri, AssetModel.group, AssetModel.extra)), + "asset_alias": rows(select(AssetAliasModel.name, AssetAliasModel.group)), + "asset_active": rows(select(AssetActive.name, AssetActive.uri)), + "schedule_ref": rows(select(DagScheduleAssetReference.dag_id, DagScheduleAssetReference.asset_id)), + "outlet_ref": rows( + select( + TaskOutletAssetReference.dag_id, + TaskOutletAssetReference.task_id, + TaskOutletAssetReference.asset_id, + ) + ), } @@ -184,6 +214,59 @@ def _two_bundles_interleaved(tmp_path): ] +def _one_asset_defined_differently_by_two_files(tmp_path): + """The case the whole ordering guarantee exists for: last write decides what the asset is.""" + return [ + _file( + tmp_path, + "first.py", + dags=[("asset_first", "first.py")], + assets=[(Asset(name="shared", uri="s3://shared", extra={"from": "first"}), None)], + ), + _file( + tmp_path, + "second.py", + dags=[("asset_second", "second.py")], + assets=[(Asset(name="shared", uri="s3://shared", extra={"from": "second"}), None)], + ), + ] + + +def _one_asset_across_two_bundles(tmp_path): + """ + Assets are not scoped to a bundle, so an interleaved sweep can reorder who wins. + + The asset is shared by the middle and last files deliberately. Were the last allowed to join + the first file's group to fill it, it would be written before the middle one and lose an asset + it should win. + """ + return [ + _file(tmp_path, "x.py", dags=[("asset_x", "x.py")]), + _file( + tmp_path, + "y.py", + dags=[("asset_y", "y.py")], + bundle=OTHER_BUNDLE, + assets=[(Asset(name="crossing", uri="s3://crossing", extra={"from": "y"}), None)], + ), + _file( + tmp_path, + "z.py", + dags=[("asset_z", "z.py")], + assets=[(Asset(name="crossing", uri="s3://crossing", extra={"from": "z"}), None)], + ), + ] + + +def _an_asset_produced_and_consumed_in_one_sweep(tmp_path): + """One file's outlet is another's schedule, so the rows land in whichever order they are written.""" + produced = Asset(name="handoff", uri="s3://handoff") + return [ + _file(tmp_path, "producer.py", dags=[("producer", "producer.py")], assets=[(None, produced)]), + _file(tmp_path, "consumer.py", dags=[("consumer", "consumer.py")], assets=[(produced, None)]), + ] + + def _warning_then_a_clean_parse(tmp_path): return [ _file( @@ -205,6 +288,9 @@ def _warning_then_a_clean_parse(tmp_path): _file_that_now_defines_nothing, _two_bundles_interleaved, _warning_then_a_clean_parse, + _one_asset_defined_differently_by_two_files, + _one_asset_across_two_bundles, + _an_asset_produced_and_consumed_in_one_sweep, ], ids=lambda fn: fn.__name__.strip("_"), ) diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index c046e894092c5..687183152cbd4 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py @@ -1539,7 +1539,9 @@ def test_persist_parsing_results_provides_its_own_session_when_caller_omits(self stat=DagFileStat(), ) - with mock.patch("airflow.dag_processing.manager.update_dag_parsing_results_in_db") as mock_update: + with mock.patch( + "airflow.dag_processing.manager.update_dag_parsing_results_in_db", autospec=True + ) as mock_update: manager.persist_parsing_results([item]) mock_update.assert_called_once() @@ -1558,7 +1560,9 @@ def test_a_failed_write_throttles_the_retry_without_claiming_success(self, sessi processor.had_callbacks = False processor.parsing_result = DagFileParsingResult(fileloc="abc.txt", serialized_dags=[]) - with mock.patch.object(manager, "persist_parsing_results", side_effect=RuntimeError("boom")): + with mock.patch.object( + manager, "persist_parsing_results", autospec=True, side_effect=RuntimeError("boom") + ): result = manager._build_parse_result(file, processor) manager._persist_sweep([result]) @@ -1667,7 +1671,6 @@ def test_files_sharing_a_dag_id_are_written_one_at_a_time(self): ) def test_files_with_distinct_dag_ids_are_written_together(self): - """The split only applies to conflicts; everything else still batches.""" manager = DagFileProcessorManager(max_runs=1) manager._bundle_versions["testing"] = "v1" @@ -1792,7 +1795,9 @@ def test_the_released_handler_still_persists_a_file_on_its_own(self, persist_fai processor.parsing_result = DagFileParsingResult(fileloc="a.py", serialized_dags=[]) side_effect = RuntimeError("boom") if persist_fails else None - with mock.patch.object(manager, "persist_parsing_result", side_effect=side_effect) as persist: + with mock.patch.object( + manager, "persist_parsing_result", autospec=True, side_effect=side_effect + ) as persist: manager.handle_parsing_result(file, processor) assert persist.call_args.kwargs["relative_fileloc"] == "a.py" @@ -1956,7 +1961,7 @@ def take_then_fail(results, **kwargs): calls.append(len(results)) raise RuntimeError("accepted, then the connection dropped") - with mock.patch.object(manager, "persist_parsing_results", side_effect=take_then_fail): + with mock.patch.object(manager, "persist_parsing_results", autospec=True, side_effect=take_then_fail): manager._persist_sweep(items) assert calls == [2], "the failed group must not be resent a file at a time" @@ -2081,6 +2086,48 @@ def take_then_fail_on_the_second(*, relative_fileloc, **kwargs): assert handed == ["file_0.py", "file_1.py"], "no file may be handed over twice" + def test_failing_to_handle_a_callback_only_run_leaves_its_timestamps_alone(self): + """ + A callback-only run is not a parse, so it never stamps a finish time. + + Stamping one when handling it fails would tell _scan_stale_dags the file had just been + parsed, and past the threshold its Dags are deactivated on the strength of it. + """ + manager = DagFileProcessorManager(max_runs=1) + manager._bundle_versions["testing"] = "v1" + file = self._ready_processor(manager, "cb.py") + manager._processors[file].had_callbacks = True + manager._processors[file].parsing_result = None + before = DagFileStat(num_dags=3, run_count=7) + manager._file_stats[file] = before + + with mock.patch.object( + manager, "_build_parse_result", autospec=True, side_effect=RuntimeError("team lookup") + ): + manager._collect_results() + + assert manager._file_stats[file] is before, "nothing about the file was learned" + + def test_cleanup_survives_a_hook_dropping_a_processor(self): + """The cleanup must not become the failure that hides whatever brought us here.""" + manager = DagFileProcessorManager(max_runs=1) + manager._bundle_versions["testing"] = "v1" + dropped = self._ready_processor(manager, "dropped.py", num_dags=1) + kept = self._ready_processor(manager, "kept.py", num_dags=1) + closed = mock.patch.object(manager._processors[kept], "close") + + def drop_one_and_fail(_): + manager._processors.pop(dropped) + raise RuntimeError("the sweep failed") + + with closed as close: + with mock.patch.object(manager, "_persist_sweep", autospec=True, side_effect=drop_one_and_fail): + with pytest.raises(RuntimeError, match="the sweep failed"): + manager._collect_results() + + close.assert_called_once() + assert manager._processors == {} + def test_a_file_that_cannot_be_handled_does_not_discard_its_neighbours(self): """ Working out what a file leaves to persist can reach the DB, for the team its bundle is in. @@ -2120,7 +2167,9 @@ def test_finished_processors_are_closed_even_when_persistence_raises(self): processor = manager._processors[file] with mock.patch.object(processor, "close") as close: - with mock.patch.object(manager, "_persist_sweep", side_effect=RuntimeError("boom")): + with mock.patch.object( + manager, "_persist_sweep", autospec=True, side_effect=RuntimeError("boom") + ): with pytest.raises(RuntimeError, match="boom"): manager._collect_results() @@ -2146,7 +2195,6 @@ def test_groups_of_different_bundles_keep_their_place_in_the_sweep(self): ], "collecting a bundle's files together must not let them overtake another bundle's" def test_collect_results_leaves_unfinished_processors_alone(self): - """A processor that has not finished must not be persisted or dropped.""" manager = DagFileProcessorManager(max_runs=1) manager._bundle_versions["testing"] = "v1" @@ -2194,7 +2242,9 @@ def persist_unless_bad_is_present(bundle_name, group, *, session): raise RuntimeError("simulated write failure") persisted.extend(item.file for item in group) - with mock.patch.object(manager, "_persist_bundle_group", side_effect=persist_unless_bad_is_present): + with mock.patch.object( + manager, "_persist_bundle_group", autospec=True, side_effect=persist_unless_bad_is_present + ): manager._persist_sweep(items) assert persisted == [good], "the healthy file must still be written by the per-file retry" @@ -2239,7 +2289,9 @@ def fail_for_a(bundle_name, group, *, session): if any(item.file == file_a for item in group): raise RuntimeError("boom") - with mock.patch.object(manager, "_persist_bundle_group", side_effect=fail_for_a) as mock_persist: + with mock.patch.object( + manager, "_persist_bundle_group", autospec=True, side_effect=fail_for_a + ) as mock_persist: manager._collect_results() # The grouped attempt, then one retry per file. @@ -2267,7 +2319,9 @@ def test_collect_results_tolerates_stale_file_handle_on_close(self): proc.logger_filehandle.close.side_effect = OSError(116, "Stale file handle") manager._processors = {file: proc} - with mock.patch("airflow.dag_processing.manager.update_dag_parsing_results_in_db") as db_write: + with mock.patch( + "airflow.dag_processing.manager.update_dag_parsing_results_in_db", autospec=True + ) as db_write: with mock.patch.object(manager, "persist_parsing_result", autospec=True) as persist: manager._collect_results() diff --git a/airflow-core/tests/unit/dag_processing/test_parse_result_query_budget.py b/airflow-core/tests/unit/dag_processing/test_parse_result_query_budget.py index 3b60355ed183b..39aa13fa2f0e9 100644 --- a/airflow-core/tests/unit/dag_processing/test_parse_result_query_budget.py +++ b/airflow-core/tests/unit/dag_processing/test_parse_result_query_budget.py @@ -343,7 +343,9 @@ def persist_parsing_result(self, *, relative_fileloc, session, **kwargs): manager._bundle_versions[BUNDLE] = None batch = [_parse_result(tmp_path, "override_a"), _parse_result(tmp_path, "override_b")] - with mock.patch("airflow.dag_processing.manager.update_dag_parsing_results_in_db") as db_write: + with mock.patch( + "airflow.dag_processing.manager.update_dag_parsing_results_in_db", autospec=True + ) as db_write: manager.persist_parsing_results(batch) assert calls == ["override_a.py", "override_b.py"] @@ -358,7 +360,9 @@ def _collect_a_two_file_sweep(manager, tmp_path: Path, sockets, name: str, sessi # The manager persists on sessions of its own, so release ours rather than contend with them. session.commit() - with mock.patch("airflow.dag_processing.manager.update_dag_parsing_results_in_db") as db_write: + with mock.patch( + "airflow.dag_processing.manager.update_dag_parsing_results_in_db", autospec=True + ) as db_write: manager._collect_results() return db_write @@ -570,7 +574,6 @@ def test_a_sweep_writes_each_bundles_files_under_its_own_version(session, testin def test_batched_sweep_records_warnings_from_every_file(session, testing_dag_bundle, tmp_path): - """Warnings are merged across the sweep, so each file's must survive the merge.""" manager = DagFileProcessorManager(max_runs=1) manager._bundle_versions[BUNDLE] = None @@ -622,7 +625,7 @@ def persist(results, **kwargs): if bundle == "bundle_b": raise OperationalError("simulated contention", None, Exception()) - with mock.patch.object(manager, "persist_parsing_results", side_effect=persist): + with mock.patch.object(manager, "persist_parsing_results", autospec=True, side_effect=persist): manager._persist_sweep([good, bad]) assert calls == ["bundle_a", "bundle_b"], "each bundle must be persisted by its own call" From 5b6a902a36b30022b055c8bb9a35e9675d9f1a06 Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Wed, 19 Aug 2026 13:48:14 +0100 Subject: [PATCH 11/11] Keep apart two files that file their Dags under one path Failing to serialize a Dag records an error against the file it is filed under, and that error stales every Dag filed under it. Two files whose Dags share a path could still be written together, so one file's failure staled the Dag its neighbour had just written clean -- where writing them apart leaves the later one active. Files were already held apart where one spoke for the other's path; this is the case where both speak for a third. A subclass that wraps the handler and sends its results elsewhere was reached only through the seam it replaced, so its results went to the metadata DB it had arranged to skip. What proved otherwise asserted only that the database had been written, which the routing it replaced satisfied just as well. --- .../src/airflow/dag_processing/manager.py | 7 +-- .../tests/unit/dag_processing/test_manager.py | 49 +++++++++++++++++++ 2 files changed, 53 insertions(+), 3 deletions(-) diff --git a/airflow-core/src/airflow/dag_processing/manager.py b/airflow-core/src/airflow/dag_processing/manager.py index 5f7564c500d79..8a26ddbf1cfd3 100644 --- a/airflow-core/src/airflow/dag_processing/manager.py +++ b/airflow-core/src/airflow/dag_processing/manager.py @@ -1517,9 +1517,9 @@ def build_persistence_groups(self, results: Sequence[FileParseResult]) -> list[l A run ends where the next file cannot join it: a different bundle, since the bundle fixes the version the write needs; a dag_id the run already claims, since merging two files defining one dag_id would lose the duplicate warning that comes from comparing an incoming - Dag against the file already recorded; a file the run and the incoming file speak - for in opposite directions, since an error recorded against a file is applied after the - Dags filed under it are written and would leave them stale; or + Dag against the file already recorded; a file the run and the incoming file both speak + for, whichever way round, since an error recorded against a file is applied after the Dags + filed under it are written and would leave every one of them stale; or ``MAX_DAGS_PER_PERSISTENCE_GROUP``. """ groups: list[list[FileParseResult]] = [] @@ -1545,6 +1545,7 @@ def build_persistence_groups(self, results: Sequence[FileParseResult]) -> list[l and claimed_dag_ids[-1].isdisjoint(dag_ids) and dag_locs_claimed[-1].isdisjoint(file_locs) and file_locs_claimed[-1].isdisjoint(dag_locs) + and dag_locs_claimed[-1].isdisjoint(dag_locs) ) if not joins_run: # A new group takes the file whatever its size; a file cannot be split. diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index 687183152cbd4..61f06b4ff1aec 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py @@ -1858,6 +1858,34 @@ def test_the_startup_check_runs_before_the_parsing_loop(self): warn.assert_called_once() + def test_an_override_delegating_to_super_reaches_the_replacement_seam(self): + """ + Asserting only that the database was written cannot tell this routing from the old. + + A deployment that wraps the handler and has adopted the batch seam must reach that seam, + not the per-file one it replaced, or its results go to the metadata DB it arranged to skip. + """ + seen: list[str] = [] + + class WrappingApiManager(DagFileProcessorManager): + def handle_parsing_result(self, file, proc, *, session=None): + super().handle_parsing_result(file, proc) + + def persist_parsing_results(self, results, *, session=None): + seen.extend(str(item.file.rel_path) for item in results) + + manager = WrappingApiManager(max_runs=1) + manager._bundle_versions["testing"] = "v1" + self._ready_processor(manager, "wrapped.py", num_dags=1) + + with mock.patch( + "airflow.dag_processing.manager.update_dag_parsing_results_in_db", autospec=True + ) as write: + manager._collect_results() + + assert seen == ["wrapped.py"] + write.assert_not_called() + @pytest.mark.parametrize( ("overrides", "warns"), [ @@ -1986,6 +2014,27 @@ def test_a_file_a_dag_is_filed_under_is_not_written_alongside_it(self): ["b.py"], ], "b.py must be written after the Dag that could blame it" + def test_two_files_filing_dags_under_one_path_are_written_apart(self): + """ + Failing to serialize either one records an error against the path they share. + + That error stales every Dag filed under it, so a Dag written clean beside it would be + staled by its neighbour's failure — where writing them apart leaves the later one active. + """ + manager = DagFileProcessorManager(max_runs=1) + + groups = manager.build_persistence_groups( + [ + self._item("file_a.py", ["dag_a"], filed_under="shared.py"), + self._item("file_c.py", ["dag_c"], filed_under="shared.py"), + ] + ) + + assert [[str(item.file.rel_path) for item in group] for group in groups] == [ + ["file_a.py"], + ["file_c.py"], + ] + def test_a_dag_is_not_written_alongside_an_error_against_the_file_it_is_filed_under(self): """ An error is applied after the Dags filed under that file are written.