Publish Dag processor results through the Execution API A processor without metadata database access needs to retain completed imports until the server accepts them, while continuing to supervise its other parsing processes. Bundle setup and parse-time credentials also need the bundle's authorization context throughout this path.
diff --git a/airflow-core/src/airflow/dag_processing/api_client.py b/airflow-core/src/airflow/dag_processing/api_client.py index 673c135..964b0ae 100644 --- a/airflow-core/src/airflow/dag_processing/api_client.py +++ b/airflow-core/src/airflow/dag_processing/api_client.py
@@ -26,12 +26,20 @@ from pathlib import Path from time import monotonic from typing import Any +from urllib.parse import quote from uuid import UUID import httpx2 import structlog from uuid6 import uuid7 +from airflow.api_fastapi.execution_api.datamodels.dag_parsing import ( + DagBundleInventoryBody, + DagBundleInventoryResponse, + DagBundleStateResponse, + DagParseResultBody, + DagParseResultResponse, +) from airflow.api_fastapi.execution_api.datamodels.job import ( DagParseTokenBody, DagParseTokenResponse, @@ -168,7 +176,7 @@ finally: self._parse_context.reset(selected) - def _get_parse_token(self, context: DagParseContext, *, retry: bool) -> str: + def _get_parse_token(self, context: DagParseContext, *, retry: bool, timeout: httpx2.Timeout) -> str: if context.token is not None and monotonic() < context.expires_at: if monotonic() < context.renew_at: return context.token @@ -186,12 +194,12 @@ ) if monotonic() < context.expires_at: return context.token - return self._exchange_parse_token(context, retry=retry) + return self._exchange_parse_token(context, retry=retry, timeout=timeout) def _exchange_parse_token( self, context: DagParseContext, *, retry: bool, timeout: httpx2.Timeout | None = None ) -> str: - self._ensure_job_token(retry=retry) + self._ensure_job_token(retry=retry, timeout=timeout) started_at = monotonic() response = super().request( "POST", @@ -304,10 +312,18 @@ } ) - def _ensure_job_token(self, *, retry: bool) -> None: + def _get_write_timeout(self) -> httpx2.Timeout: + """Fail fast while the API is unreachable, but let a slow write commit and acknowledge.""" + bounded = self._get_bounded_timeout() + return httpx2.Timeout( + connect=bounded.connect, read=self.timeout.read, write=self.timeout.write, pool=bounded.pool + ) + + def _ensure_job_token(self, *, retry: bool, timeout: httpx2.Timeout | None = None) -> None: self._require_job_id() + timeout = timeout or (None if retry else self._get_bounded_timeout()) if monotonic() >= self._expires_at: - self._register_job(retry=retry, timeout=None if retry else self._get_bounded_timeout()) + self._register_job(retry=retry, timeout=timeout) return if self.restart_required or monotonic() < self._retry_renewal_at: return @@ -324,7 +340,7 @@ restart_required=self.restart_required, ) if monotonic() >= self._expires_at: - self._register_job(retry=retry, timeout=None if retry else self._get_bounded_timeout()) + self._register_job(retry=retry, timeout=timeout) def request(self, *args, retry: bool = True, **kwargs) -> httpx2.Response: """Use a parsing credential for subprocess requests, and the Job credential for manager work.""" @@ -332,7 +348,11 @@ try: headers = httpx2.Headers(kwargs.get("headers")) if context := self._parse_context.get(): - kwargs["auth"] = BearerAuth(self._get_parse_token(context, retry=retry)) + # Avoid retry sleeps on the manager loop; required credentials use the runtime request's timeout. + retry = False + timeout = kwargs.get("timeout", httpx2.USE_CLIENT_DEFAULT) + timeout = self.timeout if timeout is httpx2.USE_CLIENT_DEFAULT else httpx2.Timeout(timeout) + kwargs["auth"] = BearerAuth(self._get_parse_token(context, retry=retry, timeout=timeout)) headers["Airflow-Dag-Bundle"] = context.request.bundle_name else: self._ensure_job_token(retry=retry) @@ -356,6 +376,39 @@ response = self.request("POST", f"jobs/{job_id}/heartbeat", retry=False, headers=_JOB_API_HEADERS) return JobHeartbeatResponse.model_validate_json(response.content).state + def publish_parse_result(self, body: DagParseResultBody) -> DagParseResultResponse: + """Publish once; the manager retains the request and schedules acknowledgment recovery.""" + response = self.request( + "POST", + f"jobs/{self._require_job_id()}/parse-results", + content=body.model_dump_json(), + headers=_JOB_API_HEADERS, + retry=False, + timeout=self._get_write_timeout(), + ) + return DagParseResultResponse.model_validate_json(response.content) + + def get_bundles(self) -> list[DagBundleStateResponse]: + response = self.request( + "GET", + f"jobs/{self._require_job_id()}/bundles", + headers=_JOB_API_HEADERS, + retry=False, + timeout=self._get_bounded_timeout(), + ) + return [DagBundleStateResponse.model_validate(item) for item in response.json()] + + def publish_inventory(self, bundle_name: str, body: DagBundleInventoryBody) -> DagBundleInventoryResponse: + response = self.request( + "POST", + f"jobs/{self._require_job_id()}/bundles/{quote(bundle_name, safe='')}/inventory", + content=body.model_dump_json(), + headers=_JOB_API_HEADERS, + retry=False, + timeout=self._get_write_timeout(), + ) + return DagBundleInventoryResponse.model_validate_json(response.content) + def complete_job(self, state: TerminalJobState) -> None: """Complete this Job, retaining its identity and outcome across lost acknowledgments.""" body = JobCompleteBody(state=state)
diff --git a/airflow-core/src/airflow/dag_processing/bundles/manager.py b/airflow-core/src/airflow/dag_processing/bundles/manager.py index fe74fff..6aa50af 100644 --- a/airflow-core/src/airflow/dag_processing/bundles/manager.py +++ b/airflow-core/src/airflow/dag_processing/bundles/manager.py
@@ -22,6 +22,7 @@ import os import warnings from collections import defaultdict +from contextlib import nullcontext from pathlib import Path from typing import TYPE_CHECKING, Any, NamedTuple, cast @@ -40,7 +41,8 @@ from airflow.utils.session import NEW_SESSION, create_session, provide_session if TYPE_CHECKING: - from collections.abc import Iterable + from collections.abc import Callable, Collection, Iterable + from contextlib import AbstractContextManager from sqlalchemy.engine import CursorResult from sqlalchemy.orm import Session @@ -312,8 +314,16 @@ class DagBundlesManager(LoggingMixin): """Manager for DAG bundles.""" - def __init__(self, *args, **kwargs): + def __init__( + self, + *args, + bundle_names: Collection[str] | None = None, + bundle_context: Callable[[str], AbstractContextManager] | None = None, + **kwargs, + ): super().__init__(*args, **kwargs) + self._bundle_names = bundle_names + self._bundle_context = bundle_context or (lambda _: nullcontext()) self._bundle_config: dict[str, _InternalBundleConfig] = {} self.parse_config() @@ -333,6 +343,8 @@ return for bundle_config in bundle_config_list: + if self._bundle_names is not None and bundle_config.name not in self._bundle_names: + continue if bundle_config.team_name and not conf.getboolean("core", "multi_team"): raise AirflowConfigException( "Section `dag_processor` key `dag_bundle_config_list` " @@ -349,7 +361,13 @@ self.log.info("DAG bundles loaded: %s", ", ".join(self._bundle_config.keys())) @provide_session - def sync_bundles_to_db(self, *, deactivate_missing: bool = True, session: Session = NEW_SESSION) -> None: + def sync_bundles_to_db( + self, + *, + deactivate_missing: bool = True, + include_bundle_urls: bool = True, + session: Session = NEW_SESSION, + ) -> None: """ Persist the configured DAG bundles into ``DagBundleModel`` rows. @@ -359,6 +377,8 @@ ``DagBag`` plus ``sync_bag_to_db`` (or, in production, the DAG processor); calling this method does not trigger that work. + :param include_bundle_urls: Sync ownership without constructing bundles when ``False``, so + API authorization can be established before bundle code requests secrets. :param deactivate_missing: When ``True`` (the default), any bundle stored in the database that is not present in this manager's config is marked inactive. This is only correct when the calling process sees the @@ -400,24 +420,27 @@ if not team: raise _bundle_item_exc(f"Team '{config.team_name}' does not exist") - try: - new_template, new_params = _extract_and_sign_template(name) - except Exception as e: - self.log.exception("Error creating bundle '%s': %s", name, e) - continue + if include_bundle_urls: + try: + with self._bundle_context(name): + new_template, new_params = _extract_and_sign_template(name) + except Exception as e: + self.log.exception("Error creating bundle '%s': %s", name, e) + continue if bundle := stored.pop(name, None): bundle.active = True - if new_template != bundle.signed_url_template: + if include_bundle_urls and new_template != bundle.signed_url_template: bundle.signed_url_template = new_template self.log.debug("Updated URL template for bundle %s", name) - if new_params != bundle.template_params: + if include_bundle_urls and new_params != bundle.template_params: bundle.template_params = new_params self.log.debug("Updated template parameters for bundle %s", name) else: bundle = DagBundleModel(name=name) - bundle.signed_url_template = new_template - bundle.template_params = new_params + if include_bundle_urls: + bundle.signed_url_template = new_template + bundle.template_params = new_params session.add(bundle) self.log.info("Added new DAG bundle %s to the database", name) @@ -694,9 +717,10 @@ cfg_bundle = self._bundle_config.get(name) if not cfg_bundle: raise ValueError(f"Requested bundle '{name}' is not configured.") - return cfg_bundle.bundle_class( - name=name, version=version, version_data=version_data, **cfg_bundle.kwargs - ) + with self._bundle_context(name): + return cfg_bundle.bundle_class( + name=name, version=version, version_data=version_data, **cfg_bundle.kwargs + ) @classmethod def is_bundle_configured(cls, name: str) -> bool: @@ -716,7 +740,9 @@ """ for name, cfg in self._bundle_config.items(): try: - yield cfg.bundle_class(name=name, version=None, **cfg.kwargs) + with self._bundle_context(name): + bundle = cfg.bundle_class(name=name, version=None, **cfg.kwargs) + yield bundle except Exception as e: self.log.exception("Error creating bundle '%s': %s", name, e) # Skip this bundle and continue with others
diff --git a/airflow-core/src/airflow/dag_processing/manager.py b/airflow-core/src/airflow/dag_processing/manager.py index 19aee04..f177e2d 100644 --- a/airflow-core/src/airflow/dag_processing/manager.py +++ b/airflow-core/src/airflow/dag_processing/manager.py
@@ -29,16 +29,20 @@ import signal import sys import time -from collections import OrderedDict, defaultdict +from collections import OrderedDict, defaultdict, deque +from contextlib import AbstractContextManager, nullcontext from copy import deepcopy from dataclasses import dataclass, field from datetime import datetime, timedelta from operator import attrgetter, itemgetter from pathlib import Path from typing import TYPE_CHECKING, Any, Literal, NamedTuple, cast +from uuid import UUID import attrs +import httpx2 import structlog +from pydantic import ValidationError from sqlalchemy import or_, select, update from sqlalchemy.exc import OperationalError from sqlalchemy.orm import load_only @@ -48,6 +52,13 @@ from airflow._shared.observability.metrics import stats from airflow._shared.observability.metrics.stats import normalize_name_for_stats from airflow._shared.timezones import timezone +from airflow.api_fastapi.execution_api.datamodels.dag_parsing import ( + MAX_PUBLICATION_BYTES, + DagBundleInventoryBody, + DagParseResultBody, + ParseSourceCode, + ParseWarning, +) from airflow.api_fastapi.execution_api.in_process import InProcessExecutionAPI from airflow.callbacks.callback_requests import DagCallbackRequest from airflow.configuration import conf @@ -102,15 +113,28 @@ from sqlalchemy.sql import Select from airflow.callbacks.callback_requests import CallbackRequest + from airflow.dag_processing.api_client import DagProcessorAPIClient from airflow.dag_processing.bundles.base import BaseDagBundle from airflow.sdk.api.client import Client +# Each send has its own short timeout; this bounds how long one loop spends on queued sends +# before it goes back to heartbeating and supervising parsing processes. +_API_SEND_BUDGET_SECONDS = 1.0 +_MAX_REJECTION_REASON_LENGTH = 1024 +_REJECTED_CONTENT_REASONS = frozenset({"dag_owned_by_another_bundle", "plugin_not_granted"}) + + +def _is_rejected_content(status_code: int, reason: str | None) -> bool: + """Whether the API refused a result for what it contains, so sending it again cannot succeed.""" + return status_code in (413, 422) or (status_code in (403, 409) and reason in _REJECTED_CONTENT_REASONS) + class BundleState(NamedTuple): """Persisted refresh state for a DAG bundle.""" last_refreshed: datetime | None version: str | None + revision: UUID | None = None @attrs.define @@ -123,6 +147,20 @@ last_duration: float | None = None run_count: int = 0 last_num_of_db_queries: int = 0 + last_attempt_time: datetime | None = None + + +@attrs.define +class PendingDagPublication: + """A completed import retained until publication succeeds or exhausts its retries.""" + + body: DagParseResultBody + stat: DagFileStat + refresh_generation: int + attempts: int = 0 + next_attempt_time: float = 0.0 + is_rejection_report: bool = False + """Whether this replaces a result the API refused, so a further refusal is final.""" @dataclass(frozen=True) @@ -264,10 +302,15 @@ _bundle_parse_sources: dict[str, DagParseSource] = attrs.field(factory=dict, init=False) _bundle_refresh_generations: dict[str, int] = attrs.field(factory=lambda: defaultdict(int), init=False) _bundles_waiting_for_refresh: set[str] = attrs.field(factory=set, init=False) + _dispatch_sequence: int = attrs.field(default=0, init=False) _multi_team: bool = attrs.field(factory=lambda: conf.getboolean("core", "multi_team"), init=False) _bundle_name_to_team_name: dict[str, str | None] = attrs.field(factory=dict, init=False) _processors: dict[DagFileInfo, BaseDagFileProcessorProcess] = attrs.field(factory=dict, init=False) + _pending_publications: dict[DagFileInfo, PendingDagPublication] = attrs.field(factory=dict, init=False) + _pending_inventories: dict[str, tuple[DagBundleInventoryBody, set[DagFileInfo], DagParseSource]] = ( + attrs.field(factory=dict, init=False) + ) _parsing_start_time: float | None = attrs.field(default=None, init=False) _num_run: int = attrs.field(default=0, init=False) @@ -302,8 +345,15 @@ ) """Resolved once per process so file discovery and the deactivation scan use the same value.""" - _api_server: InProcessExecutionAPI = attrs.field(init=False, factory=InProcessExecutionAPI) - """API server to interact with Metadata DB""" + api_client: DagProcessorAPIClient | None = None + """ + Registered HTTP client for this manager's API requests, used instead of the in-process API. + + Bundle credentials and parse-time requests then go through it, on behalf of the bundle they are for. + """ + + _api_server: InProcessExecutionAPI | None = attrs.field(init=False, default=None) + """In-process API server, created only when there is no :attr:`api_client`.""" def register_exit_signals(self): """Register signals that stop child processes.""" @@ -317,9 +367,17 @@ return {} missing = [name for name in bundle_names if name not in self._bundle_name_to_team_name] if missing: - queried = DagBundleModel.get_team_names(missing) - for name in missing: - self._bundle_name_to_team_name[name] = queried.get(name) + try: + if self.api_client is not None: + queried = {bundle.name: bundle.team_name for bundle in self.api_client.get_bundles()} + else: + queried = DagBundleModel.get_team_names(missing) + except httpx2.HTTPError: + # Team names only tag metrics, so parsing continues and the lookup is retried later. + self.log.warning("Unable to look up the teams of bundles %s; retrying later", missing) + else: + for name in missing: + self._bundle_name_to_team_name[name] = queried.get(name) return {name: self._bundle_name_to_team_name.get(name) for name in bundle_names} def _get_team_name(self, bundle_name: str) -> str | None: @@ -334,12 +392,34 @@ self.log.debug("Finished terminating DAG processors.") sys.exit(os.EX_OK) - def sync_bundles(self) -> None: + def _create_bundle_manager(self) -> DagBundlesManager: + if self.api_client is None: + return DagBundlesManager() + return DagBundlesManager( + bundle_names=self.bundle_names_to_parse, + bundle_context=self.api_client.use_bundle, + ) + + def sync_bundles(self, *, include_bundle_urls: bool = True) -> None: """Sync configured DAG bundles to the metadata database.""" + if self.api_client is not None: + provisioned = {bundle.name for bundle in self.api_client.get_bundles()} + configured = set( + self.bundle_names_to_parse or self._create_bundle_manager().get_all_bundle_names() + ) + if missing := configured - provisioned: + raise ValueError( + f"Bundles must be provisioned before starting the processor: {sorted(missing)}" + ) + return # When this processor only parses a subset of bundles, it does not see the full # bundle configuration and must not deactivate bundles owned by other processors. - dag_bundle_manager = DagBundlesManager() - dag_bundle_manager.sync_bundles_to_db(deactivate_missing=not self.bundle_names_to_parse) + dag_bundle_manager = self._create_bundle_manager() + dag_bundle_manager.sync_bundles_to_db( + deactivate_missing=not self.bundle_names_to_parse, include_bundle_urls=include_bundle_urls + ) + if not include_bundle_urls: + return # Best-effort legacy repair: a failure here must not crash DFP startup. # Affected Dags self-heal on the next successful parse. try: @@ -349,7 +429,7 @@ def get_all_bundles(self) -> list[BaseDagBundle]: """Return configured DAG bundles filtered by ``bundle_names_to_parse`` if provided.""" - return list(DagBundlesManager().get_all_dag_bundles()) + return list(self._create_bundle_manager().get_all_dag_bundles()) def run(self): """ @@ -358,15 +438,18 @@ By processing them in separate processes, we can get parallelism and isolation from potentially harmful user code. """ - self.before_run() try: + self.before_run() return self._run_parsing_loop() finally: self.after_run() def before_run(self) -> None: """Set up state required before the parsing loop starts. Default implementation; override to customize.""" - self.prepare_server_process_context() + if self.api_client is None: + self.prepare_server_process_context() + else: + self.prepare_api_secrets_context(self.api_client) self.prepare_process_context() self.register_exit_signals() self.log.info("Processing files using up to %s processes at a time ", self._parallelism) @@ -378,13 +461,21 @@ gc.freeze() def after_run(self) -> None: - """Tear down state after the parsing loop exits. Default no-op; override to customize.""" + """Tear down state after the parsing loop exits. Default implementation; override to customize.""" + if self.api_client is None: + return + from airflow.dag_processing.api_client import DagProcessorSecretsComms + from airflow.sdk.execution_time import task_runner + + if isinstance(getattr(task_runner, "SUPERVISOR_COMMS", None), DagProcessorSecretsComms): + del task_runner.SUPERVISOR_COMMS def warm_importers(self) -> None: """Build each bundle's Dag importers, so parse processes forked later share them.""" for bundle in self._dag_bundles: try: - get_importer_registry(bundle.name).warm_importers() + with self._use_bundle(bundle.name): + get_importer_registry(bundle.name).warm_importers() except Exception: # The importer fails again when the bundle is listed, which reports it per refresh. self.log.exception("Error loading Dag importers for bundle %s", bundle.name) @@ -406,6 +497,28 @@ # Related: https://github.com/apache/airflow/pull/57459 os.environ["_AIRFLOW_PROCESS_CONTEXT"] = "server" + def prepare_api_secrets_context(self, api_client: DagProcessorAPIClient) -> None: + """ + Resolve this process's own connection and variable lookups through ``api_client``. + + Bundle code runs in this process, for example a Git bundle reading its connection. Its secret + lookups use the API. Parse processes replace + ``SUPERVISOR_COMMS`` with their own channel before running any Dag code. + """ + from airflow.dag_processing.api_client import DagProcessorSecretsComms + from airflow.sdk.execution_time import task_runner + + # SDK cache keys do not carry the bundle identity used to authorize these requests. + SecretCache.reset() + os.environ["_AIRFLOW_PROCESS_CONTEXT"] = "client" + task_runner.SUPERVISOR_COMMS = DagProcessorSecretsComms(api_client) # type: ignore[assignment] + + def _use_bundle(self, bundle_name: str) -> AbstractContextManager[object]: + """Make API requests on behalf of the bundle; a no-op with the in-process API.""" + if self.api_client is None: + return nullcontext() + return self.api_client.use_bundle(bundle_name) + def prepare_process_context(self) -> None: """Initialize transport-neutral process state (selector, stats) before the parsing loop starts.""" # Initialization is delayed until here to avoid fork issues in some @@ -436,6 +549,8 @@ def _scan_stale_dags(self): """Scan and deactivate DAGs which are no longer present in files.""" + if self.api_client is not None: + return now = time.monotonic() elapsed_time_since_refresh = now - self._last_deactivate_stale_dags_time if elapsed_time_since_refresh > self.parsing_cleanup_interval: @@ -463,6 +578,12 @@ def cleanup_stale_bundle_versions(self) -> None: """Clean up stale DAG bundle version usage records.""" + if self.api_client is not None: + manager = BundleUsageTrackingManager() + for bundle in self._dag_bundles: + if bundle.supports_versioning: + manager._remove_stale_bundle_versions_for_bundle(bundle.name) + return BundleUsageTrackingManager().remove_stale_bundle_versions() @provide_session @@ -558,7 +679,8 @@ def _run_parsing_loop(self): # initialize cache to mutualize calls to Variable.get in DAGs # needs to be done before this process is forked to create the DAG parsing processes. - SecretCache.init() + if self.api_client is None: + SecretCache.init() poll_time = 0.0 @@ -589,6 +711,7 @@ self._collect_results() + self._publish_pending_results() for callback in self.fetch_callbacks(): self._add_callback_to_queue(callback) self._scan_stale_dags() @@ -649,6 +772,8 @@ Default implementation reads from the metadata DB; override to source requests from an API. """ + if self.api_client is not None: + return [] return self._claim_priority_files() def request_bundle_refresh(self, bundle_names: str | Iterable[str]) -> None: @@ -711,6 +836,8 @@ Default implementation reads from the metadata DB; override to source callbacks from an API. """ + if self.api_client is not None: + return [] return self._fetch_callbacks_from_db() @provide_session @@ -785,21 +912,18 @@ self.log.error("Bundle %s is not initialized, skipping callback", request.bundle_name) return None return loaded - try: - bundle = DagBundlesManager().get_bundle( - name=request.bundle_name, - version=request.bundle_version, - version_data=request.version_data, - ) - except ValueError: - self.log.error("Bundle %s no longer configured, skipping callback", request.bundle_name) - return None - if bundle.supports_versioning: + with self._use_bundle(request.bundle_name): try: - bundle.initialize() + bundle = self._create_bundle_manager().get_bundle( + name=request.bundle_name, + version=request.bundle_version, + version_data=request.version_data, + ) + if bundle.supports_versioning: + bundle.initialize() except Exception: self.log.exception( - "Error initializing bundle %s version %s for callback, skipping", + "Error preparing bundle %s version %s for callback, skipping", request.bundle_name, request.bundle_version, ) @@ -846,6 +970,11 @@ Returns ``None`` if the bundle has no database record. """ + if self.api_client is not None: + for bundle in self.api_client.get_bundles(): + if bundle.name == bundle_name: + return BundleState(bundle.last_refreshed, bundle.version, bundle.revision) + return None row = session.execute( select(DagBundleModel.last_refreshed, DagBundleModel.version).where( DagBundleModel.name == bundle_name @@ -876,6 +1005,44 @@ values["version"] = version session.execute(update(DagBundleModel).where(DagBundleModel.name == bundle_name).values(**values)) + def _publish_bundle_inventory(self, bundle_name: str, known_files: dict[str, set[DagFileInfo]]) -> None: + if self.api_client is None: + raise ValueError("Inventory publication requires an API client") + body, found_files, source = self._pending_inventories[bundle_name] + if len(body.model_dump_json().encode()) > MAX_PUBLICATION_BYTES: + self.log.error( + "Bundle %s inventory exceeds the %s-byte publication limit; retaining the previous inventory", + bundle_name, + MAX_PUBLICATION_BYTES, + ) + self._discard_unusable_bundle_source(bundle_name) + return + try: + receipt = self.api_client.publish_inventory(bundle_name, body) + except httpx2.HTTPStatusError as error: + if error.response.status_code == 409: + del self._pending_inventories[bundle_name] + self.request_bundle_refresh(bundle_name) + elif error.response.status_code == 413: + self._discard_unusable_bundle_source(bundle_name) + elif error.response.status_code < 500: + raise + self.log.warning("Inventory publication was rejected for %s", bundle_name) + return + except httpx2.RequestError: + self.log.warning("Unable to publish inventory for %s; retaining it for retry", bundle_name) + return + del self._pending_inventories[bundle_name] + self._bundle_parse_sources[bundle_name] = attrs.evolve(source, bundle_revision=receipt.revision) + self._bundle_versions[bundle_name] = source.bundle_version + known_files[bundle_name] = found_files + + def _discard_unusable_bundle_source(self, bundle_name: str) -> None: + self._pending_inventories.pop(bundle_name, None) + self._bundle_parse_sources.pop(bundle_name, None) + self._force_refresh_bundles.discard(bundle_name) + self._bundles_waiting_for_refresh.discard(bundle_name) + def _invalidate_bundle_parse_source(self, bundle_name: str) -> None: # Initialization or refresh may change files even when the provider raises afterwards. self._bundle_refresh_generations[bundle_name] += 1 @@ -892,6 +1059,7 @@ now_seconds < next_check and not self._force_refresh_bundles and not self._bundles_waiting_for_refresh + and not self._pending_inventories ): self.log.debug( "Not time to check if DAG Bundles need refreshed yet - skipping. Next check in %.2f seconds", @@ -901,132 +1069,202 @@ self._bundles_last_refreshed = now_seconds + catalog: dict[str, BundleState] | None = None + if self.api_client is not None and self._dag_bundles: + try: + bundles = self.api_client.get_bundles() + catalog = { + bundle.name: BundleState(bundle.last_refreshed, bundle.version, bundle.revision) + for bundle in bundles + } + self._bundle_name_to_team_name = {bundle.name: bundle.team_name for bundle in bundles} + except httpx2.HTTPError: + self.log.warning("Unable to read the bundle catalog; retrying on a later refresh") + any_refreshed = False - active_bundles = {file.bundle_name for file in self._processors if file.bundle_version is None} + active_bundles = { + file.bundle_name + for file in self._processors.keys() | self._pending_publications.keys() + if file.bundle_version is None + } for bundle in self._dag_bundles: - # TODO: AIP-66 handle errors in the case of incomplete cloning? And test this. - # What if the cloning/refreshing took too long(longer than the dag processor timeout) - if not bundle.is_initialized: + with self._use_bundle(bundle.name): + if catalog is not None and bundle.name not in catalog: + self._discard_unusable_bundle_source(bundle.name) + continue + if bundle.name in self._pending_inventories: + self._publish_bundle_inventory(bundle.name, known_files) + any_refreshed = True + continue + if self.api_client is not None and catalog is None: + continue + # TODO: AIP-66 handle errors in the case of incomplete cloning? And test this. + # What if the cloning/refreshing took too long(longer than the dag processor timeout) + if not bundle.is_initialized: + self._invalidate_bundle_parse_source(bundle.name) + try: + bundle.initialize() + any_refreshed = True + except Exception as e: + self.log.exception("Error initializing bundle %s: %s", bundle.name, e) + continue + try: + bundle_state = ( + catalog.get(bundle.name) + if catalog is not None + else self.get_bundle_state(bundle.name) + ) + except httpx2.HTTPError: + self.log.warning("Unable to read metadata for bundle %s; retrying later", bundle.name) + continue + except Exception: + self.log.exception("Error fetching state for bundle %s", bundle.name) + continue + if bundle_state is None: + self.log.warning("Bundle model not found for %s", bundle.name) + continue + elapsed_time_since_refresh = ( + now - (bundle_state.last_refreshed or utc_epoch()) + ).total_seconds() + if bundle.supports_versioning: + # we will also check the version of the bundle to see if another DAG processor has seen + # a new version + pre_refresh_version = self._bundle_versions.get(bundle.name) + # Use `is None` (not falsy) so an empty-string version is treated as a valid cached value. + if pre_refresh_version is None: + pre_refresh_version, _ = unpack_bundle_version(bundle.get_current_version(), bundle) + current_version_matches_db = pre_refresh_version == bundle_state.version + else: + # With no versioning, it always "matches" + current_version_matches_db = True + + previously_seen = bundle.name in self._bundle_versions + if ( + bundle.name not in self._bundles_waiting_for_refresh + and bundle.name in self._bundle_parse_sources + and self.should_skip_refresh( + bundle=bundle, + elapsed_time_since_refresh=elapsed_time_since_refresh, + current_version_matches_db=current_version_matches_db, + previously_seen=previously_seen, + ) + ): + self.log.debug("Not time to refresh bundle %s", bundle.name) + continue + + if bundle.name in active_bundles: + # Forced refreshes wait too: refreshing under an import would discard its result. + self._bundles_waiting_for_refresh.add(bundle.name) + continue + + self._bundles_waiting_for_refresh.discard(bundle.name) + self.log.info("Refreshing bundle %s", bundle.name) + self._invalidate_bundle_parse_source(bundle.name) try: - bundle.initialize() + bundle.refresh() any_refreshed = True - except Exception as e: - self.log.exception("Error initializing bundle %s: %s", bundle.name, e) + except Exception: + self.log.exception("Error refreshing bundle %s", bundle.name) continue - try: - bundle_state = self.get_bundle_state(bundle.name) - except Exception: - self.log.exception("Error fetching state for bundle %s", bundle.name) - continue - if bundle_state is None: - self.log.warning("Bundle model not found for %s", bundle.name) - continue - elapsed_time_since_refresh = (now - (bundle_state.last_refreshed or utc_epoch())).total_seconds() - if bundle.supports_versioning: - # we will also check the version of the bundle to see if another DAG processor has seen - # a new version - pre_refresh_version = self._bundle_versions.get(bundle.name) - # Use `is None` (not falsy) so an empty-string version is treated as a valid cached value. - if pre_refresh_version is None: - pre_refresh_version, _ = unpack_bundle_version(bundle.get_current_version(), bundle) - current_version_matches_db = pre_refresh_version == bundle_state.version - else: - # With no versioning, it always "matches" - current_version_matches_db = True - previously_seen = bundle.name in self._bundle_versions - if ( - bundle.name not in self._bundles_waiting_for_refresh - and bundle.name in self._bundle_parse_sources - and self.should_skip_refresh( - bundle=bundle, - elapsed_time_since_refresh=elapsed_time_since_refresh, - current_version_matches_db=current_version_matches_db, - previously_seen=previously_seen, - ) - ): - self.log.debug("Not time to refresh bundle %s", bundle.name) - continue + self._force_refresh_bundles.discard(bundle.name) - if bundle.name in active_bundles: - # Forced refreshes wait too: refreshing under an import would discard its result. - self._bundles_waiting_for_refresh.add(bundle.name) - continue - - self._bundles_waiting_for_refresh.discard(bundle.name) - self.log.info("Refreshing bundle %s", bundle.name) - - self._invalidate_bundle_parse_source(bundle.name) - try: - bundle.refresh() - any_refreshed = True - except Exception: - self.log.exception("Error refreshing bundle %s", bundle.name) - continue - - self._force_refresh_bundles.discard(bundle.name) - - if bundle.supports_versioning: - version_after_refresh, version_data_after_refresh = unpack_bundle_version( - bundle.get_current_version(), bundle - ) - else: - version_after_refresh = None - version_data_after_refresh = None - - self._bundle_parse_sources[bundle.name] = DagParseSource( - bundle_version=version_after_refresh, - version_data=deepcopy(version_data_after_refresh), - refresh_generation=self._bundle_refresh_generations[bundle.name], - ) - - if bundle.supports_versioning: - # We can short-circuit the rest of this if (1) bundle was seen before by - # this dag processor and (2) the version of the bundle did not change - # after refreshing it - if previously_seen and pre_refresh_version == version_after_refresh: - self.log.debug( - "Bundle %s version not changed after refresh: %s", - bundle.name, - version_after_refresh, + if bundle.supports_versioning: + version_after_refresh, version_data_after_refresh = unpack_bundle_version( + bundle.get_current_version(), bundle ) - try: - self.update_bundle_state(bundle.name, last_refreshed=now, version=None) - except Exception: - self.log.exception("Error persisting state for bundle %s", bundle.name) + else: + version_after_refresh = None + version_data_after_refresh = None + + self._bundle_parse_sources[bundle.name] = DagParseSource( + bundle_version=version_after_refresh, + version_data=deepcopy(version_data_after_refresh), + refresh_generation=self._bundle_refresh_generations[bundle.name], + ) + + if bundle.supports_versioning: + # We can short-circuit the rest of this if (1) bundle was seen before by + # this dag processor and (2) the version of the bundle did not change + # after refreshing it + if ( + self.api_client is None + and previously_seen + and pre_refresh_version == version_after_refresh + ): + self.log.debug( + "Bundle %s version not changed after refresh: %s", + bundle.name, + version_after_refresh, + ) + try: + self.update_bundle_state(bundle.name, last_refreshed=now, version=None) + except Exception: + self.log.exception("Error persisting state for bundle %s", bundle.name) + continue + + self.log.info( + "Version changed for %s, new version: %s", bundle.name, version_after_refresh + ) + try: + found_files = self._find_files_in_bundle(bundle) + except Exception: + # Keep the bundle's known files and Dags, and leave its version unadvanced so the + # next refresh lists it again. + self.log.exception("Error listing Dag definitions in bundle %s", bundle.name) + if self.api_client is not None: + self._bundle_parse_sources.pop(bundle.name, None) continue - self.log.info("Version changed for %s, new version: %s", bundle.name, version_after_refresh) - try: - found_files = self._find_files_in_bundle(bundle) - except Exception: - # Keep the bundle's known files and Dags, and leave its version unadvanced so the - # next refresh lists it again. - self.log.exception("Error listing Dag definitions in bundle %s", bundle.name) - continue + if self.api_client is not None: + self._dispatch_sequence += 1 + try: + inventory = DagBundleInventoryBody( + attempt_id=uuid7(), + dispatch_sequence=self._dispatch_sequence, + expected_revision=bundle_state.revision, + version=version_after_refresh, + files=sorted(self._get_observed_filelocs(found_files)), + failed_sources=sorted( + {loc for file in found_files for loc in file.failed_sources} + ), + ) + except ValidationError: + # Handled like a failed listing, so the API never sees a partial inventory. + self.log.exception("Bundle %s has a listing the API cannot accept", bundle.name) + self._discard_unusable_bundle_source(bundle.name) + continue + self._pending_inventories[bundle.name] = ( + inventory, + found_files, + self._bundle_parse_sources.pop(bundle.name), + ) + self._publish_bundle_inventory(bundle.name, known_files) + continue - # Persistence failure must not skip file scanning (bundle is already refreshed locally). - # _bundle_versions is only advanced on success to stay consistent with the DB. - try: - self.update_bundle_state(bundle.name, last_refreshed=now, version=version_after_refresh) - except Exception: - self.log.exception("Error persisting state for bundle %s", bundle.name) - else: - self._bundle_versions[bundle.name] = version_after_refresh + # Persistence failure must not skip file scanning (bundle is already refreshed locally). + # _bundle_versions is only advanced on success to stay consistent with the DB. + try: + self.update_bundle_state(bundle.name, last_refreshed=now, version=version_after_refresh) + except Exception: + self.log.exception("Error persisting state for bundle %s", bundle.name) + else: + self._bundle_versions[bundle.name] = version_after_refresh - known_files[bundle.name] = found_files + known_files[bundle.name] = found_files - self.deactivate_deleted_dags(bundle_name=bundle.name, present=found_files) - self.clear_orphaned_import_errors( - bundle_name=bundle.name, - observed_filelocs=self._get_observed_filelocs(found_files), - failed_sources={loc for file in found_files for loc in file.failed_sources}, - ) + self.deactivate_deleted_dags(bundle_name=bundle.name, present=found_files) + self.clear_orphaned_import_errors( + bundle_name=bundle.name, + observed_filelocs=self._get_observed_filelocs(found_files), + failed_sources={loc for file in found_files for loc in file.failed_sources}, + ) if any_refreshed: # Bundle-to-team assignments can only change on bundle refresh, so clear the cache. - self._bundle_name_to_team_name = {} + if self.api_client is None: + self._bundle_name_to_team_name = {} self.handle_removed_files(known_files=known_files) self._resort_file_queue() self._add_new_files_to_queue(known_files=known_files) @@ -1311,6 +1549,9 @@ stats_to_remove = {file for file in self._file_stats if file.presence_key not in present_keys} for file in stats_to_remove: del self._file_stats[file] + for file in list(self._pending_publications): + if file.presence_key not in present_keys: + del self._pending_publications[file] def terminate_orphan_processes(self, present: set[DagFileInfo]): """Stop processors that are working on deleted files.""" @@ -1399,6 +1640,19 @@ if proc.parsing_result is not None: try: + if self.api_client is not None: + pending = PendingDagPublication( + body=self._build_parse_result(file, proc, run_duration), + stat=next_stat, + refresh_generation=proc.parse_source.refresh_generation, + ) + if len(pending.body.model_dump_json().encode()) > MAX_PUBLICATION_BYTES: + pending = self._build_rejection_report( + pending, + f"it exceeds the {MAX_PUBLICATION_BYTES // (1024 * 1024)} MiB publication limit", + ) + self._pending_publications[file] = pending + return self.persist_parsing_result( bundle_name=file.bundle_name, bundle_version=proc.parse_source.bundle_version, @@ -1416,19 +1670,157 @@ 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, - ) + self._record_failed_publication(file, next_stat) return self._file_stats[file] = next_stat + def _record_failed_publication(self, file: DagFileInfo, stat: DagFileStat) -> None: + self._file_stats[file] = attrs.evolve( + self._file_stats[file], + last_attempt_time=timezone.utcnow(), + last_duration=stat.last_duration, + run_count=stat.run_count, + ) + + def _publish_pending_results(self, *, deadline: float | None = None) -> None: + """Publish queued results until the deadline, stopping early while the API is unavailable.""" + if self.api_client is None: + return + from airflow.dag_processing.api_client import get_error_reason + + if deadline is None: + deadline = time.monotonic() + _API_SEND_BUDGET_SECONDS + files = deque(self._pending_publications) + while files and time.monotonic() < deadline: + file = files.popleft() + pending = self._pending_publications[file] + if pending.refresh_generation != self._bundle_refresh_generations.get(file.bundle_name, 0): + del self._pending_publications[file] + self._add_files_to_queue([file], mode="front") + continue + if time.monotonic() < pending.next_attempt_time: + continue + pending.attempts += 1 + try: + self.api_client.publish_parse_result(pending.body) + except Exception as error: + status_error = error if isinstance(error, httpx2.HTTPStatusError) else None + reason = get_error_reason(status_error) if status_error is not None else None + if reason == "source_changed": + # Another processor's inventory moved the revision; refreshing fetches it. + self.log.info( + "Bundle %s changed before %s was published; reparsing", file.bundle_name, file + ) + del self._pending_publications[file] + self.request_bundle_refresh(file.bundle_name) + self._add_files_to_queue([file], mode="front") + return + retryable = isinstance(error, httpx2.RequestError) or ( + status_error is not None and status_error.response.status_code >= 500 + ) + if retryable and pending.attempts < conf.getint("workers", "execution_api_retries"): + pending.next_attempt_time = time.monotonic() + min( + conf.getfloat("workers", "execution_api_retry_wait_max"), + max( + conf.getfloat("workers", "execution_api_retry_wait_min"), + 2 ** (pending.attempts - 1), + ), + ) + self.log.warning("Unable to publish %s; retrying on a later loop", file) + del self._pending_publications[file] + self._pending_publications[file] = pending + return + if ( + status_error is not None + and not pending.is_rejection_report + and _is_rejected_content(status_error.response.status_code, reason) + ): + self.log.warning("The API server rejected the parsing result for %s", file) + status_code = status_error.response.status_code + detail = f"HTTP {status_code}: {reason}" if reason else f"HTTP {status_code}" + if reason == "invalid_dag": + message = status_error.response.json()["detail"].get("message") + if isinstance(message, str): + detail += f": {message}" + self._pending_publications[file] = self._build_rejection_report( + pending, f"the API server rejected it ({detail})" + ) + files.append(file) + continue + self.log.exception("Failed to publish parsing result for %s", file) + self._record_failed_publication(file, pending.stat) + else: + self._file_stats[file] = pending.stat + del self._pending_publications[file] + + def _build_rejection_report(self, pending: PendingDagPublication, reason: str) -> PendingDagPublication: + """Replace a result that cannot be published with an import error users can see.""" + self._dispatch_sequence += 1 + if len(reason) > _MAX_REJECTION_REASON_LENGTH: + reason = reason[:_MAX_REJECTION_REASON_LENGTH] + "..." + import_errors = { + pending.body.relative_fileloc: f"The parsing result was not published: {reason}.", + } + return PendingDagPublication( + body=DagParseResultBody( + attempt_id=uuid7(), + dispatch_sequence=self._dispatch_sequence, + bundle_name=pending.body.bundle_name, + relative_fileloc=pending.body.relative_fileloc, + bundle_version=pending.body.bundle_version, + bundle_revision=pending.body.bundle_revision, + parse_duration=pending.body.parse_duration, + serialized_dags=[], + source_codes={}, + import_errors=import_errors, + ), + stat=attrs.evolve(pending.stat, num_dags=0, import_errors=len(import_errors)), + refresh_generation=pending.refresh_generation, + is_rejection_report=True, + ) + + def _build_parse_result( + self, file: DagFileInfo, proc: BaseDagFileProcessorProcess, run_duration: float + ) -> DagParseResultBody: + result = proc.parsing_result + if result is None or self.api_client is None: + raise ValueError("API publication requires a client and completed parse result") + published_dag_ids = {dag.dag_id for dag in result.serialized_dags} + warnings = [ + ParseWarning( + **{ + key: warning[key] if isinstance(warning, dict) else getattr(warning, key) + for key in ("dag_id", "warning_type", "message") + } + ) + for warning in result.warnings or [] + ] + # The API accepts warnings only for published Dags; a Dag that failed to serialize has an import error. + warnings = [warning for warning in warnings if warning.dag_id in published_dag_ids] + source_codes = {} + for dag in result.serialized_dags: + source = result.dag_source_codes.get(dag.dag_id) + source_codes[dag.dag_id] = ParseSourceCode( + source_code=source.source_code if source else None, + language=source.language if source else "python", + ) + return DagParseResultBody( + attempt_id=proc.id, + dispatch_sequence=proc.dispatch_sequence, + bundle_name=file.bundle_name, + relative_fileloc=str(file.rel_path), + bundle_version=proc.parse_source.bundle_version, + bundle_revision=proc.parse_source.bundle_revision, + version_data=proc.parse_source.version_data, + parse_duration=run_duration, + serialized_dags=[dag.data for dag in result.serialized_dags], + import_errors=result.import_errors or {}, + parsed_definitions=result.parsed_definitions, + warnings=warnings, + source_codes=source_codes, + ) + def persist_parsing_result( self, *, @@ -1534,8 +1926,12 @@ @functools.cached_property def client(self) -> Client: + if self.api_client is not None: + return self.api_client + from airflow.sdk.api.client import Client + self._api_server = InProcessExecutionAPI() client = Client(base_url=None, token="", dry_run=True, transport=self._api_server.transport) # Mypy is wrong -- the setter accepts a string on the property setter! `URLType = URL | str` client.base_url = "http://in-process.invalid./" @@ -1608,10 +2004,12 @@ self._readmit_files_with_source() bundle_to_team = self._get_team_names({file.bundle_name for file in self._file_queue}) - while self._parallelism > len(self._processors) and self._file_queue: + while ( + self._parallelism > len(self._processors) + len(self._pending_publications) and self._file_queue + ): file, _ = self._file_queue.popitem(last=False) # Stop creating duplicate processor i.e. processor with the same filepath - if file in self._processors: + if file in self._processors or file in self._pending_publications: continue source = self._get_dispatch_source(file) @@ -1621,6 +2019,8 @@ source = deepcopy(source) processor = self._create_process(file) processor.parse_source = source + self._dispatch_sequence += 1 + processor.dispatch_sequence = self._dispatch_sequence stats.incr( "dag_processing.processes", tags=prune_dict( @@ -1646,6 +2046,7 @@ tracked_presence_keys = {file.presence_key for file in self._file_queue} tracked_presence_keys.update(file.presence_key for file in self._file_stats) tracked_presence_keys.update(file.presence_key for file in self._processors) + tracked_presence_keys.update(file.presence_key for file in self._pending_publications) tracked_presence_keys.update(file.presence_key for file in self._files_awaiting_source) for files in known_files.values(): for file in files: @@ -1685,7 +2086,7 @@ modified_datetime = datetime.fromtimestamp(modified_timestamp, tz=timezone.utc) files_with_mtime[file] = modified_timestamp stat = file_stats_by_presence_key.get(file.presence_key) - last_time = stat.last_finish_time if stat else None + last_time = (stat.last_attempt_time or stat.last_finish_time) if stat else None if not last_time: continue if modified_datetime > last_time: @@ -1712,7 +2113,7 @@ ), None, ) - last_time = stat.last_finish_time if stat else None + last_time = (stat.last_attempt_time or stat.last_finish_time) if stat else None if not last_time: return False elapsed_ss = (now - last_time).total_seconds() @@ -1738,6 +2139,7 @@ # If the file path is already being processed, or if a file was # processed recently, wait until the next batch in_progress_keys = {file.presence_key for file in self._processors} + in_progress_keys.update(file.presence_key for file in self._pending_publications) in_progress_keys.update(file.presence_key for file in self._files_awaiting_source) file_stats_by_presence_key = {file.presence_key: stat for file, stat in self._file_stats.items()} now = timezone.utcnow() @@ -1750,7 +2152,7 @@ for file in bundle_files: files.append(file) stat = file_stats_by_presence_key.get(file.presence_key) - last_time = stat.last_finish_time if stat else None + last_time = (stat.last_attempt_time or stat.last_finish_time) if stat else None if last_time and (now - last_time).total_seconds() < self._file_process_interval: recently_processed.add(file) @@ -1891,6 +2293,10 @@ return False if self._num_run < self.max_runs: return False + if self._pending_publications: + return False + if self._pending_inventories: + return False if any( file not in self._callback_to_execute for file in [*self._file_queue, *self._files_awaiting_source]
diff --git a/airflow-core/src/airflow/dag_processing/processor.py b/airflow-core/src/airflow/dag_processing/processor.py index e90466a..3adb457 100644 --- a/airflow-core/src/airflow/dag_processing/processor.py +++ b/airflow-core/src/airflow/dag_processing/processor.py
@@ -21,10 +21,12 @@ import importlib import logging import os +import sys import traceback from collections.abc import Callable, Sequence from pathlib import Path from typing import TYPE_CHECKING, Annotated, BinaryIO, ClassVar, Literal, cast +from uuid import UUID import attrs from pydantic import BaseModel, Field, TypeAdapter @@ -87,6 +89,7 @@ from structlog.typing import FilteringBoundLogger + from airflow.dag_processing.api_client import DagParseContext from airflow.sdk.api.client import Client from airflow.sdk.bases.operator import BaseOperator from airflow.sdk.definitions.context import Context @@ -577,6 +580,7 @@ bundle_version: str | None = None version_data: dict | None = None refresh_generation: int = 0 + bundle_revision: UUID | None = None @attrs.define(kw_only=True) @@ -596,6 +600,7 @@ """The file's parse log, which the process closes. Without one, the output is discarded.""" parsing_result: DagFileParsingResult | None = None parse_source: DagParseSource = attrs.field(factory=DagParseSource) + dispatch_sequence: int = 0 decoder: ClassVar[TypeAdapter[ToManager]] = TypeAdapter[ToManager](ToManager) had_callbacks: bool = False # Track if this process was started with callbacks to prevent stale DAG detection false positives @@ -608,6 +613,7 @@ bundle_name: str dag_file_rel_path: str + _api_parse_context: DagParseContext | None = attrs.field(init=False, default=None) def _get_target_loggers(self) -> tuple[FilteringBoundLogger, ...]: base = super()._get_target_loggers() @@ -638,10 +644,15 @@ self.parsing_result = msg return None, {} + def _handle_previous_successful_dag_run( + self, msg: GetPrevSuccessfulDagRun, log: FilteringBoundLogger, req_id: int + ) -> RequestResult: + # The run is looked up by task instance, and a parse process is not one, so the answer is always empty. + return PrevSuccessfulDagRunResult(), {"exclude_unset": True} + _client_request_types: ClassVar[tuple[type[BaseModel], ...]] = ( DeleteVariable, GetConnection, - GetPrevSuccessfulDagRun, GetPreviousDagRun, GetPreviousTI, GetTICount, @@ -662,7 +673,12 @@ "dict[type[BaseModel], RequestHandler[BaseDagFileProcessorProcess]]", WatchedSubprocess._get_shared_request_handlers(*_client_request_types, MaskSecret), ), - **dict([register_request_method(DagFileParsingResult, _handle_parsing_result)]), + **dict( + [ + register_request_method(DagFileParsingResult, _handle_parsing_result), + register_request_method(GetPrevSuccessfulDagRun, _handle_previous_successful_dag_run), + ] + ), } def _handle_request(self, msg, log: FilteringBoundLogger, req_id: int) -> None: @@ -675,7 +691,25 @@ ), ) return - super()._handle_request(msg, log, req_id) + # The HTTP client pulls in the Execution API versions. A DagProcessorAPIClient can only exist once + # its module is loaded, so the database-backed mode never has to import it. + api_client_module = sys.modules.get("airflow.dag_processing.api_client") + if api_client_module is None or not isinstance(self.client, api_client_module.DagProcessorAPIClient): + super()._handle_request(msg, log, req_id) + return + from airflow.api_fastapi.execution_api.datamodels.job import DagParseTokenBody + from airflow.dag_processing.api_client import DagParseContext + + if self._api_parse_context is None: + self._api_parse_context = DagParseContext( + request=DagParseTokenBody( + attempt_id=self.id, + bundle_name=self.bundle_name, + relative_fileloc=self.dag_file_rel_path, + ) + ) + with self.client.use_parse(self._api_parse_context): + super()._handle_request(msg, log, req_id) def _reject_request(self, msg, log: FilteringBoundLogger, req_id: int) -> None: log.error("Unhandled request", msg=msg)
diff --git a/airflow-core/tests/unit/dag_processing/bundles/test_dag_bundle_manager.py b/airflow-core/tests/unit/dag_processing/bundles/test_dag_bundle_manager.py index aef28b9..f859903 100644 --- a/airflow-core/tests/unit/dag_processing/bundles/test_dag_bundle_manager.py +++ b/airflow-core/tests/unit/dag_processing/bundles/test_dag_bundle_manager.py
@@ -278,6 +278,63 @@ @pytest.mark.db_test +@pytest.mark.parametrize("existing", [False, True]) +@conf_vars({("core", "load_examples"): "False"}) +@patch.object(BasicBundle, "__init__", autospec=True, side_effect=AssertionError("constructed too early")) +def test_metadata_sync_does_not_construct_bundles_or_clear_urls(constructor, clear_db, session, existing): + if existing: + bundle = DagBundleModel(name="my-test-bundle") + bundle.signed_url_template = "saved" + bundle.template_params = {"x": 1} + session.add(bundle) + session.commit() + with conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps(BASIC_BUNDLE_CONFIG)}): + DagBundlesManager().sync_bundles_to_db(include_bundle_urls=False) + session.expire_all() + bundle = session.get(DagBundleModel, "my-test-bundle") + assert bundle.active + if existing: + assert (bundle.signed_url_template, bundle.template_params) == ("saved", {"x": 1}) + constructor.assert_not_called() + + +@pytest.mark.parametrize("operation", ["get", "list", "sync"]) +@pytest.mark.db_test +@conf_vars({("core", "load_examples"): "False"}) +def test_bundle_construction_and_url_resolution_use_the_bundle_context(clear_db, operation): + active = [] + seen = [] + + @contextlib.contextmanager + def bundle_context(name): + active.append(name) + try: + yield + finally: + active.pop() + + def read_secret(*args, **kwargs): + seen.append(active[-1]) + + config = [*BASIC_BUNDLE_CONFIG, {"name": "excluded", "classpath": "missing.Bundle", "kwargs": {}}] + with ( + conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps(config)}), + patch.object(BasicBundle, "__init__", autospec=True, side_effect=read_secret), + patch.object(BasicBundle, "view_url_template", autospec=True, side_effect=read_secret), + ): + manager = DagBundlesManager(bundle_names=["my-test-bundle"], bundle_context=bundle_context) + if operation == "get": + manager.get_bundle("my-test-bundle") + elif operation == "list": + list(manager.get_all_dag_bundles()) + else: + manager.sync_bundles_to_db(deactivate_missing=False) + assert seen + assert set(seen) == {"my-test-bundle"} + assert active == [] + + +@pytest.mark.db_test @conf_vars({("core", "LOAD_EXAMPLES"): "False"}) def test_sync_bundles_to_db(clear_db, session): def _get_bundle_names_and_active():
diff --git a/airflow-core/tests/unit/dag_processing/test_api_client.py b/airflow-core/tests/unit/dag_processing/test_api_client.py index 48c947b..9e685a3 100644 --- a/airflow-core/tests/unit/dag_processing/test_api_client.py +++ b/airflow-core/tests/unit/dag_processing/test_api_client.py
@@ -25,7 +25,12 @@ import jwt import pytest from tenacity import wait_none +from uuid6 import uuid7 +from airflow.api_fastapi.execution_api.datamodels.dag_parsing import ( + DagBundleInventoryBody, + DagParseResultBody, +) from airflow.api_fastapi.execution_api.datamodels.job import DagParseTokenBody, JobState, TerminalJobState from airflow.dag_processing import api_client from airflow.dag_processing.api_client import ( @@ -767,12 +772,69 @@ ) client.register_job() with client.use_parse(make_parse_context()): + with pytest.raises(httpx2.ReadError): + client.get("variables/key") client.get("variables/key") assert requests[1].content == requests[2].content assert len(requests) == 4 +@pytest.mark.parametrize( + ("request_kwargs", "expected_timeout"), + [ + pytest.param({}, httpx2.Timeout(5), id="client-default"), + pytest.param({"timeout": 2.5}, httpx2.Timeout(2.5), id="explicit-timeout"), + pytest.param({"timeout": None}, httpx2.Timeout(None), id="disabled-timeout"), + pytest.param( + {"timeout": httpx2.Timeout(4, connect=2)}, + httpx2.Timeout(4, connect=2), + id="per-phase-timeout", + ), + ], +) +def test_parse_requests_make_one_attempt_with_the_request_timeout( + make_client, request_kwargs, expected_timeout +): + client, requests = make_client( + make_registration_response(), make_parse_token_response(), httpx2.Response(503), timeout=5 + ) + client.register_job() + + with client.use_parse(make_parse_context()), pytest.raises(httpx2.HTTPStatusError): + client.get("variables/key", **request_kwargs) + + assert len(requests) == 3 + assert requests[1].extensions["timeout"] == expected_timeout.as_dict() + assert requests[2].extensions["timeout"] == expected_timeout.as_dict() + + +@pytest.mark.parametrize(("expired", "tick"), [("parse", 60), ("job", 300)]) +def test_required_token_renewals_use_the_runtime_request_timeout(make_client, clock, expired, tick): + client, requests = make_client( + make_registration_response(), + make_parse_token_response(), + httpx2.Response(200), + *([make_registration_response()] if expired == "job" else []), + make_parse_token_response(), + httpx2.Response(200), + timeout=5, + ) + client.register_job() + + with client.use_parse(make_parse_context()): + client.get("variables/key") + clock.return_value = tick + client.get("variables/key", timeout=2.5) + + assert [request.url.path for request in requests[3:]] == [ + *(["/execution/jobs"] if expired == "job" else []), + "/execution/jobs/1/parse-token", + "/execution/variables/key", + ] + assert all(request.extensions["timeout"] == httpx2.Timeout(2.5).as_dict() for request in requests[3:]) + + def test_interleaved_parses_keep_distinct_file_credentials(make_client): first_response = make_parse_token_response() second_response = make_parse_token_response( @@ -951,3 +1013,69 @@ assert client.restart_required assert client.heartbeat() == JobState.RESTARTING assert requests[-1].headers["Authorization"] == f"Bearer {make_job_token()}" + + +@pytest.mark.parametrize("expired", [False, True]) +@pytest.mark.parametrize("failure", [httpx2.ReadTimeout("Lost response"), httpx2.Response(503)]) +def test_publication_attempt_does_not_sleep_through_heartbeats(make_client, clock, expired, failure): + body = DagParseResultBody( + attempt_id=uuid7(), + dispatch_sequence=1, + bundle_name="bundle-a", + relative_fileloc="dag.py", + parse_duration=0.1, + serialized_dags=[], + source_codes={}, + ) + responses = [make_registration_response()] + if expired: + responses.append(make_registration_response()) + responses.extend( + [ + failure, + httpx2.Response(200, json={"state": "running"}), + httpx2.Response( + 200, json={"attempt_id": str(body.attempt_id), "accepted_at": "2026-10-05T12:00:00Z"} + ), + ] + ) + client, requests = make_client(*responses) + client.register_job() + if expired: + clock.return_value = 300 + with pytest.raises(httpx2.HTTPError): + client.publish_parse_result(body) + assert len(requests) == (3 if expired else 2) + publication_timeout = requests[-1].extensions["timeout"] + assert publication_timeout["connect"] <= 1 + assert publication_timeout["pool"] <= 1 + assert publication_timeout["read"] == client.timeout.read > 1 + if expired: + assert all(value <= 1 for value in requests[1].extensions["timeout"].values()) + assert client.heartbeat() == JobState.RUNNING + assert client.publish_parse_result(body).attempt_id == body.attempt_id + assert requests[-1].content == requests[-3].content + + +def test_inventory_waits_for_a_slow_write_but_connects_quickly(make_client): + client, requests = make_client( + make_registration_response(), + httpx2.Response( + 200, + json={ + "attempt_id": str(uuid7()), + "accepted_at": "2026-10-05T12:00:00Z", + "revision": str(uuid7()), + }, + ), + ) + client.register_job() + + client.publish_inventory( + "bundle-a", + DagBundleInventoryBody(attempt_id=uuid7(), dispatch_sequence=1, expected_revision=None, files=[]), + ) + + timeout = requests[-1].extensions["timeout"] + assert timeout["connect"] <= 1 + assert timeout["read"] == client.timeout.read > 1
diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index 9418be4..8f6959b 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py
@@ -23,20 +23,25 @@ import os import random import re +import runpy import selectors import shutil import signal +import sys import textwrap import time import zipfile from collections import Counter, OrderedDict, defaultdict, namedtuple -from contextlib import contextmanager +from contextlib import contextmanager, nullcontext from datetime import datetime, timedelta from pathlib import Path from socket import socket, socketpair from unittest import mock from unittest.mock import MagicMock +from uuid import UUID +import httpx2 +import jwt import msgspec import pytest import structlog @@ -46,10 +51,21 @@ from uuid6 import uuid7 from airflow._shared.timezones import timezone +from airflow.api_fastapi.execution_api.datamodels.dag_parsing import ( + DagBundleInventoryBody, + DagBundleInventoryResponse, + DagBundleStateResponse, + DagParseResultBody, + ParseSourceCode, +) from airflow.api_fastapi.execution_api.datamodels.taskinstance import TaskInstance as TIDataModel from airflow.callbacks.callback_requests import DagCallbackRequest, TaskCallbackRequest +from airflow.dag_processing.api_client import ( + DagProcessorAPIClient, + DagProcessorSecretsComms, +) from airflow.dag_processing.bundles.base import BaseDagBundle, BundleVersion -from airflow.dag_processing.bundles.manager import DagBundlesManager +from airflow.dag_processing.bundles.manager import DagBundlesManager, _load_bundle_config_snapshot from airflow.dag_processing.collection import update_dag_parsing_results_in_db from airflow.dag_processing.dagbag import DagBag from airflow.dag_processing.lang_sdk_processor import LangSDKDagFileProcessorProcess @@ -58,6 +74,7 @@ DagFileInfo, DagFileProcessorManager, DagFileStat, + PendingDagPublication, ) from airflow.dag_processing.processor import ( DagFileParseRequest, @@ -74,7 +91,8 @@ from airflow.models.serialized_dag import SerializedDagModel from airflow.models.team import Team from airflow.providers.standard.operators.empty import EmptyOperator -from airflow.sdk import DAG as SdkDAG, BaseOperator +from airflow.sdk import DAG as SdkDAG, BaseOperator, Connection, SecretCache +from airflow.sdk.execution_time import task_runner from airflow.sdk.importers import DagDefinition, DagImporterRegistry, DagImportError, DagSourceCode from airflow.serialization.serialized_objects import DagSerialization, LazyDeserializedDAG from airflow.utils.net import get_hostname @@ -557,14 +575,18 @@ manager = DagFileProcessorManager(max_runs=1) with mock.patch("airflow.dag_processing.manager.DagBundlesManager") as mock_bundles_manager: manager.sync_bundles() - mock_bundles_manager.return_value.sync_bundles_to_db.assert_called_once_with(deactivate_missing=True) + mock_bundles_manager.return_value.sync_bundles_to_db.assert_called_once_with( + deactivate_missing=True, include_bundle_urls=True + ) def test_sync_bundles_does_not_deactivate_missing_when_filtered(self): """A processor started with ``--bundle-name`` owns a subset and must not deactivate others.""" manager = DagFileProcessorManager(max_runs=1, bundle_names_to_parse=["only-mine"]) with mock.patch("airflow.dag_processing.manager.DagBundlesManager") as mock_bundles_manager: manager.sync_bundles() - mock_bundles_manager.return_value.sync_bundles_to_db.assert_called_once_with(deactivate_missing=False) + mock_bundles_manager.return_value.sync_bundles_to_db.assert_called_once_with( + deactivate_missing=False, include_bundle_urls=True + ) @pytest.mark.parametrize( "safe_mode", @@ -2320,10 +2342,14 @@ assert not manager._file_queue assert not manager._files_awaiting_source + @pytest.mark.parametrize("api_mode", [False, True]) @mock.patch.object(DagFileProcessorManager, "_create_process", autospec=True) @mock.patch.object(DagFileProcessorManager, "persist_parsing_result", autospec=True) - def test_periodic_refresh_drains_imports_before_admitting_more(self, persist, create): - manager = DagFileProcessorManager(max_runs=1, bundle_refresh_check_interval=300) + def test_periodic_refresh_drains_imports_before_admitting_more(self, persist, create, api_mode): + client = mock.create_autospec(DagProcessorAPIClient, instance=True) if api_mode else None + if client is not None: + client.publish_inventory.return_value.revision = uuid7() + manager = DagFileProcessorManager(max_runs=1, bundle_refresh_check_interval=300, api_client=client) bundle = self._make_refresh_bundle(supports_versioning=True, current_version="v2") bundle.refresh_interval = 300 manager._bundle_versions[bundle.name] = "v1" @@ -2332,6 +2358,7 @@ queued = DagFileInfo(bundle_name=bundle.name, rel_path=Path("queued.py")) processor, _ = self.mock_processor() processor.had_callbacks = False + processor.dispatch_sequence = 1 processor.parse_source = DagParseSource(bundle_version="v1") processor.parsing_result = DagFileParsingResult(fileloc="active.py", serialized_dags=[]) manager._processors[active] = processor @@ -2344,8 +2371,15 @@ assert manager._bundles_waiting_for_refresh == {bundle.name} manager._collect_results() - persist.assert_called_once() - assert persist.call_args.kwargs["bundle_version"] == "v1" + if api_mode: + persist.assert_not_called() + self._refresh_with_mocked_state(manager, bundle, BundleState(last_refreshed=None, version="v1")) + bundle.refresh.assert_not_called() + manager._publish_pending_results() + assert client.publish_parse_result.call_args.args[0].bundle_version == "v1" + else: + persist.assert_called_once() + assert persist.call_args.kwargs["bundle_version"] == "v1" self._refresh_with_mocked_state( manager, bundle, BundleState(last_refreshed=timezone.utcnow(), version="v1") ) @@ -2393,7 +2427,8 @@ assert manager._file_stats[file].num_dags == 0 assert manager._file_stats[file].import_errors == 0 assert manager._file_stats[file].run_count == 1 - assert manager._file_stats[file].last_finish_time is not None + assert manager._file_stats[file].last_finish_time is None + assert manager._file_stats[file].last_attempt_time is not None assert manager._file_stats[file].last_duration is not None assert manager.processed_recently(timezone.utcnow(), file) is True @@ -2472,7 +2507,8 @@ 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 assert manager._file_stats[file_a].run_count == stat_a_before.run_count + 1 - assert manager._file_stats[file_a].last_finish_time is not None + assert manager._file_stats[file_a].last_finish_time == stat_a_before.last_finish_time + assert manager._file_stats[file_a].last_attempt_time is not None assert manager._file_stats[file_b] is not stat_b_before assert manager._file_stats[file_b].run_count == 2 assert len(manager._processors) == 0 @@ -4375,7 +4411,25 @@ manager._force_refresh_bundles = {bundle.name} if force else set() mock_get = mock.patch.object(manager, "get_bundle_state", return_value=initial_state) mock_update = mock.patch.object(manager, "update_bundle_state") + mock_catalog = ( + mock.patch.object( + manager.api_client, + "get_bundles", + return_value=[ + DagBundleStateResponse( + name=bundle.name, + last_refreshed=initial_state.last_refreshed, + version=initial_state.version, + revision=initial_state.revision, + team_name=None, + ) + ], + ) + if manager.api_client is not None + else nullcontext() + ) with ( + mock_catalog, mock_get as patched_get, mock_update as patched_update, mock.patch.object(manager, "_find_files_in_bundle", return_value=set()), @@ -4875,6 +4929,239 @@ f"({SWEEP_CALLS} x {fixed} fixed + {SWEEP_FILES} x {per_dag} per Dag)." ) + @pytest.mark.parametrize("publication_fails", [False, True]) + @mock.patch.object(DagFileProcessorManager, "persist_parsing_result", autospec=True) + @mock.patch.object(DagFileProcessorManager, "_create_process", autospec=True) + def test_api_publication_uses_dispatch_identity_without_local_persistence( + self, create, persist, session, publication_fails + ): + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + manager = DagFileProcessorManager(max_runs=1, api_client=client) + manager._bundle_parse_sources["testing"] = DagParseSource(bundle_version="v1") + file = _get_file_infos(["dag.py"])[0] + manager._file_queue[file] = None + proc, read_end = self.mock_processor() + proc.had_callbacks = False + proc.parsing_result = DagFileParsingResult(fileloc="dag.py", serialized_dags=[]) + create.return_value = proc + if publication_fails: + client.publish_parse_result.side_effect = RuntimeError("publication failed") + try: + manager._start_new_processes() + manager.handle_parsing_result(file, proc, session=session) + client.publish_parse_result.assert_not_called() + manager._publish_pending_results() + body = client.publish_parse_result.call_args.args[0] + assert body.attempt_id == proc.id + assert body.dispatch_sequence == proc.dispatch_sequence == 1 + assert body.bundle_version == "v1" + assert body.relative_fileloc == "dag.py" + assert body.serialized_dags == [] + assert manager._file_stats[file].run_count == 1 + persist.assert_not_called() + finally: + read_end.close() + proc.close() + + @pytest.mark.parametrize("source_available", [True, False]) + def test_api_publication_carries_diagnostics_and_explicit_source_availability(self, source_available): + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + manager = DagFileProcessorManager(max_runs=1, api_client=client) + file = _get_file_infos(["dag.py"])[0] + dag = SdkDAG("published", schedule=None) + dag.fileloc = "/worker/dag.py" + dag.relative_fileloc = "dag.py" + proc, read_end = self.mock_processor() + proc.dispatch_sequence = 1 + warning = {"dag_id": dag.dag_id, "warning_type": "python:deprecated", "message": "warning"} + unserialized_warning = {"dag_id": "failed", "warning_type": "python:deprecated", "message": "warning"} + proc.parsing_result = DagFileParsingResult( + fileloc=dag.fileloc, + serialized_dags=[LazyDeserializedDAG.from_dag(dag)], + warnings=[warning, unserialized_warning], + parsed_definitions=["dag.py"], + import_errors={"dag.py": "import error"}, + dag_source_codes={dag.dag_id: DagSourceCode(source_code="captured", language="python")} + if source_available + else {}, + ) + try: + body = manager._build_parse_result(file, proc, 0.5) + assert [published.model_dump() for published in body.warnings] == [warning] + assert body.import_errors == {"dag.py": "import error"} + assert body.parsed_definitions == ["dag.py"] + assert body.source_codes[dag.dag_id].source_code == ("captured" if source_available else None) + finally: + read_end.close() + proc.close() + + @pytest.mark.usefixtures("testing_dag_bundle") + @conf_vars({("workers", "execution_api_retries"): "1"}) + @pytest.mark.parametrize("failure", [413, 422, 403, 500, "timeout"]) + def test_rejected_publication_preserves_stale_detection(self, session, time_machine, failure): + time_machine.move_to("2026-10-05T12:00:00Z", tick=False) + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + request = httpx2.Request("POST", "http://api/jobs/1/parse-results") + client.publish_parse_result.side_effect = ( + httpx2.ReadTimeout("Lost response", request=request) + if failure == "timeout" + else httpx2.HTTPStatusError("Rejected", request=request, response=httpx2.Response(failure)) + ) + manager = DagFileProcessorManager(max_runs=1, api_client=client) + file = _get_file_infos(["dag.py"])[0] + previous = timezone.utcnow() - timedelta(seconds=manager.stale_dag_threshold + 1) + manager._file_stats[file] = DagFileStat(num_dags=1, last_finish_time=previous) + dag = SdkDAG("published", schedule=None) + dag.fileloc = "/worker/dag.py" + dag.relative_fileloc = "dag.py" + session.add( + DagModel( + dag_id=dag.dag_id, + bundle_name="testing", + relative_fileloc="dag.py", + fileloc=dag.fileloc, + last_parsed_time=previous, + is_stale=False, + ) + ) + session.commit() + proc, read_end = self.mock_processor() + proc.dispatch_sequence = 1 + proc.parsing_result = DagFileParsingResult( + fileloc=dag.fileloc, serialized_dags=[LazyDeserializedDAG.from_dag(dag)] + ) + try: + manager.handle_parsing_result(file, proc) + manager._publish_pending_results() + manager._scan_stale_dags() + session.expire_all() + assert not session.get(DagModel, dag.dag_id).is_stale + assert manager._file_stats[file].last_finish_time == previous + assert manager._file_stats[file].num_dags == 1 + assert manager.processed_recently(timezone.utcnow(), file) + assert not manager._pending_publications + finally: + read_end.close() + proc.close() + + def test_publication_retries_leave_the_loop_responsive(self, mocker): + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + client.publish_parse_result.side_effect = [httpx2.ReadTimeout("Lost acknowledgment"), None, None] + heartbeat_observations = [] + with selectors.DefaultSelector() as selector: + manager = DagFileProcessorManager( + max_runs=1, + api_client=client, + selector=selector, + heartbeat=lambda: heartbeat_observations.append(client.publish_parse_result.call_count), + ) + for method in ( + "_refresh_dag_bundles", + "_queue_requested_files_for_parsing", + "_scan_stale_dags", + "_cleanup_stale_bundle_versions", + "print_stats", + ): + mocker.patch.object(DagFileProcessorManager, method, autospec=True) + mocker.patch.object(DagFileProcessorManager, "fetch_callbacks", autospec=True, return_value=[]) + for file in _get_file_infos(["a.py", "b.py"]): + proc, read_end = self.mock_processor() + proc.dispatch_sequence = 1 + proc.parsing_result = DagFileParsingResult(fileloc=str(file.rel_path), serialized_dags=[]) + try: + manager.handle_parsing_result(file, proc) + finally: + read_end.close() + proc.close() + assert not manager.max_runs_reached() + manager._run_parsing_loop() + calls = client.publish_parse_result.call_args_list + assert [call.args[0].relative_fileloc for call in calls] == ["a.py", "b.py", "a.py"] + assert calls[0].args[0] is calls[2].args[0] + assert {0, 1}.issubset(heartbeat_observations) + assert not manager._pending_publications + assert all(stat.run_count == 1 for stat in manager._file_stats.values()) + + @pytest.mark.parametrize("change", ["refresh", "delete"]) + def test_pending_publication_is_discarded_when_source_changes(self, change): + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + manager = DagFileProcessorManager(max_runs=1, api_client=client) + file = _get_file_infos(["dag.py"])[0] + proc, read_end = self.mock_processor() + proc.dispatch_sequence = 1 + proc.parsing_result = DagFileParsingResult(fileloc="dag.py", serialized_dags=[]) + try: + manager.handle_parsing_result(file, proc) + if change == "refresh": + manager._invalidate_bundle_parse_source("testing") + else: + manager.remove_orphaned_file_stats(present=set()) + manager._publish_pending_results() + client.publish_parse_result.assert_not_called() + assert not manager._pending_publications + assert list(manager._file_queue) == ([file] if change == "refresh" else []) + finally: + read_end.close() + proc.close() + + @mock.patch.object(DagFileProcessorManager, "_create_process", autospec=True) + def test_pending_publications_bound_admission_and_prevent_reimports(self, create): + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + manager = DagFileProcessorManager(max_runs=1, api_client=client, parallelism=1) + manager._file_parsing_sort_mode = "alphabetical" + file, other = _get_file_infos(["dag.py", "other.py"]) + proc, read_end = self.mock_processor() + proc.dispatch_sequence = 1 + proc.parsing_result = DagFileParsingResult(fileloc="dag.py", serialized_dags=[]) + manager._bundle_parse_sources["testing"] = DagParseSource() + try: + manager.handle_parsing_result(file, proc) + manager.prepare_file_queue(known_files={"testing": {file, other}}) + assert list(manager._file_queue) == [other] + manager._start_new_processes() + create.assert_not_called() + assert list(manager._file_queue) == [other] + manager._parallelism = 2 + manager._file_queue = OrderedDict.fromkeys([file]) + manager._start_new_processes() + create.assert_not_called() + finally: + read_end.close() + proc.close() + + @pytest.mark.parametrize("failure", ["invalid_path", "too_many_files"]) + @mock.patch.object(DagFileProcessorManager, "_get_observed_filelocs", autospec=True) + def test_unusable_inventory_preserves_version_and_recovers_after_refresh(self, observed, failure): + observed.return_value = ( + {r"bad\dag.py"} if failure == "invalid_path" else {f"{index}.py" for index in range(100001)} + ) + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + manager = DagFileProcessorManager( + max_runs=1, api_client=client, parallelism=1, bundle_refresh_check_interval=0 + ) + bundle = self._make_refresh_bundle() + manager._bundles_waiting_for_refresh.add(bundle.name) + manager._bundle_versions[bundle.name] = "previous" + + self._refresh_with_mocked_state( + manager, bundle, BundleState(last_refreshed=None, version=None), force=True + ) + + bundle.refresh.assert_called_once() + client.publish_inventory.assert_not_called() + assert bundle.name not in manager._bundle_parse_sources + assert bundle.name not in manager._pending_inventories + assert manager._bundle_versions[bundle.name] == "previous" + assert not manager._force_refresh_bundles + assert not manager._bundles_waiting_for_refresh + + observed.return_value = set() + client.publish_inventory.return_value.revision = uuid7() + self._refresh_with_mocked_state(manager, bundle, BundleState(last_refreshed=None, version=None)) + assert bundle.name in manager._bundle_parse_sources + assert manager._bundle_versions[bundle.name] is None + assert not manager._pending_inventories + class TestMultiTeamMetrics: """Tests for team_name tag on dag processing metrics in multi-team mode.""" @@ -5250,3 +5537,774 @@ assert result == "dags_test_test_dag.py" assert caplog.entries == [] + + +API_CLIENT_BUNDLES = [ + { + "name": name, + "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", + "kwargs": {"path": "/"}, + } + for name in ("bundle_a", "bundle_b") +] + + +def _read_bundle_connection() -> None: + Connection.get("bundle_conn") + + +class TestDagFileProcessorManagerWithAPIClient: + """With an API client, bundle code resolves its connections through it, for the bundle it belongs to.""" + + @pytest.fixture + def api_requests(self) -> list[httpx2.Request]: + return [] + + @pytest.fixture + def api_client(self, tmp_path, api_requests): + token_file = tmp_path / "processor.jwt" + token_file.write_text("session-token") + job_token = jwt.encode({"iat": 0, "exp": 300, "scope": "dag_processor", "job_id": 1}, "secret") + + def handle(request: httpx2.Request) -> httpx2.Response: + api_requests.append(request) + if request.url.path.endswith("/jobs"): + return httpx2.Response(201, json={"job_id": 1, "token": job_token, "expires_in": 300}) + if request.url.path.endswith("/bundles"): + return httpx2.Response( + 200, + json=[ + { + "name": bundle["name"], + "version": None, + "last_refreshed": None, + "revision": None, + "team_name": None, + } + for bundle in API_CLIENT_BUNDLES + ], + ) + return httpx2.Response( + 200, + json={ + "conn_id": "bundle_conn", + "conn_type": "http", + "host": None, + "schema": None, + "login": None, + "password": None, + "port": None, + "extra": None, + }, + ) + + with DagProcessorAPIClient( + base_url="http://api/execution/", + token_file=token_file, + hostname="processor-1", + transport=httpx2.MockTransport(handle), + ) as client: + client.register_job() + yield client + + @pytest.fixture + def make_manager(self, api_client): + managers = [] + + def make(**kwargs) -> DagFileProcessorManager: + manager = DagFileProcessorManager(max_runs=1, api_client=api_client, **kwargs) + manager.prepare_api_secrets_context(api_client) + managers.append(manager) + return manager + + yield make + for manager in managers: + manager.after_run() + + @pytest.fixture(autouse=True) + def configured_bundles(self): + with conf_vars( + { + ("dag_processor", "dag_bundle_config_list"): json.dumps(API_CLIENT_BUNDLES), + ("core", "load_examples"): "False", + } + ): + _load_bundle_config_snapshot.cache_clear() + yield + _load_bundle_config_snapshot.cache_clear() + + @staticmethod + def _get_connection_bundles(api_requests: list[httpx2.Request]) -> list[str]: + return [ + request.headers["Airflow-Dag-Bundle"] + for request in api_requests + if request.url.path.startswith("/execution/connections/") + ] + + def test_uses_the_api_client_instead_of_the_in_process_api(self, make_manager, api_client): + manager = make_manager() + + assert manager.client is api_client + assert manager._api_server is None + + @conf_vars({("secrets", "use_cache"): "True"}) + def test_api_mode_does_not_share_sdk_secret_cache_between_bundles(self, make_manager, monkeypatch): + monkeypatch.setattr(SecretCache, "_cache", {"inherited": "secret"}) + manager = make_manager() + manager.heartbeat = mock.Mock(spec=lambda: None, side_effect=SystemExit) + with pytest.raises(SystemExit): + manager._run_parsing_loop() + assert SecretCache._cache is None + + @mock.patch("airflow.dag_processing.manager.get_importer_registry", autospec=True) + def test_importer_warmup_selects_its_bundle(self, registry, make_manager, api_requests): + manager = make_manager(bundle_names_to_parse=["bundle_b"]) + manager._dag_bundles = manager.get_all_bundles() + registry.return_value.warm_importers.side_effect = _read_bundle_connection + manager.warm_importers() + assert self._get_connection_bundles(api_requests) == ["bundle_b"] + + @mock.patch.object( + DagFileProcessorManager, + "prepare_process_context", + autospec=True, + side_effect=ValueError("startup failed"), + ) + def test_startup_failure_clears_the_api_secrets_context(self, prepare, api_client): + manager = DagFileProcessorManager(max_runs=1, api_client=api_client) + with pytest.raises(ValueError, match="startup failed"): + manager.run() + assert not hasattr(task_runner, "SUPERVISOR_COMMS") + + @mock.patch("airflow.dag_processing.manager.gc", autospec=True) + def test_run_resolves_secrets_through_the_api_client_instead_of_the_database(self, _, api_client): + manager = DagFileProcessorManager(max_runs=1, api_client=api_client) + with ( + mock.patch.object(manager, "prepare_server_process_context") as server_context, + mock.patch.object(manager, "prepare_process_context"), + mock.patch.object(manager, "register_exit_signals"), + mock.patch.object(manager, "prepare_bundles"), + mock.patch.object(manager, "_symlink_latest_log_directory"), + mock.patch.object(manager, "warm_importers"), + ): + manager.before_run() + server_context.assert_not_called() + assert isinstance(task_runner.SUPERVISOR_COMMS, DagProcessorSecretsComms) + + manager.after_run() + + assert not hasattr(task_runner, "SUPERVISOR_COMMS") + + @pytest.mark.parametrize( + ("bundle_names_to_parse", "expected"), + [(None, ["bundle_a", "bundle_b"]), (["bundle_b"], ["bundle_b"])], + ) + @mock.patch("airflow.dag_processing.bundles.local.LocalDagBundle.__init__", autospec=True) + def test_creates_each_bundle_for_its_bundle( + self, mock_bundle, make_manager, api_requests, bundle_names_to_parse, expected + ): + def create(bundle, *, name, **kwargs): + _read_bundle_connection() + bundle.name = name + + mock_bundle.side_effect = create + + bundles = make_manager(bundle_names_to_parse=bundle_names_to_parse).get_all_bundles() + + assert [bundle.name for bundle in bundles] == expected + assert self._get_connection_bundles(api_requests) == expected + + def test_refreshes_each_bundle_for_its_bundle(self, make_manager, api_requests): + manager = make_manager() + bundle = MagicMock(spec=BaseDagBundle) + bundle.name = "bundle_b" + bundle.is_initialized = False + + def initialize(): + _read_bundle_connection() + raise OSError("clone failed") + + bundle.initialize.side_effect = initialize + manager._dag_bundles = [bundle] + + manager._refresh_dag_bundles(known_files={}) + + assert self._get_connection_bundles(api_requests) == ["bundle_b"] + + @mock.patch.object(DagBundlesManager, "get_bundle", autospec=True) + def test_prepares_a_callback_bundle_for_its_bundle(self, mock_get_bundle, make_manager, api_requests): + def create(_, name, version=None, version_data=None): + _read_bundle_connection() + bundle = MagicMock(spec=BaseDagBundle) + bundle.supports_versioning = False + return bundle + + mock_get_bundle.side_effect = create + request = DagCallbackRequest( + filepath="dag.py", + dag_id="dag_id", + run_id="run_id", + bundle_name="bundle_b", + bundle_version="v1", + context_from_server=None, + is_failure_callback=False, + ) + + assert make_manager().prepare_callback_bundle(request) is not None + assert self._get_connection_bundles(api_requests) == ["bundle_b"] + + +@mock.patch.object(DagFileProcessorManager, "_find_files_in_bundle", autospec=True, return_value=set()) +@conf_vars({("core", "multi_team"): "True"}) +def test_refresh_reads_one_fresh_catalog_per_pass(find_files, tmp_path): + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + manager = DagFileProcessorManager(max_runs=1, api_client=client, bundle_refresh_check_interval=0) + states = [] + known_files = {} + for name in ("a", "b", "c"): + bundle = _make_bundle(tmp_path) + bundle.name = name + bundle.is_initialized = True + bundle.supports_versioning = False + bundle.refresh_interval = 300 + manager._dag_bundles.append(bundle) + manager._bundle_versions[name] = None + manager._bundle_parse_sources[name] = DagParseSource() + known_files[name] = set() + states.append( + DagBundleStateResponse( + name=name, version=None, last_refreshed=timezone.utcnow(), revision=uuid7(), team_name=None + ) + ) + client.get_bundles.return_value = states + client.publish_inventory.return_value.revision = uuid7() + manager._refresh_dag_bundles(known_files) + client.get_bundles.assert_called_once_with() + find_files.assert_not_called() + + states[0] = states[0].model_copy( + update={"last_refreshed": None, "team_name": "new_team", "revision": uuid7()} + ) + states.pop() + manager._refresh_dag_bundles(known_files) + assert client.get_bundles.call_count == 2 + find_files.assert_called_once_with(manager, manager._dag_bundles[0]) + assert client.publish_inventory.call_args.args[1].expected_revision == states[0].revision + assert manager._bundle_name_to_team_name == {"a": "new_team", "b": None} + manager._dag_bundles[2].refresh.assert_not_called() + + file = DagFileInfo(bundle_name="b", rel_path=Path("a.py")) + known_files["b"] = {file} + pending_body = DagBundleInventoryBody( + attempt_id=uuid7(), dispatch_sequence=3, expected_revision=states[0].revision, files=[] + ) + manager._pending_inventories["a"] = (pending_body, set(), DagParseSource()) + client.get_bundles.side_effect = httpx2.ConnectError("API unavailable") + manager._refresh_dag_bundles(known_files) + assert client.get_bundles.call_count == 3 + assert known_files["b"] == {file} + assert client.publish_inventory.call_args.args[1] is pending_body + assert find_files.call_count == 1 + + +@pytest.mark.parametrize("catalog_failed", [False, True]) +@pytest.mark.parametrize("has_pending_inventory", [False, True]) +@mock.patch("airflow.dag_processing.manager.time", autospec=True) +@mock.patch.object(DagFileProcessorManager, "_find_files_in_bundle", autospec=True) +def test_catalog_omission_releases_source_and_allows_recovery( + find_files, clock, catalog_failed, has_pending_inventory, tmp_path +): + clock.monotonic.return_value = 100 + (tmp_path / "a.py").touch() + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + manager = DagFileProcessorManager(max_runs=1, api_client=client, bundle_refresh_check_interval=10) + bundles = {} + files = {} + for name in ("missing", "ready"): + bundle = _make_bundle(tmp_path) + bundle.name = name + bundle.is_initialized = True + bundle.supports_versioning = False + bundle.refresh_interval = 300 + bundles[name] = bundle + file = DagFileInfo(bundle_name=name, rel_path=Path("a.py"), bundle_path=tmp_path) + files[name] = {file} + manager._bundle_versions[name] = None + manager._bundle_parse_sources[name] = DagParseSource() + manager._dag_bundles = list(bundles.values()) + previous_files = files.copy() + previous_versions = manager._bundle_versions.copy() + ready_source = manager._bundle_parse_sources["ready"] + missing_source = manager._bundle_parse_sources["missing"] + pending_body = DagBundleInventoryBody( + attempt_id=uuid7(), dispatch_sequence=1, expected_revision=None, files=["a.py"] + ) + if has_pending_inventory: + manager._pending_inventories["missing"] = (pending_body, files["missing"], missing_source) + manager._bundle_parse_sources.pop("missing") + manager._force_refresh_bundles.add("missing") + manager._bundles_waiting_for_refresh.add("missing") + ready_state = DagBundleStateResponse( + name="ready", version=None, last_refreshed=timezone.utcnow(), revision=None, team_name=None + ) + client.get_bundles.return_value = [ready_state] + if catalog_failed: + client.get_bundles.side_effect = httpx2.ConnectError("Catalog unavailable") + client.publish_inventory.side_effect = _make_status_error(503) + + manager._refresh_dag_bundles(files) + + client.get_bundles.assert_called_once_with() + find_files.assert_not_called() + assert files == previous_files + assert manager._bundle_versions == previous_versions + assert manager._bundle_parse_sources["ready"] is ready_source + if catalog_failed: + assert manager._force_refresh_bundles == {"missing"} + assert manager._bundles_waiting_for_refresh == {"missing"} + if has_pending_inventory: + client.publish_inventory.assert_called_once_with("missing", pending_body) + assert manager._pending_inventories["missing"][0] is pending_body + else: + client.publish_inventory.assert_not_called() + assert manager._bundle_parse_sources["missing"] is missing_source + return + + client.publish_inventory.assert_not_called() + assert not manager._pending_inventories + assert not manager._force_refresh_bundles + assert not manager._bundles_waiting_for_refresh + assert manager._get_dispatch_source(next(iter(files["missing"]))) is None + manager._refresh_dag_bundles(files) + client.get_bundles.assert_called_once_with() + + clock.monotonic.return_value = 110 + client.get_bundles.return_value = [ready_state, ready_state.model_copy(update={"name": "missing"})] + client.publish_inventory.side_effect = None + revision = uuid7() + client.publish_inventory.return_value.revision = revision + find_files.return_value = files["missing"] + manager._refresh_dag_bundles(files) + + bundles["missing"].refresh.assert_called_once() + bundles["ready"].refresh.assert_not_called() + assert manager._bundle_parse_sources["missing"].bundle_revision == revision + assert files == previous_files + assert manager._bundle_versions == previous_versions + + +@pytest.mark.parametrize("status_code", [409, 503]) +def test_api_inventory_failure_preserves_known_files(status_code): + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + manager = DagFileProcessorManager(max_runs=1, api_client=client) + body = DagBundleInventoryBody(attempt_id=uuid7(), dispatch_sequence=1, expected_revision=None, files=[]) + manager._pending_inventories["testing"] = (body, set(), DagParseSource()) + known_files = {"testing": set(_get_file_infos(["previous.py"]))} + previous = known_files.copy() + response = httpx2.Response(status_code, request=httpx2.Request("POST", "http://api/inventory")) + client.publish_inventory.side_effect = httpx2.HTTPStatusError( + "Unavailable", request=response.request, response=response + ) + manager._publish_bundle_inventory("testing", known_files) + assert known_files == previous + assert not manager._bundle_parse_sources + if status_code == 409: + assert not manager._pending_inventories + assert "testing" in manager._force_refresh_bundles + else: + assert manager._pending_inventories["testing"][0] is body + revision = uuid7() + client.publish_inventory.side_effect = None + client.publish_inventory.return_value = DagBundleInventoryResponse( + attempt_id=body.attempt_id, accepted_at=timezone.utcnow(), revision=revision + ) + manager._publish_bundle_inventory("testing", known_files) + assert client.publish_inventory.call_args.args[1] is body + assert known_files == {"testing": set()} + assert manager._bundle_parse_sources["testing"].bundle_revision == revision + + +@pytest.mark.parametrize("rejection", ["local_limit", "server_limit"]) +@mock.patch("airflow.dag_processing.manager.MAX_PUBLICATION_BYTES", 1024) +@mock.patch("airflow.dag_processing.manager.time", autospec=True) +@mock.patch.object(DagFileProcessorManager, "_create_process", autospec=True) +def test_oversized_inventory_preserves_previous_inventory_and_other_bundle_progress( + create, clock, rejection, tmp_path +): + clock.monotonic.return_value = 100 + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + client.get_bundles.return_value = [ + DagBundleStateResponse(name=name, version=None, last_refreshed=None, revision=None, team_name=None) + for name in ("oversized", "ready") + ] + manager = DagFileProcessorManager(max_runs=1, api_client=client, bundle_refresh_check_interval=10) + old_file, ready_file = [ + DagFileInfo(bundle_name=name, rel_path=Path("a.py"), bundle_path=tmp_path) + for name in ("oversized", "ready") + ] + (tmp_path / "a.py").touch() + known_files = {"oversized": {old_file}} + manager._bundle_versions["oversized"] = "previous" + manager._force_refresh_bundles.add("oversized") + for name, files in ( + ("oversized", [f"{index}/{'é' * 80}.py" for index in range(6)] if rejection == "local_limit" else []), + ("ready", ["a.py"]), + ): + bundle = _make_bundle(tmp_path) + bundle.name = name + manager._dag_bundles.append(bundle) + body = DagBundleInventoryBody( + attempt_id=uuid7(), dispatch_sequence=1, expected_revision=None, version="new", files=files + ) + if name == "oversized" and rejection == "local_limit": + assert len(body.model_dump_json()) < 1024 < len(body.model_dump_json().encode()) + manager._pending_inventories[name] = ( + body, + {ready_file} if name == "ready" else set(), + DagParseSource(bundle_version="new"), + ) + + def publish(bundle_name, body): + if bundle_name == "oversized": + assert rejection == "server_limit", "The oversized inventory must be rejected before sending" + raise _make_status_error(413) + return DagBundleInventoryResponse( + attempt_id=body.attempt_id, accepted_at=timezone.utcnow(), revision=uuid7() + ) + + client.publish_inventory.side_effect = publish + manager._refresh_dag_bundles(known_files) + manager._start_new_processes() + + assert known_files == {"oversized": {old_file}, "ready": {ready_file}} + assert manager._bundle_versions == {"oversized": "previous", "ready": "new"} + assert not manager._pending_inventories + assert not manager._force_refresh_bundles + assert set(manager._bundle_parse_sources) == {"ready"} + assert list(manager._files_awaiting_source) == [old_file] + create.assert_called_once_with(manager, ready_file) + calls = client.get_bundles.call_count + manager._refresh_dag_bundles(known_files) + assert client.get_bundles.call_count == calls + + +@pytest.mark.parametrize("auth_status", [401, 403]) +def test_oversized_inventory_byte_boundary_preserves_auth_failures(auth_status): + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + manager = DagFileProcessorManager(max_runs=1, api_client=client) + body = DagBundleInventoryBody( + attempt_id=uuid7(), dispatch_sequence=1, expected_revision=None, files=["é.py"] + ) + pending = (body, set(), DagParseSource()) + manager._pending_inventories["testing"] = pending + client.publish_inventory.return_value = DagBundleInventoryResponse( + attempt_id=body.attempt_id, accepted_at=timezone.utcnow(), revision=uuid7() + ) + size = len(body.model_dump_json().encode()) + with mock.patch("airflow.dag_processing.manager.MAX_PUBLICATION_BYTES", size): + manager._publish_bundle_inventory("testing", {}) + client.publish_inventory.assert_called_once() + manager._pending_inventories["testing"] = pending + with mock.patch("airflow.dag_processing.manager.MAX_PUBLICATION_BYTES", size - 1): + manager._publish_bundle_inventory("testing", {}) + client.publish_inventory.assert_called_once() + assert not manager._pending_inventories + assert not manager._bundle_parse_sources + manager._pending_inventories["testing"] = pending + client.publish_inventory.side_effect = _make_status_error(auth_status) + with pytest.raises(httpx2.HTTPStatusError): + manager._publish_bundle_inventory("testing", {}) + + +def test_publication_for_a_changed_source_refreshes_and_reparses(): + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + manager = DagFileProcessorManager(max_runs=1, api_client=client) + file = _get_file_infos(["a.py"])[0] + manager._pending_publications[file] = PendingDagPublication( + body=mock.create_autospec(DagParseResultBody, instance=True), stat=DagFileStat(), refresh_generation=0 + ) + manager._pending_publications[file].body.attempt_id = uuid7() + response = httpx2.Response( + 409, + json={"detail": {"reason": "source_changed"}}, + request=httpx2.Request("POST", "http://api/parse-results"), + ) + client.publish_parse_result.side_effect = httpx2.HTTPStatusError( + "Rejected", request=response.request, response=response + ) + + manager._publish_pending_results() + + assert not manager._pending_publications + assert manager._force_refresh_bundles == {"testing"} + assert list(manager._file_queue) == [file] + + +@conf_vars({("core", "multi_team"): "True"}) +def test_api_team_lookup_failure_is_retried_without_stopping_the_processor(): + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + client.get_bundles.side_effect = [ + httpx2.ConnectError("unavailable"), + [ + DagBundleStateResponse( + name="testing", version=None, last_refreshed=None, revision=None, team_name="team_a" + ) + ], + ] + manager = DagFileProcessorManager(max_runs=1, api_client=client) + + assert manager._get_team_names({"testing"}) == {"testing": None} + assert manager._get_team_names({"testing"}) == {"testing": "team_a"} + + +def _make_pending_publication( + file: DagFileInfo, *, attempt_id: UUID | None = None, import_errors: dict[str, str] | None = None +) -> PendingDagPublication: + dag = SdkDAG(f"dag_in_{file.rel_path.stem}", schedule=None) + dag.fileloc = f"/worker/{file.rel_path}" + dag.relative_fileloc = str(file.rel_path) + return PendingDagPublication( + body=DagParseResultBody( + attempt_id=attempt_id or uuid7(), + dispatch_sequence=1, + bundle_name=file.bundle_name, + relative_fileloc=str(file.rel_path), + parse_duration=0.1, + serialized_dags=[LazyDeserializedDAG.from_dag(dag).data], + source_codes={dag.dag_id: ParseSourceCode(source_code="# source", language="python")}, + import_errors=import_errors or {}, + ), + stat=DagFileStat(num_dags=1, run_count=1, last_finish_time=timezone.utcnow()), + refresh_generation=0, + ) + + +def _make_status_error( + status_code: int, reason: str | None = None, *, message: str | None = None +) -> httpx2.HTTPStatusError: + response = httpx2.Response( + status_code, + json={"detail": {"reason": reason, "message": message}} if reason else None, + request=httpx2.Request("POST", "http://api/parse-results"), + ) + return httpx2.HTTPStatusError("Rejected", request=response.request, response=response) + + +def test_publication_sends_every_ready_result_within_the_loop_budget(): + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + manager = DagFileProcessorManager(max_runs=1, api_client=client) + files = _get_file_infos(["a.py", "b.py"]) + for file in files: + manager._pending_publications[file] = _make_pending_publication(file) + + manager._publish_pending_results(deadline=time.monotonic()) + client.publish_parse_result.assert_not_called() + + manager._publish_pending_results() + sent = [call.args[0].relative_fileloc for call in client.publish_parse_result.call_args_list] + assert sent == ["a.py", "b.py"] + assert not manager._pending_publications + assert all(manager._file_stats[file].num_dags == 1 for file in files) + + +@conf_vars({("workers", "execution_api_retries"): "1"}) +@pytest.mark.parametrize( + ("error", "message"), + [ + (_make_status_error(409, "dag_owned_by_another_bundle"), "HTTP 409: dag_owned_by_another_bundle"), + (_make_status_error(403, "plugin_not_granted"), "HTTP 403: plugin_not_granted"), + (_make_status_error(413), "HTTP 413"), + (_make_status_error(422), "HTTP 422"), + ( + _make_status_error(422, "invalid_dag", message="Unknown role: missing_role"), + "HTTP 422: invalid_dag: Unknown role: missing_role", + ), + ], +) +def test_result_the_api_refuses_is_replaced_by_an_import_error(error, message): + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + client.publish_parse_result.side_effect = [error, None] + manager = DagFileProcessorManager(max_runs=1, api_client=client) + manager._dispatch_sequence = 7 + file = _get_file_infos(["a.py"])[0] + original = _make_pending_publication( + file, + attempt_id=uuid7(), + import_errors={"a.py/member.py": "member failed"}, + ) + manager._pending_publications[file] = original + + manager._publish_pending_results() + + report = client.publish_parse_result.call_args_list[1].args[0] + assert report.attempt_id != original.body.attempt_id + assert report.dispatch_sequence == 8 + assert report.serialized_dags == [] + assert report.source_codes == {} + assert report.import_errors == { + "a.py": f"The parsing result was not published: the API server rejected it ({message}).", + } + assert not manager._pending_publications + assert manager._file_stats[file].num_dags == 0 + assert manager._file_stats[file].import_errors == 1 + + +@conf_vars({("workers", "execution_api_retries"): "1"}) +def test_server_failure_after_publication_retries_preserves_previous_result(): + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + client.publish_parse_result.side_effect = [_make_status_error(500), None] + manager = DagFileProcessorManager(max_runs=1, api_client=client) + file = _get_file_infos(["a.py"])[0] + previous = DagFileStat(num_dags=2, last_finish_time=DEFAULT_DATE) + manager._file_stats[file] = previous + pending = _make_pending_publication(file, attempt_id=uuid7()) + manager._pending_publications[file] = pending + + manager._publish_pending_results() + + client.publish_parse_result.assert_called_once_with(pending.body) + assert not manager._pending_publications + assert manager._file_stats[file].num_dags == previous.num_dags + assert manager._file_stats[file].last_finish_time == previous.last_finish_time + assert manager._file_stats[file].last_attempt_time is not None + + +@mock.patch("airflow.dag_processing.manager.MAX_PUBLICATION_BYTES", 32 * 1024) +@pytest.mark.parametrize("oversized_field", ["import_errors", "version_data", "parsed_definitions"]) +def test_rejection_report_discards_oversized_parse_metadata(oversized_field): + limit = 32 * 1024 + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + client.publish_parse_result.side_effect = [_make_status_error(413), None] + manager = DagFileProcessorManager(max_runs=1, api_client=client) + file = _get_file_infos(["a.py"])[0] + original = _make_pending_publication(file) + original.body.bundle_version = "v1" + original.body.bundle_revision = uuid7() + oversized = { + "import_errors": {"a.py/member.py": "é" * limit}, + "version_data": {"metadata": "é" * limit}, + "parsed_definitions": ["a.py/member.py"] * limit, + } + setattr(original.body, oversized_field, oversized[oversized_field]) + assert len(original.body.model_dump_json().encode()) > limit + manager._pending_publications[file] = original + + manager._publish_pending_results() + + report = client.publish_parse_result.call_args_list[1].args[0] + assert len(report.model_dump_json().encode()) < limit + assert DagParseResultBody.model_validate_json(report.model_dump_json()) == report + assert report.bundle_version == original.body.bundle_version + assert report.bundle_revision == original.body.bundle_revision + assert report.version_data is None + assert report.parsed_definitions == [] + assert list(report.import_errors) == ["a.py"] + assert manager._file_stats[file].import_errors == 1 + assert not manager._pending_publications + + +@mock.patch("airflow.dag_processing.manager.MAX_PUBLICATION_BYTES", 32 * 1024) +def test_rejection_report_bounds_server_validation_message(): + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + message = "Unknown role: " + "é" * (32 * 1024) + client.publish_parse_result.side_effect = [_make_status_error(422, "invalid_dag", message=message), None] + manager = DagFileProcessorManager(max_runs=1, api_client=client) + file = _get_file_infos(["a.py"])[0] + manager._pending_publications[file] = _make_pending_publication(file) + + manager._publish_pending_results() + + report = client.publish_parse_result.call_args_list[1].args[0] + assert "Unknown role: é" in report.import_errors["a.py"] + assert len(report.model_dump_json().encode()) < 32 * 1024 + + +def test_direct_manager_does_not_import_api_client(monkeypatch, mocker): + monkeypatch.setitem(sys.modules, "airflow.dag_processing.api_client", None) + namespace = runpy.run_path(sys.modules[DagFileProcessorManager.__module__].__file__) + manager_class = namespace["DagFileProcessorManager"] + bundle_manager = mocker.patch.object(manager_class, "_create_bundle_manager", autospec=True) + bundle_manager.return_value.get_bundle.side_effect = ValueError("missing bundle") + manager = manager_class(max_runs=1) + + manager._publish_pending_results() + manager.after_run() + request = DagCallbackRequest( + filepath="a.py", + dag_id="dag1", + run_id="run1", + is_failure_callback=False, + bundle_name="testing", + bundle_version="v1", + ) + assert manager.prepare_callback_bundle(request) is None + + +@pytest.mark.parametrize("reason", ["bundle_inactive", "publication_superseded", None]) +def test_refusal_unrelated_to_content_is_not_reported_as_an_import_error(reason): + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + client.publish_parse_result.side_effect = _make_status_error(409, reason) + manager = DagFileProcessorManager(max_runs=1, api_client=client) + file = _get_file_infos(["a.py"])[0] + manager._pending_publications[file] = _make_pending_publication(file) + + manager._publish_pending_results() + + client.publish_parse_result.assert_called_once() + assert not manager._pending_publications + + +def test_refused_import_error_report_is_not_replaced_again(): + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + client.publish_parse_result.side_effect = _make_status_error(413) + manager = DagFileProcessorManager(max_runs=1, api_client=client) + file = _get_file_infos(["a.py"])[0] + manager._pending_publications[file] = _make_pending_publication(file) + + manager._publish_pending_results() + + assert client.publish_parse_result.call_count == 2 + assert not manager._pending_publications + assert manager._file_stats[file].num_dags == 0 + + +@mock.patch("airflow.dag_processing.manager.MAX_PUBLICATION_BYTES", 1024 * 1024) +def test_result_over_the_publication_limit_is_reported_without_sending_it(): + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + manager = DagFileProcessorManager(max_runs=1, api_client=client) + file = _get_file_infos(["a.py"])[0] + dag = SdkDAG("large", schedule=None, doc_md="x" * 1024 * 1024) + dag.fileloc = "/worker/a.py" + dag.relative_fileloc = "a.py" + proc = mock.create_autospec(DagFileProcessorProcess, instance=True) + proc.id = uuid7() + proc.start_time = time.monotonic() + proc.had_callbacks = False + proc.dispatch_sequence = 1 + proc.parse_source = DagParseSource() + proc.parsing_result = DagFileParsingResult( + fileloc="a.py", serialized_dags=[LazyDeserializedDAG.from_dag(dag)] + ) + + manager.handle_parsing_result(file, proc) + + pending = manager._pending_publications[file] + assert pending.is_rejection_report + assert pending.body.serialized_dags == [] + assert pending.body.import_errors == { + "a.py": "The parsing result was not published: it exceeds the 1 MiB publication limit." + } + client.publish_parse_result.assert_not_called() + + +@pytest.mark.parametrize("method", ["claim_priority_files", "fetch_callbacks"]) +def test_internal_api_manager_leaves_requested_work_for_activation(method, mocker): + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + manager = DagFileProcessorManager(max_runs=1, api_client=client) + database_method = ( + "_claim_priority_files" if method == "claim_priority_files" else "_fetch_callbacks_from_db" + ) + read_database = mocker.patch.object(manager, database_method, autospec=True) + + assert getattr(manager, method)() == [] + read_database.assert_not_called()
diff --git a/airflow-core/tests/unit/dag_processing/test_processor.py b/airflow-core/tests/unit/dag_processing/test_processor.py index 629cb6b..d00fe91 100644 --- a/airflow-core/tests/unit/dag_processing/test_processor.py +++ b/airflow-core/tests/unit/dag_processing/test_processor.py
@@ -18,6 +18,7 @@ from __future__ import annotations import inspect +import json import logging import os import pathlib @@ -29,8 +30,10 @@ from collections.abc import Callable, Iterable from socket import socket, socketpair from typing import TYPE_CHECKING, Any, BinaryIO -from unittest.mock import MagicMock, PropertyMock, patch +from unittest.mock import ANY, MagicMock, PropertyMock, patch +import httpx2 +import jwt import pytest import structlog from pydantic import TypeAdapter @@ -51,6 +54,7 @@ EmailRequest, TaskCallbackRequest, ) +from airflow.dag_processing.api_client import DagProcessorAPIClient from airflow.dag_processing.dagbag import DagBag from airflow.dag_processing.manager import DagFileProcessorManager, process_parse_results from airflow.dag_processing.processor import ( @@ -85,6 +89,7 @@ TICount, ToSupervisor, ToTask, + VariableResult, XComResult, XComSequenceSliceResult, ) @@ -2414,7 +2419,11 @@ @pytest.mark.parametrize( "message_type", - sorted(set(typing.get_args(typing.get_args(ToManager)[0])) - {DagFileParsingResult}, key=str), + sorted( + set(typing.get_args(typing.get_args(ToManager)[0])) + - {DagFileParsingResult, comms.GetPrevSuccessfulDagRun}, + key=str, + ), ids=lambda message_type: message_type.__name__, ) def test_reuses_shared_request_handlers(self, message_type): @@ -2480,15 +2489,92 @@ send_msg.assert_called_once_with(proc, None, request_id=42, error=None) @patch.object(DagFileProcessorProcess, "send_msg", autospec=True) - def test_previous_successful_run_uses_process_id(self, send_msg, proc): - proc.client.task_instances.get_previous_successful_dagrun.return_value = ( - comms.PrevSuccessfulDagRunResult() - ) + def test_previous_successful_run_is_answered_locally(self, send_msg, proc): proc._handle_request( comms.GetPrevSuccessfulDagRun(ti_id=uuid.uuid4()), structlog.get_logger(), req_id=42 ) - proc.client.task_instances.get_previous_successful_dagrun.assert_called_once_with(proc.id) + assert not proc.client.mock_calls + send_msg.assert_called_once_with( + proc, comms.PrevSuccessfulDagRunResult(), request_id=42, error=None, exclude_unset=True + ) + + @patch.object(supervisor.WatchedSubprocess, "_handle_request", autospec=True) + def test_default_mode_requests_do_not_load_the_api_client(self, handle_request, proc, monkeypatch): + monkeypatch.setitem(sys.modules, "airflow.dag_processing.api_client", None) + msg = GetVariable(key="my_key") + + proc._handle_request(msg, structlog.get_logger(), req_id=42) + + handle_request.assert_called_once_with(proc, msg, ANY, 42) + + @patch.object(DagFileProcessorProcess, "send_msg", autospec=True) + def test_api_client_requests_are_made_for_the_files_bundle(self, send_msg, proc, tmp_path): + token_file = tmp_path / "processor.jwt" + token_file.write_text("session-token") + job_token = jwt.encode({"iat": 0, "exp": 300, "scope": "dag_processor", "job_id": 1}, "secret") + parse_token = jwt.encode( + { + "iat": 0, + "exp": 300, + "sub": str(proc.id), + "scope": "dag_parse", + "job_id": 1, + "dag_bundles": [proc.bundle_name], + "relative_fileloc": proc.dag_file_rel_path, + }, + "secret", + ) + requests = [] + + def handle(request): + requests.append(request) + if request.url.path.endswith("/jobs"): + return httpx2.Response(201, json={"job_id": 1, "token": job_token, "expires_in": 300}) + if request.url.path.endswith("/parse-token"): + return httpx2.Response(200, json={"token": parse_token, "expires_in": 300}) + return httpx2.Response(200, json={"key": "my_key", "value": "my_value"}) + + with DagProcessorAPIClient( + base_url="http://api/execution/", + token_file=token_file, + hostname="processor-1", + transport=httpx2.MockTransport(handle), + ) as client: + client.register_job() + proc.client = client + proc._handle_request(GetVariable(key="my_key"), structlog.get_logger(), req_id=42) + proc._handle_request(GetVariable(key="my_key"), structlog.get_logger(), req_id=43) + + assert requests[-1].url.path == "/execution/variables/my_key" + assert requests[-1].headers["Airflow-Dag-Bundle"] == "mybundle" + assert requests[-1].headers["Authorization"] == f"Bearer {parse_token}" + assert json.loads(requests[1].content) == { + "attempt_id": str(proc.id), + "bundle_name": proc.bundle_name, + "relative_fileloc": proc.dag_file_rel_path, + } + assert sum(request.url.path.endswith("/parse-token") for request in requests) == 1 + assert send_msg.call_count == 2 + send_msg.assert_any_call( + proc, + VariableResult(key="my_key", value="my_value"), + request_id=42, + error=None, + exclude_unset=True, + ) + + @patch.object(DagFileProcessorProcess, "send_msg", autospec=True) + def test_local_parse_messages_do_not_require_a_token_exchange(self, send_msg, proc, tmp_path): + with DagProcessorAPIClient( + base_url="http://unused/", token_file=tmp_path / "missing-token", hostname="processor-1" + ) as client: + proc.client = client + proc._handle_request( + comms.GetPrevSuccessfulDagRun(ti_id=uuid.uuid4()), structlog.get_logger(), req_id=42 + ) + + assert client.job_id is None send_msg.assert_called_once_with( proc, comms.PrevSuccessfulDagRunResult(), request_id=42, error=None, exclude_unset=True )
diff --git a/generated/known_sdk_imports_in_core.txt b/generated/known_sdk_imports_in_core.txt index 91d224d..9c761d5 100644 --- a/generated/known_sdk_imports_in_core.txt +++ b/generated/known_sdk_imports_in_core.txt
@@ -7,7 +7,7 @@ airflow-core/src/airflow/dag_processing/dagbag.py::3 airflow-core/src/airflow/dag_processing/importer_routing.py::1 airflow-core/src/airflow/dag_processing/lang_sdk_processor.py::6 -airflow-core/src/airflow/dag_processing/manager.py::5 +airflow-core/src/airflow/dag_processing/manager.py::7 airflow-core/src/airflow/dag_processing/processor.py::14 airflow-core/src/airflow/exceptions.py::1 airflow-core/src/airflow/executors/base_executor.py::3
diff --git a/task-sdk/src/airflow/sdk/execution_time/cache.py b/task-sdk/src/airflow/sdk/execution_time/cache.py index 0c220fd..1d1910b 100644 --- a/task-sdk/src/airflow/sdk/execution_time/cache.py +++ b/task-sdk/src/airflow/sdk/execution_time/cache.py
@@ -97,7 +97,7 @@ @classmethod def reset(cls): - """Use for test purposes only.""" + """Detach this process from the shared cache, including entries inherited across a fork.""" cls._cache = None @classmethod