You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by ep...@apache.org on 2026/10/15 10:47:06 UTC
(airflow) branch dag-processor-activation created (now 05a73311d22)
This is an automated email from the ASF dual-hosted git repository.
ephraimbuddy pushed a change to branch dag-processor-activation
in repository https://gitbox.apache.org/repos/asf/airflow.git
at 05a73311d22 Run Dag processors without metadata database access
This branch includes the following new commits:
new 05a73311d22 Run Dag processors without metadata database access
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
(airflow) 01/01: Run Dag processors without metadata database access
Posted by ep...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
ephraimbuddy pushed a commit to branch dag-processor-activation
in repository https://gitbox.apache.org/repos/asf/airflow.git
commit 05a73311d22ce06258761336b75e4a161a550a1a
Author: Ephraim Anierobi <sp...@gmail.com>
AuthorDate: Thu Oct 15 11:23:24 2026 +0100
Run Dag processors without metadata database access
Standalone processors need to resolve bundle credentials and publish
parsed metadata without inheriting database or signing credentials.
API-managed Jobs also need lifecycle and requested-work recovery so
transient outages do not lose callbacks or priority parses.
---
.../security/dag-processor-api/docker-compose.yaml | 53 ++
.../docs/security/jwt_token_authentication.rst | 196 ++++-
airflow-core/docs/security/security_model.rst | 15 +-
airflow-core/src/airflow/cli/cli_config.py | 15 +
.../airflow/cli/commands/dag_processor_command.py | 66 +-
.../src/airflow/config_templates/config.yml | 40 +-
.../src/airflow/dag_processing/api_client.py | 44 +-
airflow-core/src/airflow/dag_processing/health.py | 56 ++
airflow-core/src/airflow/dag_processing/manager.py | 295 ++++++-
.../src/airflow/dag_processing/processor.py | 7 +-
.../src/airflow/jobs/dag_processor_job_runner.py | 94 +-
airflow-core/src/airflow/utils/cli.py | 16 +-
.../dag_processing/test_api_isolation.py | 306 +++++++
.../execution_api/test_dag_processor_client.py | 150 +++-
.../cli/commands/test_dag_processor_command.py | 136 +++
.../tests/unit/dag_processing/test_health.py | 64 ++
.../tests/unit/dag_processing/test_manager.py | 947 ++++++++++++++++++++-
.../unit/jobs/test_dag_processor_job_runner.py | 275 ++++++
.../observability/metrics/metrics_template.yaml | 32 +-
19 files changed, 2723 insertions(+), 84 deletions(-)
diff --git a/airflow-core/docs/security/dag-processor-api/docker-compose.yaml b/airflow-core/docs/security/dag-processor-api/docker-compose.yaml
new file mode 100644
index 00000000000..a73b114243f
--- /dev/null
+++ b/airflow-core/docs/security/dag-processor-api/docker-compose.yaml
@@ -0,0 +1,53 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+---
+# Run independently of the quick-start Compose file so its credentials are not inherited.
+services:
+ dag-processor:
+ image: ${AIRFLOW_IMAGE_NAME:?Set an Airflow image with processor API support}
+ command: [dag-processor]
+ user: "${AIRFLOW_UID:-50000}:0"
+ environment:
+ USER: airflow
+ AIRFLOW_CONFIG: /tmp/airflow.cfg
+ AIRFLOW__CORE__LOAD_EXAMPLES: "False"
+ AIRFLOW__CORE__EXECUTION_API_SERVER_URL: http://airflow-apiserver:8080/execution/
+ AIRFLOW__DAG_PROCESSOR__EXECUTION_API_TOKEN_FILE: /run/dag-processor/processor.jwt
+ AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_STORAGE_PATH: /tmp/dag-bundles
+ AIRFLOW__DAG_PROCESSOR__HEALTH_CHECK_FILE: /tmp/dag-processor-health.json
+ CONNECTION_CHECK_MAX_COUNT: "0"
+ volumes:
+ - ./tokens:/run/dag-processor:ro
+ - ./dags:/opt/airflow/dags:ro
+ - ./logs:/opt/airflow/logs
+ read_only: true
+ tmpfs:
+ - /tmp
+ cap_drop: [ALL]
+ security_opt: [no-new-privileges:true]
+ healthcheck:
+ test: [CMD, airflow, dag-processor, --check-health]
+ interval: 30s
+ timeout: 10s
+ start_period: 30s
+ retries: 5
+ networks: [processor-api]
+ restart: unless-stopped
+networks:
+ processor-api:
+ name: ${DAG_PROCESSOR_API_NETWORK:-dag-processor-api}
+ external: true
diff --git a/airflow-core/docs/security/jwt_token_authentication.rst b/airflow-core/docs/security/jwt_token_authentication.rst
index e72f9d3faa6..e978e054bf2 100644
--- a/airflow-core/docs/security/jwt_token_authentication.rst
+++ b/airflow-core/docs/security/jwt_token_authentication.rst
@@ -471,8 +471,10 @@ Dag File Processor and Triggerer
---------------------------------
The **Dag File Processor** and **Triggerer** are internal Airflow components that also
-interact with the Execution API, but they do so via an **in-process** transport
-(``InProcessExecutionAPI``) rather than over the network. This in-process API:
+interact with the Execution API, but by default they do so via an **in-process** transport
+(``InProcessExecutionAPI``) rather than over the network. A Dag File Processor can instead use
+authenticated HTTP without database access; see `Dag processor HTTP client credentials`_. This
+in-process API:
- Runs the Execution API application directly within the same process, using an ASGI/WSGI
bridge.
@@ -514,16 +516,18 @@ guidance, and the planned strategic and tactical improvements.
Dag processor HTTP client credentials
-------------------------------------
-The core ``DagProcessorAPIClient`` lets processor integrations authenticate to the Execution
-API without holding its signing key. The standard ``airflow dag-processor`` command does
-not yet use this client. Provisioning a token file alone does not switch that command to
-HTTP or remove its need for database access.
+Set ``[dag_processor] execution_api_token_file`` to run ``airflow dag-processor`` with
+authenticated HTTP requests. Job registration, heartbeats, completion, bundle secret
+lookups, bundle metadata, result publication, and callback and priority delivery then use
+the Execution API. The API client is created
+in the processor process, after any daemon or hot-reload fork. Leaving the setting
+unset preserves the in-process API and database-backed Job lifecycle.
-On a trusted host with the API signing key, provision a session token for the bundles
-the client may access::
+On a trusted host with database access and the server keys, synchronize the configured
+bundle catalog and provision a separate token file for each processor::
- airflow dag-processor-token --token-file /run/airflow/processor.jwt \
- --bundle-name dags-folder --rotate
+ airflow dag-processor --sync-bundles-only
+ airflow dag-processor-token --token-file /run/airflow/processor.jwt --bundle-name dags-folder --rotate
The command writes the token atomically. A new file is readable only by its owner, and replacing
an existing file keeps its mode and group, so grant the processor read access once. ACLs are not
@@ -532,14 +536,30 @@ token replacement stays visible. Each processor process needs its own session. R
provisioner with the same token file and bundles keeps that session, so running processors keep
their Jobs.
+Configure ``[core] execution_api_server_url`` to reach the API server, and start the
+processor::
+
+ AIRFLOW__DAG_PROCESSOR__EXECUTION_API_TOKEN_FILE=/run/airflow/processor.jwt \
+ airflow dag-processor --bundle-name dags-folder
+
+The processor does not need metadata database credentials, the Fernet key, or API signing
+keys. Remove them from its environment and mounted configuration, and block its database
+network path. Enabling the setting alone does not remove existing credentials or isolate
+code from other secrets on the processor host. Catalog synchronization is a trusted
+deployment operation: repeat it when bundle configuration changes. It loads provider
+classes for bundle URLs; the API server does not construct bundles or access their sources.
+The processor does not use the shared SDK secret cache in this mode, even when
+``[secrets] use_cache`` is enabled, because SDK cache lookups do not include the bundle
+identity that the API uses for authorization.
+
The three credentials have different issuers and permissions:
- ``dag_processor_session`` is issued by trusted provisioning. It can only register a
Job or renew that Job's credential through ``POST /jobs``.
- ``dag_processor`` is returned by Job registration. It can heartbeat and complete that
- Job, read Connections and Variables for a granted bundle, and exchange a parsing
- credential through ``POST /jobs/{job_id}/parse-token``. It cannot outlive the session
- token used to obtain it.
+ Job, read Connections and Variables for a granted bundle, publish parsing results and
+ bundle inventories, claim requested work, and exchange a parsing credential through
+ ``POST /jobs/{job_id}/parse-token``. It cannot outlive the session token used to obtain it.
- ``dag_parse`` is returned by the parsing-token exchange for one attempt, bundle, and
relative file location. It permits parse-time requests, including callback-context reads,
but cannot manage Jobs or exchange tokens. It cannot outlive its parent Job credential.
@@ -557,14 +577,148 @@ parsing subprocess can without this client. ``Variable.set`` upserts on the key
multi-team mode a write can overwrite another team's Variable with the same key, as it can with task
tokens.
-The client attempts renewal at 80% of token lifetime. Transient early-renewal failures leave
-the existing token usable until expiry. A retired registration signals that its owner must
-restart. Completion retains its original identity and outcome across retries.
-
-File and attempt identity are carried in the token but are not currently used for access
-checks or request attribution. Archives retain one file identity under the current processor
-model. These credentials do not isolate hostile code from other credentials accessible on
-the processor host; the process and deployment limitations above still apply.
+For each ``dag_parse`` token, the manager chooses the bundle, bundle-relative file location,
+and parsing-attempt UUID.
+
+The manager publishes completed imports through ``POST /jobs/{job_id}/parse-results``
+using its Job credential. Parsing tokens cannot publish. A request contains one file or
+container's serialized Dags, diagnostics, source text (or an explicit unavailable marker),
+and the bundle version captured at dispatch. The API never reads the processor's source
+paths. It checks bundle access, existing Dag ownership and team-scoped plugin classes
+before deserializing results and rechecks ownership on the rows used for persistence.
+A Dag can move between files within its bundle. Recovery from a removed bundle is
+limited to stale legacy rows with neither a relative file location nor a Dag version.
+
+Each request carries an attempt UUID and a dispatch sequence within its Job. A checkpoint
+and the metadata writes commit together. Replaying the latest identical request returns
+the saved receipt; conflicting payloads and superseded sequences are rejected. Checkpoints
+retain one entry per Job and source and are removed with Job cleanup. There is no total
+ordering between different Jobs publishing a source without versioning; an already-accepted
+request's retry does not overwrite another Job's result. Listener side effects are not
+transactional and may repeat after a failed transaction.
+
+Publication and inventory request bodies are limited to 16 MiB, and the API refuses a
+larger body before reading the rest of it. A rejected publication does not fall back
+to direct database writes, and does not prevent other files from being processed. Transport
+failures and server errors are retried; exhausting those retries preserves the last accepted
+result. When the API refuses a result for its content, or the result is over the limit, the
+processor publishes an import error for the file instead, so the problem is visible.
+Dags whose definition path matches an import error are marked stale until a later parse
+publishes them again. For a container-level error, missing member Dags are marked stale
+only when their last successful parse is older than ``[dag_processor] stale_dag_threshold``,
+using the server's clock. Invalid ``access_control`` produces an import error with both
+current and supported older FAB releases. Permission backend failures remain retryable
+server errors. Empty imports still publish a receipt and clear the file's earlier import
+errors. A bundle refresh, including one forced by a parse request, waits for that bundle's
+active imports instead of discarding them; only initializing a bundle again invalidates and
+requeues older in-flight results from that manager.
+
+Inventories report discovered definitions and sources whose discovery failed separately.
+Failed sources preserve their existing definitions and diagnostics, including archive members;
+other sources still undergo deletion reconciliation. Each inventory carries the previously
+observed server revision; conflicting refreshes must
+rediscover before retrying. Imports carry the accepted revision, so an older snapshot cannot
+publish after a changed inventory. Only accepted publications establish missing Dags within
+a file; rejected requests and incomplete discovery do not establish absence. Every
+``[scheduler] parsing_cleanup_interval``, the scheduler marks the Dags of inactive or unassigned
+bundles stale, removes stale warnings, and ends expired processor Jobs, in bounded batches and in
+both modes.
+
+Callback and priority requests remain stored while claimed. A lost claim response or
+acknowledgment can be retried, and requests from a retired Job become available to another
+processor. A callback acknowledgment confirms delivery, not successful execution of user
+code. As with other at-least-once delivery, a crash after a callback's side effect and before
+acknowledgment can repeat that side effect. A graceful shutdown acknowledges finished work
+before the processor's Job ends. When a callback's bundle version cannot be
+prepared, the processor retries with backoff for up to ``[dag_processor] job_heartbeat_timeout``
+and then acknowledges the callback as failed. Priority requests are assigned to an import
+started after their requested refresh. Only that import's publication or terminal failure
+acknowledges the request; an older in-flight import cannot consume it. Discovery after the
+requested refresh can also acknowledge a request if it confirms the definition is absent.
+
+The manager retains a bounded queue of pending publications. Each loop it sends queued
+publications and acknowledgments for up to a second, stopping at the first failure that may
+be transient. Connecting is limited to a second, but a publication or inventory waits up to
+``[workers] execution_api_timeout`` for its response, so a slow write can still be
+acknowledged. Retries keep the same payload and identity while
+heartbeats and subprocess supervision continue between attempts. Pending results
+count against parsing capacity. Failed publication attempts throttle reparsing but
+do not advance the timestamps used to detect stale Dags.
+
+The returned ``dag_parse`` token binds these identifiers to the processor's session and
+Job. It authorizes parse-time operations within that bundle and has no Job-management or
+token-exchange access. A bundle header cannot override its signed bundle. The token expires
+no later than the management credential; ending or replacing the Job invalidates it.
+The supervisor caches it for the attempt and attempts renewal at 80% of its lifetime.
+Transient early-renewal failures leave the existing token usable until expiry.
+
+Parse-time requests use ``[workers] execution_api_timeout`` (five seconds by default),
+unless the request specifies a timeout. Required parsing-token exchanges and expired
+Job-token renewals use the same timeout. Each makes one attempt without retry sleeps in
+the manager loop; optional early token renewal caps each HTTP timeout phase at one second.
+
+Heartbeats use ``[scheduler] job_heartbeat_sec`` and are also sent between bundle
+operations and before each callback's bundle is prepared. Transient API failures are
+retried on the next heartbeat interval. Once a Job has gone
+``[dag_processor] job_heartbeat_timeout`` without a recorded heartbeat, the scheduler's next cleanup (every
+``[scheduler] parsing_cleanup_interval``) ends it and releases its claimed work, and a processor
+that could not reach the API for that long stops. Set the timeout above the longest single
+bundle initialization, refresh, or listing, and set it for the scheduler as well as the
+processor: the scheduler's value decides when the Job ends. The processor does not update its
+local health file during one of those operations either, so ``health_check_threshold`` and the
+liveness probe's failure budget must also cover it. Likewise, the API server reads
+``[dag_processor] health_check_threshold`` to decide whether a session's previous Job is still
+alive, so keep it the same for the API server and the processor. A restart request, replaced Job, or retired
+registration also stops parsing and terminates its children. The deployment's process
+supervisor must restart the command, which creates a new registration. After a crash,
+registration retries ``job_running`` with backoff for up to the health-check threshold
+plus one heartbeat interval, retaining the same registration ID while waiting for the old
+Job to expire.
+Failed runs attempt to record failure without masking the original error; completion
+failures after a successful run are reported to the caller.
+
+Use ``airflow dag-processor --check-health`` for local liveness and ``--check-ready`` for
+local liveness plus a recent API heartbeat. Both read ``[dag_processor] health_check_file``
+without querying the database. Give each processor its own writable health file. A paused
+manager fails liveness even while the API remains healthy; a transient API outage affects
+readiness without immediately failing liveness.
+
+Set ``CONNECTION_CHECK_MAX_COUNT=0`` in the processor container too: the standard image
+entrypoint otherwise checks the database before starting the command. Do not set
+``_AIRFLOW_DB_MIGRATE`` or ``_AIRFLOW_WWW_USER_CREATE`` in that container.
+
+Isolated deployment examples
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+The following examples run only the processor, alongside an existing trusted deployment.
+Use an image containing processor API support; this feature targets Airflow 3.5.
+Synchronize the bundle catalog and run token provisioning in the trusted deployment as
+described above. Keep the token provisioner running, and mount its directory rather than
+a single file so atomic token replacement remains visible.
+
+For Compose, use the :download:`standalone processor definition <dag-processor-api/docker-compose.yaml>`
+as its own project, rather than merging it with the quick-start file. It has no shared
+``env_file``, server configuration mount, or dependency on a database/migration service.
+Create its ``tokens``, ``dags``, and ``logs`` directories with permissions for the configured
+container user; provision ``tokens/processor.jwt``. Set ``AIRFLOW_IMAGE_NAME`` to the image
+with API support. Create an internal network and attach only the API server and processor::
+
+ docker network create --internal dag-processor-api
+ docker network connect --alias airflow-apiserver dag-processor-api <api-server-container>
+ docker compose -p dag-processor-api -f docker-compose.yaml up -d
+
+The API server keeps its existing database network. The processor joins only
+``dag-processor-api``. Verify readiness with ``docker compose -p dag-processor-api
+-f docker-compose.yaml exec dag-processor airflow dag-processor --check-ready``.
+The examples do not grant provider credentials through environment variables; bundle
+Connection lookups use the API. Review any image customizations, added mounts, or environment
+settings for server credentials before deployment.
+
+The token carries file and attempt identity, but the API does not currently use them
+for access checks or request attribution. Connection and Variable permissions remain
+bundle/team based. Archives retain one file identity under the current processor model.
+This exchange does not isolate hostile code from credentials accessible on the processor
+host. The process and deployment limitations described above still apply.
Workload Isolation and Current Limitations
diff --git a/airflow-core/docs/security/security_model.rst b/airflow-core/docs/security/security_model.rst
index 2ae6144f409..926407720a1 100644
--- a/airflow-core/docs/security/security_model.rst
+++ b/airflow-core/docs/security/security_model.rst
@@ -496,6 +496,14 @@ potentially still executes with direct database access in the Dag File Processor
the database directly, including the JWT signing key configuration if it is available in the
process environment. If a Dag author obtains the JWT signing key, they could forge arbitrary tokens.
+ Since Airflow 3.5 the Dag File Processor can instead run without metadata database access, by
+ setting ``[dag_processor] execution_api_token_file``. It then uses short-lived Execution API
+ credentials limited to its granted bundles, and each parsing attempt gets a bundle-scoped token
+ carrying its file and attempt identity. This mode removes the processor's need for database
+ credentials, the Fernet key and signing keys, but does not remove them: Deployment Managers must
+ take them out of the processor's environment and block its database network path. The Triggerer
+ still needs database access. See :doc:`/security/jwt_token_authentication`.
+
**Dag File Processor and Triggerer are shared across teams**
In the default deployment, a **single Dag File Processor instance** parses all Dag files and a
**single Triggerer instance** handles all triggers — regardless of team assignment. There is no
@@ -772,6 +780,10 @@ model — Airflow does not enforce these natively.
value itself. However, access to the file it points to should be restricted to the Scheduler
(which generates tokens) and the API Server (which validates them).
+ A Dag File Processor running with ``[dag_processor] execution_api_token_file`` needs none of the
+ metadata database variables, the Fernet key or the JWT signing secret listed for it above; it
+ authenticates with the token file instead.
+
.. START AUTOGENERATED PROVIDER SENSITIVE VARS
**Provider-specific sensitive configuration variables:**
@@ -845,7 +857,8 @@ model — Airflow does not enforce these natively.
the metadata database exclusively through the API server (similar to how workers use the
Execution API today). This would eliminate the need for these components to have database
credentials at all, providing security by design rather than relying on deployment-level
- measures.
+ measures. Airflow 3.5 offers this for the Dag File Processor as an opt-in mode
+ (``[dag_processor] execution_api_token_file``); the Triggerer still uses the database.
* **Tactical (shorter-term)**: Native support for Unix user impersonation in the Dag File
Processor and Triggerer child processes, so that Dag author code runs as a different, low-
privilege user that cannot access the parent's credentials or the database.
diff --git a/airflow-core/src/airflow/cli/cli_config.py b/airflow-core/src/airflow/cli/cli_config.py
index aec3ec81dfa..b6e9f2e6a66 100644
--- a/airflow-core/src/airflow/cli/cli_config.py
+++ b/airflow-core/src/airflow/cli/cli_config.py
@@ -2304,6 +2304,21 @@ core_commands: list[CLICommand] = [
help="Start a dag processor instance",
func=lazy_load_command("airflow.cli.commands.dag_processor_command.dag_processor"),
args=(
+ Arg(
+ ("--check-health",),
+ action="store_true",
+ help="Check the API-mode processor's local loop heartbeat and exit",
+ ),
+ Arg(
+ ("--check-ready",),
+ action="store_true",
+ help="Check the API-mode processor's loop and API heartbeat and exit",
+ ),
+ Arg(
+ ("--sync-bundles-only",),
+ action="store_true",
+ help="Provision the configured bundle catalog using trusted database and signing credentials, then exit",
+ ),
ARG_PID,
ARG_DAEMON,
ARG_BUNDLE_NAME,
diff --git a/airflow-core/src/airflow/cli/commands/dag_processor_command.py b/airflow-core/src/airflow/cli/commands/dag_processor_command.py
index b8195bead28..b49a4d66623 100644
--- a/airflow-core/src/airflow/cli/commands/dag_processor_command.py
+++ b/airflow-core/src/airflow/cli/commands/dag_processor_command.py
@@ -19,19 +19,28 @@
from __future__ import annotations
import logging
-from typing import Any
+from typing import TYPE_CHECKING, Any
from airflow.cli.commands.daemon_utils import run_command_with_daemon_option
from airflow.configuration import conf
-from airflow.dag_processing.bundles.manager import get_configured_bundle_team_names
+from airflow.dag_processing.bundles.manager import (
+ _load_bundle_config_snapshot,
+ get_configured_bundle_team_names,
+)
+from airflow.dag_processing.health import check_processor_health
from airflow.dag_processing.manager import DagFileProcessorManager
+from airflow.executors.base_executor import get_execution_api_server_url
from airflow.jobs.dag_processor_job_runner import DagProcessorJobRunner
from airflow.jobs.job import Job, run_job
from airflow.utils import cli as cli_utils
from airflow.utils.memray_utils import MemrayTraceComponents, enable_memray_trace
+from airflow.utils.net import get_hostname
from airflow.utils.process_utils import set_component_mp_start_method
from airflow.utils.providers_configuration_loader import providers_configuration_loaded
+if TYPE_CHECKING:
+ from collections.abc import Callable
+
log = logging.getLogger(__name__)
@@ -60,7 +69,12 @@ def _get_team_names(bundle_names: list[str] | None) -> list[str]:
def _create_dag_processor_job_runner(args: Any) -> DagProcessorJobRunner:
"""Create DagFileProcessorProcess instance."""
if args.bundle_name:
- cli_utils.validate_dag_bundle_arg(args.bundle_name)
+ if conf.get("dag_processor", "execution_api_token_file", fallback=None):
+ # Bundle constructors can read secrets; defer them until the API client is registered.
+ if unknown := set(args.bundle_name) - _load_bundle_config_snapshot().names:
+ raise SystemExit(f"Bundles not found: {', '.join(sorted(unknown))}")
+ else:
+ cli_utils.validate_dag_bundle_arg(args.bundle_name)
return DagProcessorJobRunner(
job=Job(bundle_names=args.bundle_name, team_names=_get_team_names(args.bundle_name)),
processor=DagFileProcessorManager(
@@ -72,13 +86,55 @@ def _create_dag_processor_job_runner(args: Any) -> DagProcessorJobRunner:
@enable_memray_trace(component=MemrayTraceComponents.dag_processor)
def _run_dag_processor_job(job_runner: DagProcessorJobRunner) -> None:
- run_job(job=job_runner.job, execute_callable=job_runner._execute)
+ if not (token_file := conf.get("dag_processor", "execution_api_token_file", fallback=None)):
+ run_job(job=job_runner.job, execute_callable=job_runner._execute)
+ return
+
+ # Create the connection pool after the daemon or hot-reload fork.
+ from airflow.dag_processing.api_client import DagProcessorAPIClient
+
+ with DagProcessorAPIClient(
+ base_url=get_execution_api_server_url(),
+ token_file=token_file,
+ hostname=get_hostname(),
+ unixname=job_runner.job.unixname,
+ bundle_names=job_runner.processor.bundle_names_to_parse
+ or sorted(_load_bundle_config_snapshot().names),
+ ) as client:
+ job_runner.run_with_api(client)
-@cli_utils.action_cli
@providers_configuration_loaded
def dag_processor(args):
"""Start Airflow Dag Processor Job."""
+ if getattr(args, "check_health", False) or getattr(args, "check_ready", False):
+ if not conf.get("dag_processor", "execution_api_token_file", fallback=None):
+ raise SystemExit("Local health probes require Dag processor API mode")
+ return check_processor_health(readiness=getattr(args, "check_ready", False))
+ if getattr(args, "sync_bundles_only", False):
+ return _run_with_action_log(_sync_bundle_catalog, args)
+ if not conf.get("dag_processor", "execution_api_token_file", fallback=None):
+ return _run_with_action_log(_start_dag_processor, args)
+ if getattr(args, "verbose", False):
+ cli_utils.set_verbose_logging()
+ return _start_dag_processor(args)
+
+
+def _run_with_action_log(func: Callable[[Any], Any], args: Any) -> Any:
+ # The action log names its event after the decorated function; keep the command's own name.
+ def dag_processor(args):
+ return func(args)
+
+ return cli_utils.action_cli(dag_processor)(args)
+
+
+def _sync_bundle_catalog(args):
+ if args.bundle_name:
+ cli_utils.validate_dag_bundle_arg(args.bundle_name)
+ DagFileProcessorManager(max_runs=0, bundle_names_to_parse=args.bundle_name).sync_bundles()
+
+
+def _start_dag_processor(args):
set_component_mp_start_method("dag_processor")
job_runner = _create_dag_processor_job_runner(args)
diff --git a/airflow-core/src/airflow/config_templates/config.yml b/airflow-core/src/airflow/config_templates/config.yml
index 5eb6a37879e..c4aa143271c 100644
--- a/airflow-core/src/airflow/config_templates/config.yml
+++ b/airflow-core/src/airflow/config_templates/config.yml
@@ -3269,6 +3269,32 @@ dag_processor:
- how often to refresh DAGs from those sources
- how many files to parse concurrently
options:
+ execution_api_token_file:
+ description: |
+ Path to a session token written by ``airflow dag-processor-token``. When set,
+ ``airflow dag-processor`` uses the authenticated Execution API for its Job lifecycle,
+ bundle metadata and secrets, parse-time requests, result publication, and callback
+ and priority-request delivery. The API URL comes from
+ ``[core] execution_api_server_url``. Each processor needs its own session token.
+
+ Provision the bundle catalog on a trusted host before starting the processor.
+ The processor does not need metadata database credentials, the Fernet key, or
+ API signing keys. Remove those credentials from its deployment and block database
+ network access to enforce isolation; setting this option does not remove them.
+ Leave unset to use the in-process API and database-backed Job lifecycle.
+ version_added: 3.5.0
+ type: string
+ default: ~
+ example: "/run/airflow/dag-processor.jwt"
+ health_check_file:
+ description: |
+ Local heartbeat file written by a Dag processor using the Execution API.
+ Use a separate writable file for each processor. ``airflow dag-processor --check-health``
+ checks its loop heartbeat; ``--check-ready`` also checks the last successful API heartbeat.
+ version_added: 3.5.0
+ type: string
+ default: "{AIRFLOW_HOME}/dag_processor_health.json"
+ example: ~
dag_bundle_storage_path:
description: |
String path to folder where Airflow bundles can store files locally. Not templated.
@@ -3432,6 +3458,10 @@ dag_processor:
dag processor is considered unhealthy.
This is used by the health check in the **/health** endpoint and in ``airflow jobs check``
CLI for DagProcessorJob.
+ For a Dag processor using the Execution API, the API server also reads this value to decide
+ whether a session's previous Job is still alive, and the processor reads it for its local
+ ``--check-health`` and ``--check-ready`` probes and to bound how long registration waits for
+ a previous Job to expire. Set the same value for the API server and the processor.
version_added: 3.2.0
type: integer
example: ~
@@ -3440,7 +3470,15 @@ dag_processor:
description: |
Seconds a Dag processor using the Execution API may go without a recorded Job heartbeat.
After this, the scheduler's next cleanup (every ``[scheduler] parsing_cleanup_interval``)
- ends the Job and releases the callbacks and parse requests it claimed.
+ ends the Job and releases the callbacks and parse requests it claimed, and a processor that
+ could not reach the API stops. Processors do not heartbeat while initializing, refreshing,
+ or listing a single bundle, so set this above the longest of those operations; the local
+ health check must also tolerate them, through ``health_check_threshold`` and the probe's
+ failure budget. It also bounds how long a processor retries preparing the bundle
+ version of a claimed callback before acknowledging the callback as failed.
+ The scheduler reads this value to end Jobs and the processor reads it to stop, so set the
+ same value for both; changing it only in the processor's configuration has no effect on
+ when the scheduler ends the Job.
version_added: 3.5.0
type: integer
example: ~
diff --git a/airflow-core/src/airflow/dag_processing/api_client.py b/airflow-core/src/airflow/dag_processing/api_client.py
index 964b0aead80..14f6e6db4c7 100644
--- a/airflow-core/src/airflow/dag_processing/api_client.py
+++ b/airflow-core/src/airflow/dag_processing/api_client.py
@@ -25,7 +25,7 @@ from contextvars import ContextVar
from dataclasses import dataclass, field
from pathlib import Path
from time import monotonic
-from typing import Any
+from typing import Any, Literal
from urllib.parse import quote
from uuid import UUID
@@ -39,6 +39,9 @@ from airflow.api_fastapi.execution_api.datamodels.dag_parsing import (
DagBundleStateResponse,
DagParseResultBody,
DagParseResultResponse,
+ ProcessorWorkAckBody,
+ ProcessorWorkClaimBody,
+ ProcessorWorkItem,
)
from airflow.api_fastapi.execution_api.datamodels.job import (
DagParseTokenBody,
@@ -142,6 +145,7 @@ class DagProcessorAPIClient(Client):
self._job_id: int | None = None
self._completion_state: TerminalJobState | None = None
self._completed = False
+ self._pending_claims: dict[str, ProcessorWorkClaimBody] = {}
super().__init__(base_url=base_url, token="", **kwargs)
@property
@@ -398,6 +402,44 @@ class DagProcessorAPIClient(Client):
)
return [DagBundleStateResponse.model_validate(item) for item in response.json()]
+ def claim_work(
+ self, kind: Literal["callbacks", "priority"], bundle_names: list[str], limit: int
+ ) -> list[ProcessorWorkItem]:
+ if not bundle_names:
+ return []
+ body = self._pending_claims.setdefault(
+ kind,
+ ProcessorWorkClaimBody(
+ claim_id=uuid7(),
+ bundle_names=bundle_names,
+ limit=min(limit, 100),
+ ),
+ )
+ response = self.request(
+ "POST",
+ f"jobs/{self._require_job_id()}/requested-work/{kind}/claim",
+ content=body.model_dump_json(),
+ headers=_JOB_API_HEADERS,
+ retry=False,
+ timeout=self._get_bounded_timeout(),
+ )
+ work = [ProcessorWorkItem.model_validate(item) for item in response.json()]
+ del self._pending_claims[kind]
+ return work
+
+ def acknowledge_work(
+ self, kind: Literal["callbacks", "priority"], work: ProcessorWorkItem, *, failed: bool = False
+ ) -> None:
+ body = ProcessorWorkAckBody(claim_id=work.claim_id, state="failed" if failed else "success")
+ self.request(
+ "POST",
+ f"jobs/{self._require_job_id()}/requested-work/{kind}/{work.id}/ack",
+ content=body.model_dump_json(),
+ headers=_JOB_API_HEADERS,
+ retry=False,
+ timeout=self._get_bounded_timeout(),
+ )
+
def publish_inventory(self, bundle_name: str, body: DagBundleInventoryBody) -> DagBundleInventoryResponse:
response = self.request(
"POST",
diff --git a/airflow-core/src/airflow/dag_processing/health.py b/airflow-core/src/airflow/dag_processing/health.py
new file mode 100644
index 00000000000..640724a85ef
--- /dev/null
+++ b/airflow-core/src/airflow/dag_processing/health.py
@@ -0,0 +1,56 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""Local probes for a processor that does not have metadata database credentials."""
+
+from __future__ import annotations
+
+import json
+import os
+from pathlib import Path
+from time import monotonic
+
+from airflow.configuration import conf
+
+
+def write_processor_health(last_api_heartbeat: float) -> None:
+ path = Path(conf.get("dag_processor", "health_check_file"))
+ path.parent.mkdir(parents=True, exist_ok=True)
+ temporary = path.with_name(f".{path.name}.{os.getpid()}")
+ temporary.write_text(
+ json.dumps(
+ {
+ "pid": os.getpid(),
+ "loop": monotonic(),
+ "api": last_api_heartbeat,
+ }
+ )
+ )
+ temporary.replace(path)
+
+
+def check_processor_health(*, readiness: bool = False) -> None:
+ try:
+ record = json.loads(Path(conf.get("dag_processor", "health_check_file")).read_text())
+ os.kill(record["pid"], 0)
+ now = monotonic()
+ maximum_age = conf.getint("dag_processor", "health_check_threshold")
+ fields = ("loop", "api") if readiness else ("loop",)
+ if any(not 0 <= now - record[field] <= maximum_age for field in fields):
+ raise ValueError("Processor heartbeat has expired")
+ except (OSError, ValueError, KeyError, TypeError) as error:
+ raise SystemExit("Dag processor is not healthy") from error
diff --git a/airflow-core/src/airflow/dag_processing/manager.py b/airflow-core/src/airflow/dag_processing/manager.py
index f177e2dad38..8021e81ef4b 100644
--- a/airflow-core/src/airflow/dag_processing/manager.py
+++ b/airflow-core/src/airflow/dag_processing/manager.py
@@ -23,6 +23,7 @@ import functools
import gc
import inspect
import logging
+import math
import os
import random
import selectors
@@ -42,7 +43,7 @@ from uuid import UUID
import attrs
import httpx2
import structlog
-from pydantic import ValidationError
+from pydantic import TypeAdapter, ValidationError
from sqlalchemy import or_, select, update
from sqlalchemy.exc import OperationalError
from sqlalchemy.orm import load_only
@@ -58,9 +59,10 @@ from airflow.api_fastapi.execution_api.datamodels.dag_parsing import (
DagParseResultBody,
ParseSourceCode,
ParseWarning,
+ ProcessorWorkItem,
)
from airflow.api_fastapi.execution_api.in_process import InProcessExecutionAPI
-from airflow.callbacks.callback_requests import DagCallbackRequest
+from airflow.callbacks.callback_requests import CallbackRequest, DagCallbackRequest
from airflow.configuration import conf
from airflow.dag_processing.bundles.base import (
BundleUsageTrackingManager,
@@ -112,7 +114,6 @@ if TYPE_CHECKING:
from sqlalchemy.orm import Session
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
@@ -163,6 +164,27 @@ class PendingDagPublication:
"""Whether this replaces a result the API refused, so a further refusal is final."""
+@attrs.define
+class DeferredCallback:
+ """A claimed callback whose bundle could not be prepared, retried until it expires."""
+
+ expires_at: float
+ attempts: int = 0
+ next_attempt_time: float = 0.0
+
+
+@attrs.define
+class PriorityParseClaim:
+ """A priority request served only by an import after its requested refresh."""
+
+ work: ProcessorWorkItem
+ refresh_generation: int
+ attempt_id: UUID | None = None
+
+
+_MAX_CALLBACK_RETRY_DELAY = 60.0
+
+
@dataclass(frozen=True)
class DagFileInfo:
"""Information about a DAG file."""
@@ -313,6 +335,14 @@ class DagFileProcessorManager(LoggingMixin):
)
_parsing_start_time: float | None = attrs.field(default=None, init=False)
+ _callback_claims: dict[int, ProcessorWorkItem] = attrs.field(factory=dict, init=False)
+ _deferred_api_callbacks: list[CallbackRequest] = attrs.field(factory=list, init=False)
+ _deferred_callback_retries: dict[int, DeferredCallback] = attrs.field(factory=dict, init=False)
+ _callback_attempts: dict[UUID, list[ProcessorWorkItem]] = attrs.field(factory=dict, init=False)
+ _priority_claims: dict[str, PriorityParseClaim] = attrs.field(factory=dict, init=False)
+ _pending_work_acks: list[tuple[Literal["callbacks", "priority"], ProcessorWorkItem, bool]] = attrs.field(
+ factory=list, init=False
+ )
_num_run: int = attrs.field(default=0, init=False)
_callback_to_execute: dict[DagFileInfo, list[CallbackRequest]] = attrs.field(
@@ -712,6 +742,7 @@ class DagFileProcessorManager(LoggingMixin):
self._collect_results()
self._publish_pending_results()
+ self._acknowledge_requested_work()
for callback in self.fetch_callbacks():
self._add_callback_to_queue(callback)
self._scan_stale_dags()
@@ -772,9 +803,36 @@ class DagFileProcessorManager(LoggingMixin):
Default implementation reads from the metadata DB; override to source requests from an API.
"""
- if self.api_client is not None:
+ if self.api_client is None:
+ return self._claim_priority_files()
+ available = (
+ self._parallelism
+ - len(self._priority_claims)
+ - sum(kind == "priority" for kind, _, _ in self._pending_work_acks)
+ )
+ if available <= 0:
+ return []
+ bundles = {bundle.name: bundle for bundle in self._dag_bundles}
+ try:
+ work = self.api_client.claim_work("priority", list(bundles), available)
+ except httpx2.HTTPError:
+ self.log.warning("Unable to claim priority parses; retrying on a later loop")
return []
- return self._claim_priority_files()
+ files = []
+ for item in work:
+ self._priority_claims[item.id] = PriorityParseClaim(
+ work=item, refresh_generation=self._bundle_refresh_generations[item.bundle_name] + 1
+ )
+ bundle = bundles[item.bundle_name]
+ path = find_enclosing_file(bundle.path / item.relative_fileloc)
+ files.append(
+ DagFileInfo(
+ rel_path=path.relative_to(bundle.path) if path else Path(item.relative_fileloc),
+ bundle_name=bundle.name,
+ bundle_path=bundle.path,
+ )
+ )
+ return files
def request_bundle_refresh(self, bundle_names: str | Iterable[str]) -> None:
"""
@@ -836,9 +894,80 @@ class DagFileProcessorManager(LoggingMixin):
Default implementation reads from the metadata DB; override to source callbacks from an API.
"""
- if self.api_client is not None:
+ if self.api_client is None:
+ return self._fetch_callbacks_from_db()
+ available = (
+ self.max_callbacks_per_loop
+ - len(self._callback_claims)
+ - sum(len(work) for work in self._callback_attempts.values())
+ - sum(kind == "callbacks" for kind, _, _ in self._pending_work_acks)
+ )
+ if available <= 0:
+ return list(self._deferred_api_callbacks)
+ bundles = [b.name for b in self._dag_bundles if not b.supports_versioning or b.is_initialized]
+ try:
+ work = self.api_client.claim_work("callbacks", bundles, available)
+ except httpx2.HTTPError:
+ self.log.warning("Unable to claim callbacks; retrying on a later loop")
+ return list(self._deferred_api_callbacks)
+ for item in work:
+ if item.callback is None:
+ raise ValueError("Callback claim returned no request")
+ request: CallbackRequest = TypeAdapter(CallbackRequest).validate_json(item.callback)
+ self._callback_claims[id(request)] = item
+ self._deferred_api_callbacks.append(request)
+ return list(self._deferred_api_callbacks)
+
+ def _acknowledge_requested_work(self, *, deadline: float | None = None) -> None:
+ """Send queued acknowledgments in order until the deadline, keeping claims after a failure."""
+ if self.api_client is None:
+ return
+ if deadline is None:
+ deadline = time.monotonic() + _API_SEND_BUDGET_SECONDS
+ while self._pending_work_acks and time.monotonic() < deadline:
+ kind, item, failed = self._pending_work_acks[0]
+ try:
+ self.api_client.acknowledge_work(kind, item, failed=failed)
+ except httpx2.HTTPStatusError as error:
+ if error.response.status_code >= 500:
+ self.log.warning("Unable to acknowledge requested work; retaining the claim")
+ return
+ if error.response.status_code != 409:
+ # Retrying cannot succeed; the server releases the claim when this Job ends.
+ self.log.warning(
+ "The API refused to acknowledge %s work %s (HTTP %s)",
+ kind,
+ item.id,
+ error.response.status_code,
+ )
+ except httpx2.RequestError:
+ self.log.warning("Unable to acknowledge requested work; retaining the claim")
+ return
+ self._pending_work_acks.pop(0)
+
+ def _finish_callback_attempt(self, proc: BaseDagFileProcessorProcess) -> None:
+ for item in self._callback_attempts.pop(proc.id, []):
+ self._pending_work_acks.append(("callbacks", item, proc._exit_code != 0))
+
+ def _finish_priority_claims(self, attempt_id: UUID, *, failed: bool) -> None:
+ """Acknowledge only requests assigned to this import, never newer requests for its file."""
+ for key, claim in list(self._priority_claims.items()):
+ if claim.attempt_id == attempt_id:
+ self._pending_work_acks.append(("priority", claim.work, failed))
+ del self._priority_claims[key]
+
+ def _get_priority_claims(self, file: DagFileInfo) -> list[PriorityParseClaim]:
+ if file.bundle_version is not None:
return []
- return self._fetch_callbacks_from_db()
+ return [
+ claim
+ for claim in self._priority_claims.values()
+ if claim.work.bundle_name == file.bundle_name
+ and (
+ claim.work.relative_fileloc == str(file.rel_path)
+ or claim.work.relative_fileloc.startswith(str(file.rel_path) + "/")
+ )
+ ]
@provide_session
@retry_db_transaction
@@ -921,7 +1050,16 @@ class DagFileProcessorManager(LoggingMixin):
)
if bundle.supports_versioning:
bundle.initialize()
- except Exception:
+ except Exception as error:
+ if self.api_client is not None:
+ from airflow.dag_processing.api_client import DagProcessorRegistrationRetired
+
+ if isinstance(error, DagProcessorRegistrationRetired):
+ raise
+ if self.api_client.restart_required:
+ raise DagProcessorRegistrationRetired(
+ "The Dag processor registration has ended; restart required"
+ ) from error
self.log.exception(
"Error preparing bundle %s version %s for callback, skipping",
request.bundle_name,
@@ -934,11 +1072,19 @@ class DagFileProcessorManager(LoggingMixin):
self.log.debug("Queuing %s CallbackRequest: %s", type(request).__name__, request)
if get_claiming_importer(request.filepath, request.bundle_name) is not None:
self._log_dropped_lang_sdk_callback(request)
+ if item := self._callback_claims.pop(id(request), None):
+ self._pending_work_acks.append(("callbacks", item, True))
+ self._remove_deferred_callback(request)
return
+ retry = self._deferred_callback_retries.get(id(request))
+ if retry is not None and time.monotonic() < retry.next_attempt_time:
+ return
+ self.heartbeat()
bundle = self.prepare_callback_bundle(request)
if bundle is None:
+ if id(request) in self._callback_claims:
+ self._defer_claimed_callback(request)
return
-
file_info = DagFileInfo(
rel_path=Path(request.filepath),
bundle_path=bundle.path,
@@ -946,10 +1092,41 @@ class DagFileProcessorManager(LoggingMixin):
bundle_version=request.bundle_version,
)
self._callback_to_execute[file_info].append(request)
+ if id(request) in self._callback_claims:
+ self._remove_deferred_callback(request)
self._add_files_to_queue([file_info], mode="front")
team_name = self._get_team_name(file_info.bundle_name)
stats.incr("dag_processing.other_callback_count", tags=prune_dict({"team_name": team_name}))
+ def _remove_deferred_callback(self, request: CallbackRequest) -> None:
+ # Equal callback payloads can belong to different claims.
+ self._deferred_api_callbacks = [
+ deferred for deferred in self._deferred_api_callbacks if deferred is not request
+ ]
+ self._deferred_callback_retries.pop(id(request), None)
+
+ def _defer_claimed_callback(self, request: CallbackRequest) -> None:
+ """Retry preparing a claimed callback's bundle with backoff, then acknowledge it as failed."""
+ now = time.monotonic()
+ timeout = conf.getint("dag_processor", "job_heartbeat_timeout")
+ retry = self._deferred_callback_retries.setdefault(id(request), DeferredCallback(now + timeout))
+ if now >= retry.expires_at:
+ self.log.error(
+ "Dropping %s for %s in bundle %s: its bundle version %s could not be prepared within %s seconds",
+ type(request).__name__,
+ request.filepath,
+ request.bundle_name,
+ request.bundle_version,
+ timeout,
+ )
+ self._pending_work_acks.append(("callbacks", self._callback_claims.pop(id(request)), True))
+ self._remove_deferred_callback(request)
+ return
+ retry.attempts += 1
+ retry.next_attempt_time = min(
+ now + min(2 ** (retry.attempts - 1), _MAX_CALLBACK_RETRY_DELAY), retry.expires_at
+ )
+
def _log_dropped_lang_sdk_callback(self, request: CallbackRequest) -> None:
if isinstance(request, DagCallbackRequest):
target = f"dag_id={request.dag_id} run_id={request.run_id}"
@@ -1036,12 +1213,35 @@ class DagFileProcessorManager(LoggingMixin):
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
+ failed_sources = set(body.failed_sources)
+ for key, claim in list(self._priority_claims.items()):
+ item = claim.work
+ if (
+ item.bundle_name == bundle_name
+ and source.refresh_generation >= claim.refresh_generation
+ and item.relative_fileloc not in body.files
+ and not is_source_discovery_failed(item.relative_fileloc, failed_sources)
+ ):
+ self._pending_work_acks.append(("priority", item, True))
+ del self._priority_claims[key]
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)
+ active_attempts = {
+ processor.id for file, processor in self._processors.items() if file.bundle_name == bundle_name
+ }
+ active_attempts.update(
+ pending.body.attempt_id
+ for file, pending in self._pending_publications.items()
+ if file.bundle_name == bundle_name
+ )
+ for key, claim in list(self._priority_claims.items()):
+ if claim.work.bundle_name == bundle_name and claim.attempt_id not in active_attempts:
+ self._pending_work_acks.append(("priority", claim.work, True))
+ del self._priority_claims[key]
def _invalidate_bundle_parse_source(self, bundle_name: str) -> None:
# Initialization or refresh may change files even when the provider raises afterwards.
@@ -1088,6 +1288,8 @@ class DagFileProcessorManager(LoggingMixin):
if file.bundle_version is None
}
for bundle in self._dag_bundles:
+ # Bundle operations can block for minutes; an API-mode Job ends after job_heartbeat_timeout.
+ self.heartbeat()
with self._use_bundle(bundle.name):
if catalog is not None and bundle.name not in catalog:
self._discard_unusable_bundle_source(bundle.name)
@@ -1537,10 +1739,14 @@ class DagFileProcessorManager(LoggingMixin):
def purge_removed_files_from_queue(self, present: set[DagFileInfo]):
"""Remove from queue any files no longer observed locally."""
present_keys = {file.presence_key for file in present}
- self._file_queue = OrderedDict((x, None) for x in self._file_queue if x.presence_key in present_keys)
- self._files_awaiting_source = {
- file: None for file in self._files_awaiting_source if file.presence_key in present_keys
- }
+
+ def is_retained(file: DagFileInfo) -> bool:
+ return file.presence_key in present_keys or bool(
+ self.api_client is not None and self._callback_to_execute.get(file)
+ )
+
+ self._file_queue = OrderedDict((x, None) for x in self._file_queue if is_retained(x))
+ self._files_awaiting_source = {x: None for x in self._files_awaiting_source if is_retained(x)}
stats.gauge("dag_processing.file_path_queue_size", len(self._file_queue))
def remove_orphaned_file_stats(self, present: set[DagFileInfo]):
@@ -1560,6 +1766,8 @@ class DagFileProcessorManager(LoggingMixin):
bundle_to_team = self._get_team_names({file.bundle_name for file in self._processors})
for file in list(self._processors.keys()):
+ if self.api_client is not None and self._processors[file].id in self._callback_attempts:
+ continue
if file.presence_key not in present_keys:
processor = self._processors.pop(file, None)
if not processor:
@@ -1671,8 +1879,11 @@ class DagFileProcessorManager(LoggingMixin):
file.bundle_name,
)
self._record_failed_publication(file, next_stat)
+ self._finish_priority_claims(proc.id, failed=True)
return
+ if proc.parsing_result is None and not is_callback_only:
+ self._finish_priority_claims(proc.id, failed=True)
self._file_stats[file] = next_stat
def _record_failed_publication(self, file: DagFileInfo, stat: DagFileStat) -> None:
@@ -1750,12 +1961,15 @@ class DagFileProcessorManager(LoggingMixin):
continue
self.log.exception("Failed to publish parsing result for %s", file)
self._record_failed_publication(file, pending.stat)
+ self._finish_priority_claims(pending.body.attempt_id, failed=True)
else:
self._file_stats[file] = pending.stat
+ self._finish_priority_claims(pending.body.attempt_id, failed=False)
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._finish_priority_claims(pending.body.attempt_id, failed=True)
self._dispatch_sequence += 1
if len(reason) > _MAX_REJECTION_REASON_LENGTH:
reason = reason[:_MAX_REJECTION_REASON_LENGTH] + "..."
@@ -1875,6 +2089,7 @@ class DagFileProcessorManager(LoggingMixin):
for file in finished:
processor = self._processors.pop(file)
+ self._finish_callback_attempt(processor)
processor.close()
def _get_log_dir(self) -> str:
@@ -1938,15 +2153,22 @@ class DagFileProcessorManager(LoggingMixin):
return client
def _create_process(self, dag_file: DagFileInfo) -> BaseDagFileProcessorProcess:
- id = uuid7()
+ process_id = uuid7()
callback_to_execute_for_file = self._callback_to_execute.pop(dag_file, [])
+ claims = [
+ self._callback_claims.pop(id(request))
+ for request in callback_to_execute_for_file
+ if id(request) in self._callback_claims
+ ]
+ if claims:
+ self._callback_attempts[process_id] = claims
logger, logger_filehandle = self._get_logger_for_dag_file(dag_file)
subprocess_logs_to_stdout = conf.get("logging", "dag_processor_log_target") == "stdout"
if get_claiming_importer(dag_file.absolute_path, dag_file.bundle_name) is not None:
return LangSDKDagFileProcessorProcess.start(
- id=id,
+ id=process_id,
path=dag_file.absolute_path,
bundle_path=cast("Path", dag_file.bundle_path),
bundle_name=dag_file.bundle_name,
@@ -1959,7 +2181,7 @@ class DagFileProcessorManager(LoggingMixin):
)
return DagFileProcessorProcess.start(
- id=id,
+ id=process_id,
path=dag_file.absolute_path,
bundle_path=cast("Path", dag_file.bundle_path),
bundle_name=dag_file.bundle_name,
@@ -1984,13 +2206,16 @@ class DagFileProcessorManager(LoggingMixin):
return None
return self._bundle_parse_sources.get(file.bundle_name)
+ def _has_requested_work(self, file: DagFileInfo) -> bool:
+ return file in self._callback_to_execute or bool(self._get_priority_claims(file))
+
def _readmit_files_with_source(self) -> None:
ready = [file for file in self._files_awaiting_source if self._get_dispatch_source(file) is not None]
requested: list[DagFileInfo] = []
regular: list[DagFileInfo] = []
for file in ready:
del self._files_awaiting_source[file]
- if file in self._callback_to_execute:
+ if self._has_requested_work(file):
requested.append(file)
else:
regular.append(file)
@@ -2004,12 +2229,17 @@ class DagFileProcessorManager(LoggingMixin):
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) + len(self._pending_publications) and self._file_queue
- ):
+ # Requested work skipped in this pass goes back to the front once the pass is over.
+ requested_later: dict[DagFileInfo, None] = {}
+ for _ in range(len(self._file_queue)):
+ if len(self._processors) + len(self._pending_publications) >= self._parallelism:
+ break
file, _ = self._file_queue.popitem(last=False)
+ priority_claims = self._get_priority_claims(file)
# Stop creating duplicate processor i.e. processor with the same filepath
if file in self._processors or file in self._pending_publications:
+ if priority_claims:
+ requested_later[file] = None
continue
source = self._get_dispatch_source(file)
@@ -2021,6 +2251,12 @@ class DagFileProcessorManager(LoggingMixin):
processor.parse_source = source
self._dispatch_sequence += 1
processor.dispatch_sequence = self._dispatch_sequence
+ if file.bundle_version is None:
+ for claim in priority_claims:
+ if processor.had_callbacks or source.refresh_generation < claim.refresh_generation:
+ requested_later[file] = None
+ else:
+ claim.attempt_id = processor.id
stats.incr(
"dag_processing.processes",
tags=prune_dict(
@@ -2035,6 +2271,8 @@ class DagFileProcessorManager(LoggingMixin):
self._processors[file] = processor
stats.gauge("dag_processing.file_path_queue_size", len(self._file_queue))
+ if requested_later:
+ self._add_files_to_queue(list(reversed(requested_later)), mode="front")
def _add_new_files_to_queue(self, known_files: dict[str, set[DagFileInfo]]):
"""
@@ -2250,10 +2488,13 @@ class DagFileProcessorManager(LoggingMixin):
last_num_of_db_queries=0,
)
self._file_stats[file] = stat
+ if not processor.had_callbacks:
+ self._finish_priority_claims(processor.id, failed=True)
# Clean up `self._processors` after iterating over it
for proc in processors_to_remove:
processor = self._processors.pop(proc)
+ self._finish_callback_attempt(processor)
processor.close()
def _add_files_to_queue(
@@ -2297,6 +2538,13 @@ class DagFileProcessorManager(LoggingMixin):
return False
if self._pending_inventories:
return False
+ if (
+ self._pending_work_acks
+ or self._callback_claims
+ or self._callback_attempts
+ or self._priority_claims
+ ):
+ return False
if any(
file not in self._callback_to_execute
for file in [*self._file_queue, *self._files_awaiting_source]
@@ -2324,10 +2572,15 @@ class DagFileProcessorManager(LoggingMixin):
processor.kill(signal.SIGTERM, escalation_delay=5.0)
def end(self):
- """Kill all child processes on exit since we don't want to leave them as orphaned."""
+ """Kill child processes so none are orphaned, then acknowledge work that already finished."""
pids_to_kill = [p.pid for p in self._processors.values()]
if pids_to_kill:
kill_child_processes_by_pids(pids_to_kill)
+ # Claims the server releases after this Job ends are delivered again, repeating their callbacks.
+ try:
+ self._acknowledge_requested_work(deadline=math.inf)
+ except Exception:
+ self.log.exception("Unable to acknowledge requested work before exiting")
def emit_metrics(*, parse_time: float, dag_file_stats: Sequence[DagFileStat]):
diff --git a/airflow-core/src/airflow/dag_processing/processor.py b/airflow-core/src/airflow/dag_processing/processor.py
index 3adb457b425..95f6cb2d082 100644
--- a/airflow-core/src/airflow/dag_processing/processor.py
+++ b/airflow-core/src/airflow/dag_processing/processor.py
@@ -40,6 +40,7 @@ from airflow.callbacks.callback_requests import (
)
from airflow.configuration import conf
from airflow.dag_processing.bundles.base import BundleVersionLock
+from airflow.dag_processing.bundles.manager import get_configured_bundle_team_names
from airflow.dag_processing.dagbag import BundleDagBag, DagBag, LangSDKSerializedDAG
from airflow.models.dag import DagModel
from airflow.sdk.exceptions import TaskNotFound
@@ -426,7 +427,11 @@ def _execute_dag_callbacks(dagbag: DagBag, request: DagCallbackRequest, log: Fil
{
"dag_id": request.dag_id,
"team_name": (
- DagModel.get_team_name(request.dag_id)
+ (
+ get_configured_bundle_team_names().get(request.bundle_name)
+ if conf.get("dag_processor", "execution_api_token_file", fallback=None)
+ else DagModel.get_team_name(request.dag_id)
+ )
if conf.getboolean("core", "multi_team")
else None
),
diff --git a/airflow-core/src/airflow/jobs/dag_processor_job_runner.py b/airflow-core/src/airflow/jobs/dag_processor_job_runner.py
index 3b25d32d4a5..3af40ae3f91 100644
--- a/airflow-core/src/airflow/jobs/dag_processor_job_runner.py
+++ b/airflow-core/src/airflow/jobs/dag_processor_job_runner.py
@@ -17,19 +17,31 @@
from __future__ import annotations
+from time import monotonic, sleep
from typing import TYPE_CHECKING
+import httpx2
+
from airflow._shared.observability.metrics import stats
+from airflow._shared.timezones import timezone
+from airflow.configuration import conf
+from airflow.dag_processing.health import write_processor_health
from airflow.jobs.base_job_runner import BaseJobRunner
-from airflow.jobs.job import Job, perform_heartbeat
+from airflow.jobs.job import Job, JobState, execute_job, perform_heartbeat
+from airflow.listeners.listener import get_listener_manager
from airflow.utils.log.logging_mixin import LoggingMixin
if TYPE_CHECKING:
from sqlalchemy.orm import Session
+ from airflow.dag_processing.api_client import DagProcessorAPIClient
from airflow.dag_processing.manager import DagFileProcessorManager
+class DagProcessorHeartbeatTimeout(RuntimeError):
+ """The processor could not confirm its Job heartbeat before the scheduler may end the Job."""
+
+
class DagProcessorJobRunner(BaseJobRunner, LoggingMixin):
"""
DagProcessorJobRunner is a job runner that runs a DagFileProcessorManager processor.
@@ -49,12 +61,92 @@ class DagProcessorJobRunner(BaseJobRunner, LoggingMixin):
):
super().__init__(job)
self.processor = processor
+ self._last_api_heartbeat = 0.0
+ self._next_api_heartbeat = 0.0
self.processor.heartbeat = lambda: perform_heartbeat(
job=self.job,
heartbeat_callback=self.heartbeat_callback,
only_if_necessary=True,
)
+ def run_with_api(self, client: DagProcessorAPIClient) -> int | None:
+ """Let the API own the Job row while retaining the normal processor cleanup and listeners."""
+ # The client loads versioned routes which import this runner.
+ from airflow.api_fastapi.execution_api.datamodels.job import TerminalJobState
+
+ self.processor.api_client = client
+ self.job.id = self._register_api_job(client)
+ self.job.state = JobState.RUNNING
+ self.job.start_date = self.job.latest_heartbeat = timezone.utcnow()
+ self._last_api_heartbeat = monotonic()
+ self._next_api_heartbeat = self._last_api_heartbeat + self.job.heartrate
+ self.processor.heartbeat = lambda: self._heartbeat_api(client)
+ stats.incr("job_start", 1, 1)
+ failed = True
+ try:
+ result = execute_job(self.job, execute_callable=self._execute)
+ failed = False
+ return result
+ finally:
+ self.job.end_date = timezone.utcnow()
+ self.job.state = JobState.FAILED if failed else JobState.SUCCESS
+ try:
+ get_listener_manager().hook.before_stopping(component=self.job)
+ except Exception:
+ self.log.exception("Error calling Dag processor shutdown listener")
+ try:
+ client.complete_job(TerminalJobState(self.job.state))
+ except Exception:
+ if not failed:
+ raise
+ self.log.exception("Unable to record the failed Dag processor Job")
+ finally:
+ stats.incr("job_end", 1, 1)
+
+ def _register_api_job(self, client: DagProcessorAPIClient) -> int:
+ from airflow.dag_processing.api_client import DagProcessorJobAlreadyRunning
+
+ deadline = monotonic() + conf.getint("dag_processor", "health_check_threshold") + self.job.heartrate
+ delay = 1.0
+ while True:
+ try:
+ return client.register_job()
+ except DagProcessorJobAlreadyRunning:
+ remaining = deadline - monotonic()
+ if remaining <= 0:
+ raise
+ self.log.info("Waiting for the previous Dag processor Job heartbeat to expire")
+ sleep(min(delay, remaining))
+ delay = min(delay * 2, 5.0)
+
+ def _heartbeat_api(self, client: DagProcessorAPIClient) -> None:
+ from airflow.dag_processing.api_client import DagProcessorRegistrationRetired
+
+ write_processor_health(self._last_api_heartbeat)
+ if client.restart_required:
+ raise DagProcessorRegistrationRetired("The Dag processor registration requires a restart")
+ now = monotonic()
+ if now < self._next_api_heartbeat:
+ return
+ self._next_api_heartbeat = now + self.job.heartrate
+ try:
+ state = client.heartbeat()
+ except httpx2.HTTPError as error:
+ if isinstance(error, httpx2.HTTPStatusError) and error.response.status_code < 500:
+ raise
+ stats.incr("dag_processor_heartbeat_failure", 1, 1)
+ if monotonic() - self._last_api_heartbeat >= conf.getint(
+ "dag_processor", "job_heartbeat_timeout"
+ ):
+ raise DagProcessorHeartbeatTimeout("Dag processor API heartbeat timed out") from error
+ self.log.warning("Unable to heartbeat the Dag processor Job; retrying on the next interval")
+ return
+ if client.restart_required or state != JobState.RUNNING:
+ raise DagProcessorRegistrationRetired("The API requested a Dag processor restart")
+ self._last_api_heartbeat = monotonic()
+ self.job.latest_heartbeat = timezone.utcnow()
+ self.heartbeat_callback()
+
def _execute(self) -> int | None:
self.log.info("Starting the Dag Processor Job")
try:
diff --git a/airflow-core/src/airflow/utils/cli.py b/airflow-core/src/airflow/utils/cli.py
index d5af192b350..b2ade41fda9 100644
--- a/airflow-core/src/airflow/utils/cli.py
+++ b/airflow-core/src/airflow/utils/cli.py
@@ -60,6 +60,14 @@ def _check_cli_args(args):
)
+def set_verbose_logging() -> None:
+ """Enable debug output for the root logger and its handlers."""
+ root_logger = logging.getLogger()
+ root_logger.setLevel(logging.DEBUG)
+ for handler in root_logger.handlers:
+ handler.setLevel(logging.DEBUG)
+
+
def action_cli(func=None, check_db=True):
def action_logging(f: T) -> T:
"""
@@ -97,12 +105,8 @@ def action_cli(func=None, check_db=True):
_check_cli_args(args)
metrics = _build_metrics(f.__name__, args[0])
cli_action_loggers.on_pre_execution(**metrics)
- verbose = getattr(args[0], "verbose", False)
- root_logger = logging.getLogger()
- if verbose:
- root_logger.setLevel(logging.DEBUG)
- for handler in root_logger.handlers:
- handler.setLevel(logging.DEBUG)
+ if getattr(args[0], "verbose", False):
+ set_verbose_logging()
try:
# Check and run migrations if necessary
if check_db:
diff --git a/airflow-core/tests/integration/dag_processing/test_api_isolation.py b/airflow-core/tests/integration/dag_processing/test_api_isolation.py
new file mode 100644
index 00000000000..6701c7f4a26
--- /dev/null
+++ b/airflow-core/tests/integration/dag_processing/test_api_isolation.py
@@ -0,0 +1,306 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""Run with Breeze and PostgreSQL; the processor joins an API-only internal network."""
+
+from __future__ import annotations
+
+import json
+import secrets
+import shutil
+import socket
+import subprocess
+import time
+from pathlib import Path
+from urllib.parse import urlsplit
+
+import pytest
+import yaml
+from sqlalchemy import delete, select
+from sqlalchemy.engine import make_url
+from uuid6 import uuid7
+
+from airflow.callbacks.callback_requests import DagCallbackRequest
+from airflow.configuration import conf
+from airflow.jobs.job import Job
+from airflow.models.callback import DagProcessorCallback
+from airflow.models.dag import DagModel
+from airflow.models.dagbag import DagPriorityParsingRequest
+from airflow.models.dagbundle import DagBundleModel
+from airflow.models.dagcode import DagCode
+from airflow.models.variable import Variable
+
+from tests_common.test_utils.db import clear_db_callbacks, clear_db_dag_bundles, clear_db_dags
+from unit.api_fastapi.execution_api.conftest import async_db_engine as async_db_engine
+from unit.api_fastapi.execution_api.test_dag_processor_client import (
+ api_requests as api_requests,
+ api_url as api_url,
+ clean_db as clean_db,
+ freeze_time as freeze_time,
+ provision_token as provision_token,
+)
+
+pytestmark = [
+ pytest.mark.db_test,
+ pytest.mark.skipif(
+ not Path("/var/run/docker.sock").exists() or not shutil.which("docker"),
+ reason="Run in Breeze with its Docker socket mounted",
+ ),
+]
+
+
+@pytest.fixture
+def api_bind_host():
+ return "0.0.0.0"
+
+
+@pytest.fixture
+def api_secret():
+ return secrets.token_hex(32)
+
+
+def run_docker(*args, input=None, check=True, timeout=30):
+ return subprocess.run(
+ ["docker", *args],
+ input=input,
+ capture_output=True,
+ text=True,
+ check=check,
+ timeout=timeout,
+ )
+
+
+@pytest.mark.parametrize("shutdown", ["signal", "api_restart"])
+def test_processor_without_database_network_or_server_credentials(
+ api_url,
+ provision_token,
+ session,
+ tmp_path,
+ shutdown,
+):
+ database = make_url(conf.get("database", "sql_alchemy_conn"))
+ if not database.host:
+ pytest.skip("Requires PostgreSQL or MySQL in Breeze")
+ database_ip = socket.gethostbyname(database.host)
+ database_port = database.port or (5432 if database.drivername.startswith("postgres") else 3306)
+ current = json.loads(run_docker("inspect", socket.gethostname()).stdout)[0]
+ mounts = {mount["Destination"]: mount["Source"] for mount in current["Mounts"]}
+ suffix = uuid7().hex
+ network = f"processor-isolation-{suffix}"
+ worker = f"processor-isolation-worker-{suffix}"
+ clear_db_callbacks()
+ clear_db_dags()
+ clear_db_dag_bundles()
+ session.execute(delete(DagPriorityParsingRequest))
+ session.add(DagBundleModel(name="bundle-a"))
+ session.add(DagPriorityParsingRequest(bundle_name="bundle-a", relative_fileloc="a.py"))
+ callback = DagProcessorCallback(
+ priority_weight=1,
+ callback=DagCallbackRequest(
+ filepath="a.py",
+ bundle_name="bundle-a",
+ bundle_version=None,
+ dag_id="isolated",
+ run_id="run",
+ ),
+ )
+ session.add(callback)
+ session.commit()
+ callback_id = callback.id
+ Variable.set("isolation-value", "from-api")
+ Variable.delete("isolation-callback")
+ source_code = (
+ "from airflow.sdk import DAG, Variable\n"
+ "def callback(context):\n"
+ " Variable.set('isolation-callback', 'done')\n"
+ "dag = DAG('isolated', schedule=None, description=Variable.get('isolation-value'), "
+ "on_failure_callback=callback)\n"
+ )
+ run_docker("network", "create", "--internal", network)
+ try:
+ run_docker("network", "connect", network, current["Id"])
+ inspected = json.loads(run_docker("inspect", current["Id"]).stdout)[0]
+ api_ip = inspected["NetworkSettings"]["Networks"][network]["IPAddress"]
+ worker_api_url = f"http://{api_ip}:{urlsplit(api_url).port}/execution/"
+ example = Path(__file__).resolve().parents[3] / "docs/security/dag-processor-api/docker-compose.yaml"
+ processor_environment = yaml.safe_load(example.read_text())["services"]["dag-processor"][
+ "environment"
+ ]
+ environment = {
+ "USER": processor_environment["USER"],
+ "AIRFLOW_USER_HOME_DIR": "/tmp/airflow-user",
+ "AIRFLOW_HOME": "/tmp/isolated-airflow",
+ "AIRFLOW_CONFIG": "/tmp/isolated-airflow/airflow.cfg",
+ "AIRFLOW__DATABASE__SQL_ALCHEMY_CONN": f"postgresql+psycopg://blocked:blocked@{database_ip}:{database_port}/blocked",
+ "AIRFLOW__CORE__FERNET_KEY": "",
+ "AIRFLOW__API_AUTH__JWT_SECRET": "",
+ "AIRFLOW__API_AUTH__JWT_PRIVATE_KEY_PATH": "",
+ "AIRFLOW__CORE__LOAD_EXAMPLES": "False",
+ "AIRFLOW__CORE__EXECUTION_API_SERVER_URL": worker_api_url,
+ "AIRFLOW__DAG_PROCESSOR__EXECUTION_API_TOKEN_FILE": "/tmp/processor.jwt",
+ "AIRFLOW__DAG_PROCESSOR__HEALTH_CHECK_THRESHOLD": "5",
+ "AIRFLOW__SCHEDULER__JOB_HEARTBEAT_SEC": "1",
+ "CONNECTION_CHECK_MAX_COUNT": "0",
+ "AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST": json.dumps(
+ [
+ {
+ "name": "bundle-a",
+ "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle",
+ "kwargs": {"path": "/tmp/worker-only-dags"},
+ }
+ ]
+ ),
+ }
+ arguments = [
+ "run",
+ "--detach",
+ "--name",
+ worker,
+ "--network",
+ network,
+ "--cap-drop",
+ "ALL",
+ "--security-opt",
+ "no-new-privileges",
+ "--user",
+ "54321:0",
+ "--read-only",
+ "--tmpfs",
+ "/tmp:rw,exec,size=256m",
+ "--entrypoint",
+ "sleep",
+ ]
+ for package in ("airflow-core", "task-sdk", "shared", "providers", "devel-common"):
+ arguments += [
+ "--mount",
+ f"type=bind,source={mounts[f'/opt/airflow/{package}']},target=/opt/airflow/{package},readonly",
+ ]
+ for key, value in environment.items():
+ arguments += ["--env", f"{key}={value}"]
+ arguments += [current["Image"], "600"]
+ run_docker(*arguments)
+ run_docker(
+ "exec",
+ "-i",
+ worker,
+ "sh",
+ "-c",
+ "umask 077; cat > /tmp/processor.jwt",
+ input=provision_token().read_text(),
+ )
+ run_docker(
+ "exec",
+ "-i",
+ worker,
+ "sh",
+ "-c",
+ "mkdir -p /tmp/worker-only-dags; cat > /tmp/worker-only-dags/a.py",
+ input=source_code,
+ )
+ entrypoint = Path(__file__).resolve().parents[4] / "scripts/docker/entrypoint_prod.sh"
+ run_docker(
+ "exec",
+ "-i",
+ worker,
+ "sh",
+ "-c",
+ "cat > /tmp/processor-entrypoint.sh",
+ input=entrypoint.read_text(),
+ )
+ assert not Path("/tmp/worker-only-dags/a.py").exists()
+ denied = run_docker(
+ "exec",
+ worker,
+ "python",
+ "-c",
+ (
+ "import socket,sys\n"
+ f"try: socket.create_connection(({database_ip!r}, {database_port}), timeout=2)\n"
+ "except OSError: sys.exit(0)\n"
+ "sys.exit(1)\n"
+ ),
+ )
+ assert denied.returncode == 0
+ with (tmp_path / "processor.log").open("w+") as output:
+ process = subprocess.Popen(
+ [
+ "docker",
+ "exec",
+ worker,
+ "bash",
+ "/tmp/processor-entrypoint.sh",
+ "airflow",
+ "dag-processor",
+ "--bundle-name",
+ "bundle-a",
+ ],
+ stdout=output,
+ stderr=subprocess.STDOUT,
+ )
+ deadline = time.monotonic() + 90
+ while time.monotonic() < deadline:
+ session.expire_all()
+ dag = session.get(DagModel, "isolated")
+ delivered = session.get(DagProcessorCallback, callback_id)
+ if dag and delivered is None:
+ break
+ if process.poll() is not None:
+ output.seek(0)
+ pytest.fail(output.read())
+ time.sleep(0.2)
+ else:
+ output.seek(0)
+ pytest.fail(output.read())
+ assert Variable.get("isolation-callback") == "done"
+ assert session.scalar(select(DagPriorityParsingRequest)) is None
+ assert (
+ session.scalar(select(DagCode.source_code).where(DagCode.dag_id == "isolated")) == source_code
+ )
+ run_docker("exec", worker, "airflow", "dag-processor", "--check-ready")
+ pause = (
+ "import json,os,signal; "
+ "pid=json.load(open('/tmp/isolated-airflow/dag_processor_health.json'))['pid']; "
+ "os.kill(pid, signal.SIGSTOP)"
+ )
+ run_docker("exec", worker, "python", "-c", pause)
+ time.sleep(6)
+ assert (
+ run_docker(
+ "exec", worker, "airflow", "dag-processor", "--check-health", check=False
+ ).returncode
+ != 0
+ )
+ run_docker("exec", worker, "python", "-c", pause.replace("SIGSTOP", "SIGCONT"))
+ if shutdown == "signal":
+ run_docker("exec", worker, "python", "-c", pause.replace("SIGSTOP", "SIGTERM"))
+ else:
+ session.scalar(select(Job)).state = "restarting"
+ session.commit()
+ assert (process.wait(timeout=30) == 0) == (shutdown == "signal")
+ session.expire_all()
+ assert session.scalar(select(Job.state)) == ("success" if shutdown == "signal" else "failed")
+ finally:
+ run_docker("rm", "--force", worker, check=False)
+ run_docker("network", "disconnect", network, current["Id"], check=False)
+ run_docker("network", "rm", network, check=False)
+ clear_db_callbacks()
+ clear_db_dags()
+ clear_db_dag_bundles()
+ session.execute(delete(DagPriorityParsingRequest))
+ session.commit()
+ Variable.delete("isolation-value")
+ Variable.delete("isolation-callback")
diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/test_dag_processor_client.py b/airflow-core/tests/unit/api_fastapi/execution_api/test_dag_processor_client.py
index 0c7484acffd..977cb2ec5c5 100644
--- a/airflow-core/tests/unit/api_fastapi/execution_api/test_dag_processor_client.py
+++ b/airflow-core/tests/unit/api_fastapi/execution_api/test_dag_processor_client.py
@@ -16,9 +16,13 @@
# under the License.
from __future__ import annotations
+import json
+import os
import socket
+import subprocess
import threading
from datetime import UTC, datetime
+from pathlib import Path
from unittest import mock
import httpx2
@@ -32,20 +36,32 @@ from uuid6 import uuid7
from airflow.api_fastapi.app import create_app
from airflow.api_fastapi.auth.tokens import JWTGenerator, JWTValidator
from airflow.api_fastapi.execution_api.app import lifespan
+from airflow.api_fastapi.execution_api.datamodels.dag_parsing import DagParseResultBody
from airflow.api_fastapi.execution_api.datamodels.job import DagParseTokenBody, JobState, TerminalJobState
+from airflow.configuration import conf
from airflow.dag_processing.api_client import (
DagParseContext,
DagProcessorAPIClient,
DagProcessorRegistrationRetired,
)
+from airflow.dag_processing.bundles.local import LocalDagBundle
+from airflow.dag_processing.bundles.manager import DagBundlesManager
from airflow.jobs.job import Job
+from airflow.models.dag import DagModel
+from airflow.models.dag_parse_checkpoint import DagParseCheckpoint
from airflow.models.dagbundle import DagBundleModel
+from airflow.models.serialized_dag import SerializedDagModel
+from airflow.models.team import Team
from airflow.models.variable import Variable
+from airflow.sdk import Variable as SDKVariable
from airflow.sdk.api.client import Client
from tests_common.test_utils.config import conf_vars
from tests_common.test_utils.db import (
+ clear_db_dag_bundles,
+ clear_db_dags,
clear_db_jobs,
+ clear_db_teams,
clear_db_variables,
)
@@ -218,7 +234,7 @@ class LostAcknowledgmentTransport(httpx2.BaseTransport):
def handle_request(self, request):
response = self.transport.handle_request(request)
if (
- request.url.path.endswith(("/jobs", "/complete", "/parse-token"))
+ request.url.path.endswith(("/jobs", "/complete", "/parse-token", "/parse-results"))
and request.url.path not in self.lost
and response.is_success
):
@@ -247,16 +263,38 @@ def test_lost_registration_and_completion_responses_over_http(api_url, provision
request=DagParseTokenBody(attempt_id=uuid7(), bundle_name="bundle-a", relative_fileloc="dag.py")
)
with processor.use_parse(context):
+ with pytest.raises(httpx2.ReadError):
+ processor.variables.get("parse-token-key")
assert processor.variables.get("parse-token-key").value == "value"
with pytest.raises(httpx2.HTTPStatusError) as error:
processor.post(f"jobs/{job_id}/heartbeat")
assert error.value.response.status_code == 403
+ with pytest.raises(httpx2.HTTPStatusError) as publication_error:
+ processor.post(f"jobs/{job_id}/parse-results", json={})
+ assert publication_error.value.response.status_code == 403
+ body = DagParseResultBody(
+ attempt_id=uuid7(),
+ dispatch_sequence=1,
+ bundle_name="bundle-a",
+ relative_fileloc="dag.py",
+ parse_duration=0.1,
+ serialized_dags=[],
+ source_codes={},
+ )
+ with pytest.raises(httpx2.ReadError):
+ processor.publish_parse_result(body)
+ processor.heartbeat()
+ receipt = processor.publish_parse_result(body)
+ assert receipt.attempt_id == body.attempt_id
+ assert session.scalar(select(DagParseCheckpoint)).attempt_id == body.attempt_id
+ assert processor.publish_parse_result(body) == receipt
processor.complete_job(TerminalJobState.SUCCESS)
assert session.get(Job, job_id).state == JobState.SUCCESS.value
assert transport.lost == {
"/execution/jobs",
f"/execution/jobs/{job_id}/parse-token",
f"/execution/jobs/{job_id}/complete",
+ f"/execution/jobs/{job_id}/parse-results",
}
@@ -371,3 +409,113 @@ def test_replaced_job_signals_restart(clock, api_url, provision_token, session):
assert replaced.restart_required
assert replacement.heartbeat() == JobState.RUNNING
+
+
+class SecretReadingLocalBundle(LocalDagBundle):
+ def __init__(self, **kwargs):
+ super().__init__(**kwargs)
+ if SDKVariable.get("processor-client-key") != "from-api":
+ raise ValueError("Bundle constructor did not resolve its Variable")
+
+
+@pytest.mark.parametrize("multi_team", [False, True])
+@pytest.mark.parametrize("database_access", [False, True])
+def test_normal_processor_command_uses_authenticated_api(
+ api_url,
+ api_requests,
+ provision_token,
+ tmp_path,
+ session,
+ multi_team,
+ database_access,
+):
+ clear_db_dags()
+ clear_db_dag_bundles()
+ clear_db_teams()
+ team_name = "processor-team" if multi_team else None
+ if team_name:
+ session.add(Team(name=team_name))
+ session.commit()
+ dag_file = tmp_path / "api_dag.py"
+ dag_file.write_text(
+ "from airflow.sdk import DAG, Variable\n"
+ 'dag = DAG("authenticated_processor", schedule=None, description=Variable.get("processor-client-key"))\n'
+ )
+ bundle_config = [
+ {
+ "name": "bundle-a",
+ "classpath": f"{__name__}.SecretReadingLocalBundle",
+ "kwargs": {"path": str(tmp_path)},
+ **({"team_name": team_name} if team_name else {}),
+ }
+ ]
+ config = {
+ ("dag_processor", "execution_api_token_file"): str(provision_token()),
+ ("core", "execution_api_server_url"): api_url,
+ ("core", "load_examples"): "False",
+ ("core", "multi_team"): str(multi_team),
+ ("dag_processor", "dag_bundle_config_list"): json.dumps(bundle_config),
+ ("scheduler", "job_heartbeat_sec"): "0",
+ ("database", "sql_alchemy_conn"): conf.get("database", "sql_alchemy_conn"),
+ }
+ try:
+ with conf_vars(config):
+ Variable.set("processor-client-key", "from-api", team_name=team_name)
+ DagBundlesManager().sync_bundles_to_db(include_bundle_urls=False)
+ processor_config = dict(config)
+ if not database_access:
+ processor_config.update(
+ {
+ (
+ "database",
+ "sql_alchemy_conn",
+ ): "postgresql+psycopg://blocked:blocked@127.0.0.1:1/blocked",
+ ("core", "fernet_key"): "",
+ ("api_auth", "jwt_secret"): "",
+ ("api_auth", "jwt_private_key_path"): "",
+ }
+ )
+ result = subprocess.run(
+ ["airflow", "dag-processor", "--num-runs", "1", "--bundle-name", "bundle-a"],
+ env={
+ **os.environ,
+ **{
+ f"AIRFLOW__{section.upper()}__{key.upper()}": value
+ for (section, key), value in processor_config.items()
+ },
+ "PYTHONPATH": f"{Path(__file__).parents[3]}:{os.environ.get('PYTHONPATH', '')}",
+ },
+ capture_output=True,
+ text=True,
+ timeout=90,
+ check=False,
+ )
+ assert result.returncode == 0, result.stdout + result.stderr
+
+ session.expire_all()
+ job = session.scalars(select(Job)).one()
+ assert job.state == JobState.SUCCESS
+ assert job.end_date is not None
+ assert job.team_names == ([team_name] if team_name else [])
+ assert session.get(DagModel, "authenticated_processor").description == "from-api"
+ assert (
+ session.scalar(
+ select(SerializedDagModel).where(SerializedDagModel.dag_id == "authenticated_processor")
+ )
+ is not None
+ )
+ requests = {(path, claims.get("scope")) for path, claims in api_requests}
+ assert ("/execution/jobs", "dag_processor_session") in requests
+ assert (f"/execution/jobs/{job.id}/heartbeat", "dag_processor") in requests
+ assert (f"/execution/jobs/{job.id}/complete", "dag_processor") in requests
+ assert (f"/execution/jobs/{job.id}/parse-results", "dag_processor") in requests
+ assert ("/execution/variables/processor-client-key", "dag_processor") in requests
+ assert ("/execution/variables/processor-client-key", "dag_parse") in requests
+ parsing = [claims for _, claims in api_requests if claims.get("scope") == "dag_parse"]
+ assert all(claims["relative_fileloc"] == "api_dag.py" for claims in parsing)
+ finally:
+ clear_db_jobs()
+ clear_db_dags()
+ clear_db_dag_bundles()
+ clear_db_variables()
+ clear_db_teams()
diff --git a/airflow-core/tests/unit/cli/commands/test_dag_processor_command.py b/airflow-core/tests/unit/cli/commands/test_dag_processor_command.py
index 77de6f928d3..b0a69dd9378 100644
--- a/airflow-core/tests/unit/cli/commands/test_dag_processor_command.py
+++ b/airflow-core/tests/unit/cli/commands/test_dag_processor_command.py
@@ -17,6 +17,7 @@
# under the License.
from __future__ import annotations
+import logging
import os
from unittest import mock
@@ -24,6 +25,9 @@ import pytest
from airflow.cli import cli_parser
from airflow.cli.commands import dag_processor_command
+from airflow.dag_processing.manager import DagFileProcessorManager
+from airflow.jobs.dag_processor_job_runner import DagProcessorJobRunner
+from airflow.jobs.job import Job
from tests_common.test_utils.config import conf_vars
@@ -166,3 +170,135 @@ class TestDagProcessorCommand:
mock_daemon_option.call_args.kwargs["callback"]()
memray.Tracker.assert_called_once()
+
+ @conf_vars({("dag_processor", "execution_api_token_file"): "/run/processor.jwt"})
+ @mock.patch("airflow.cli.commands.dag_processor_command.run_command_with_daemon_option", autospec=True)
+ @mock.patch("airflow.dag_processing.api_client.DagProcessorAPIClient", autospec=True)
+ @mock.patch.object(DagProcessorJobRunner, "run_with_api", autospec=True)
+ @mock.patch("airflow.cli.commands.dag_processor_command.run_job", autospec=True)
+ def test_api_client_is_created_in_the_daemon_callback(
+ self, run_job, run_with_api, client_class, daemon, configure_testing_dag_bundle
+ ):
+ with configure_testing_dag_bundle(os.devnull):
+ args = self.parser.parse_args(["dag-processor", "--bundle-name", "testing"])
+ dag_processor_command.dag_processor(args)
+ client_class.assert_not_called()
+ daemon.call_args.kwargs["callback"]()
+ assert client_class.call_args.kwargs["token_file"] == "/run/processor.jwt"
+ assert client_class.call_args.kwargs["bundle_names"] == ["testing"]
+ run_with_api.assert_called_once()
+ assert run_with_api.call_args.args[1] is client_class.return_value.__enter__.return_value
+ client_class.return_value.__exit__.assert_called_once()
+ run_job.assert_not_called()
+
+ @conf_vars({("dag_processor", "execution_api_token_file"): "/run/processor.jwt"})
+ @mock.patch("airflow.utils.cli.validate_dag_bundle_arg", autospec=True)
+ def test_api_bundle_validation_does_not_construct_bundles(self, validate, configure_testing_dag_bundle):
+ with configure_testing_dag_bundle(os.devnull):
+ runner = dag_processor_command._create_dag_processor_job_runner(
+ self.parser.parse_args(["dag-processor", "--bundle-name", "testing"])
+ )
+ assert runner.processor.bundle_names_to_parse == ["testing"]
+ with pytest.raises(SystemExit, match="Bundles not found: missing"):
+ dag_processor_command._create_dag_processor_job_runner(
+ self.parser.parse_args(["dag-processor", "--bundle-name", "missing"])
+ )
+ validate.assert_not_called()
+
+ @conf_vars(
+ {
+ ("dag_processor", "execution_api_token_file"): "/run/processor.jwt",
+ ("core", "load_examples"): "False",
+ }
+ )
+ @mock.patch("airflow.dag_processing.api_client.DagProcessorAPIClient", autospec=True)
+ def test_api_client_closes_on_failure(self, client_class, configure_testing_dag_bundle):
+ runner = mock.create_autospec(DagProcessorJobRunner, instance=True)
+ runner.job = Job()
+ runner.processor = mock.create_autospec(DagFileProcessorManager, instance=True)
+ runner.processor.bundle_names_to_parse = None
+ runner.run_with_api.side_effect = RuntimeError("failed")
+ with configure_testing_dag_bundle(os.devnull), pytest.raises(RuntimeError, match="failed"):
+ dag_processor_command._run_dag_processor_job(runner)
+ assert client_class.call_args.kwargs["bundle_names"] == ["testing"]
+ client_class.return_value.__exit__.assert_called_once()
+
+
+@pytest.mark.parametrize("api_mode", [False, True])
+@pytest.mark.parametrize("verbose", [False, True])
+@mock.patch("airflow.cli.commands.dag_processor_command._start_dag_processor", autospec=True)
+@mock.patch("airflow.utils.cli.cli_action_loggers.on_pre_execution", autospec=True)
+@mock.patch("airflow.utils.cli.cli_action_loggers.on_post_execution", autospec=True)
+@mock.patch("airflow.utils.db.check_and_run_migrations", autospec=True)
+@mock.patch("airflow.utils.db.synchronize_log_template", autospec=True)
+def test_verbose_logging_in_both_modes(
+ sync_log_template,
+ check_migrations,
+ on_post_execution,
+ on_pre_execution,
+ start,
+ api_mode,
+ verbose,
+ monkeypatch,
+):
+ root_logger = logging.getLogger()
+ original_level = root_logger.level
+ handler = logging.NullHandler(level=logging.INFO)
+ monkeypatch.setattr(root_logger, "handlers", [handler])
+ root_logger.setLevel(logging.INFO)
+ try:
+ with conf_vars(
+ {("dag_processor", "execution_api_token_file"): "/run/processor.jwt" if api_mode else ""}
+ ):
+ args = cli_parser.get_parser().parse_args(["dag-processor", *(["--verbose"] if verbose else [])])
+ dag_processor_command.dag_processor(args)
+
+ expected_level = logging.DEBUG if verbose else logging.INFO
+ assert root_logger.level == expected_level
+ assert handler.level == expected_level
+ start.assert_called_once_with(args)
+ if api_mode:
+ on_pre_execution.assert_not_called()
+ on_post_execution.assert_not_called()
+ check_migrations.assert_not_called()
+ sync_log_template.assert_not_called()
+ else:
+ on_pre_execution.assert_called_once()
+ on_post_execution.assert_called_once()
+ finally:
+ root_logger.setLevel(original_level)
+
+
+@mock.patch.object(DagFileProcessorManager, "sync_bundles", autospec=True)
+@mock.patch.object(DagProcessorJobRunner, "_execute", autospec=True)
+def test_catalog_provisioning_does_not_start_a_processor(execute, sync):
+ args = cli_parser.get_parser().parse_args(["dag-processor", "--sync-bundles-only"])
+ dag_processor_command.dag_processor(args)
+ sync.assert_called_once()
+ execute.assert_not_called()
+
+
+@pytest.mark.parametrize("command", [["dag-processor"], ["dag-processor", "--sync-bundles-only"]])
+@mock.patch("airflow.utils.cli.cli_action_loggers.on_pre_execution", autospec=True)
+@mock.patch.object(DagFileProcessorManager, "sync_bundles", autospec=True)
+@mock.patch("airflow.cli.commands.dag_processor_command.DagProcessorJobRunner")
+def test_action_log_keeps_the_command_name(mock_runner, sync, on_pre_execution, command):
+ mock_runner.return_value.job_type = "DagProcessorJob"
+ dag_processor_command.dag_processor(cli_parser.get_parser().parse_args(command))
+ assert on_pre_execution.call_args.kwargs["sub_command"] == "dag_processor"
+
+
+@mock.patch.object(DagFileProcessorManager, "sync_bundles", autospec=True)
+def test_catalog_provisioning_for_named_bundles_keeps_other_bundles(sync, configure_testing_dag_bundle):
+ parser = cli_parser.get_parser()
+ with configure_testing_dag_bundle(os.devnull):
+ dag_processor_command.dag_processor(
+ parser.parse_args(["dag-processor", "--sync-bundles-only", "--bundle-name", "testing"])
+ )
+ with pytest.raises(SystemExit, match="Bundles not found: unknown"):
+ dag_processor_command.dag_processor(
+ parser.parse_args(["dag-processor", "--sync-bundles-only", "--bundle-name", "unknown"])
+ )
+ manager = sync.call_args.args[0]
+ assert manager.bundle_names_to_parse == ["testing"]
+ sync.assert_called_once()
diff --git a/airflow-core/tests/unit/dag_processing/test_health.py b/airflow-core/tests/unit/dag_processing/test_health.py
new file mode 100644
index 00000000000..ffd118fdb31
--- /dev/null
+++ b/airflow-core/tests/unit/dag_processing/test_health.py
@@ -0,0 +1,64 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+from unittest import mock
+
+import pytest
+
+from airflow.dag_processing.health import check_processor_health, write_processor_health
+
+from tests_common.test_utils.config import conf_vars
+
+
+@pytest.mark.parametrize(
+ ("loop_age", "api_age", "readiness", "healthy"),
+ [
+ (0, 0, False, True),
+ (31, 0, False, False),
+ (0, 31, False, True),
+ (0, 31, True, False),
+ (0, 0, True, True),
+ (-1, 0, False, False),
+ ],
+)
+@mock.patch("airflow.dag_processing.health.monotonic", autospec=True)
+def test_local_health_requires_a_fresh_loop(clock, tmp_path, loop_age, api_age, readiness, healthy):
+ with conf_vars(
+ {
+ ("dag_processor", "health_check_file"): str(tmp_path / "health.json"),
+ ("dag_processor", "health_check_threshold"): "30",
+ }
+ ):
+ clock.return_value = 100 - loop_age
+ write_processor_health(100 - api_age)
+ clock.return_value = 100
+ if healthy:
+ check_processor_health(readiness=readiness)
+ else:
+ with pytest.raises(SystemExit):
+ check_processor_health(readiness=readiness)
+
+
+@pytest.mark.parametrize("contents", [None, "broken", "{}"])
+def test_missing_or_invalid_health_file_fails_closed(tmp_path, contents):
+ path = tmp_path / "health.json"
+ if contents is not None:
+ path.write_text(contents)
+ with conf_vars({("dag_processor", "health_check_file"): str(path)}), pytest.raises(SystemExit):
+ check_processor_health()
diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py
index 8f6959bb4fb..c6d55d96a68 100644
--- a/airflow-core/tests/unit/dag_processing/test_manager.py
+++ b/airflow-core/tests/unit/dag_processing/test_manager.py
@@ -57,11 +57,13 @@ from airflow.api_fastapi.execution_api.datamodels.dag_parsing import (
DagBundleStateResponse,
DagParseResultBody,
ParseSourceCode,
+ ProcessorWorkItem,
)
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,
+ DagProcessorRegistrationRetired,
DagProcessorSecretsComms,
)
from airflow.dag_processing.bundles.base import BaseDagBundle, BundleVersion
@@ -74,7 +76,9 @@ from airflow.dag_processing.manager import (
DagFileInfo,
DagFileProcessorManager,
DagFileStat,
+ DeferredCallback,
PendingDagPublication,
+ PriorityParseClaim,
)
from airflow.dag_processing.processor import (
DagFileParseRequest,
@@ -476,10 +480,11 @@ class TestDagFileProcessorManager:
"test_zip.zip/broken_dag.py",
}
+ @pytest.mark.parametrize("api_mode", [False, True], ids=["direct", "api"])
@pytest.mark.parametrize("absolute", [False, True], ids=["relative", "absolute"])
@mock.patch("airflow.dag_processing.manager.get_importer_registry", autospec=True)
def test_find_files_in_bundle_makes_discovery_errors_bundle_relative(
- self, mock_get_registry, tmp_path, absolute
+ self, mock_get_registry, tmp_path, absolute, api_mode
):
(tmp_path / "broken.dag").write_text("")
source_reference = os.fspath(tmp_path / "broken.dag") if absolute else "broken.dag"
@@ -487,8 +492,9 @@ class TestDagFileProcessorManager:
(mock.sentinel.importer, DagImportError(source_reference=source_reference, message="boom"))
]
bundle = _make_bundle(tmp_path)
+ client = mock.create_autospec(DagProcessorAPIClient, instance=True) if api_mode else None
- found_files = DagFileProcessorManager(max_runs=1)._find_files_in_bundle(bundle)
+ found_files = DagFileProcessorManager(max_runs=1, api_client=client)._find_files_in_bundle(bundle)
assert [file.definition_locs for file in found_files] == [frozenset({"broken.dag"})]
assert [file.failed_sources for file in found_files] == [frozenset({"broken.dag"})]
@@ -3354,6 +3360,7 @@ class TestDagFileProcessorManager:
with mock.patch.object(
DagFileProcessorProcess, "start", side_effect=lambda *args, **kwargs: self.mock_processor()[0]
) as start:
+ manager._bundle_parse_sources["testing"] = DagParseSource()
manager._start_new_processes()
# Callbacks passed to processor
assert start.call_args_list == [
@@ -4762,7 +4769,7 @@ class TestDagFileProcessorManager:
@pytest.mark.usefixtures("testing_dag_bundle")
def test_bundle_version_data_stored_after_refresh(self, session):
- """Capture version metadata alongside the refreshed source."""
+ """Keep the local source metadata independently of its persisted bundle state."""
from airflow.dag_processing.bundles.base import BundleVersion
manager = DagFileProcessorManager(max_runs=1)
@@ -5129,9 +5136,59 @@ class TestDagFileProcessorManager:
read_end.close()
proc.close()
+ def test_requested_work_is_readmitted_ahead_of_queued_files(self):
+ manager = DagFileProcessorManager(max_runs=1)
+ requested = DagFileInfo(bundle_name="testing", rel_path=Path("requested.py"))
+ parked = DagFileInfo(bundle_name="testing", rel_path=Path("parked.py"))
+ queued = DagFileInfo(bundle_name="other", rel_path=Path("queued.py"))
+ _claim_priority_parse(manager, requested)
+ manager._files_awaiting_source = dict.fromkeys([parked, requested])
+ manager._file_queue[queued] = None
+ manager._bundle_parse_sources["testing"] = DagParseSource()
+
+ manager._readmit_files_with_source()
+
+ assert list(manager._file_queue) == [requested, queued, parked]
+
+ @mock.patch.object(DagFileProcessorManager, "_create_process", autospec=True)
+ def test_requested_file_still_parsing_stays_ahead_of_the_queue(self, create):
+ manager = DagFileProcessorManager(max_runs=1)
+ manager._parallelism = 2
+ requested, other, later = _get_file_infos(["requested.py", "other.py", "later.py"])
+ _claim_priority_parse(manager, requested)
+ manager._bundle_parse_sources["testing"] = DagParseSource()
+ manager._processors[requested] = self.mock_processor()[0]
+ manager._file_queue = OrderedDict.fromkeys([requested, other, later])
+ create.return_value = self.mock_processor()[0]
+
+ manager._start_new_processes()
+
+ assert other in manager._processors
+ assert list(manager._file_queue) == [requested, later]
+
+ @mock.patch.object(DagFileProcessorManager, "prepare_callback_bundle", autospec=True)
+ def test_queuing_a_callback_heartbeats_before_preparing_its_bundle(self, mock_prepare_callback_bundle):
+ events = []
+ manager = DagFileProcessorManager(max_runs=1)
+ manager.heartbeat = lambda: events.append("heartbeat")
+ mock_prepare_callback_bundle.side_effect = lambda manager, request: events.append("prepare")
+ request = DagCallbackRequest(
+ filepath="file1.py",
+ dag_id="dag1",
+ run_id="run1",
+ is_failure_callback=False,
+ bundle_name="testing",
+ bundle_version="some_commit_hash",
+ msg=None,
+ )
+
+ manager._add_callback_to_queue(request)
+
+ assert events == ["heartbeat", "prepare"]
+
@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):
+ def test_listing_the_api_cannot_accept_releases_priority_capacity(self, observed, failure):
observed.return_value = (
{r"bad\dag.py"} if failure == "invalid_path" else {f"{index}.py" for index in range(100001)}
)
@@ -5140,6 +5197,10 @@ class TestDagFileProcessorManager:
max_runs=1, api_client=client, parallelism=1, bundle_refresh_check_interval=0
)
bundle = self._make_refresh_bundle()
+ work = ProcessorWorkItem(
+ id=str(uuid7()), claim_id=uuid7(), bundle_name=bundle.name, relative_fileloc="a.py"
+ )
+ manager._priority_claims[work.id] = PriorityParseClaim(work=work, refresh_generation=1)
manager._bundles_waiting_for_refresh.add(bundle.name)
manager._bundle_versions[bundle.name] = "previous"
@@ -5154,6 +5215,12 @@ class TestDagFileProcessorManager:
assert manager._bundle_versions[bundle.name] == "previous"
assert not manager._force_refresh_bundles
assert not manager._bundles_waiting_for_refresh
+ assert not manager._priority_claims
+ assert manager._pending_work_acks == [("priority", work, True)]
+ manager._acknowledge_requested_work()
+ client.claim_work.return_value = []
+ manager.claim_priority_files()
+ client.claim_work.assert_called_once_with("priority", [bundle.name], 1)
observed.return_value = set()
client.publish_inventory.return_value.revision = uuid7()
@@ -5162,6 +5229,34 @@ class TestDagFileProcessorManager:
assert manager._bundle_versions[bundle.name] is None
assert not manager._pending_inventories
+ def test_refresh_dag_bundles_heartbeats_before_each_bundle(self):
+ events = []
+ manager = DagFileProcessorManager(max_runs=1)
+ manager.heartbeat = lambda: events.append("heartbeat")
+ manager._dag_bundles = []
+ for name in ("first", "second"):
+ bundle = self._make_refresh_bundle()
+ bundle.name = name
+ bundle.refresh.side_effect = lambda name=name: events.append(name)
+ manager._dag_bundles.append(bundle)
+ manager._force_refresh_bundles = set()
+
+ with (
+ mock.patch.object(
+ manager, "get_bundle_state", return_value=BundleState(last_refreshed=None, version=None)
+ ),
+ mock.patch.object(manager, "update_bundle_state"),
+ mock.patch.object(manager, "_find_files_in_bundle", return_value=set()),
+ mock.patch.object(manager, "deactivate_deleted_dags"),
+ mock.patch.object(manager, "clear_orphaned_import_errors"),
+ mock.patch.object(manager, "handle_removed_files"),
+ mock.patch.object(manager, "_resort_file_queue"),
+ mock.patch.object(manager, "_add_new_files_to_queue"),
+ ):
+ manager._refresh_dag_bundles({})
+
+ assert events == ["heartbeat", "first", "heartbeat", "second"]
+
class TestMultiTeamMetrics:
"""Tests for team_name tag on dag processing metrics in multi-team mode."""
@@ -5251,13 +5346,13 @@ class TestMultiTeamMetrics:
@mock.patch("airflow.dag_processing.manager.stats.incr")
def test_start_new_processes_includes_team_name(self, mock_incr, mock_get_team_name):
manager = DagFileProcessorManager(max_runs=1)
+ manager._bundle_parse_sources["testing"] = DagParseSource()
dag_file = DagFileInfo(
bundle_name="testing",
rel_path=Path("dag_file.py"),
bundle_path=TEST_DAGS_FOLDER,
)
manager._file_queue[dag_file] = None
- manager._bundle_parse_sources["testing"] = DagParseSource()
manager._create_process = MagicMock(return_value=self.mock_processor())
manager._start_new_processes()
@@ -5753,6 +5848,247 @@ class TestDagFileProcessorManagerWithAPIClient:
assert make_manager().prepare_callback_bundle(request) is not None
assert self._get_connection_bundles(api_requests) == ["bundle_b"]
+ @mock.patch("airflow.sdk.api.client.ConnectionOperations.get", autospec=True)
+ @mock.patch.object(DagBundlesManager, "get_bundle", autospec=True)
+ def test_git_callback_connection_outage_is_deferred(self, get_bundle, get_connection, make_manager):
+ git_bundle = pytest.importorskip("airflow.providers.git.bundles.git").GitDagBundle
+ get_bundle.side_effect = lambda _, **kwargs: git_bundle(tracking_ref="main", **kwargs)
+ get_connection.side_effect = httpx2.ConnectError("API unavailable")
+ manager = make_manager()
+ request = DagCallbackRequest(
+ filepath="a.py", bundle_name="bundle_b", bundle_version="v1", dag_id="dag", run_id="run"
+ )
+ work = ProcessorWorkItem(
+ id=str(uuid7()), claim_id=uuid7(), bundle_name="bundle_b", relative_fileloc="a.py"
+ )
+ manager._callback_claims[id(request)] = work
+ manager._deferred_api_callbacks.append(request)
+
+ manager._add_callback_to_queue(request)
+
+ get_connection.assert_called_once()
+ assert id(request) in manager._deferred_callback_retries
+ assert not manager._pending_work_acks
+
+
+@pytest.mark.parametrize(
+ ("failure", "stage"),
+ [
+ ("transport", "construct"),
+ ("server", "construct"),
+ ("retired", "construct"),
+ ("wrapped_retired", "construct"),
+ ("retired", "initialize"),
+ ],
+)
+@mock.patch.object(DagFileProcessorManager, "_create_bundle_manager", autospec=True)
+def test_callback_construction_failure_uses_bounded_retry(create, failure, stage):
+ client = mock.create_autospec(DagProcessorAPIClient, instance=True)
+ client.restart_required = failure == "wrapped_retired"
+ manager = DagFileProcessorManager(max_runs=1, api_client=client)
+ request = DagCallbackRequest(
+ filepath="a.py", bundle_name="testing", bundle_version="v1", dag_id="dag", run_id="run"
+ )
+ work = ProcessorWorkItem(
+ id=str(uuid7()), claim_id=uuid7(), bundle_name="testing", relative_fileloc="a.py"
+ )
+ manager._callback_claims[id(request)] = work
+ manager._deferred_api_callbacks.append(request)
+ response = httpx2.Response(503, request=httpx2.Request("GET", "http://api/connections/git"))
+ errors = {
+ "transport": httpx2.ConnectError("API unavailable"),
+ "server": httpx2.HTTPStatusError("API unavailable", request=response.request, response=response),
+ "retired": DagProcessorRegistrationRetired("retired"),
+ "wrapped_retired": RuntimeError("provider wrapped a retired registration"),
+ }
+ if stage == "construct":
+ create.return_value.get_bundle.side_effect = errors[failure]
+ else:
+ create.return_value.get_bundle.return_value.supports_versioning = True
+ create.return_value.get_bundle.return_value.initialize.side_effect = errors[failure]
+ if "retired" in failure:
+ with pytest.raises(DagProcessorRegistrationRetired):
+ manager._add_callback_to_queue(request)
+ assert not manager._deferred_callback_retries
+ return
+
+ manager._add_callback_to_queue(request)
+ assert id(request) in manager._deferred_callback_retries
+ assert not manager._pending_work_acks
+ manager._deferred_callback_retries[id(request)].next_attempt_time = 0
+ create.return_value.get_bundle.side_effect = None
+ create.return_value.get_bundle.return_value = _make_bundle(TEST_DAGS_FOLDER)
+ manager._add_callback_to_queue(request)
+ assert not manager._deferred_callback_retries
+ assert not manager._deferred_api_callbacks
+ assert list(manager._callback_to_execute.values()) == [[request]]
+
+
+@conf_vars({("dag_processor", "job_heartbeat_timeout"): "10"})
+@mock.patch("airflow.dag_processing.manager.time", autospec=True)
+@mock.patch.object(DagFileProcessorManager, "_create_process", autospec=True)
+@mock.patch.object(DagFileProcessorManager, "_create_bundle_manager", autospec=True)
+def test_callback_constructor_outage_allows_other_work_until_retry_expires(
+ create_bundle_manager, create_process, mock_time, tmp_path
+):
+ mock_time.monotonic.return_value = 0
+ client = mock.create_autospec(DagProcessorAPIClient, instance=True)
+ client.restart_required = False
+ manager = DagFileProcessorManager(max_runs=1, api_client=client, parallelism=2)
+ bundle = _make_bundle(tmp_path)
+ bundle.is_initialized = True
+ manager._dag_bundles = [bundle]
+ manager._bundle_parse_sources[bundle.name] = DagParseSource()
+
+ def get_bundle(*, name, version, version_data):
+ if version == "unavailable":
+ raise httpx2.ConnectError("Bundle connection lookup unavailable")
+ return bundle
+
+ create_bundle_manager.return_value.get_bundle.side_effect = get_bundle
+ work = []
+ for version in ("unavailable", "available"):
+ request = DagCallbackRequest(
+ filepath=f"{version}.py",
+ bundle_name=bundle.name,
+ bundle_version=version,
+ dag_id=version,
+ run_id="run",
+ )
+ work.append(
+ ProcessorWorkItem(
+ id=str(uuid7()),
+ claim_id=uuid7(),
+ bundle_name=bundle.name,
+ relative_fileloc=request.filepath,
+ callback=request.model_dump_json(),
+ )
+ )
+ client.claim_work.return_value = work
+ for request in manager.fetch_callbacks():
+ manager._add_callback_to_queue(request)
+ assert len(manager._deferred_api_callbacks) == 1
+ normal = DagFileInfo(bundle_name=bundle.name, rel_path=Path("normal.py"), bundle_path=tmp_path)
+ manager._add_files_to_queue([normal], mode="back")
+
+ def create(manager, file):
+ proc = mock.create_autospec(DagFileProcessorProcess, instance=True)
+ proc.id = uuid7()
+ proc.had_callbacks = file.bundle_version is not None
+ return proc
+
+ create_process.side_effect = create
+ manager._start_new_processes()
+ assert {file.rel_path for file in manager._processors} == {Path("available.py"), Path("normal.py")}
+ assert not manager._pending_work_acks
+
+ mock_time.monotonic.return_value = 10
+ for request in list(manager._deferred_api_callbacks):
+ manager._add_callback_to_queue(request)
+ assert not manager._deferred_api_callbacks
+ assert not manager._deferred_callback_retries
+ assert manager._pending_work_acks == [("callbacks", work[0], True)]
+ manager._acknowledge_requested_work()
+ manager._acknowledge_requested_work()
+ client.acknowledge_work.assert_called_once_with("callbacks", work[0], failed=True)
+
+
+@pytest.mark.parametrize("outcome", ["queued", "expired", "unsupported"])
+@mock.patch("airflow.dag_processing.manager.time", autospec=True)
+@mock.patch("airflow.dag_processing.manager.get_claiming_importer", autospec=True)
+@mock.patch.object(DagFileProcessorManager, "prepare_callback_bundle", autospec=True)
+def test_equal_claimed_callbacks_are_removed_by_identity(prepare, importer, clock, outcome):
+ client = mock.create_autospec(DagProcessorAPIClient, instance=True)
+ manager = DagFileProcessorManager(max_runs=1, api_client=client)
+ first = DagCallbackRequest(
+ filepath="a.py", bundle_name="testing", bundle_version="v1", dag_id="dag", run_id="run"
+ )
+ second = first.model_copy(deep=True)
+ work = [
+ ProcessorWorkItem(id=str(uuid7()), claim_id=uuid7(), bundle_name="testing", relative_fileloc="a.py")
+ for _ in range(2)
+ ]
+ manager._callback_claims = {id(first): work[0], id(second): work[1]}
+ manager._deferred_api_callbacks = [first, second]
+ manager._deferred_callback_retries = {
+ id(first): DeferredCallback(expires_at=20, next_attempt_time=5),
+ id(second): DeferredCallback(expires_at=0),
+ }
+ clock.monotonic.return_value = 0
+ importer.return_value = object() if outcome == "unsupported" else None
+ prepare.return_value = None if outcome == "expired" else _make_bundle(TEST_DAGS_FOLDER)
+
+ manager._add_callback_to_queue(second)
+
+ assert len(manager._deferred_api_callbacks) == 1
+ assert manager._deferred_api_callbacks[0] is first
+ assert set(manager._deferred_callback_retries) == {id(first)}
+ assert manager._callback_claims[id(first)] is work[0]
+ if outcome == "queued":
+ assert next(iter(manager._callback_to_execute.values()))[0] is second
+ assert manager._callback_claims[id(second)] is work[1]
+ assert not manager._pending_work_acks
+ else:
+ assert id(second) not in manager._callback_claims
+ assert manager._pending_work_acks == [("callbacks", work[1], True)]
+
+
+@mock.patch("airflow.dag_processing.manager.get_claiming_importer", autospec=True, return_value=None)
+@mock.patch.object(DagFileProcessorProcess, "start", autospec=True)
+@mock.patch.object(DagFileProcessorManager, "_get_logger_for_dag_file", autospec=True)
+@mock.patch.object(DagFileProcessorManager, "prepare_callback_bundle", autospec=True)
+def test_equal_claimed_callbacks_dispatch_once_each_after_preparation_recovers(
+ prepare, logger, start, importer
+):
+ client = mock.create_autospec(DagProcessorAPIClient, instance=True)
+ client.claim_work.return_value = []
+ manager = DagFileProcessorManager(max_runs=1, api_client=client)
+ first = DagCallbackRequest(
+ filepath="a.py", bundle_name="testing", bundle_version="v1", dag_id="dag", run_id="run"
+ )
+ second = first.model_copy(deep=True)
+ work = [
+ ProcessorWorkItem(id=str(uuid7()), claim_id=uuid7(), bundle_name="testing", relative_fileloc="a.py")
+ for _ in range(2)
+ ]
+ manager._callback_claims = {id(first): work[0], id(second): work[1]}
+ manager._deferred_api_callbacks = [first, second]
+ bundle = _make_bundle(TEST_DAGS_FOLDER)
+ prepare.side_effect = [None, bundle, bundle]
+ logger.return_value = (mock.MagicMock(spec=structlog.stdlib.BoundLogger), mock.mock_open()())
+
+ def create_process(**kwargs):
+ process = mock.create_autospec(DagFileProcessorProcess, instance=True)
+ process.id = kwargs["id"]
+ process._exit_code = 0
+ return process
+
+ start.side_effect = create_process
+ for request in manager.fetch_callbacks():
+ manager._add_callback_to_queue(request)
+
+ for expected in (second, first):
+ file = next(iter(manager._callback_to_execute))
+ process = manager._create_process(file)
+ assert start.call_args.kwargs["callbacks"][0] is expected
+ manager._finish_callback_attempt(process)
+ manager._acknowledge_requested_work()
+ if expected is second:
+ manager._deferred_callback_retries[id(first)].next_attempt_time = 0
+ for request in manager.fetch_callbacks():
+ manager._add_callback_to_queue(request)
+
+ assert start.call_count == 2
+ assert manager.fetch_callbacks() == []
+ assert not manager._callback_claims
+ assert not manager._callback_attempts
+ assert not manager._deferred_callback_retries
+ assert not manager._callback_to_execute
+ assert client.acknowledge_work.call_args_list == [
+ mock.call("callbacks", work[1], failed=False),
+ mock.call("callbacks", work[0], failed=False),
+ ]
+
@mock.patch.object(DagFileProcessorManager, "_find_files_in_bundle", autospec=True, return_value=set())
@conf_vars({("core", "multi_team"): "True"})
@@ -5811,7 +6147,7 @@ def test_refresh_reads_one_fresh_catalog_per_pass(find_files, tmp_path):
@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(
+def test_catalog_omission_releases_priority_capacity_and_allows_recovery(
find_files, clock, catalog_failed, has_pending_inventory, tmp_path
):
clock.monotonic.return_value = 100
@@ -5820,6 +6156,7 @@ def test_catalog_omission_releases_source_and_allows_recovery(
manager = DagFileProcessorManager(max_runs=1, api_client=client, bundle_refresh_check_interval=10)
bundles = {}
files = {}
+ claims = {}
for name in ("missing", "ready"):
bundle = _make_bundle(tmp_path)
bundle.name = name
@@ -5831,6 +6168,9 @@ def test_catalog_omission_releases_source_and_allows_recovery(
files[name] = {file}
manager._bundle_versions[name] = None
manager._bundle_parse_sources[name] = DagParseSource()
+ work = ProcessorWorkItem(id=str(uuid7()), claim_id=uuid7(), bundle_name=name, relative_fileloc="a.py")
+ claims[name] = PriorityParseClaim(work=work, refresh_generation=1)
+ manager._priority_claims[work.id] = claims[name]
manager._dag_bundles = list(bundles.values())
previous_files = files.copy()
previous_versions = manager._bundle_versions.copy()
@@ -5859,7 +6199,10 @@ def test_catalog_omission_releases_source_and_allows_recovery(
assert files == previous_files
assert manager._bundle_versions == previous_versions
assert manager._bundle_parse_sources["ready"] is ready_source
+ assert manager._priority_claims[claims["ready"].work.id] is claims["ready"]
if catalog_failed:
+ assert manager._priority_claims[claims["missing"].work.id] is claims["missing"]
+ assert not manager._pending_work_acks
assert manager._force_refresh_bundles == {"missing"}
assert manager._bundles_waiting_for_refresh == {"missing"}
if has_pending_inventory:
@@ -5871,12 +6214,18 @@ def test_catalog_omission_releases_source_and_allows_recovery(
return
client.publish_inventory.assert_not_called()
+ assert claims["missing"].work.id not in manager._priority_claims
+ assert manager._pending_work_acks == [("priority", claims["missing"].work, True)]
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()
+ manager._acknowledge_requested_work()
+ client.claim_work.return_value = []
+ manager.claim_priority_files()
+ client.claim_work.assert_called_once_with("priority", ["missing", "ready"], 1)
clock.monotonic.return_value = 110
client.get_bundles.return_value = [ready_state, ready_state.model_copy(update={"name": "missing"})]
@@ -5988,6 +6337,106 @@ def test_oversized_inventory_preserves_previous_inventory_and_other_bundle_progr
assert client.get_bundles.call_count == calls
+@pytest.mark.parametrize("rejection", ["local_limit", "server_limit"])
+@pytest.mark.parametrize("retained_attempt", [None, "processor", "publication"])
+@mock.patch("airflow.dag_processing.manager.MAX_PUBLICATION_BYTES", 1024)
+@mock.patch.object(DagFileProcessorManager, "_create_process", autospec=True)
+def test_oversized_inventory_releases_unassigned_priority_capacity(
+ create, retained_attempt, rejection, tmp_path
+):
+ client = mock.create_autospec(DagProcessorAPIClient, instance=True)
+ manager = DagFileProcessorManager(max_runs=1, api_client=client, parallelism=3 if retained_attempt else 1)
+ for name in ("oversized", "ready"):
+ bundle = _make_bundle(tmp_path)
+ bundle.name = name
+ manager._dag_bundles.append(bundle)
+ (tmp_path / "a.py").touch()
+ rejected_work, ready_work = [
+ ProcessorWorkItem(id=str(uuid7()), claim_id=uuid7(), bundle_name=name, relative_fileloc="a.py")
+ for name in ("oversized", "ready")
+ ]
+ retained_claims = {}
+ if retained_attempt:
+ for name, attempt_id in (("oversized", uuid7()), ("ready", None)):
+ work = ProcessorWorkItem(
+ id=str(uuid7()), claim_id=uuid7(), bundle_name=name, relative_fileloc="b.py"
+ )
+ retained_claims[work.id] = PriorityParseClaim(
+ work=work, refresh_generation=1, attempt_id=attempt_id
+ )
+ if name == "oversized":
+ active_file = DagFileInfo(bundle_name=name, rel_path=Path("b.py"), bundle_path=tmp_path)
+ if retained_attempt == "processor":
+ process = mock.create_autospec(DagFileProcessorProcess, instance=True)
+ process.id = attempt_id
+ manager._processors[active_file] = process
+ else:
+ pending_body = mock.create_autospec(DagParseResultBody, instance=True)
+ pending_body.attempt_id = attempt_id
+ manager._pending_publications[active_file] = PendingDagPublication(
+ body=pending_body, stat=DagFileStat(), refresh_generation=1
+ )
+ manager._priority_claims.update(retained_claims)
+ client.claim_work.side_effect = [[rejected_work], [ready_work]]
+ manager._queue_requested_files_for_parsing()
+ assert manager._priority_claims[rejected_work.id].attempt_id is None
+ body = DagBundleInventoryBody(
+ attempt_id=uuid7(),
+ dispatch_sequence=1,
+ expected_revision=None,
+ files=[f"{index}/{'é' * 80}.py" for index in range(6)] if rejection == "local_limit" else [],
+ )
+ manager._pending_inventories["oversized"] = (body, set(), DagParseSource(refresh_generation=1))
+ client.publish_inventory.side_effect = _make_status_error(413)
+
+ manager._publish_bundle_inventory("oversized", {})
+
+ assert manager._priority_claims == retained_claims
+ assert manager._pending_work_acks == [("priority", rejected_work, True)]
+ if rejection == "local_limit":
+ client.publish_inventory.assert_not_called()
+ else:
+ client.publish_inventory.assert_called_once_with("oversized", body)
+ assert manager.claim_priority_files() == []
+ client.claim_work.assert_called_once_with("priority", ["oversized", "ready"], 1)
+ manager._acknowledge_requested_work()
+ client.acknowledge_work.assert_called_once_with("priority", rejected_work, failed=True)
+ manager._queue_requested_files_for_parsing()
+ assert client.claim_work.call_args_list == [
+ mock.call("priority", ["oversized", "ready"], 1),
+ mock.call("priority", ["oversized", "ready"], 1),
+ ]
+ manager._force_refresh_bundles.remove("ready")
+ manager._invalidate_bundle_parse_source("ready")
+ ready_file = DagFileInfo(bundle_name="ready", rel_path=Path("a.py"), bundle_path=tmp_path)
+ body = DagBundleInventoryBody(
+ attempt_id=uuid7(), dispatch_sequence=2, expected_revision=None, files=["a.py", "b.py"]
+ )
+ manager._pending_inventories["ready"] = (
+ body,
+ {ready_file},
+ DagParseSource(refresh_generation=1),
+ )
+ client.publish_inventory.side_effect = None
+ client.publish_inventory.return_value = DagBundleInventoryResponse(
+ attempt_id=body.attempt_id, accepted_at=timezone.utcnow(), revision=uuid7()
+ )
+ manager._publish_bundle_inventory("ready", {})
+ process = mock.create_autospec(DagFileProcessorProcess, instance=True)
+ process.id = uuid7()
+ process.had_callbacks = False
+ create.return_value = process
+
+ manager._start_new_processes()
+
+ create.assert_called_once_with(manager, ready_file)
+ assert manager._priority_claims[ready_work.id].attempt_id == process.id
+ assert {key: claim for key, claim in manager._priority_claims.items() if key != ready_work.id} == (
+ retained_claims
+ )
+ assert not manager._pending_work_acks
+
+
@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)
@@ -6016,14 +6465,318 @@ def test_oversized_inventory_byte_boundary_preserves_auth_failures(auth_status):
manager._publish_bundle_inventory("testing", {})
+@pytest.mark.parametrize("status_code", [403, 404, 409, 503])
+def test_requested_work_ack_retains_only_retryable_failures(status_code):
+ client = mock.create_autospec(DagProcessorAPIClient, instance=True)
+ manager = DagFileProcessorManager(max_runs=1, api_client=client)
+ work = ProcessorWorkItem(
+ id=str(uuid7()), claim_id=uuid7(), bundle_name="testing", relative_fileloc="a.py"
+ )
+ manager._pending_work_acks.append(("callbacks", work, False))
+ response = httpx2.Response(status_code, request=httpx2.Request("POST", "http://api/ack"))
+ client.acknowledge_work.side_effect = httpx2.HTTPStatusError(
+ "Unavailable", request=response.request, response=response
+ )
+ manager._acknowledge_requested_work()
+ assert bool(manager._pending_work_acks) == (status_code == 503)
+ if status_code == 503:
+ client.acknowledge_work.side_effect = None
+ manager._acknowledge_requested_work()
+ assert not manager._pending_work_acks
+ assert client.acknowledge_work.call_args.args == ("callbacks", work)
+
+
+def test_api_inventory_cleanup_preserves_claimed_callbacks():
+ client = mock.create_autospec(DagProcessorAPIClient, instance=True)
+ manager = DagFileProcessorManager(max_runs=1, api_client=client)
+ file = _get_file_infos(["removed.py"])[0]
+ request = DagCallbackRequest(
+ filepath="removed.py", bundle_name="testing", bundle_version="older", dag_id="dag", run_id="run"
+ )
+ manager._callback_to_execute[file].append(request)
+ manager._file_queue[file] = None
+ processor = mock.create_autospec(DagFileProcessorProcess, instance=True)
+ processor.id = uuid7()
+ manager._processors[file] = processor
+ manager._callback_attempts[processor.id] = [
+ ProcessorWorkItem(
+ id=str(uuid7()), claim_id=uuid7(), bundle_name="testing", relative_fileloc="removed.py"
+ )
+ ]
+ manager.purge_removed_files_from_queue(set())
+ manager.terminate_orphan_processes(set())
+ assert list(manager._file_queue) == [file]
+ processor.kill.assert_not_called()
+ processor.close.assert_not_called()
+
+
+@pytest.mark.parametrize("generation", [0, 1])
+@pytest.mark.parametrize("discovery_failed", [False, True])
+def test_inventory_only_fails_priority_requests_after_authoritative_refresh(generation, discovery_failed):
+ client = mock.create_autospec(DagProcessorAPIClient, instance=True)
+ client.publish_inventory.return_value.revision = uuid7()
+ manager = DagFileProcessorManager(max_runs=1, api_client=client)
+ work = ProcessorWorkItem(
+ id=str(uuid7()), claim_id=uuid7(), bundle_name="testing", relative_fileloc="archive.zip/a.py"
+ )
+ manager._priority_claims[work.id] = PriorityParseClaim(work=work, refresh_generation=1)
+ body = DagBundleInventoryBody(
+ attempt_id=uuid7(),
+ dispatch_sequence=1,
+ expected_revision=None,
+ files=["archive.zip"],
+ failed_sources=["archive.zip"] if discovery_failed else [],
+ )
+ manager._pending_inventories["testing"] = (body, set(), DagParseSource(refresh_generation=generation))
+
+ manager._publish_bundle_inventory("testing", {})
+
+ failed = generation == 1 and not discovery_failed
+ assert manager._pending_work_acks == ([("priority", work, True)] if failed else [])
+ assert (work.id in manager._priority_claims) is not failed
+
+
+@pytest.mark.parametrize("kind", ["callbacks", "priority"])
+def test_pending_requested_work_does_not_block_other_available_capacity(kind):
+ client = mock.create_autospec(DagProcessorAPIClient, instance=True)
+ client.claim_work.return_value = []
+ manager = DagFileProcessorManager(max_runs=1, api_client=client, parallelism=2, max_callbacks_per_loop=2)
+ bundle = mock.create_autospec(BaseDagBundle, instance=True)
+ bundle.name = "testing"
+ bundle.supports_versioning = False
+ manager._dag_bundles = [bundle]
+ work = ProcessorWorkItem(
+ id=str(uuid7()), claim_id=uuid7(), bundle_name="testing", relative_fileloc="missing.py"
+ )
+ if kind == "priority":
+ manager._priority_claims[work.id] = PriorityParseClaim(work=work, refresh_generation=1)
+ manager.claim_priority_files()
+ else:
+ manager._callback_claims[1] = work
+ manager.fetch_callbacks()
+ client.claim_work.assert_called_once_with(kind, ["testing"], 1)
+
+
+def _claim_priority_parse(manager: DagFileProcessorManager, file: DagFileInfo) -> ProcessorWorkItem:
+ work = ProcessorWorkItem(
+ id=str(uuid7()), claim_id=uuid7(), bundle_name=file.bundle_name, relative_fileloc=str(file.rel_path)
+ )
+ manager._priority_claims[work.id] = PriorityParseClaim(
+ work=work, refresh_generation=0, attempt_id=uuid7()
+ )
+ manager._file_stats[file] = DagFileStat()
+ return work
+
+
+@pytest.mark.parametrize("parallelism", [1, 2])
+@pytest.mark.parametrize("older_outcome", ["published", "completed", "timeout"])
+@mock.patch.object(DagFileProcessorManager, "_create_process", autospec=True)
+def test_priority_request_waits_for_its_own_refreshed_import(create, tmp_path, parallelism, older_outcome):
+ (tmp_path / "a.py").write_text("from airflow.sdk import DAG\ndag = DAG('test')\n")
+ client = mock.create_autospec(DagProcessorAPIClient, instance=True)
+ manager = DagFileProcessorManager(
+ max_runs=1,
+ api_client=client,
+ parallelism=parallelism,
+ processor_timeout=1,
+ bundle_refresh_check_interval=0,
+ )
+ bundle = _make_bundle(tmp_path)
+ bundle.is_initialized = True
+ bundle.supports_versioning = False
+ bundle.refresh_interval = 300
+ manager._dag_bundles = [bundle]
+ file = DagFileInfo(bundle_name=bundle.name, rel_path=Path("a.py"), bundle_path=tmp_path)
+ manager._bundle_parse_sources[bundle.name] = DagParseSource()
+ old = mock.create_autospec(DagFileProcessorProcess, instance=True)
+ old.id = uuid7()
+ old.start_time = time.monotonic() - 5
+ old.had_callbacks = False
+ old.parse_source = DagParseSource()
+ old.dispatch_sequence = 1
+ old.parsing_result = DagFileParsingResult(fileloc="a.py", serialized_dags=[])
+ client.get_bundles.return_value = [
+ DagBundleStateResponse(
+ name=bundle.name, version=None, last_refreshed=None, revision=None, team_name=None
+ )
+ ]
+ client.publish_inventory.side_effect = lambda name, body: DagBundleInventoryResponse(
+ attempt_id=body.attempt_id, revision=uuid7(), accepted_at=timezone.utcnow()
+ )
+ known_files = {bundle.name: {file}}
+ if older_outcome == "published":
+ body = mock.create_autospec(DagParseResultBody, instance=True)
+ body.attempt_id = old.id
+ manager._pending_publications[file] = PendingDagPublication(
+ body=body, stat=DagFileStat(), refresh_generation=0
+ )
+ else:
+ manager._processors[file] = old
+ work = ProcessorWorkItem(
+ id=str(uuid7()), claim_id=uuid7(), bundle_name=bundle.name, relative_fileloc="a.py"
+ )
+ client.claim_work.return_value = [work]
+ manager._queue_requested_files_for_parsing()
+ manager._refresh_dag_bundles(known_files)
+ bundle.refresh.assert_not_called()
+ manager._start_new_processes()
+ if older_outcome == "published":
+ manager._publish_pending_results()
+ elif older_outcome == "completed":
+ manager._processors.pop(file)
+ manager.handle_parsing_result(file, old)
+ manager._publish_pending_results()
+ else:
+ manager._kill_timed_out_processors()
+ assert not manager._pending_work_acks
+ assert work.id in manager._priority_claims
+ assert file in manager._file_queue
+ create.assert_not_called()
+
+ manager._start_new_processes()
+ create.assert_not_called()
+ manager._refresh_dag_bundles(known_files)
+ bundle.refresh.assert_called_once_with()
+ assert manager._bundle_parse_sources[bundle.name].refresh_generation == 1
+ new = mock.create_autospec(DagFileProcessorProcess, instance=True)
+ new.id = uuid7()
+ new.start_time = time.monotonic()
+ new.had_callbacks = False
+ new.parsing_result = DagFileParsingResult(fileloc="a.py", serialized_dags=[])
+ create.return_value = new
+ manager._start_new_processes()
+ create.assert_called_once_with(manager, file)
+
+ later = work.model_copy(update={"id": str(uuid7()), "claim_id": uuid7()})
+ manager._parallelism = 2
+ client.claim_work.return_value = [later]
+ manager._queue_requested_files_for_parsing()
+ manager.handle_parsing_result(file, new)
+ manager._publish_pending_results()
+ assert manager._pending_work_acks == [("priority", work, False)]
+ assert later.id in manager._priority_claims
+ assert file in manager._file_queue
+
+
+@pytest.mark.parametrize("failure", ["response_lost", "server_error", "source_changed"])
+@mock.patch.object(DagFileProcessorManager, "_create_process", autospec=True)
+def test_priority_request_during_publication_retry_completes_after_refresh(create, tmp_path, failure):
+ (tmp_path / "a.py").write_text("from airflow.sdk import DAG\ndag = DAG('test')\n")
+ client = mock.create_autospec(DagProcessorAPIClient, instance=True)
+ manager = DagFileProcessorManager(max_runs=1, api_client=client, bundle_refresh_check_interval=0)
+ bundle = _make_bundle(tmp_path)
+ bundle.is_initialized = True
+ bundle.supports_versioning = False
+ bundle.refresh_interval = 300
+ manager._dag_bundles = [bundle]
+ file = DagFileInfo(bundle_name=bundle.name, rel_path=Path("a.py"), bundle_path=tmp_path)
+ client.get_bundles.return_value = [
+ DagBundleStateResponse(
+ name=bundle.name, version=None, last_refreshed=None, revision=None, team_name=None
+ )
+ ]
+ client.publish_inventory.side_effect = lambda name, body: DagBundleInventoryResponse(
+ attempt_id=body.attempt_id, revision=uuid7(), accepted_at=timezone.utcnow()
+ )
+ known_files = {bundle.name: {file}}
+ first = ProcessorWorkItem(
+ id=str(uuid7()), claim_id=uuid7(), bundle_name=bundle.name, relative_fileloc="a.py"
+ )
+ client.claim_work.return_value = [first]
+ manager._queue_requested_files_for_parsing()
+ manager._refresh_dag_bundles(known_files)
+ proc = mock.create_autospec(DagFileProcessorProcess, instance=True)
+ proc.id = uuid7()
+ proc.start_time = time.monotonic()
+ proc.had_callbacks = False
+ proc.parsing_result = DagFileParsingResult(fileloc="a.py", serialized_dags=[])
+ create.return_value = proc
+ manager._start_new_processes()
+ manager._processors.pop(file)
+ manager.handle_parsing_result(file, proc)
+ if failure == "response_lost":
+ error = httpx2.ReadError("Lost publication acknowledgment")
+ else:
+ response = httpx2.Response(
+ 409 if failure == "source_changed" else 503,
+ json={"detail": {"reason": failure}},
+ request=httpx2.Request("POST", "http://api/parse-results"),
+ )
+ error = httpx2.HTTPStatusError("Publication failed", request=response.request, response=response)
+ client.publish_parse_result.side_effect = error
+ manager._publish_pending_results()
+ assert not manager._pending_work_acks
+
+ later = first.model_copy(update={"id": str(uuid7()), "claim_id": uuid7()})
+ client.claim_work.return_value = [later]
+ manager._queue_requested_files_for_parsing()
+ client.publish_parse_result.side_effect = None
+ if failure != "source_changed":
+ manager._refresh_dag_bundles(known_files)
+ assert bundle.refresh.call_count == 1
+ manager._start_new_processes()
+ assert create.call_count == 1
+ manager._pending_publications[file].next_attempt_time = 0
+ manager._publish_pending_results()
+ assert manager._pending_work_acks == [("priority", first, False)]
+ assert later.id in manager._priority_claims
+ manager._refresh_dag_bundles(known_files)
+ assert bundle.refresh.call_count == 2
+ replacement = mock.create_autospec(DagFileProcessorProcess, instance=True)
+ replacement.id = uuid7()
+ replacement.start_time = time.monotonic()
+ replacement.had_callbacks = False
+ replacement.parsing_result = DagFileParsingResult(fileloc="a.py", serialized_dags=[])
+ create.return_value = replacement
+ manager._start_new_processes()
+ assert manager._processors[file] is replacement
+ manager._processors.pop(file)
+ manager.handle_parsing_result(file, replacement)
+ manager._publish_pending_results()
+ assert manager._pending_work_acks == [("priority", first, False), ("priority", later, False)]
+ assert not manager._priority_claims
+ assert not manager._pending_publications
+ manager._acknowledge_requested_work()
+ manager._acknowledge_requested_work()
+ assert client.acknowledge_work.call_args_list == [
+ mock.call("priority", first, failed=False),
+ mock.call("priority", later, failed=False),
+ ]
+
+
+@pytest.mark.parametrize(
+ ("status_code", "failed"), [(None, False), (409, True)], ids=["published", "rejected"]
+)
+def test_publication_outcome_acknowledges_priority_claims(status_code, failed):
+ client = mock.create_autospec(DagProcessorAPIClient, instance=True)
+ manager = DagFileProcessorManager(max_runs=1, api_client=client)
+ file = _get_file_infos(["a.py"])[0]
+ work = _claim_priority_parse(manager, file)
+ manager._pending_publications[file] = PendingDagPublication(
+ body=mock.create_autospec(DagParseResultBody, instance=True), stat=DagFileStat(), refresh_generation=0
+ )
+ manager._pending_publications[file].body.attempt_id = manager._priority_claims[work.id].attempt_id
+ if status_code:
+ response = httpx2.Response(status_code, 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 manager._pending_work_acks == [("priority", work, failed)]
+ assert not manager._priority_claims
+
+
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]
+ work = _claim_priority_parse(manager, file)
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()
+ manager._pending_publications[file].body.attempt_id = manager._priority_claims[work.id].attempt_id
response = httpx2.Response(
409,
json={"detail": {"reason": "source_changed"}},
@@ -6038,6 +6791,46 @@ def test_publication_for_a_changed_source_refreshes_and_reparses():
assert not manager._pending_publications
assert manager._force_refresh_bundles == {"testing"}
assert list(manager._file_queue) == [file]
+ assert manager._priority_claims[work.id].work == work
+ assert not manager._pending_work_acks
+
+
+def test_rejected_inventory_releases_priority_attempt_discarded_after_source_change():
+ client = mock.create_autospec(DagProcessorAPIClient, instance=True)
+ manager = DagFileProcessorManager(max_runs=1, api_client=client, parallelism=1)
+ file = _get_file_infos(["a.py"])[0]
+ work = _claim_priority_parse(manager, file)
+ body = mock.create_autospec(DagParseResultBody, instance=True)
+ body.attempt_id = manager._priority_claims[work.id].attempt_id
+ manager._pending_publications[file] = PendingDagPublication(
+ body=body, stat=DagFileStat(), refresh_generation=0
+ )
+ client.publish_parse_result.side_effect = _make_status_error(409, "source_changed")
+
+ manager._publish_pending_results()
+
+ assert not manager._pending_publications
+ assert manager._priority_claims[work.id].attempt_id is not None
+ inventory = DagBundleInventoryBody(
+ attempt_id=uuid7(), dispatch_sequence=2, expected_revision=None, files=["a.py"]
+ )
+ manager._pending_inventories[file.bundle_name] = (
+ inventory,
+ {file},
+ DagParseSource(refresh_generation=1),
+ )
+ client.publish_inventory.side_effect = _make_status_error(413)
+
+ manager._publish_bundle_inventory(file.bundle_name, {})
+
+ assert not manager._priority_claims
+ assert manager._pending_work_acks == [("priority", work, True)]
+ assert not manager._force_refresh_bundles
+ manager._acknowledge_requested_work()
+ client.acknowledge_work.assert_called_once_with("priority", work, failed=True)
+ client.claim_work.return_value = []
+ manager.claim_priority_files()
+ client.claim_work.assert_called_once_with("priority", [], 1)
@conf_vars({("core", "multi_team"): "True"})
@@ -6057,6 +6850,91 @@ def test_api_team_lookup_failure_is_retried_without_stopping_the_processor():
assert manager._get_team_names({"testing"}) == {"testing": "team_a"}
+@pytest.mark.parametrize(
+ ("outcome", "acknowledged"),
+ [("no_result", True), ("build_failed", True), ("callback_only", False)],
+)
+def test_parse_without_publication_acknowledges_priority_claims(outcome, acknowledged):
+ client = mock.create_autospec(DagProcessorAPIClient, instance=True)
+ manager = DagFileProcessorManager(max_runs=1, api_client=client)
+ file = _get_file_infos(["a.py"])[0]
+ work = _claim_priority_parse(manager, file)
+ proc = mock.create_autospec(DagFileProcessorProcess, instance=True)
+ proc.id = manager._priority_claims[work.id].attempt_id
+ proc.start_time = time.monotonic()
+ proc.had_callbacks = outcome == "callback_only"
+ proc.parse_source = DagParseSource()
+ proc.parsing_result = (
+ DagFileParsingResult(fileloc="a.py", serialized_dags=[]) if outcome == "build_failed" else None
+ )
+
+ with mock.patch.object(manager, "_build_parse_result", side_effect=ValueError("invalid result")):
+ manager.handle_parsing_result(file, proc)
+
+ assert manager._pending_work_acks == ([("priority", work, True)] if acknowledged else [])
+ assert bool(manager._priority_claims) is not acknowledged
+
+
+@pytest.mark.parametrize("had_callbacks", [False, True])
+def test_timed_out_parse_acknowledges_priority_claims(had_callbacks):
+ client = mock.create_autospec(DagProcessorAPIClient, instance=True)
+ manager = DagFileProcessorManager(max_runs=1, api_client=client, processor_timeout=1)
+ file = _get_file_infos(["a.py"])[0]
+ work = _claim_priority_parse(manager, file)
+ processor = mock.create_autospec(DagFileProcessorProcess, instance=True)
+ processor.id = manager._priority_claims[work.id].attempt_id
+ processor.start_time = time.monotonic() - 5
+ processor.had_callbacks = had_callbacks
+ manager._processors[file] = processor
+
+ manager._kill_timed_out_processors()
+
+ processor.kill.assert_called_once_with(signal.SIGKILL)
+ assert manager._pending_work_acks == ([] if had_callbacks else [("priority", work, True)])
+
+
+@pytest.mark.parametrize("recovers", [False, True], ids=["dropped", "recovered"])
+@conf_vars({("dag_processor", "job_heartbeat_timeout"): "10"})
+@mock.patch("airflow.dag_processing.manager.time", autospec=True)
+@mock.patch.object(DagFileProcessorManager, "prepare_callback_bundle", autospec=True, return_value=None)
+def test_unpreparable_callback_is_retried_with_backoff_until_the_job_heartbeat_timeout(
+ prepare, mock_time, recovers
+):
+ client = mock.create_autospec(DagProcessorAPIClient, instance=True)
+ manager = DagFileProcessorManager(max_runs=1, api_client=client)
+ request = DagCallbackRequest(
+ filepath="a.py", bundle_name="testing", bundle_version="v1", dag_id="dag", run_id="run"
+ )
+ work = ProcessorWorkItem(
+ id=str(uuid7()), claim_id=uuid7(), bundle_name="testing", relative_fileloc="a.py"
+ )
+ manager._callback_claims[id(request)] = work
+ manager._deferred_api_callbacks.append(request)
+
+ for now in (0, 0.5, 1, 2):
+ mock_time.monotonic.return_value = now
+ manager._add_callback_to_queue(request)
+ assert prepare.call_count == 2
+
+ if recovers:
+ bundle = mock.create_autospec(BaseDagBundle, instance=True)
+ bundle.path = TEST_DAGS_FOLDER
+ prepare.return_value = bundle
+ mock_time.monotonic.return_value = 10
+ manager._add_callback_to_queue(request)
+
+ assert prepare.call_count == 3
+ assert not manager._deferred_callback_retries
+ assert not manager._deferred_api_callbacks
+ if recovers:
+ assert manager._callback_claims == {id(request): work}
+ assert list(manager._callback_to_execute.values()) == [[request]]
+ assert not manager._pending_work_acks
+ else:
+ assert not manager._callback_claims
+ assert manager._pending_work_acks == [("callbacks", work, True)]
+
+
def _make_pending_publication(
file: DagFileInfo, *, attempt_id: UUID | None = None, import_errors: dict[str, str] | None = None
) -> PendingDagPublication:
@@ -6127,9 +7005,10 @@ def test_result_the_api_refuses_is_replaced_by_an_import_error(error, message):
manager = DagFileProcessorManager(max_runs=1, api_client=client)
manager._dispatch_sequence = 7
file = _get_file_infos(["a.py"])[0]
+ work = _claim_priority_parse(manager, file)
original = _make_pending_publication(
file,
- attempt_id=uuid7(),
+ attempt_id=manager._priority_claims[work.id].attempt_id,
import_errors={"a.py/member.py": "member failed"},
)
manager._pending_publications[file] = original
@@ -6147,6 +7026,7 @@ def test_result_the_api_refuses_is_replaced_by_an_import_error(error, message):
assert not manager._pending_publications
assert manager._file_stats[file].num_dags == 0
assert manager._file_stats[file].import_errors == 1
+ assert manager._pending_work_acks == [("priority", work, True)]
@conf_vars({("workers", "execution_api_retries"): "1"})
@@ -6155,9 +7035,10 @@ def test_server_failure_after_publication_retries_preserves_previous_result():
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]
+ work = _claim_priority_parse(manager, file)
previous = DagFileStat(num_dags=2, last_finish_time=DEFAULT_DATE)
manager._file_stats[file] = previous
- pending = _make_pending_publication(file, attempt_id=uuid7())
+ pending = _make_pending_publication(file, attempt_id=manager._priority_claims[work.id].attempt_id)
manager._pending_publications[file] = pending
manager._publish_pending_results()
@@ -6167,6 +7048,7 @@ def test_server_failure_after_publication_retries_preserves_previous_result():
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
+ assert manager._pending_work_acks == [("priority", work, True)]
@mock.patch("airflow.dag_processing.manager.MAX_PUBLICATION_BYTES", 32 * 1024)
@@ -6273,11 +7155,12 @@ 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]
+ work = _claim_priority_parse(manager, file)
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.id = manager._priority_claims[work.id].attempt_id
proc.start_time = time.monotonic()
proc.had_callbacks = False
proc.dispatch_sequence = 1
@@ -6294,17 +7177,45 @@ def test_result_over_the_publication_limit_is_reported_without_sending_it():
assert pending.body.import_errors == {
"a.py": "The parsing result was not published: it exceeds the 1 MiB publication limit."
}
+ assert manager._pending_work_acks == [("priority", work, True)]
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):
+def test_acknowledgments_are_sent_within_the_loop_budget():
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)
+ work = [
+ ProcessorWorkItem(id=str(uuid7()), claim_id=uuid7(), bundle_name="testing", relative_fileloc="a.py")
+ for _ in range(3)
+ ]
+ manager._pending_work_acks = [("priority", item, False) for item in work]
+
+ manager._acknowledge_requested_work(deadline=time.monotonic())
+ client.acknowledge_work.assert_not_called()
- assert getattr(manager, method)() == []
- read_database.assert_not_called()
+ manager._acknowledge_requested_work()
+ assert client.acknowledge_work.call_args_list == [
+ mock.call("priority", item, failed=False) for item in work
+ ]
+ assert not manager._pending_work_acks
+
+
+@pytest.mark.parametrize("failure", [None, httpx2.ConnectError("unavailable"), RuntimeError("closed")])
+def test_shutdown_acknowledges_finished_work(failure):
+ client = mock.create_autospec(DagProcessorAPIClient, instance=True)
+ client.acknowledge_work.side_effect = failure
+ manager = DagFileProcessorManager(max_runs=1, api_client=client)
+ work = [
+ ProcessorWorkItem(id=str(uuid7()), claim_id=uuid7(), bundle_name="testing", relative_fileloc="a.py")
+ for _ in range(2)
+ ]
+ manager._pending_work_acks = [("callbacks", item, False) for item in work]
+
+ manager.end()
+
+ if failure is None:
+ assert client.acknowledge_work.call_count == 2
+ assert not manager._pending_work_acks
+ else:
+ client.acknowledge_work.assert_called_once()
+ assert len(manager._pending_work_acks) == 2
diff --git a/airflow-core/tests/unit/jobs/test_dag_processor_job_runner.py b/airflow-core/tests/unit/jobs/test_dag_processor_job_runner.py
new file mode 100644
index 00000000000..f55ed3ca524
--- /dev/null
+++ b/airflow-core/tests/unit/jobs/test_dag_processor_job_runner.py
@@ -0,0 +1,275 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from unittest import mock
+
+import httpx2
+import pytest
+
+from airflow.api_fastapi.execution_api.datamodels.job import TerminalJobState
+from airflow.dag_processing.api_client import (
+ DagProcessorAPIClient,
+ DagProcessorJobAlreadyRunning,
+ DagProcessorRegistrationRetired,
+)
+from airflow.dag_processing.manager import DagFileProcessorManager
+from airflow.jobs.dag_processor_job_runner import DagProcessorHeartbeatTimeout, DagProcessorJobRunner
+from airflow.jobs.job import Job, JobState
+
+from tests_common.test_utils.config import conf_vars
+
+
+@pytest.fixture
+def client():
+ client = mock.create_autospec(DagProcessorAPIClient, instance=True)
+ client.register_job.return_value = 42
+ client.restart_required = False
+ client.heartbeat.return_value = JobState.RUNNING
+ return client
+
+
+@pytest.fixture
+def runner():
+ return DagProcessorJobRunner(
+ job=Job(heartrate=5), processor=mock.create_autospec(DagFileProcessorManager, instance=True)
+ )
+
+
+@pytest.fixture
+def clock():
+ with mock.patch(
+ "airflow.jobs.dag_processor_job_runner.monotonic", autospec=True, return_value=0
+ ) as clock:
+ yield clock
+
+
+@pytest.mark.parametrize(
+ ("error", "outcome"),
+ [
+ (None, TerminalJobState.SUCCESS),
+ (SystemExit(0), TerminalJobState.SUCCESS),
+ (ValueError("parse failed"), TerminalJobState.FAILED),
+ (KeyboardInterrupt(), TerminalJobState.FAILED),
+ ],
+)
+@mock.patch.object(Job, "prepare_for_execution", autospec=True)
+@mock.patch.object(Job, "complete_execution", autospec=True)
+def test_api_lifecycle_cleans_up_before_completion(complete, prepare, runner, client, error, outcome):
+ events = []
+ runner.processor.run.side_effect = error
+ runner.processor.terminate.side_effect = lambda: events.append("terminate")
+ runner.processor.end.side_effect = lambda: events.append("end")
+ client.complete_job.side_effect = lambda state: events.append(state)
+
+ if outcome == TerminalJobState.FAILED:
+ with pytest.raises(type(error)):
+ runner.run_with_api(client)
+ else:
+ runner.run_with_api(client)
+
+ assert runner.job.id == 42
+ assert runner.job.state == outcome.value
+ assert runner.processor.api_client is client
+ runner.processor.sync_bundles.assert_not_called()
+ client.register_job.assert_called_once_with()
+ assert events == ["terminate", "end", outcome]
+ prepare.assert_not_called()
+ complete.assert_not_called()
+
+
+def test_failed_registration_never_starts_parsing(runner, client):
+ client.register_job.side_effect = httpx2.ConnectError("unavailable")
+ with pytest.raises(httpx2.ConnectError):
+ runner.run_with_api(client)
+ runner.processor.run.assert_not_called()
+ client.complete_job.assert_not_called()
+
+
+@pytest.mark.parametrize("available_at", [12, 15, None])
+@conf_vars({("dag_processor", "health_check_threshold"): "10"})
+@mock.patch("airflow.jobs.dag_processor_job_runner.sleep", autospec=True)
+def test_registration_waits_for_a_crashed_job_with_a_bounded_backoff(
+ sleep, runner, client, clock, available_at
+):
+ def wait(seconds):
+ clock.return_value += seconds
+
+ def register():
+ runner.processor.run.assert_not_called()
+ if available_at is None or clock.return_value < available_at:
+ raise DagProcessorJobAlreadyRunning("Still alive")
+ return 42
+
+ sleep.side_effect = wait
+ client.register_job.side_effect = register
+
+ if available_at is None:
+ with pytest.raises(DagProcessorJobAlreadyRunning):
+ runner.run_with_api(client)
+ runner.processor.run.assert_not_called()
+ client.complete_job.assert_not_called()
+ else:
+ runner.run_with_api(client)
+ runner.processor.run.assert_called_once_with()
+ assert runner.job.id == 42
+
+ assert sleep.call_args_list == [mock.call(1), mock.call(2), mock.call(4), mock.call(5)] + (
+ [mock.call(3)] if available_at != 12 else []
+ )
+ assert clock.return_value == (available_at or 15)
+ assert client.register_job.call_count == sleep.call_count + 1
+
+
+@pytest.mark.parametrize("reason", ["registration_conflict", "registration_retired", None])
+@mock.patch("airflow.jobs.dag_processor_job_runner.sleep", autospec=True)
+def test_registration_does_not_wait_on_other_conflicts(sleep, runner, client, reason):
+ error = httpx2.HTTPStatusError(
+ "conflict",
+ request=httpx2.Request("POST", "http://api/jobs"),
+ response=httpx2.Response(409, json={"detail": {"reason": reason}}),
+ )
+ client.register_job.side_effect = error
+
+ with pytest.raises(httpx2.HTTPStatusError):
+ runner.run_with_api(client)
+
+ sleep.assert_not_called()
+ client.register_job.assert_called_once_with()
+ runner.processor.run.assert_not_called()
+
+
+@mock.patch("airflow.jobs.dag_processor_job_runner.get_listener_manager", autospec=True)
+def test_shutdown_listener_failure_does_not_prevent_completion(listeners, runner, client):
+ listeners.return_value.hook.before_stopping.side_effect = ValueError("listener failed")
+ runner.run_with_api(client)
+ listeners.return_value.hook.before_stopping.assert_called_once_with(component=runner.job)
+ client.complete_job.assert_called_once_with(TerminalJobState.SUCCESS)
+
+
+@pytest.mark.parametrize("parse_error", [None, ValueError("original error")])
+def test_completion_failure_preserves_the_original_error(runner, client, parse_error):
+ runner.processor.run.side_effect = parse_error
+ client.complete_job.side_effect = httpx2.ConnectError("completion unavailable")
+ with pytest.raises(type(parse_error) if parse_error else httpx2.ConnectError) as raised:
+ runner.run_with_api(client)
+ if parse_error:
+ assert raised.value is parse_error
+
+
+def test_heartbeats_are_throttled_and_use_the_api(runner, client, clock):
+ def run():
+ runner.processor.heartbeat()
+ client.heartbeat.assert_not_called()
+ clock.return_value = 5
+ runner.processor.heartbeat()
+ runner.processor.heartbeat()
+ client.heartbeat.assert_called_once_with()
+ clock.return_value = 10
+ runner.processor.heartbeat()
+ assert client.heartbeat.call_count == 2
+
+ runner.processor.run.side_effect = run
+ runner.run_with_api(client)
+
+
+@pytest.mark.parametrize("restart", ["already_retired", "renewal_retired", "server_restart", "job_closed"])
+def test_restart_stops_parsing_and_releases_children(runner, client, clock, restart):
+ def heartbeat():
+ if restart == "job_closed":
+ raise DagProcessorRegistrationRetired("replaced")
+ if restart == "renewal_retired":
+ client.restart_required = True
+ return JobState.RESTARTING if restart == "server_restart" else JobState.RUNNING
+
+ def run():
+ client.restart_required = restart == "already_retired"
+ clock.return_value = 1 if client.restart_required else 5
+ runner.processor.heartbeat()
+ pytest.fail("Parsing continued after a restart request")
+
+ client.heartbeat.side_effect = heartbeat
+ runner.processor.run.side_effect = run
+ with pytest.raises(DagProcessorRegistrationRetired):
+ runner.run_with_api(client)
+ runner.processor.terminate.assert_called_once_with()
+ runner.processor.end.assert_called_once_with()
+ client.complete_job.assert_called_once_with(TerminalJobState.FAILED)
+
+
+@pytest.mark.parametrize("failure", ["transport", 503, 403])
+@conf_vars({("dag_processor", "job_heartbeat_timeout"): "10"})
+def test_heartbeat_failure_is_bounded(runner, client, clock, failure):
+ error = (
+ httpx2.ConnectError("unavailable")
+ if failure == "transport"
+ else httpx2.HTTPStatusError(
+ "refused",
+ request=httpx2.Request("POST", "http://api/jobs/42/heartbeat"),
+ response=httpx2.Response(failure),
+ )
+ )
+ client.heartbeat.side_effect = error
+
+ def run():
+ clock.return_value = 5
+ runner.processor.heartbeat()
+ clock.return_value = 10
+ runner.processor.heartbeat()
+
+ runner.processor.run.side_effect = run
+ with pytest.raises(httpx2.HTTPStatusError if failure == 403 else DagProcessorHeartbeatTimeout):
+ runner.run_with_api(client)
+ assert client.heartbeat.call_count == (1 if failure == 403 else 2)
+
+
+@conf_vars({("dag_processor", "job_heartbeat_timeout"): "10"})
+def test_recovered_heartbeat_resets_the_lease(runner, client, clock):
+ client.heartbeat.side_effect = [
+ httpx2.ConnectError("offline"),
+ JobState.RUNNING,
+ httpx2.ConnectError("offline"),
+ ]
+
+ def run():
+ for tick in (5, 10, 15):
+ clock.return_value = tick
+ runner.processor.heartbeat()
+
+ runner.processor.run.side_effect = run
+ runner.run_with_api(client)
+ client.complete_job.assert_called_once_with(TerminalJobState.SUCCESS)
+
+
+@conf_vars(
+ {
+ ("dag_processor", "health_check_threshold"): "30",
+ ("dag_processor", "job_heartbeat_timeout"): "300",
+ }
+)
+def test_a_stall_past_the_health_threshold_keeps_the_lease(runner, client, clock):
+ client.heartbeat.side_effect = [httpx2.ConnectError("offline"), JobState.RUNNING]
+
+ def run():
+ for tick in (120, 125):
+ clock.return_value = tick
+ runner.processor.heartbeat()
+
+ runner.processor.run.side_effect = run
+ runner.run_with_api(client)
+ assert client.heartbeat.call_count == 2
+ client.complete_job.assert_called_once_with(TerminalJobState.SUCCESS)
diff --git a/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml b/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml
index 827462d8c3b..4647a1d2722 100644
--- a/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml
+++ b/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml
@@ -39,6 +39,24 @@ metrics:
legacy_name: "-"
name_variables: ["job_name"]
+ - name: "job_start"
+ description: "Number of started jobs."
+ type: "counter"
+ legacy_name: "-"
+ name_variables: []
+
+ - name: "job_end"
+ description: "Number of ended jobs."
+ type: "counter"
+ legacy_name: "-"
+ name_variables: []
+
+ - name: "dag_processor_heartbeat_failure"
+ description: "Number of failed Dag processor heartbeats through the Execution API."
+ type: "counter"
+ legacy_name: "-"
+ name_variables: []
+
- name: "operator_failures"
description: "Operator ``{operator_name}`` failures."
type: "counter"
@@ -122,13 +140,6 @@ metrics:
legacy_name: "-"
name_variables: []
- - name: "dag_processing.results_discarded_on_refresh"
- description: "Parse results discarded because bundle initialization or refresh overlapped the import.
- Metric with bundle_name tagging. Repeated discards can indicate that refreshes are delaying parsing."
- type: "counter"
- legacy_name: "-"
- name_variables: []
-
- name: "dag_processing.processes"
description: "Relative number of currently running Dag parsing processes (ie this delta is negative when,
since the last metric was sent, processes have completed). Metric with file_path, bundle_name, and action
@@ -150,6 +161,13 @@ metrics:
legacy_name: "-"
name_variables: []
+ - name: "dag_processing.results_discarded_on_refresh"
+ description: "Parse results discarded because bundle initialization or refresh overlapped the import.
+ Metric with bundle_name tagging. Repeated discards can indicate that refreshes are delaying parsing."
+ type: "counter"
+ legacy_name: "-"
+ name_variables: []
+
- name: "dag_processing.callback_only_count"
description: "Number of DAG file processing runs that processed callbacks only, without full DAG parsing"
type: "counter"