diff --git a/airflow-core/newsfragments/71771.misc.rst b/airflow-core/newsfragments/71771.misc.rst new file mode 100644 index 0000000000000..6956495512533 --- /dev/null +++ b/airflow-core/newsfragments/71771.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..da01e6553612b 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 @@ -94,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: @@ -106,6 +108,15 @@ def _make_execution_api() -> InProcessExecutionAPI: return InProcessExecutionAPI() +MAX_DAGS_PER_PERSISTENCE_GROUP = 32 +""" +Dags at most per persistence call, bounding what one transaction holds locked. + +Counted in Dags rather than files because that is what grows: a few files generating Dags +dynamically can be thousands. +""" + + class DagParsingStat(NamedTuple): """Information on processing progress.""" @@ -158,6 +169,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 +393,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 +1274,12 @@ 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. + 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,37 @@ 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. + + 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: self.log.debug("Detected callback-only processing for %s", file) @@ -1285,37 +1340,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 +1362,13 @@ 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. 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: import_errors = { @@ -1342,9 +1382,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, @@ -1353,23 +1393,273 @@ 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, + ) + + def _overrides(self, name: str) -> bool: + """ + Report whether this manager has its own version of one of the hooks. + + Reads the instance as well as the class, so replacing a hook on one manager counts. + """ + 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 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 + 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 a sweep at a time. Override persist_parsing_results instead, " + "which is handed every file that finished together.", + DeprecationWarning, + stacklevel=2, + ) + + def _overrides_handle_parsing_result(self) -> bool: + """ + Report whether the per-file result handler has been replaced. + + 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") + + @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 is paid + once per call, so a whole sweep together costs far fewer statements than file by file. + + 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. + + 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: + 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, 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] = [] + # Last group to claim each dag_id: 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 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 + + 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[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] = {} + dag_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] + 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), + dags=dags, + import_errors=import_errors, + parse_duration=parse_durations, + warnings=dag_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] = [] + # 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(): + 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: + # 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() + + 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 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 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] + 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 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; " + "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 +2058,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..c7ac8e7146e32 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_DAGS_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): + """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() 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,372 @@ 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_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}_{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] == [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"), + [ + 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): + + 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 +1980,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 @@ -1653,10 +2014,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 new file mode 100644 index 0000000000000..ba031205c5186 --- /dev/null +++ b/airflow-core/tests/unit/dag_processing/test_parse_result_query_budget.py @@ -0,0 +1,560 @@ +# 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 when it persists parse results, 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 derives the call count from the Dag budget a group carries and measures +both prices, so it +carries no dialect-specific number and runs anywhere. +""" + +from __future__ import annotations + +import math +import re +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, 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 ( + MAX_DAGS_PER_PERSISTENCE_GROUP, + DagFileInfo, + DagFileProcessorManager, + DagFileStat, + FileParseResult, +) +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 +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.db import ( + clear_db_dags, + clear_db_import_errors, + clear_db_serialized_dags, +) + +pytestmark = pytest.mark.db_test + +BUNDLE = "testing" +OTHER_BUNDLE = "testing-other" +DAG_FILE = "budget_dags.py" + +# Per persistence call, and per Dag written by it. +FIXED_PER_CALL = 10 +PER_DAG = 3 + +SWEEP_FILES = 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_import_errors() + 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_CALL + PER_DAG, id="one-dag"), + pytest.param(5, FIXED_PER_CALL + 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_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() + sockets += [read_end, write_end] + process = MagicMock(spec=ProcessTracker) + process.wait.return_value = 0 + processor = DagFileProcessorProcess( + process_log=MagicMock(), + id=uuid7(), + pid=1234, + process=process, + stdin=write_end, + logger_filehandle=BytesIO(), + client=MagicMock(spec=Client), + 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()) + + +@pytest.mark.parametrize( + ("n_files", "dags_per_file"), + [ + 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, 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 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") + 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, n_files, sockets, f"sweep_{n_files}x{dags_per_file}", dags_per_file=dags_per_file + ) + + 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 of {total_dags} Dags costs {sweep} statements, expected {expected} " + f"({calls} x {fixed} fixed + {total_dags} x {per_dag} per Dag)." + ) + + +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_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) + ), + 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 _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([str(item.file.rel_path) for item in results]) + + manager = BatchApiManager(max_runs=1) + manager._bundle_versions[BUNDLE] = 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" + db_write.assert_not_called() + + +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( + tmp_path: Path, + dag_id: str, + *, + 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 + 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] if dag_ids is None else dag_ids, 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_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) + 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