diff --git a/airflow-core/docs/howto/docker-compose/docker-compose.yaml b/airflow-core/docs/howto/docker-compose/docker-compose.yaml index b68f1afd3cc9f..54f4dfe603f3b 100644 --- a/airflow-core/docs/howto/docker-compose/docker-compose.yaml +++ b/airflow-core/docs/howto/docker-compose/docker-compose.yaml @@ -155,7 +155,14 @@ services: <<: *airflow-common command: dag-processor healthcheck: - test: ["CMD-SHELL", 'airflow jobs check --job-type DagProcessorJob --hostname "$${HOSTNAME}"'] + test: + - CMD-SHELL + - >- + if [ -n "$${AIRFLOW__DAG_PROCESSOR__EXECUTION_API_TOKEN_FILE:-}" ]; then + airflow dag-processor --check-health; + else + airflow jobs check --job-type DagProcessorJob --hostname "$${HOSTNAME}"; + fi interval: 30s timeout: 10s retries: 5 diff --git a/airflow-core/docs/migrations-ref.rst b/airflow-core/docs/migrations-ref.rst index 316eb4bd8ca0d..ec17da204975e 100644 --- a/airflow-core/docs/migrations-ref.rst +++ b/airflow-core/docs/migrations-ref.rst @@ -39,7 +39,11 @@ Here's the list of all the Database Migrations that are executed via when you ru +-------------------------+------------------+-------------------+--------------------------------------------------------------+ | Revision ID | Revises ID | Airflow Version | Description | +=========================+==================+===================+==============================================================+ -| ``c71065a1b14f`` (head) | ``a5d7b9c13e40`` | ``3.5.0`` | Add session_id and registration_id to Job. | +| ``687eb253de24`` (head) | ``e941ab8243b7`` | ``3.5.0`` | Add source revisions for Dag processor reconciliation. | ++-------------------------+------------------+-------------------+--------------------------------------------------------------+ +| ``e941ab8243b7`` | ``c71065a1b14f`` | ``3.5.0`` | Add checkpoints for Dag result publication. | ++-------------------------+------------------+-------------------+--------------------------------------------------------------+ +| ``c71065a1b14f`` | ``a5d7b9c13e40`` | ``3.5.0`` | Add session_id and registration_id to Job. | +-------------------------+------------------+-------------------+--------------------------------------------------------------+ | ``a5d7b9c13e40`` | ``c3e7a9182f64`` | ``3.4.0`` | Preserve exact task execution attribution in asset state. | +-------------------------+------------------+-------------------+--------------------------------------------------------------+ 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 0000000000000..a73b114243f2f --- /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/dag-processor-api/helm-values.yaml b/airflow-core/docs/security/dag-processor-api/helm-values.yaml new file mode 100644 index 0000000000000..a4ecfcf877de7 --- /dev/null +++ b/airflow-core/docs/security/dag-processor-api/helm-values.yaml @@ -0,0 +1,74 @@ +# 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. +--- +# Install as a separate processor-only release, using an image with API mode. +airflowVersion: "3.5.0" +defaultAirflowTag: "3.5.0" +executor: LocalExecutor +scheduler: + enabled: false +apiServer: + enabled: false +triggerer: + enabled: false +createUserJob: + enabled: false +migrateDatabaseJob: + enabled: false +postgresql: + enabled: false +redis: + enabled: false +statsd: + enabled: false +rbac: + create: false +enableBuiltInSecretEnvVars: + AIRFLOW__CORE__FERNET_KEY: false + AIRFLOW__DATABASE__SQL_ALCHEMY_CONN: false + AIRFLOW_CONN_AIRFLOW_DB: false + AIRFLOW__API__SECRET_KEY: false + AIRFLOW__API_AUTH__JWT_SECRET: false + AIRFLOW__CELERY__RESULT_BACKEND: false + AIRFLOW__CELERY__BROKER_URL: false +config: + core: + execution_api_server_url: http://airflow-api-server.airflow.svc.cluster.local:8080/execution/ + dag_processor: + execution_api_token_file: /run/dag-processor/processor.jwt +dagProcessor: + replicas: 1 + waitForMigrations: + enabled: false + serviceAccount: + automountServiceAccountToken: false + env: + - name: CONNECTION_CHECK_MAX_COUNT + value: "0" + extraVolumes: + - name: processor-token + secret: + secretName: dag-processor-session + defaultMode: 0440 + extraVolumeMounts: + - name: processor-token + mountPath: /run/dag-processor + readOnly: true +dags: + persistence: + enabled: true + existingClaim: processor-dags diff --git a/airflow-core/docs/security/dag-processor-api/network-policy.yaml b/airflow-core/docs/security/dag-processor-api/network-policy.yaml new file mode 100644 index 0000000000000..ae021ce64da2b --- /dev/null +++ b/airflow-core/docs/security/dag-processor-api/network-policy.yaml @@ -0,0 +1,51 @@ +# 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. +--- +apiVersion: networking.k8s.io/v1 +kind: NetworkPolicy +metadata: + name: dag-processor-api-only +spec: + podSelector: + matchLabels: + component: dag-processor + policyTypes: [Ingress, Egress] + ingress: [] + egress: + - to: + - namespaceSelector: + matchLabels: + kubernetes.io/metadata.name: airflow + podSelector: + matchLabels: + component: api-server + release: airflow + ports: + - protocol: TCP + port: 8080 + - to: + - namespaceSelector: + matchLabels: + kubernetes.io/metadata.name: kube-system + podSelector: + matchLabels: + k8s-app: kube-dns + ports: + - protocol: UDP + port: 53 + - protocol: TCP + port: 53 diff --git a/airflow-core/docs/security/jwt_token_authentication.rst b/airflow-core/docs/security/jwt_token_authentication.rst index e72f9d3faa6a5..c0c5f01e9038f 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,170 @@ 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. + +For Helm, set the token-file option under ``config.dag_processor`` to select the local +probe. The chart then replaces processor pods with the ``Recreate`` strategy, because pods +that share a session token cannot run their Jobs side by side, and rejects more than one +replica for the same reason. Disable the processor's migration-wait init container, provision the catalog +separately, and remove server credentials from the pod. The Compose example selects the +local probe when ``AIRFLOW__DAG_PROCESSOR__EXECUTION_API_TOKEN_FILE`` is set. Neither the +main chart release nor the quick-start Compose file removes shared credentials or adds +database-denying network policies automatically; the isolated examples below do. +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 Kubernetes, use the :download:`processor-only Helm values ` +in a **separate release and namespace**. These values disable other components and built-in +server credential injection. Do not apply them to the release running the scheduler/API server. +Create the ``dag-processor-session`` Secret containing ``processor.jwt`` and the +``processor-dags`` PVC in that namespace, then configure the API service address and image. +The processor service account has no mounted API token; a trusted provisioner updates the Secret. + +Apply the :download:`egress policy ` in the processor +namespace before starting its pods. It permits only DNS and the trusted API server. Adjust +its namespace/pod selectors and ports to match the deployment. NetworkPolicy enforcement +requires a supporting cluster network plugin; other policies selecting these pods must not +grant broader egress. This example uses local Dag files. Remote bundles need additional, +explicit egress to their source service. + +For Compose, use the :download:`standalone processor definition ` +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 + 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 2ae6144f4093b..ffb0feeb5ad2a 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 token limited to one + bundle and file. 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/pyproject.toml b/airflow-core/pyproject.toml index fa9cd83c5869e..521b523744b94 100644 --- a/airflow-core/pyproject.toml +++ b/airflow-core/pyproject.toml @@ -103,8 +103,9 @@ dependencies = [ # parsed with no field-count or field-size bound at all. "starlette>=1.3.1", "httpx>=0.25.0", - # Only for InProcessExecutionAPI's transports, which feed httpx2-based Task SDK clients. - # The rest of airflow-core is still on httpx; tracked at https://github.com/apache/airflow/issues/70522 + # Only for code that talks to the httpx2-based Task SDK client: InProcessExecutionAPI's transports + # and the Dag processor's API client. The rest of airflow-core is still on httpx; tracked at + # https://github.com/apache/airflow/issues/70522 "httpx2>=2.0.0", 'importlib_metadata>=6.5;python_version<"3.12"', 'importlib_metadata>=7.0;python_version>="3.12"', @@ -332,6 +333,21 @@ required-version = ">=0.11.8" apache-airflow-core = {workspace = true} apache-airflow-devel-common = { workspace = true } +[tool.ruff] +extend = "../pyproject.toml" + +# Not yet migrated to httpx2; tracked at https://github.com/apache/airflow/issues/70522 +[tool.ruff.lint.extend-per-file-ignores] +"src/airflow/api/client/local_client.py" = ["TID251"] +"src/airflow/api_fastapi/auth/tokens.py" = ["TID251"] +"src/airflow/example_dags/example_dag_decorator.py" = ["TID251"] +"tests/unit/api_fastapi/auth/test_tokens.py" = ["TID251"] +"tests/unit/api_fastapi/core_api/routes/public/test_dag_sources.py" = ["TID251"] +"tests/unit/api_fastapi/execution_api/versions/head/test_hitl.py" = ["TID251"] + +[tool.ruff.lint.flake8-tidy-imports.extend-banned-api] +"httpx".msg = "Use httpx2; httpx objects do not interoperate with the httpx2-based Task SDK client. See https://github.com/apache/airflow/issues/70522" + [tool.airflow] shared_distributions = [ "apache-airflow-shared-configuration", diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/dag_parsing.py b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/dag_parsing.py new file mode 100644 index 0000000000000..cbd003f4dc815 --- /dev/null +++ b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/dag_parsing.py @@ -0,0 +1,170 @@ +# 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 datetime import datetime +from typing import Any, Literal +from uuid import UUID + +from pydantic import Field, model_validator + +from airflow.api_fastapi.core_api.base import StrictBaseModel +from airflow.api_fastapi.execution_api.datamodels.job import DagParseTokenBody +from airflow.models.dag import TAG_MAX_LEN +from airflow.models.dagwarning import DagWarningTypeValue +from airflow.utils.helpers import validate_key + +MAX_PUBLICATION_BYTES = 16 * 1024 * 1024 +"""Largest request body the Dag parsing endpoints read, measured as the bytes sent.""" + + +class ParseWarning(StrictBaseModel): + """A warning attached to a Dag returned by this import.""" + + dag_id: str = Field(min_length=1, max_length=250) + warning_type: DagWarningTypeValue + message: str + + +class ParseSourceCode(StrictBaseModel): + """Source captured by the importer; null explicitly means unavailable.""" + + source_code: str | None + language: str = Field(min_length=1, max_length=64) + + +class DagParseResultBody(DagParseTokenBody): + """One completed file or container import, including an empty result.""" + + dispatch_sequence: int = Field(ge=1, le=2**53 - 1) + bundle_revision: UUID | None = None + bundle_version: str | None = Field(default=None, max_length=200) + version_data: dict[str, Any] | None = None + parse_duration: float = Field(ge=0, allow_inf_nan=False) + serialized_dags: list[dict[str, Any]] + import_errors: dict[str, str] = Field(default_factory=dict) + warnings: list[ParseWarning] = Field(default_factory=list) + parsed_definitions: list[str] = Field(default_factory=list) + source_codes: dict[str, ParseSourceCode] + + @model_validator(mode="after") + def validate_file_scope(self) -> DagParseResultBody: + paths = [*self.parsed_definitions, *self.import_errors] + dag_ids: set[str] = set() + for document in self.serialized_dags: + dag = document.get("dag") + if not isinstance(dag, dict) or not isinstance(dag.get("dag_id"), str): + raise ValueError("A serialized Dag must contain a dag_id") + # The processor runs Dag code that can bypass the SDK's checks; these columns are bounded. + validate_key(dag["dag_id"]) + tags = dag.get("tags") or [] + if not isinstance(tags, list) or any( + not isinstance(tag, str) or len(tag) > TAG_MAX_LEN for tag in tags + ): + raise ValueError(f"Dag tags must be strings of at most {TAG_MAX_LEN} characters") + if dag["dag_id"] in dag_ids: + raise ValueError("Duplicate Dag id in publication") + dag_ids.add(dag["dag_id"]) + relative_fileloc = dag.get("relative_fileloc") + if not isinstance(relative_fileloc, str): + raise ValueError("Every definition must have a bundle-relative location") + paths.append(relative_fileloc) + fileloc = dag.get("fileloc") + if not isinstance(fileloc, str) or len(fileloc) > 2000: + raise ValueError("A serialized Dag must contain a fileloc of at most 2000 characters") + for path in paths: + if len(path) > 2000: + raise ValueError("Definition location exceeds 2000 characters") + self.validate_relative_fileloc(path) + if path != self.relative_fileloc and not path.startswith(self.relative_fileloc + "/"): + raise ValueError("Definition or diagnostic is outside the published file or container") + if set(self.source_codes) != dag_ids: + raise ValueError("Supply source code, or explicit null, for every serialized Dag") + if any(warning.dag_id not in dag_ids for warning in self.warnings): + raise ValueError("Warnings must refer to Dags in this publication") + return self + + +class DagParseResultResponse(StrictBaseModel): + """Receipt returned unchanged when the latest accepted publication is replayed.""" + + attempt_id: UUID + accepted_at: datetime + + +class DagBundleStateResponse(StrictBaseModel): + """Server-owned metadata for an authorized bundle.""" + + name: str + version: str | None + last_refreshed: datetime | None + revision: UUID | None + team_name: str | None + + +class DagBundleInventoryBody(StrictBaseModel): + """A discovery snapshot with explicit failures, conditional on the previously observed revision.""" + + attempt_id: UUID + dispatch_sequence: int = Field(ge=1, le=2**53 - 1) + expected_revision: UUID | None + version: str | None = Field(default=None, max_length=200) + files: list[str] = Field(max_length=100000) + failed_sources: list[str] = Field(default_factory=list, max_length=100000) + + @model_validator(mode="after") + def validate_files(self) -> DagBundleInventoryBody: + for paths in (self.files, self.failed_sources): + for path in paths: + if len(path) > 2000: + raise ValueError("Definition location exceeds 2000 characters") + DagParseTokenBody.validate_relative_fileloc(path) + if len(set(paths)) != len(paths): + raise ValueError("Inventory contains duplicate definitions") + return self + + +class DagBundleInventoryResponse(DagParseResultResponse): + """The source revision to attach to imports dispatched from this inventory.""" + + revision: UUID + + +class ProcessorWorkClaimBody(StrictBaseModel): + """A retryable bounded claim for requested work in ready bundles.""" + + claim_id: UUID + bundle_names: list[str] = Field(min_length=1, max_length=1000) + limit: int = Field(default=20, ge=1, le=100) + + +class ProcessorWorkItem(StrictBaseModel): + """A claimed callback or priority parse request.""" + + id: str + claim_id: UUID + bundle_name: str + relative_fileloc: str + callback: str | None = None + + +class ProcessorWorkAckBody(StrictBaseModel): + """Acknowledge delivery, retaining the claim identity across retries.""" + + claim_id: UUID + state: Literal["success", "failed"] diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/__init__.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/__init__.py index 4653945ef4277..b3f9a4e1d4c37 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/__init__.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/__init__.py @@ -26,6 +26,7 @@ callbacks, connection_tests, connections, + dag_parsing, dag_runs, dags, health, @@ -66,6 +67,7 @@ authenticated_router.include_router(xcoms.router, prefix="/xcoms", tags=["XComs"]) authenticated_router.include_router(hitl.router, prefix="/hitlDetails", tags=["Human in the Loop"]) authenticated_router.include_router(jobs.router, tags=["Jobs"]) +authenticated_router.include_router(dag_parsing.router, prefix="/jobs", tags=["Dag Parsing"]) authenticated_router.include_router(task_state_store.router, prefix="/store/ti", tags=["Task State Store"]) authenticated_router.include_router( asset_state_store.router, prefix="/store/asset", tags=["Asset State Store"] diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/dag_parsing.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/dag_parsing.py new file mode 100644 index 0000000000000..d0c3fe3ad3897 --- /dev/null +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/dag_parsing.py @@ -0,0 +1,523 @@ +# 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 + +import hashlib +import json +from datetime import timedelta +from typing import TYPE_CHECKING, Any, Literal, cast +from uuid import UUID + +from cadwyn import VersionedAPIRouter +from fastapi import HTTPException, Request, Security, status +from sqlalchemy import delete, or_, select +from sqlalchemy.orm import selectinload +from starlette.requests import ClientDisconnect +from uuid6 import uuid7 + +from airflow._shared.timezones import timezone +from airflow.api_fastapi.common.db.common import SessionDep +from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc +from airflow.api_fastapi.execution_api.datamodels.dag_parsing import ( + MAX_PUBLICATION_BYTES, + DagBundleInventoryBody, + DagBundleInventoryResponse, + DagBundleStateResponse, + DagParseResultBody, + DagParseResultResponse, + ProcessorWorkAckBody, + ProcessorWorkClaimBody, + ProcessorWorkItem, +) +from airflow.api_fastapi.execution_api.security import CurrentExecutionToken, ExecutionAPIRoute, require_auth +from airflow.configuration import conf +from airflow.dag_processing.collection import ( + DagBundleOwnershipError, + DagPublicationValidationError, + _reject_other_teams_plugin_classes, + is_source_discovery_failed, + update_dag_parsing_results_in_db, + validate_dag_bundle_ownership, +) +from airflow.exceptions import DeserializationError +from airflow.jobs.job import Job +from airflow.models.asset import remove_references_to_deleted_dags +from airflow.models.callback import DagProcessorCallback +from airflow.models.dag import DagModel +from airflow.models.dag_parse_checkpoint import DagParseCheckpoint +from airflow.models.dagbag import DagPriorityParsingRequest +from airflow.models.dagbundle import DagBundleModel +from airflow.models.dagwarning import DagWarning +from airflow.models.errors import ParseImportError +from airflow.sdk.importers.base import DagSourceCode # noqa: SDK001 +from airflow.serialization.serialized_objects import DagSerialization, LazyDeserializedDAG +from airflow.utils.sqlalchemy import with_row_locks +from airflow.utils.state import CallbackState + +if TYPE_CHECKING: + from collections.abc import Callable, Coroutine + + from starlette.responses import Response + from starlette.types import Message + + +class DagParsingRoute(ExecutionAPIRoute): + """Refuse request bodies over the publication limit before they are read into memory.""" + + def get_route_handler(self) -> Callable[[Request], Coroutine[Any, Any, Response]]: + handler = super().get_route_handler() + + async def handle_bounded_request(request: Request) -> Response: + too_large = HTTPException( + status.HTTP_413_CONTENT_TOO_LARGE, + detail=f"Request body exceeds {MAX_PUBLICATION_BYTES // (1024 * 1024)} MiB", + ) + declared = request.headers.get("content-length", "") + if declared.isdigit() and int(declared) > MAX_PUBLICATION_BYTES: + raise too_large + body = bytearray() + try: + async for chunk in request.stream(): + body += chunk + if len(body) > MAX_PUBLICATION_BYTES: + raise too_large + except ClientDisconnect as error: + # FastAPI reports a body it could not read as a bad request rather than a server error. + raise HTTPException( + status.HTTP_400_BAD_REQUEST, detail="The client disconnected while sending the body" + ) from error + replayed = False + + async def receive() -> Message: + nonlocal replayed + if replayed: + return await request.receive() + replayed = True + return {"type": "http.request", "body": bytes(body), "more_body": False} + + return await handler(Request(request.scope, receive)) + + return handle_bounded_request + + +router = VersionedAPIRouter(route_class=DagParsingRoute) + + +@router.post( + "/{job_id}/parse-results", + responses=create_openapi_http_exception_doc( + [ + status.HTTP_404_NOT_FOUND, + status.HTTP_409_CONFLICT, + status.HTTP_413_CONTENT_TOO_LARGE, + status.HTTP_422_UNPROCESSABLE_CONTENT, + ] + ), + dependencies=[Security(require_auth, scopes=["token:dag_processor"])], +) +def publish_parse_result( + job_id: int, body: DagParseResultBody, *, session: SessionDep, token=CurrentExecutionToken +) -> DagParseResultResponse: + """Atomically publish a file's results and a receipt, without reading the processor's filesystem.""" + if job_id != token.claims.job_id: + raise HTTPException(status.HTTP_404_NOT_FOUND, detail={"reason": "not_found"}) + if body.bundle_name not in token.claims.dag_bundles: + raise HTTPException(status.HTTP_403_FORBIDDEN, detail={"reason": "bundle_not_granted"}) + try: + payload = json.dumps(body.model_dump(mode="json"), sort_keys=True, allow_nan=False).encode() + except ValueError: + raise HTTPException(status.HTTP_422_UNPROCESSABLE_CONTENT, detail="Invalid JSON value") + payload_hash = hashlib.sha256(payload).hexdigest() + source_key = hashlib.sha256( + json.dumps(["parse", body.bundle_name, body.relative_fileloc]).encode() + ).hexdigest() + + # The Job lock serializes checkpoint creation and fences completion/replacement in the same transaction. + job = session.scalar(select(Job).where(Job.id == job_id).with_for_update()) + if job is None or job.session_id != token.id or job.end_date is not None: + raise HTTPException(status.HTTP_403_FORBIDDEN, detail={"reason": "job_closed"}) + checkpoint = session.get(DagParseCheckpoint, (job_id, source_key)) + if checkpoint is not None: + if checkpoint.attempt_id == body.attempt_id: + if checkpoint.payload_hash != payload_hash: + raise HTTPException(status.HTTP_409_CONFLICT, detail={"reason": "publication_conflict"}) + return DagParseResultResponse( + attempt_id=checkpoint.attempt_id, accepted_at=checkpoint.accepted_at + ) + if body.dispatch_sequence <= checkpoint.dispatch_sequence: + raise HTTPException(status.HTTP_409_CONFLICT, detail={"reason": "publication_superseded"}) + + bundle = session.scalar( + select(DagBundleModel).where(DagBundleModel.name == body.bundle_name).with_for_update(key_share=True) + ) + if bundle is None or not bundle.active: + raise HTTPException(status.HTTP_409_CONFLICT, detail={"reason": "bundle_inactive"}) + if body.bundle_revision != bundle.parse_revision or ( + bundle.parse_revision is not None and body.bundle_version != bundle.version + ): + raise HTTPException(status.HTTP_409_CONFLICT, detail={"reason": "source_changed"}) + dags = [LazyDeserializedDAG(data=document) for document in body.serialized_dags] + dag_ids = [dag.dag_id for dag in dags] + existing = session.scalars( + select(DagModel) + .where(DagModel.dag_id.in_(dag_ids)) + .order_by(DagModel.dag_id) + .with_for_update(key_share=True) + ).all() + try: + validate_dag_bundle_ownership(existing, body.bundle_name, session=session) + except DagBundleOwnershipError as error: + raise HTTPException( + status.HTTP_409_CONFLICT, detail={"reason": "dag_owned_by_another_bundle"} + ) from error + # A complete accepted import proves absence; an import error only ages out prior definitions. + cutoff = timezone.utcnow() - timedelta(seconds=conf.getint("dag_processor", "stale_dag_threshold")) + source_filter = or_( + DagModel.relative_fileloc == body.relative_fileloc, + DagModel.relative_fileloc.startswith(body.relative_fileloc + "/", autoescape=True), + ) + missing = session.scalars( + select(DagModel) + .where(DagModel.bundle_name == body.bundle_name, source_filter, DagModel.dag_id.not_in(dag_ids)) + .order_by(DagModel.dag_id) + .with_for_update(key_share=True) + ) + for stored_dag in missing: + failed = any( + path == body.relative_fileloc or path == stored_dag.relative_fileloc + for path in body.import_errors + ) + if not failed or (stored_dag.last_parsed_time is not None and stored_dag.last_parsed_time < cutoff): + stored_dag.is_stale = True + if len(_reject_other_teams_plugin_classes(body.bundle_name, dags, {}, session=session)) != len(dags): + raise HTTPException(status.HTTP_403_FORBIDDEN, detail={"reason": "plugin_not_granted"}) + try: + for dag in dags: + DagSerialization.validate_serialized_dag(dag.data) + except (DeserializationError, ValueError, TypeError, KeyError) as error: + raise HTTPException(status.HTTP_422_UNPROCESSABLE_CONTENT, detail="Invalid serialized Dag") from error + + import_errors = {(body.bundle_name, path): error for path, error in body.import_errors.items()} + files_parsed = {(body.bundle_name, path) for path in [body.relative_fileloc, *body.parsed_definitions]} + files_parsed.update(import_errors) + try: + update_dag_parsing_results_in_db( + bundle_name=body.bundle_name, + bundle_version=body.bundle_version, + version_data=body.version_data, + dags=dags, + import_errors=import_errors, + parse_duration=body.parse_duration, + warnings={DagWarning(**warning.model_dump()) for warning in body.warnings}, + files_parsed=files_parsed, + dag_source_codes={ + dag_id: DagSourceCode( + source_code=source.source_code + if source.source_code is not None + else "Source unavailable", + language=source.language, + ) + for dag_id, source in body.source_codes.items() + }, + atomic=True, + enforce_bundle_ownership=True, + resolve_bundle_paths=False, + session=session, + ) + except DagPublicationValidationError as error: + raise HTTPException( + status.HTTP_422_UNPROCESSABLE_CONTENT, + detail={"reason": "invalid_dag", "message": str(error)}, + ) from error + except DagBundleOwnershipError as error: + raise HTTPException( + status.HTTP_409_CONFLICT, detail={"reason": "dag_owned_by_another_bundle"} + ) from error + if checkpoint is None: + checkpoint = DagParseCheckpoint(job_id=job_id, source_key=source_key) + session.add(checkpoint) + checkpoint.attempt_id = body.attempt_id + checkpoint.dispatch_sequence = body.dispatch_sequence + checkpoint.payload_hash = payload_hash + checkpoint.accepted_at = timezone.utcnow() + checkpoint.bundle_revision = body.bundle_revision + session.flush() + return DagParseResultResponse(attempt_id=checkpoint.attempt_id, accepted_at=checkpoint.accepted_at) + + +@router.get( + "/{job_id}/bundles", + responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), + dependencies=[Security(require_auth, scopes=["token:dag_processor"])], +) +def get_bundles( + job_id: int, *, session: SessionDep, token=CurrentExecutionToken +) -> list[DagBundleStateResponse]: + """Read the provisioned catalog without constructing providers on the API server.""" + _validate_job(job_id, token, session=session) + bundles = session.scalars( + select(DagBundleModel) + .options(selectinload(DagBundleModel.teams)) + .where(DagBundleModel.name.in_(token.claims.dag_bundles), DagBundleModel.active.is_(True)) + .order_by(DagBundleModel.name) + ) + return [ + DagBundleStateResponse( + name=bundle.name, + version=bundle.version, + last_refreshed=bundle.last_refreshed, + revision=bundle.parse_revision, + team_name=bundle.team_name, + ) + for bundle in bundles + ] + + +def _validate_job(job_id, token, *, session): + if job_id != token.claims.job_id: + raise HTTPException(status.HTTP_404_NOT_FOUND, detail={"reason": "not_found"}) + job = session.scalar(select(Job).where(Job.id == job_id).with_for_update()) + if job is None or job.session_id != token.id or job.end_date is not None: + raise HTTPException(status.HTTP_403_FORBIDDEN, detail={"reason": "job_closed"}) + + +@router.post( + "/{job_id}/bundles/{bundle_name:path}/inventory", + responses=create_openapi_http_exception_doc( + [status.HTTP_404_NOT_FOUND, status.HTTP_409_CONFLICT, status.HTTP_413_CONTENT_TOO_LARGE] + ), + dependencies=[Security(require_auth, scopes=["token:dag_processor"])], +) +def publish_inventory( + job_id: int, + bundle_name: str, + body: DagBundleInventoryBody, + *, + session: SessionDep, + token=CurrentExecutionToken, +) -> DagBundleInventoryResponse: + """Reconcile observed sources in one bundle, preserving sources whose discovery failed.""" + if bundle_name not in token.claims.dag_bundles: + raise HTTPException(status.HTTP_403_FORBIDDEN, detail={"reason": "bundle_not_granted"}) + _validate_job(job_id, token, session=session) + source_key = hashlib.sha256(json.dumps(["inventory", bundle_name]).encode()).hexdigest() + payload_hash = hashlib.sha256(body.model_dump_json().encode()).hexdigest() + checkpoint = session.get(DagParseCheckpoint, (job_id, source_key)) + if checkpoint is not None: + if checkpoint.attempt_id == body.attempt_id: + if checkpoint.payload_hash != payload_hash or checkpoint.bundle_revision is None: + raise HTTPException(status.HTTP_409_CONFLICT, detail={"reason": "publication_conflict"}) + return DagBundleInventoryResponse( + attempt_id=checkpoint.attempt_id, + accepted_at=checkpoint.accepted_at, + revision=checkpoint.bundle_revision, + ) + if body.dispatch_sequence <= checkpoint.dispatch_sequence: + raise HTTPException(status.HTTP_409_CONFLICT, detail={"reason": "publication_superseded"}) + bundle = session.scalar( + select(DagBundleModel).where(DagBundleModel.name == bundle_name).with_for_update(key_share=True) + ) + if bundle is None or not bundle.active: + raise HTTPException(status.HTTP_409_CONFLICT, detail={"reason": "bundle_inactive"}) + if body.expected_revision != bundle.parse_revision: + raise HTTPException(status.HTTP_409_CONFLICT, detail={"reason": "source_changed"}) + inventory_hash = hashlib.sha256( + json.dumps([body.version, sorted(body.files), sorted(body.failed_sources)]).encode() + ).hexdigest() + if bundle.inventory_hash != inventory_hash or bundle.parse_revision is None: + bundle.parse_revision = uuid7() + bundle.inventory_hash = inventory_hash + bundle.version = body.version + bundle.last_refreshed = timezone.utcnow() + present = set(body.files) + failed_sources = set(body.failed_sources) + + def is_removed(relative_fileloc: str | None) -> bool: + return ( + relative_fileloc is not None + and relative_fileloc not in present + and not is_source_discovery_failed(relative_fileloc, failed_sources) + ) + + # Lock only the Dags being deactivated, so the rest of the bundle stays available to the scheduler. + removed = sorted( + dag_id + for dag_id, relative_fileloc in session.execute( + select(DagModel.dag_id, DagModel.relative_fileloc).where( + DagModel.bundle_name == bundle_name, DagModel.is_stale.is_(False) + ) + ) + if is_removed(relative_fileloc) + ) + deactivated = False + for start in range(0, len(removed), 500): + for dag in session.scalars( + select(DagModel) + .where( + DagModel.dag_id.in_(removed[start : start + 500]), + DagModel.bundle_name == bundle_name, + DagModel.is_stale.is_(False), + ) + .order_by(DagModel.dag_id) + .with_for_update(key_share=True) + ): + if is_removed(dag.relative_fileloc): + dag.is_stale = True + deactivated = True + if deactivated: + session.flush() + remove_references_to_deleted_dags(session=session) + error_ids = [ + error.id + for error in session.scalars( + select(ParseImportError).where(ParseImportError.bundle_name == bundle_name) + ) + if error.filename not in present and not is_source_discovery_failed(error.filename, failed_sources) + ] + for start in range(0, len(error_ids), 500): + session.execute( + delete(ParseImportError).where(ParseImportError.id.in_(error_ids[start : start + 500])) + ) + if checkpoint is None: + checkpoint = DagParseCheckpoint(job_id=job_id, source_key=source_key) + session.add(checkpoint) + checkpoint.attempt_id = body.attempt_id + checkpoint.dispatch_sequence = body.dispatch_sequence + checkpoint.payload_hash = payload_hash + checkpoint.accepted_at = bundle.last_refreshed + checkpoint.bundle_revision = bundle.parse_revision + session.flush() + return DagBundleInventoryResponse( + attempt_id=checkpoint.attempt_id, + accepted_at=checkpoint.accepted_at, + revision=bundle.parse_revision, + ) + + +@router.post( + "/{job_id}/requested-work/{kind}/claim", + responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), + dependencies=[Security(require_auth, scopes=["token:dag_processor"])], +) +def claim_work( + job_id: int, + kind: Literal["callbacks", "priority"], + body: ProcessorWorkClaimBody, + *, + session: SessionDep, + token=CurrentExecutionToken, +) -> list[ProcessorWorkItem]: + """Retain requests until acknowledged, and replay a lost claim response.""" + if not set(body.bundle_names) <= token.claims.dag_bundles: + raise HTTPException(status.HTTP_403_FORBIDDEN, detail={"reason": "bundle_not_granted"}) + _validate_job(job_id, token, session=session) + model = DagProcessorCallback if kind == "callbacks" else DagPriorityParsingRequest + replay = list( + cast( + "list[DagProcessorCallback | DagPriorityParsingRequest]", + session.scalars( + select(model) + .where(model.processor_job_id == job_id, model.processor_claim_id == body.claim_id) + .order_by(model.id) + ).all(), + ) + ) + if not replay: + ready = select(DagBundleModel.name).where(DagBundleModel.active.is_(True)) + retired = select(Job.id).where(Job.end_date.is_not(None)) + query = select(model).where( + model.bundle_name.in_(body.bundle_names), + model.bundle_name.in_(ready), + or_(model.processor_job_id.is_(None), model.processor_job_id.in_(retired)), + ) + if kind == "callbacks": + query = query.where( + or_(DagProcessorCallback.state.is_(None), DagProcessorCallback.state == CallbackState.QUEUED) + ).order_by(DagProcessorCallback.priority_weight.desc(), DagProcessorCallback.id) + else: + query = query.order_by(DagPriorityParsingRequest.id) + replay = list( + cast( + "list[DagProcessorCallback | DagPriorityParsingRequest]", + session.scalars( + with_row_locks(query.limit(body.limit), session=session, skip_locked=True) + ).all(), + ) + ) + for work in replay: + work.processor_job_id = job_id + work.processor_claim_id = body.claim_id + if isinstance(work, DagProcessorCallback): + work.state = CallbackState.QUEUED + result = [] + for work in replay: + if work.bundle_name is None or work.bundle_name not in token.claims.dag_bundles: + raise HTTPException(status.HTTP_403_FORBIDDEN, detail={"reason": "bundle_not_granted"}) + if isinstance(work, DagProcessorCallback): + request = work.get_callback_request() + path = request.filepath + callback = request.to_json() + else: + path, callback = work.relative_fileloc, None + result.append( + ProcessorWorkItem( + id=str(work.id), + claim_id=body.claim_id, + bundle_name=work.bundle_name, + relative_fileloc=path, + callback=callback, + ) + ) + return result + + +@router.post( + "/{job_id}/requested-work/{kind}/{work_id}/ack", + responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND, status.HTTP_409_CONFLICT]), + status_code=status.HTTP_204_NO_CONTENT, + dependencies=[Security(require_auth, scopes=["token:dag_processor"])], +) +def acknowledge_work( + job_id: int, + kind: Literal["callbacks", "priority"], + work_id: str, + body: ProcessorWorkAckBody, + *, + session: SessionDep, + token=CurrentExecutionToken, +) -> None: + """Reject acknowledgments that could consume a replacement owner's claim.""" + _validate_job(job_id, token, session=session) + model = DagProcessorCallback if kind == "callbacks" else DagPriorityParsingRequest + try: + key = UUID(work_id) if kind == "callbacks" else work_id + except ValueError as error: + raise HTTPException(status.HTTP_422_UNPROCESSABLE_CONTENT, detail="Invalid callback ID") from error + work = cast( + "DagProcessorCallback | DagPriorityParsingRequest | None", + session.scalar(select(model).where(model.id == key).with_for_update()), + ) + if work is None: + return + if work.bundle_name not in token.claims.dag_bundles: + raise HTTPException(status.HTTP_403_FORBIDDEN, detail={"reason": "bundle_not_granted"}) + if (work.processor_job_id, work.processor_claim_id) != (job_id, body.claim_id): + raise HTTPException(status.HTTP_409_CONFLICT, detail={"reason": "claim_replaced"}) + # Callback runners log individual failures without failing the child process. This + # acknowledges delivery, preserving the existing consume-after-processing semantics. + session.delete(work) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/jobs.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/jobs.py index f0b5f4f82eecf..018e5faaab42c 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/jobs.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/jobs.py @@ -172,7 +172,7 @@ def register_job( A registration creates one Job. Repeating it while that Job is open returns the Job with a fresh token; once the Job completes or is replaced, the registration is refused for good. A new registration is refused while the session's Job is alive, and otherwise ends and replaces it, which ends every token - issued for the replaced Job. + issued for the replaced Job and releases its claims. """ requested = token.claims.dag_bundles if body.bundle_names is None else set(body.bundle_names) if ungranted := requested - token.claims.dag_bundles: @@ -195,6 +195,7 @@ def register_job( detail={"reason": "job_running", "message": f"Session already has running Job {previous.id}"}, ) if previous.end_date is None: + # Its claims are released only once it ends, and maintenance skips Jobs without a session. previous.state = JobState.FAILED previous.end_date = now previous.session_id = None diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2027_02_28.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2027_02_28.py index b54a50236284a..a02678ac36237 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2027_02_28.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2027_02_28.py @@ -21,7 +21,7 @@ class AddDagProcessorJobEndpoints(VersionChange): - """Add Dag processor lifecycle and parse-token endpoints.""" + """Add Dag processor lifecycle, discovery inventory, publication, and requested-work endpoints.""" description = __doc__ @@ -30,4 +30,9 @@ class AddDagProcessorJobEndpoints(VersionChange): endpoint("/jobs/{job_id}/heartbeat", ["POST"]).didnt_exist, endpoint("/jobs/{job_id}/complete", ["POST"]).didnt_exist, endpoint("/jobs/{job_id}/parse-token", ["POST"]).didnt_exist, + endpoint("/jobs/{job_id}/parse-results", ["POST"]).didnt_exist, + endpoint("/jobs/{job_id}/bundles", ["GET"]).didnt_exist, + endpoint("/jobs/{job_id}/bundles/{bundle_name:path}/inventory", ["POST"]).didnt_exist, + endpoint("/jobs/{job_id}/requested-work/{kind}/claim", ["POST"]).didnt_exist, + endpoint("/jobs/{job_id}/requested-work/{kind}/{work_id}/ack", ["POST"]).didnt_exist, ) diff --git a/airflow-core/src/airflow/cli/cli_config.py b/airflow-core/src/airflow/cli/cli_config.py index aec3ec81dfa47..b6e9f2e6a666e 100644 --- a/airflow-core/src/airflow/cli/cli_config.py +++ b/airflow-core/src/airflow/cli/cli_config.py @@ -2304,6 +2304,21 @@ class GroupCommand(NamedTuple): 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 b8195bead2895..b49a4d6662393 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 873b519576fab..c4aa143271cc4 100644 --- a/airflow-core/src/airflow/config_templates/config.yml +++ b/airflow-core/src/airflow/config_templates/config.yml @@ -2775,7 +2775,9 @@ scheduler: description: | How often (in seconds) to check for stale DAGs (DAGs which are no longer present in the expected files) which should be deactivated, as well as assets that are no longer - referenced and should be marked as orphaned. + referenced and should be marked as orphaned. The same cleanup marks the Dags of inactive + or unassigned bundles stale, removes stale Dag warnings, and ends Dag processor Jobs that + exceeded ``[dag_processor] job_heartbeat_timeout``. version_added: 2.5.0 type: integer example: ~ @@ -3267,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. @@ -3430,10 +3458,31 @@ 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: ~ default: "30" + job_heartbeat_timeout: + 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, 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: ~ + default: "300" disable_bundle_versioning: description: | Always run tasks with the latest code. If set to True, the bundle version will not diff --git a/airflow-core/src/airflow/dag_processing/api_client.py b/airflow-core/src/airflow/dag_processing/api_client.py index 10e44a7ad206f..14f6e6db4c750 100644 --- a/airflow-core/src/airflow/dag_processing/api_client.py +++ b/airflow-core/src/airflow/dag_processing/api_client.py @@ -25,13 +25,24 @@ 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 -import httpx +import httpx2 import structlog from uuid6 import uuid7 +from airflow.api_fastapi.execution_api.datamodels.dag_parsing import ( + DagBundleInventoryBody, + DagBundleInventoryResponse, + DagBundleStateResponse, + DagParseResultBody, + DagParseResultResponse, + ProcessorWorkAckBody, + ProcessorWorkClaimBody, + ProcessorWorkItem, +) from airflow.api_fastapi.execution_api.datamodels.job import ( DagParseTokenBody, DagParseTokenResponse, @@ -73,10 +84,10 @@ class DagProcessorJobAlreadyRunning(RuntimeError): """The session's previous Job has not completed or stopped heartbeating yet.""" -def get_error_reason(error: httpx.HTTPStatusError) -> str | None: +def get_error_reason(error: httpx2.HTTPStatusError) -> str | None: try: payload = error.response.json() - except (ValueError, httpx.ResponseNotRead): + except (ValueError, httpx2.ResponseNotRead): return None detail = payload.get("detail") if isinstance(payload, dict) else None return detail.get("reason") if isinstance(detail, dict) else None @@ -134,6 +145,7 @@ def __init__( 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 @@ -168,14 +180,14 @@ def use_parse(self, context: DagParseContext) -> Iterator[DagProcessorAPIClient] finally: self._parse_context.reset(selected) - def _get_parse_token(self, context: DagParseContext, *, retry: bool) -> str: + def _get_parse_token(self, context: DagParseContext, *, retry: bool, timeout: httpx2.Timeout) -> str: if context.token is not None and monotonic() < context.expires_at: if monotonic() < context.renew_at: return context.token try: self._exchange_parse_token(context, retry=False, timeout=self._get_bounded_timeout()) - except (httpx.HTTPError, ValueError) as error: - if isinstance(error, httpx.HTTPStatusError) and error.response.status_code < 500: + except (httpx2.HTTPError, ValueError) as error: + if isinstance(error, httpx2.HTTPStatusError) and error.response.status_code < 500: raise context.renew_at = monotonic() + 30 log.warning( @@ -186,12 +198,12 @@ def _get_parse_token(self, context: DagParseContext, *, retry: bool) -> str: ) if monotonic() < context.expires_at: return context.token - return self._exchange_parse_token(context, retry=retry) + return self._exchange_parse_token(context, retry=retry, timeout=timeout) def _exchange_parse_token( - self, context: DagParseContext, *, retry: bool, timeout: httpx.Timeout | None = None + self, context: DagParseContext, *, retry: bool, timeout: httpx2.Timeout | None = None ) -> str: - self._ensure_job_token(retry=retry) + self._ensure_job_token(retry=retry, timeout=timeout) started_at = monotonic() response = super().request( "POST", @@ -208,7 +220,7 @@ def _exchange_parse_token( context.renew_at = started_at + lifetime * 0.8 return parsed.token - def _update_auth(self, response: httpx.Response) -> None: + def _update_auth(self, response: httpx2.Response) -> None: # Task-token refresh headers cannot replace a provisioned or Job-bound credential. pass @@ -241,14 +253,14 @@ def register_job(self, *, retry: bool = True) -> int: self._check_can_run() try: return self._register_job(retry=retry) - except httpx.HTTPStatusError as error: + except httpx2.HTTPStatusError as error: if error.response.status_code == 409 and get_error_reason(error) == "job_running": raise DagProcessorJobAlreadyRunning( "The previous Dag processor Job is still alive" ) from error raise - def _register_job(self, *, retry: bool, timeout: httpx.Timeout | None = None) -> int: + def _register_job(self, *, retry: bool, timeout: httpx2.Timeout | None = None) -> int: if self.restart_required: raise DagProcessorRegistrationRetired( "The Dag processor registration has ended; restart required" @@ -269,7 +281,7 @@ def _register_job(self, *, retry: bool, timeout: httpx.Timeout | None = None) -> timeout=timeout or self.timeout, ) break - except httpx.HTTPStatusError as error: + except httpx2.HTTPStatusError as error: if error.response.status_code == 409 and get_error_reason(error) == "registration_retired": self._restart_required = True raise DagProcessorRegistrationRetired( @@ -296,18 +308,26 @@ def _register_job(self, *, retry: bool, timeout: httpx.Timeout | None = None) -> self._retry_renewal_at = 0.0 return registered.job_id - def _get_bounded_timeout(self) -> httpx.Timeout: - return httpx.Timeout( + def _get_bounded_timeout(self) -> httpx2.Timeout: + return httpx2.Timeout( **{ key: min(value if value is not None else 1.0, 1.0) for key, value in self.timeout.as_dict().items() } ) - def _ensure_job_token(self, *, retry: bool) -> None: + def _get_write_timeout(self) -> httpx2.Timeout: + """Fail fast while the API is unreachable, but let a slow write commit and acknowledge.""" + bounded = self._get_bounded_timeout() + return httpx2.Timeout( + connect=bounded.connect, read=self.timeout.read, write=self.timeout.write, pool=bounded.pool + ) + + def _ensure_job_token(self, *, retry: bool, timeout: httpx2.Timeout | None = None) -> None: self._require_job_id() + timeout = timeout or (None if retry else self._get_bounded_timeout()) if monotonic() >= self._expires_at: - self._register_job(retry=retry, timeout=None if retry else self._get_bounded_timeout()) + self._register_job(retry=retry, timeout=timeout) return if self.restart_required or monotonic() < self._retry_renewal_at: return @@ -315,7 +335,7 @@ def _ensure_job_token(self, *, retry: bool) -> None: if self._read_session_token() != self._registered_with or monotonic() >= self._renew_at: # Do not spend the runtime request's retry budget on an optional early renewal. self._register_job(retry=False, timeout=self._get_bounded_timeout()) - except (httpx.HTTPError, OSError, ValueError, DagProcessorRegistrationRetired) as error: + except (httpx2.HTTPError, OSError, ValueError, DagProcessorRegistrationRetired) as error: self._retry_renewal_at = monotonic() + 30 log.warning( "Unable to renew Dag processor Job token", @@ -324,15 +344,19 @@ def _ensure_job_token(self, *, retry: bool) -> None: restart_required=self.restart_required, ) if monotonic() >= self._expires_at: - self._register_job(retry=retry, timeout=None if retry else self._get_bounded_timeout()) + self._register_job(retry=retry, timeout=timeout) - def request(self, *args, retry: bool = True, **kwargs) -> httpx.Response: + def request(self, *args, retry: bool = True, **kwargs) -> httpx2.Response: """Use a parsing credential for subprocess requests, and the Job credential for manager work.""" self._check_can_run() try: - headers = httpx.Headers(kwargs.get("headers")) + headers = httpx2.Headers(kwargs.get("headers")) if context := self._parse_context.get(): - kwargs["auth"] = BearerAuth(self._get_parse_token(context, retry=retry)) + # Avoid retry sleeps on the manager loop; required credentials use the runtime request's timeout. + retry = False + timeout = kwargs.get("timeout", httpx2.USE_CLIENT_DEFAULT) + timeout = self.timeout if timeout is httpx2.USE_CLIENT_DEFAULT else httpx2.Timeout(timeout) + kwargs["auth"] = BearerAuth(self._get_parse_token(context, retry=retry, timeout=timeout)) headers["Airflow-Dag-Bundle"] = context.request.bundle_name else: self._ensure_job_token(retry=retry) @@ -342,7 +366,7 @@ def request(self, *args, retry: bool = True, **kwargs) -> httpx.Response: headers.setdefault("Content-Type", "application/json") kwargs["headers"] = headers return super().request(*args, retry=retry, **kwargs) - except httpx.HTTPStatusError as error: + except httpx2.HTTPStatusError as error: if error.response.status_code == 403 and get_error_reason(error) == "job_closed": self._restart_required = True raise DagProcessorRegistrationRetired( @@ -356,6 +380,77 @@ def heartbeat(self) -> JobState: response = self.request("POST", f"jobs/{job_id}/heartbeat", retry=False, headers=_JOB_API_HEADERS) return JobHeartbeatResponse.model_validate_json(response.content).state + def publish_parse_result(self, body: DagParseResultBody) -> DagParseResultResponse: + """Publish once; the manager retains the request and schedules acknowledgment recovery.""" + response = self.request( + "POST", + f"jobs/{self._require_job_id()}/parse-results", + content=body.model_dump_json(), + headers=_JOB_API_HEADERS, + retry=False, + timeout=self._get_write_timeout(), + ) + return DagParseResultResponse.model_validate_json(response.content) + + def get_bundles(self) -> list[DagBundleStateResponse]: + response = self.request( + "GET", + f"jobs/{self._require_job_id()}/bundles", + headers=_JOB_API_HEADERS, + retry=False, + timeout=self._get_bounded_timeout(), + ) + return [DagBundleStateResponse.model_validate(item) for item in response.json()] + + def 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", + f"jobs/{self._require_job_id()}/bundles/{quote(bundle_name, safe='')}/inventory", + content=body.model_dump_json(), + headers=_JOB_API_HEADERS, + retry=False, + timeout=self._get_write_timeout(), + ) + return DagBundleInventoryResponse.model_validate_json(response.content) + def complete_job(self, state: TerminalJobState) -> None: """Complete this Job, retaining its identity and outcome across lost acknowledgments.""" body = JobCompleteBody(state=state) @@ -377,7 +472,7 @@ def complete_job(self, state: TerminalJobState) -> None: json=body.model_dump(mode="json"), headers=_JOB_API_HEADERS, ) - except httpx.HTTPStatusError as error: + except httpx2.HTTPStatusError as error: if ( retried_expiry or error.response.status_code not in (401, 403) diff --git a/airflow-core/src/airflow/dag_processing/bundles/manager.py b/airflow-core/src/airflow/dag_processing/bundles/manager.py index fe74ffff7ab9f..6aa50af2ec81a 100644 --- a/airflow-core/src/airflow/dag_processing/bundles/manager.py +++ b/airflow-core/src/airflow/dag_processing/bundles/manager.py @@ -22,6 +22,7 @@ import os import warnings from collections import defaultdict +from contextlib import nullcontext from pathlib import Path from typing import TYPE_CHECKING, Any, NamedTuple, cast @@ -40,7 +41,8 @@ from airflow.utils.session import NEW_SESSION, create_session, provide_session if TYPE_CHECKING: - from collections.abc import Iterable + from collections.abc import Callable, Collection, Iterable + from contextlib import AbstractContextManager from sqlalchemy.engine import CursorResult from sqlalchemy.orm import Session @@ -312,8 +314,16 @@ def _sign_bundle_url(url: str, bundle_name: str) -> str: class DagBundlesManager(LoggingMixin): """Manager for DAG bundles.""" - def __init__(self, *args, **kwargs): + def __init__( + self, + *args, + bundle_names: Collection[str] | None = None, + bundle_context: Callable[[str], AbstractContextManager] | None = None, + **kwargs, + ): super().__init__(*args, **kwargs) + self._bundle_names = bundle_names + self._bundle_context = bundle_context or (lambda _: nullcontext()) self._bundle_config: dict[str, _InternalBundleConfig] = {} self.parse_config() @@ -333,6 +343,8 @@ def parse_config(self) -> None: return for bundle_config in bundle_config_list: + if self._bundle_names is not None and bundle_config.name not in self._bundle_names: + continue if bundle_config.team_name and not conf.getboolean("core", "multi_team"): raise AirflowConfigException( "Section `dag_processor` key `dag_bundle_config_list` " @@ -349,7 +361,13 @@ def parse_config(self) -> None: self.log.info("DAG bundles loaded: %s", ", ".join(self._bundle_config.keys())) @provide_session - def sync_bundles_to_db(self, *, deactivate_missing: bool = True, session: Session = NEW_SESSION) -> None: + def sync_bundles_to_db( + self, + *, + deactivate_missing: bool = True, + include_bundle_urls: bool = True, + session: Session = NEW_SESSION, + ) -> None: """ Persist the configured DAG bundles into ``DagBundleModel`` rows. @@ -359,6 +377,8 @@ def sync_bundles_to_db(self, *, deactivate_missing: bool = True, session: Sessio ``DagBag`` plus ``sync_bag_to_db`` (or, in production, the DAG processor); calling this method does not trigger that work. + :param include_bundle_urls: Sync ownership without constructing bundles when ``False``, so + API authorization can be established before bundle code requests secrets. :param deactivate_missing: When ``True`` (the default), any bundle stored in the database that is not present in this manager's config is marked inactive. This is only correct when the calling process sees the @@ -400,24 +420,27 @@ def _extract_and_sign_template(bundle_name: str) -> tuple[str | None, dict]: if not team: raise _bundle_item_exc(f"Team '{config.team_name}' does not exist") - try: - new_template, new_params = _extract_and_sign_template(name) - except Exception as e: - self.log.exception("Error creating bundle '%s': %s", name, e) - continue + if include_bundle_urls: + try: + with self._bundle_context(name): + new_template, new_params = _extract_and_sign_template(name) + except Exception as e: + self.log.exception("Error creating bundle '%s': %s", name, e) + continue if bundle := stored.pop(name, None): bundle.active = True - if new_template != bundle.signed_url_template: + if include_bundle_urls and new_template != bundle.signed_url_template: bundle.signed_url_template = new_template self.log.debug("Updated URL template for bundle %s", name) - if new_params != bundle.template_params: + if include_bundle_urls and new_params != bundle.template_params: bundle.template_params = new_params self.log.debug("Updated template parameters for bundle %s", name) else: bundle = DagBundleModel(name=name) - bundle.signed_url_template = new_template - bundle.template_params = new_params + if include_bundle_urls: + bundle.signed_url_template = new_template + bundle.template_params = new_params session.add(bundle) self.log.info("Added new DAG bundle %s to the database", name) @@ -694,9 +717,10 @@ def get_bundle( cfg_bundle = self._bundle_config.get(name) if not cfg_bundle: raise ValueError(f"Requested bundle '{name}' is not configured.") - return cfg_bundle.bundle_class( - name=name, version=version, version_data=version_data, **cfg_bundle.kwargs - ) + with self._bundle_context(name): + return cfg_bundle.bundle_class( + name=name, version=version, version_data=version_data, **cfg_bundle.kwargs + ) @classmethod def is_bundle_configured(cls, name: str) -> bool: @@ -716,7 +740,9 @@ def get_all_dag_bundles(self) -> Iterable[BaseDagBundle]: """ for name, cfg in self._bundle_config.items(): try: - yield cfg.bundle_class(name=name, version=None, **cfg.kwargs) + with self._bundle_context(name): + bundle = cfg.bundle_class(name=name, version=None, **cfg.kwargs) + yield bundle except Exception as e: self.log.exception("Error creating bundle '%s': %s", name, e) # Skip this bundle and continue with others diff --git a/airflow-core/src/airflow/dag_processing/collection.py b/airflow-core/src/airflow/dag_processing/collection.py index cd7228d210f44..c18a05282171f 100644 --- a/airflow-core/src/airflow/dag_processing/collection.py +++ b/airflow-core/src/airflow/dag_processing/collection.py @@ -28,12 +28,13 @@ from __future__ import annotations import traceback +from pathlib import PurePosixPath from typing import TYPE_CHECKING, Any, NamedTuple, TypeVar import structlog from sqlalchemy import delete, false, func, insert, or_, select, tuple_, update from sqlalchemy.exc import OperationalError -from sqlalchemy.orm import joinedload, load_only +from sqlalchemy.orm import Session, joinedload, load_only from airflow._shared.timezones.timezone import utcnow from airflow.assets.manager import asset_manager @@ -50,6 +51,7 @@ TaskOutletAssetReference, ) from airflow.models.dag import DagModel, DagOwnerAttributes, DagTag +from airflow.models.dag_version import DagVersion from airflow.models.dagbundle import DagBundleModel from airflow.models.dagrun import DagRun from airflow.models.dagwarning import DagWarning, DagWarningType @@ -68,13 +70,12 @@ from airflow.serialization.serialized_objects import BaseSerialization, LazyDeserializedDAG from airflow.triggers.base import BaseEventTrigger from airflow.utils.retries import MAX_DB_RETRIES, run_with_db_retries -from airflow.utils.sqlalchemy import get_dialect_name, with_row_locks +from airflow.utils.sqlalchemy import build_upsert_stmt, get_dialect_name, with_row_locks from airflow.utils.types import DagRunType if TYPE_CHECKING: from collections.abc import Collection, Iterable, Iterator - from sqlalchemy.orm import Session from sqlalchemy.sql import Select from airflow.models.serialized_dag import DagWriteMetadata @@ -86,6 +87,65 @@ log = structlog.get_logger(__name__) +class DagBundleOwnershipError(ValueError): + """A publication would replace a Dag owned by another bundle.""" + + +class DagPublicationValidationError(ValueError): + """A publication contains invalid Dag metadata.""" + + +def _is_dag_permission_validation_error(error: Exception) -> bool: + if isinstance(error, ValueError): + return True + # Supported older FAB releases use FabException for both ACL validation and backend failures. + # Only the validation method's own raises have the former meaning; no FAB import is needed. + exception_type = type(error) + module = "airflow.providers.fab.auth_manager.security_manager.override" + if exception_type.__module__ != module or exception_type.__name__ != "FabException": + return False + origin = error.__traceback__ + while origin is not None and origin.tb_next is not None: + origin = origin.tb_next + return ( + origin is not None + and origin.tb_frame.f_globals.get("__name__") == module + and origin.tb_frame.f_code.co_name == "_sync_dag_view_permissions" + ) + + +def is_source_discovery_failed(relative_fileloc: str | None, failed_sources: Collection[str]) -> bool: + """Check the source and its enclosing containers without matching sibling path prefixes.""" + if relative_fileloc is None or not failed_sources: + return False + return relative_fileloc in failed_sources or any( + str(parent) in failed_sources for parent in PurePosixPath(relative_fileloc).parents + ) + + +def validate_dag_bundle_ownership(dags: Iterable[DagModel], bundle_name: str, *, session: Session) -> None: + """Check locked Dag rows, allowing recovery of orphaned, unversioned legacy rows.""" + for dag in dags: + if dag.bundle_name == bundle_name: + continue + if dag.is_stale and dag.relative_fileloc is None: + previous_bundle = session.scalar( + select(DagBundleModel) + .where(DagBundleModel.name == dag.bundle_name) + .with_for_update(key_share=True) + ) + has_version = session.scalar( + select(DagVersion.id).where(DagVersion.dag_id == dag.dag_id).limit(1) + ) + if previous_bundle is not None and not previous_bundle.active and has_version is None: + # Adopting a row also adopts its run history, which must not cross a team boundary. + if not conf.getboolean("core", "multi_team") or previous_bundle.team_name == ( + DagBundleModel.get_team_name(bundle_name, session=session) + ): + continue + raise DagBundleOwnershipError(f"Dag {dag.dag_id!r} belongs to bundle {dag.bundle_name!r}") + + def _create_orm_dags( bundle_name: str, bundle_version: str | None, @@ -256,11 +316,13 @@ def _update_dag_owner_links(dag_owner_links: dict[str, str], dm: DagModel, *, se def _serialize_dag_capturing_errors( dag: LazyDeserializedDAG, bundle_name, + *, session: Session, bundle_version: str | None, version_data: dict | None = None, _prefetched: DagWriteMetadata | None = None, dag_source_code: DagSourceCode | None = None, + atomic: bool = False, ): """ Try to serialize the dag to the DB, but make a note of any errors. @@ -295,12 +357,36 @@ def _serialize_dag_capturing_errors( sync_dag_perms = getattr(get_auth_manager_cls(), "sync_dag_perms", None) if sync_dag_perms: - sync_dag_perms(dag, session=session) + if atomic: + # Auth managers such as FAB commit internally. Savepoints keep those commits inside + # the publication transaction. + session.flush() + connection = session.connection() + driver_connection = connection.connection.driver_connection + if ( + connection.dialect.name == "sqlite" + and driver_connection is not None + and not driver_connection.in_transaction + ): + # SQLite's legacy transaction mode does not start a transaction for SELECTs. + connection.exec_driver_sql("BEGIN") + with Session(bind=connection, join_transaction_mode="create_savepoint") as perms: + try: + sync_dag_perms(dag, session=perms) + except Exception as error: + if _is_dag_permission_validation_error(error): + raise DagPublicationValidationError(str(error)) from error + raise + perms.commit() + else: + sync_dag_perms(dag, session=session) return [] except OperationalError: raise except Exception: + if atomic: + raise log.exception("Failed to write serialized DAG dag_id=%s fileloc=%s", dag.dag_id, dag.fileloc) dagbag_import_error_traceback_depth = conf.getint("core", "dagbag_import_error_traceback_depth") return [ @@ -387,6 +473,8 @@ def _update_import_errors( files_parsed: set[tuple[str, str]], import_errors: dict[tuple[str, str], str], session: Session, + *, + resolve_bundle_paths: bool = True, ): from airflow.listeners.listener import get_listener_manager @@ -444,7 +532,8 @@ def _update_import_errors( ) if import_error is not None: get_listener_manager().hook.on_existing_dag_import_error( - filename=import_error.full_file_path(), stacktrace=stacktrace + filename=import_error.full_file_path() if resolve_bundle_paths else relative_fileloc, + stacktrace=stacktrace, ) except Exception: log.exception("error calling listener") @@ -460,7 +549,8 @@ def _update_import_errors( # sending notification when a new dag import error occurs try: get_listener_manager().hook.on_new_dag_import_error( - filename=import_error.full_file_path(), stacktrace=stacktrace + filename=import_error.full_file_path() if resolve_bundle_paths else relative_fileloc, + stacktrace=stacktrace, ) except Exception: log.exception("error calling listener") @@ -580,8 +670,11 @@ def update_dag_parsing_results_in_db( import_errors: dict[tuple[str, str], str], parse_duration: float | None, warnings: set[DagWarning], - session: Session, *, + session: Session, + atomic: bool = False, + enforce_bundle_ownership: bool = False, + resolve_bundle_paths: bool = True, version_data: dict | None = None, warning_types: tuple[DagWarningType, ...] = ( DagWarningType.DUPLICATE_DAG_ID, @@ -614,6 +707,12 @@ def update_dag_parsing_results_in_db( import errors are cleared for files that were parsed but no longer contain DAGs. :param dag_source_codes: Source code read by the Dag importers, keyed by dag_id. Dags without an entry have their source read from ``fileloc``. + :param atomic: Let the caller roll back and retry the whole transaction, including authorization + and a publication receipt. Persistence and diagnostic failures propagate in this mode. + :param enforce_bundle_ownership: Check ownership on the locked rows used for persistence. + :param resolve_bundle_paths: Construct bundles to give import error listeners absolute paths. The + API server passes False: it does not host the processor's bundles, so listeners receive + bundle-relative paths instead. """ accepted = _reject_other_teams_plugin_classes(bundle_name, dags, import_errors, session=session) if len(accepted) != len(dags): @@ -629,11 +728,13 @@ def update_dag_parsing_results_in_db( try: duplicate_warnings = _build_duplicate_dag_id_warnings(dags, bundle_name, session) except Exception: + if atomic: + raise log.exception("Error building duplicate dag_id warnings.") else: warnings = set(warnings) | duplicate_warnings - for attempt in run_with_db_retries(logger=log): + for attempt in run_with_db_retries(max_retries=1 if atomic else MAX_DB_RETRIES, logger=log): with attempt: serialize_errors = [] log.debug( @@ -644,7 +745,12 @@ def update_dag_parsing_results_in_db( log.debug("Calling the DAG.bulk_sync_to_db method") try: SerializedDAG.bulk_write_to_db( - bundle_name, bundle_version, dags, parse_duration, session=session + bundle_name, + bundle_version, + dags, + parse_duration, + enforce_bundle_ownership=enforce_bundle_ownership, + session=session, ) # Bulk prefetch metadata for all DAGs to avoid the standard per-DAG # metadata lookups in write_dag. This replaces the update-interval, @@ -664,10 +770,12 @@ def update_dag_parsing_results_in_db( session=session, _prefetched=prefetched_metadata.get(dag.dag_id), dag_source_code=dag_source_codes.get(dag.dag_id), + atomic=atomic, ) ) except OperationalError: - session.rollback() + if not atomic: + session.rollback() raise # Only now we are "complete" do we update import_errors - don't want to record errors from # previous failed attempts @@ -678,14 +786,19 @@ def update_dag_parsing_results_in_db( files_parsed=files_parsed if files_parsed is not None else set(), import_errors=import_errors, session=session, + resolve_bundle_paths=resolve_bundle_paths, ) except Exception: + if atomic: + raise log.exception("Error logging import errors!") # Record DAG warnings in the metadatabase. try: _update_dag_warnings([dag.dag_id for dag in dags], warnings, warning_types, session) except Exception: + if atomic: + raise log.exception("Error logging DAG warnings.") session.flush() @@ -697,26 +810,56 @@ class DagModelOperation(NamedTuple): dags: dict[str, LazyDeserializedDAG] bundle_name: str bundle_version: str | None + enforce_bundle_ownership: bool = False def find_orm_dags(self, *, session: Session) -> dict[str, DagModel]: """Find existing DagModel objects from DAG objects.""" - stmt: Select[*tuple[DagModel]] = with_row_locks( - ( - select(DagModel) - .options(joinedload(DagModel.tags, innerjoin=False)) - .where(DagModel.dag_id.in_(self.dags)) - .options(joinedload(DagModel.schedule_asset_references)) - .options(joinedload(DagModel.schedule_asset_alias_references)) - .options(joinedload(DagModel.task_outlet_asset_references)) - .options(joinedload(DagModel.dag_owner_links)) - ), - of=DagModel, - session=session, + stmt = ( + select(DagModel) + .options(joinedload(DagModel.tags, innerjoin=False)) + .where(DagModel.dag_id.in_(self.dags)) + .order_by(DagModel.dag_id) + .options(joinedload(DagModel.schedule_asset_references)) + .options(joinedload(DagModel.schedule_asset_alias_references)) + .options(joinedload(DagModel.task_outlet_asset_references)) + .options(joinedload(DagModel.dag_owner_links)) ) + if self.enforce_bundle_ownership: + stmt = stmt.with_for_update(of=DagModel, key_share=True).execution_options(populate_existing=True) + else: + stmt = with_row_locks(stmt, of=DagModel, session=session) return {dm.dag_id: dm for dm in session.scalars(stmt).unique()} def add_dags(self, *, session: Session) -> dict[str, DagModel]: orm_dags = self.find_orm_dags(session=session) + if self.enforce_bundle_ownership: + for dag_id in sorted(self.dags.keys() - orm_dags.keys()): + model = DagModel( + dag_id=dag_id, bundle_name=self.bundle_name, bundle_version=self.bundle_version + ) + values = { + key: getattr(model, key) + for key in ( + "dag_id", + "bundle_name", + "bundle_version", + "max_active_tasks", + "max_active_runs", + "max_consecutive_failed_dag_runs", + "has_task_concurrency_limits", + ) + } + if self.dags[dag_id].is_paused_upon_creation is not None: + values["is_paused"] = self.dags[dag_id].is_paused_upon_creation + # A competing insert must retain its ownership until we inspect the locked row. + session.execute( + build_upsert_stmt( + get_dialect_name(session), DagModel, ["dag_id"], values, {"dag_id": DagModel.dag_id} + ) + ) + orm_dags = self.find_orm_dags(session=session) + validate_dag_bundle_ownership(orm_dags.values(), self.bundle_name, session=session) + return orm_dags orm_dags.update( (model.dag_id, model) for model in _create_orm_dags( 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 0000000000000..640724a85ef00 --- /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/maintenance.py b/airflow-core/src/airflow/dag_processing/maintenance.py new file mode 100644 index 0000000000000..c5ab9877634bd --- /dev/null +++ b/airflow-core/src/airflow/dag_processing/maintenance.py @@ -0,0 +1,91 @@ +# 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. + +"""Bounded metadata maintenance owned by the scheduler, including absent bundles.""" + +from __future__ import annotations + +import logging +from datetime import timedelta + +from sqlalchemy import delete, or_, select, tuple_ + +from airflow._shared.timezones import timezone +from airflow.configuration import conf +from airflow.jobs.job import Job, JobState +from airflow.models.dag import DagModel +from airflow.models.dagbundle import DagBundleModel +from airflow.models.dagwarning import DagWarning +from airflow.utils.session import create_session +from airflow.utils.sqlalchemy import with_row_locks + +log = logging.getLogger(__name__) + + +def cleanup_processor_metadata(batch_size: int = 500) -> None: + """Commit each bounded batch; a later interval picks up remaining rows.""" + with create_session() as session: + inactive = select(DagBundleModel.name).where(DagBundleModel.active.is_(False)) + query = ( + select(DagModel) + .where( + DagModel.is_stale.is_(False), + or_(DagModel.bundle_name.is_(None), DagModel.bundle_name.in_(inactive)), + ) + .order_by(DagModel.dag_id) + .limit(batch_size) + ) + for dag in session.scalars(with_row_locks(query, session=session, skip_locked=True)): + log.info( + "Deactivating Dag %s. Its bundle %s is no longer active or is NULL.", + dag.dag_id, + dag.bundle_name, + ) + dag.is_stale = True + with create_session() as session: + warning_query = ( + select(DagWarning.dag_id, DagWarning.warning_type) + .join(DagModel) + .where( + DagModel.is_stale.is_(True), + ) + .order_by(DagWarning.dag_id, DagWarning.warning_type) + .limit(batch_size) + ) + keys = session.execute( + with_row_locks(warning_query, of=DagWarning, session=session, skip_locked=True) + ).all() + if keys: + session.execute( + delete(DagWarning).where(tuple_(DagWarning.dag_id, DagWarning.warning_type).in_(keys)) + ) + with create_session() as session: + cutoff = timezone.utcnow() - timedelta(seconds=conf.getint("dag_processor", "job_heartbeat_timeout")) + job_query = ( + select(Job) + .where( + Job.job_type == "DagProcessorJob", + Job.session_id.is_not(None), + Job.end_date.is_(None), + Job.latest_heartbeat < cutoff, + ) + .order_by(Job.id) + .limit(batch_size) + ) + for job in session.scalars(with_row_locks(job_query, session=session, skip_locked=True)): + job.state = JobState.FAILED + job.end_date = timezone.utcnow() diff --git a/airflow-core/src/airflow/dag_processing/manager.py b/airflow-core/src/airflow/dag_processing/manager.py index 343d72c5e002d..ff0e91b3b6b8a 100644 --- a/airflow-core/src/airflow/dag_processing/manager.py +++ b/airflow-core/src/airflow/dag_processing/manager.py @@ -23,22 +23,28 @@ import gc import inspect import logging +import math import os import random import selectors import signal import sys import time -from collections import OrderedDict, defaultdict +from collections import OrderedDict, defaultdict, deque +from contextlib import AbstractContextManager, nullcontext +from copy import deepcopy from dataclasses import dataclass, field from datetime import datetime, timedelta from operator import attrgetter, itemgetter from pathlib import Path from typing import TYPE_CHECKING, Any, Literal, NamedTuple, cast +from uuid import UUID import attrs +import httpx2 import structlog -from sqlalchemy import select, update +from pydantic import TypeAdapter, ValidationError +from sqlalchemy import or_, select, update from sqlalchemy.exc import OperationalError from sqlalchemy.orm import load_only from tabulate import tabulate @@ -47,22 +53,32 @@ from airflow._shared.observability.metrics import stats from airflow._shared.observability.metrics.stats import normalize_name_for_stats from airflow._shared.timezones import timezone +from airflow.api_fastapi.execution_api.datamodels.dag_parsing import ( + MAX_PUBLICATION_BYTES, + DagBundleInventoryBody, + DagParseResultBody, + ParseSourceCode, + ParseWarning, + 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, unpack_bundle_version, ) from airflow.dag_processing.bundles.manager import DagBundlesManager -from airflow.dag_processing.collection import update_dag_parsing_results_in_db +from airflow.dag_processing.collection import is_source_discovery_failed, update_dag_parsing_results_in_db from airflow.dag_processing.importer_routing import get_claiming_importer from airflow.dag_processing.lang_sdk_processor import LangSDKDagFileProcessorProcess from airflow.dag_processing.processor import ( BaseDagFileProcessorProcess, DagFileParsingResult, DagFileProcessorProcess, + DagParseSource, ) +from airflow.jobs.job import Job from airflow.models.asset import remove_references_to_deleted_dags from airflow.models.dag import DagModel from airflow.models.dagbag import DagPriorityParsingRequest @@ -98,16 +114,28 @@ 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 +# Each send has its own short timeout; this bounds how long one loop spends on queued sends +# before it goes back to heartbeating and supervising parsing processes. +_API_SEND_BUDGET_SECONDS = 1.0 +_MAX_REJECTION_REASON_LENGTH = 1024 +_REJECTED_CONTENT_REASONS = frozenset({"dag_owned_by_another_bundle", "plugin_not_granted"}) + + +def _is_rejected_content(status_code: int, reason: str | None) -> bool: + """Whether the API refused a result for what it contains, so sending it again cannot succeed.""" + return status_code in (413, 422) or (status_code in (403, 409) and reason in _REJECTED_CONTENT_REASONS) + class BundleState(NamedTuple): """Persisted refresh state for a DAG bundle.""" last_refreshed: datetime | None version: str | None + revision: UUID | None = None @attrs.define @@ -120,6 +148,41 @@ class DagFileStat: last_duration: float | None = None run_count: int = 0 last_num_of_db_queries: int = 0 + last_attempt_time: datetime | None = None + + +@attrs.define +class PendingDagPublication: + """A completed import retained until publication succeeds or exhausts its retries.""" + + body: DagParseResultBody + stat: DagFileStat + refresh_generation: int + attempts: int = 0 + next_attempt_time: float = 0.0 + is_rejection_report: bool = False + """Whether this replaces a result the API refused, so a further refusal is final.""" + + +@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) @@ -131,6 +194,7 @@ class DagFileInfo: bundle_path: Path | None = field(compare=False, default=None) bundle_version: str | None = None definition_locs: frozenset[str] = field(compare=False, default=frozenset()) + failed_sources: frozenset[str] = field(compare=False, default=frozenset()) @property def absolute_path(self) -> Path: @@ -249,19 +313,36 @@ class DagFileProcessorManager(LoggingMixin): """An overridable heartbeat called once every time around the loop""" _file_queue: OrderedDict[DagFileInfo, None] = attrs.field(factory=OrderedDict, init=False) + _files_awaiting_source: dict[DagFileInfo, None] = attrs.field(factory=dict, init=False) + """Files held out of the queue until their bundle has a usable source, so they never stop it emptying.""" _file_stats: dict[DagFileInfo, DagFileStat] = attrs.field( factory=lambda: defaultdict(DagFileStat), init=False ) _dag_bundles: list[BaseDagBundle] = attrs.field(factory=list, init=False) _bundle_versions: dict[str, str | None] = attrs.field(factory=dict, init=False) - _bundle_version_data: dict[str, dict | None] = attrs.field(factory=dict, init=False) + _bundle_parse_sources: dict[str, DagParseSource] = attrs.field(factory=dict, init=False) + _bundle_refresh_generations: dict[str, int] = attrs.field(factory=lambda: defaultdict(int), init=False) + _bundles_waiting_for_refresh: set[str] = attrs.field(factory=set, init=False) + _dispatch_sequence: int = attrs.field(default=0, init=False) _multi_team: bool = attrs.field(factory=lambda: conf.getboolean("core", "multi_team"), init=False) _bundle_name_to_team_name: dict[str, str | None] = attrs.field(factory=dict, init=False) _processors: dict[DagFileInfo, BaseDagFileProcessorProcess] = attrs.field(factory=dict, init=False) + _pending_publications: dict[DagFileInfo, PendingDagPublication] = attrs.field(factory=dict, init=False) + _pending_inventories: dict[str, tuple[DagBundleInventoryBody, set[DagFileInfo], DagParseSource]] = ( + attrs.field(factory=dict, init=False) + ) _parsing_start_time: float | None = attrs.field(default=None, init=False) + _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( @@ -294,8 +375,15 @@ class DagFileProcessorManager(LoggingMixin): ) """Resolved once per process so file discovery and the deactivation scan use the same value.""" - _api_server: InProcessExecutionAPI = attrs.field(init=False, factory=InProcessExecutionAPI) - """API server to interact with Metadata DB""" + api_client: DagProcessorAPIClient | None = None + """ + Registered HTTP client for this manager's API requests, used instead of the in-process API. + + Bundle credentials and parse-time requests then go through it, on behalf of the bundle they are for. + """ + + _api_server: InProcessExecutionAPI | None = attrs.field(init=False, default=None) + """In-process API server, created only when there is no :attr:`api_client`.""" def register_exit_signals(self): """Register signals that stop child processes.""" @@ -309,9 +397,17 @@ def _get_team_names(self, bundle_names: Collection[str]) -> dict[str, str | None return {} missing = [name for name in bundle_names if name not in self._bundle_name_to_team_name] if missing: - queried = DagBundleModel.get_team_names(missing) - for name in missing: - self._bundle_name_to_team_name[name] = queried.get(name) + try: + if self.api_client is not None: + queried = {bundle.name: bundle.team_name for bundle in self.api_client.get_bundles()} + else: + queried = DagBundleModel.get_team_names(missing) + except httpx2.HTTPError: + # Team names only tag metrics, so parsing continues and the lookup is retried later. + self.log.warning("Unable to look up the teams of bundles %s; retrying later", missing) + else: + for name in missing: + self._bundle_name_to_team_name[name] = queried.get(name) return {name: self._bundle_name_to_team_name.get(name) for name in bundle_names} def _get_team_name(self, bundle_name: str) -> str | None: @@ -326,12 +422,34 @@ def _exit_gracefully(self, signum, frame): self.log.debug("Finished terminating DAG processors.") sys.exit(os.EX_OK) - def sync_bundles(self) -> None: + def _create_bundle_manager(self) -> DagBundlesManager: + if self.api_client is None: + return DagBundlesManager() + return DagBundlesManager( + bundle_names=self.bundle_names_to_parse, + bundle_context=self.api_client.use_bundle, + ) + + def sync_bundles(self, *, include_bundle_urls: bool = True) -> None: """Sync configured DAG bundles to the metadata database.""" + if self.api_client is not None: + provisioned = {bundle.name for bundle in self.api_client.get_bundles()} + configured = set( + self.bundle_names_to_parse or self._create_bundle_manager().get_all_bundle_names() + ) + if missing := configured - provisioned: + raise ValueError( + f"Bundles must be provisioned before starting the processor: {sorted(missing)}" + ) + return # When this processor only parses a subset of bundles, it does not see the full # bundle configuration and must not deactivate bundles owned by other processors. - dag_bundle_manager = DagBundlesManager() - dag_bundle_manager.sync_bundles_to_db(deactivate_missing=not self.bundle_names_to_parse) + dag_bundle_manager = self._create_bundle_manager() + dag_bundle_manager.sync_bundles_to_db( + deactivate_missing=not self.bundle_names_to_parse, include_bundle_urls=include_bundle_urls + ) + if not include_bundle_urls: + return # Best-effort legacy repair: a failure here must not crash DFP startup. # Affected Dags self-heal on the next successful parse. try: @@ -341,7 +459,7 @@ def sync_bundles(self) -> None: def get_all_bundles(self) -> list[BaseDagBundle]: """Return configured DAG bundles filtered by ``bundle_names_to_parse`` if provided.""" - return list(DagBundlesManager().get_all_dag_bundles()) + return list(self._create_bundle_manager().get_all_dag_bundles()) def run(self): """ @@ -350,15 +468,18 @@ def run(self): By processing them in separate processes, we can get parallelism and isolation from potentially harmful user code. """ - self.before_run() try: + self.before_run() return self._run_parsing_loop() finally: self.after_run() def before_run(self) -> None: """Set up state required before the parsing loop starts. Default implementation; override to customize.""" - self.prepare_server_process_context() + if self.api_client is None: + self.prepare_server_process_context() + else: + self.prepare_api_secrets_context(self.api_client) self.prepare_process_context() self.register_exit_signals() self.log.info("Processing files using up to %s processes at a time ", self._parallelism) @@ -370,13 +491,21 @@ def before_run(self) -> None: gc.freeze() def after_run(self) -> None: - """Tear down state after the parsing loop exits. Default no-op; override to customize.""" + """Tear down state after the parsing loop exits. Default implementation; override to customize.""" + if self.api_client is None: + return + from airflow.dag_processing.api_client import DagProcessorSecretsComms + from airflow.sdk.execution_time import task_runner + + if isinstance(getattr(task_runner, "SUPERVISOR_COMMS", None), DagProcessorSecretsComms): + del task_runner.SUPERVISOR_COMMS def warm_importers(self) -> None: """Build each bundle's Dag importers, so parse processes forked later share them.""" for bundle in self._dag_bundles: try: - get_importer_registry(bundle.name).warm_importers() + with self._use_bundle(bundle.name): + get_importer_registry(bundle.name).warm_importers() except Exception: # The importer fails again when the bundle is listed, which reports it per refresh. self.log.exception("Error loading Dag importers for bundle %s", bundle.name) @@ -398,6 +527,28 @@ def prepare_server_process_context(self) -> None: # Related: https://github.com/apache/airflow/pull/57459 os.environ["_AIRFLOW_PROCESS_CONTEXT"] = "server" + def prepare_api_secrets_context(self, api_client: DagProcessorAPIClient) -> None: + """ + Resolve this process's own connection and variable lookups through ``api_client``. + + Bundle code runs in this process, for example a Git bundle reading its connection. Its secret + lookups use the API. Parse processes replace + ``SUPERVISOR_COMMS`` with their own channel before running any Dag code. + """ + from airflow.dag_processing.api_client import DagProcessorSecretsComms + from airflow.sdk.execution_time import task_runner + + # SDK cache keys do not carry the bundle identity used to authorize these requests. + SecretCache.reset() + os.environ["_AIRFLOW_PROCESS_CONTEXT"] = "client" + task_runner.SUPERVISOR_COMMS = DagProcessorSecretsComms(api_client) # type: ignore[assignment] + + def _use_bundle(self, bundle_name: str) -> AbstractContextManager[object]: + """Make API requests on behalf of the bundle; a no-op with the in-process API.""" + if self.api_client is None: + return nullcontext() + return self.api_client.use_bundle(bundle_name) + def prepare_process_context(self) -> None: """Initialize transport-neutral process state (selector, stats) before the parsing loop starts.""" # Initialization is delayed until here to avoid fork issues in some @@ -428,6 +579,8 @@ def load_dag_bundles(self) -> None: def _scan_stale_dags(self): """Scan and deactivate DAGs which are no longer present in files.""" + if self.api_client is not None: + return now = time.monotonic() elapsed_time_since_refresh = now - self._last_deactivate_stale_dags_time if elapsed_time_since_refresh > self.parsing_cleanup_interval: @@ -455,6 +608,12 @@ def _cleanup_stale_bundle_versions(self): def cleanup_stale_bundle_versions(self) -> None: """Clean up stale DAG bundle version usage records.""" + if self.api_client is not None: + manager = BundleUsageTrackingManager() + for bundle in self._dag_bundles: + if bundle.supports_versioning: + manager._remove_stale_bundle_versions_for_bundle(bundle.name) + return BundleUsageTrackingManager().remove_stale_bundle_versions() @provide_session @@ -480,23 +639,7 @@ def deactivate_stale_dags( stuck_legacy_rows = 0 for dag in dags_parsed: - # Dags whose bundle has been removed from config (bundle no longer active) are stale — - # the processor has stopped parsing their files, so the time-based check below would never fire. - # - # A NULL bundle_name means the row predates bundles (carried over from Airflow 2.x) and has not - # been parsed since the upgrade — parsing is what fills bundle_name in. If the file was removed - # as part of the upgrade, no parse will ever happen, so bundle_name stays NULL forever. Such a - # row can never hit the time-based check below either, because that matches on - # (bundle_name, relative_fileloc) and there is no bundle to match against, so without this - # branch the Dag stays active in the UI indefinitely. If the file does still exist, the next - # parse fills in bundle_name and clears is_stale, so a Dag deactivated here is reactivated. if dag.bundle_name is None or dag.bundle_name in inactive_bundles: - self.log.info( - "Deactivating Dag %s. Its bundle %s is no longer active or is NULL.", - dag.dag_id, - dag.bundle_name, - ) - to_deactivate.add(dag.dag_id) continue # A Dag upgraded from Airflow 2.x can still have a NULL relative_fileloc: # the 0082 migration adds the column as nullable, and the startup repair @@ -566,7 +709,8 @@ def deactivate_stale_dags( def _run_parsing_loop(self): # initialize cache to mutualize calls to Variable.get in DAGs # needs to be done before this process is forked to create the DAG parsing processes. - SecretCache.init() + if self.api_client is None: + SecretCache.init() poll_time = 0.0 @@ -581,6 +725,8 @@ def _run_parsing_loop(self): self._queue_requested_files_for_parsing() + self._service_processor_sockets(timeout=0) + self._collect_results() self._refresh_dag_bundles(known_files=known_files) if not self._file_queue: @@ -595,11 +741,12 @@ def _run_parsing_loop(self): 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() self._cleanup_stale_bundle_versions() - self.purge_inactive_dag_warnings() # Update number of loop iteration. self._num_run += 1 @@ -656,7 +803,36 @@ def claim_priority_files(self) -> list[DagFileInfo]: Default implementation reads from the metadata DB; override to source requests from an API. """ - return self._claim_priority_files() + 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 [] + 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: """ @@ -693,7 +869,15 @@ def _claim_priority_files(self, *, session: Session = NEW_SESSION) -> list[DagFi files: list[DagFileInfo] = [] bundles = {b.name: b for b in self._dag_bundles} requests = session.scalars( - select(DagPriorityParsingRequest).where(DagPriorityParsingRequest.bundle_name.in_(bundles.keys())) + select(DagPriorityParsingRequest).where( + DagPriorityParsingRequest.bundle_name.in_(bundles.keys()), + or_( + DagPriorityParsingRequest.processor_job_id.is_(None), + DagPriorityParsingRequest.processor_job_id.in_( + select(Job.id).where(Job.end_date.is_not(None)) + ), + ), + ) ) for request in requests: bundle = bundles[request.bundle_name] @@ -711,7 +895,80 @@ def fetch_callbacks(self) -> list[CallbackRequest]: Default implementation reads from the metadata DB; override to source callbacks from an API. """ - return self._fetch_callbacks_from_db() + 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 [ + 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 @@ -739,7 +996,15 @@ def _fetch_callbacks_from_db( self.log.debug("Skipping callback fetch for uninitialized bundles: %s", unready_bundles) query: Select[tuple[DbCallbackRequest]] = with_row_locks( select(DbCallbackRequest) - .where(DbCallbackRequest.bundle_name.in_(bundle_names)) + .where( + DbCallbackRequest.bundle_name.in_(bundle_names), + or_( + DbCallbackRequest.processor_job_id.is_(None), + DbCallbackRequest.processor_job_id.in_( + select(Job.id).where(Job.end_date.is_not(None)) + ), + ), + ) .order_by(DbCallbackRequest.priority_weight.desc()) .limit(self.max_callbacks_per_loop), of=DbCallbackRequest, @@ -777,21 +1042,27 @@ def prepare_callback_bundle(self, request: CallbackRequest) -> BaseDagBundle | N self.log.error("Bundle %s is not initialized, skipping callback", request.bundle_name) return None return loaded - try: - bundle = DagBundlesManager().get_bundle( - name=request.bundle_name, - version=request.bundle_version, - version_data=request.version_data, - ) - except ValueError: - self.log.error("Bundle %s no longer configured, skipping callback", request.bundle_name) - return None - if bundle.supports_versioning: + with self._use_bundle(request.bundle_name): try: - bundle.initialize() - except Exception: + bundle = self._create_bundle_manager().get_bundle( + name=request.bundle_name, + version=request.bundle_version, + version_data=request.version_data, + ) + if bundle.supports_versioning: + bundle.initialize() + except Exception 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 initializing bundle %s version %s for callback, skipping", + "Error preparing bundle %s version %s for callback, skipping", request.bundle_name, request.bundle_version, ) @@ -802,11 +1073,19 @@ def _add_callback_to_queue(self, request: CallbackRequest) -> None: 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, @@ -814,10 +1093,41 @@ def _add_callback_to_queue(self, request: CallbackRequest) -> None: 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}" @@ -838,6 +1148,11 @@ def get_bundle_state(self, bundle_name: str, *, session: Session = NEW_SESSION) Returns ``None`` if the bundle has no database record. """ + if self.api_client is not None: + for bundle in self.api_client.get_bundles(): + if bundle.name == bundle_name: + return BundleState(bundle.last_refreshed, bundle.version, bundle.revision) + return None row = session.execute( select(DagBundleModel.last_refreshed, DagBundleModel.version).where( DagBundleModel.name == bundle_name @@ -868,14 +1183,71 @@ def update_bundle_state( values["version"] = version session.execute(update(DagBundleModel).where(DagBundleModel.name == bundle_name).values(**values)) - def purge_inactive_dag_warnings(self) -> None: - """ - Purge warnings for inactive/stale DAGs. + def _publish_bundle_inventory(self, bundle_name: str, known_files: dict[str, set[DagFileInfo]]) -> None: + if self.api_client is None: + raise ValueError("Inventory publication requires an API client") + body, found_files, source = self._pending_inventories[bundle_name] + if len(body.model_dump_json().encode()) > MAX_PUBLICATION_BYTES: + self.log.error( + "Bundle %s inventory exceeds the %s-byte publication limit; retaining the previous inventory", + bundle_name, + MAX_PUBLICATION_BYTES, + ) + self._discard_unusable_bundle_source(bundle_name) + return + try: + receipt = self.api_client.publish_inventory(bundle_name, body) + except httpx2.HTTPStatusError as error: + if error.response.status_code == 409: + del self._pending_inventories[bundle_name] + self.request_bundle_refresh(bundle_name) + elif error.response.status_code == 413: + self._discard_unusable_bundle_source(bundle_name) + elif error.response.status_code < 500: + raise + self.log.warning("Inventory publication was rejected for %s", bundle_name) + return + except httpx2.RequestError: + self.log.warning("Unable to publish inventory for %s; retaining it for retry", bundle_name) + return + del self._pending_inventories[bundle_name] + self._bundle_parse_sources[bundle_name] = attrs.evolve(source, bundle_revision=receipt.revision) + self._bundle_versions[bundle_name] = source.bundle_version + known_files[bundle_name] = found_files + 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] - Default implementation deletes records from the metadata DB; override to - source warnings from an API or skip the cleanup entirely. - """ - DagWarning.purge_inactive_dag_warnings() + def _invalidate_bundle_parse_source(self, bundle_name: str) -> None: + # Initialization or refresh may change files even when the provider raises afterwards. + self._bundle_refresh_generations[bundle_name] += 1 + self._bundle_parse_sources.pop(bundle_name, None) def _refresh_dag_bundles(self, known_files: dict[str, set[DagFileInfo]]): """Refresh DAG bundles, if required.""" @@ -884,7 +1256,12 @@ def _refresh_dag_bundles(self, known_files: dict[str, set[DagFileInfo]]): # we don't need to check if it's time to refresh every loop - that is way too often next_check = self._bundles_last_refreshed + self.bundle_refresh_check_interval now_seconds = time.monotonic() - if now_seconds < next_check and not self._force_refresh_bundles: + if ( + now_seconds < next_check + and not self._force_refresh_bundles + and not self._bundles_waiting_for_refresh + and not self._pending_inventories + ): self.log.debug( "Not time to check if DAG Bundles need refreshed yet - skipping. Next check in %.2f seconds", next_check - now_seconds, @@ -893,112 +1270,204 @@ def _refresh_dag_bundles(self, known_files: dict[str, set[DagFileInfo]]): self._bundles_last_refreshed = now_seconds + catalog: dict[str, BundleState] | None = None + if self.api_client is not None and self._dag_bundles: + try: + bundles = self.api_client.get_bundles() + catalog = { + bundle.name: BundleState(bundle.last_refreshed, bundle.version, bundle.revision) + for bundle in bundles + } + self._bundle_name_to_team_name = {bundle.name: bundle.team_name for bundle in bundles} + except httpx2.HTTPError: + self.log.warning("Unable to read the bundle catalog; retrying on a later refresh") + any_refreshed = False + active_bundles = { + file.bundle_name + for file in self._processors.keys() | self._pending_publications.keys() + if file.bundle_version is None + } for bundle in self._dag_bundles: - # TODO: AIP-66 handle errors in the case of incomplete cloning? And test this. - # What if the cloning/refreshing took too long(longer than the dag processor timeout) - if not bundle.is_initialized: - try: - bundle.initialize() + # 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) + continue + if bundle.name in self._pending_inventories: + self._publish_bundle_inventory(bundle.name, known_files) any_refreshed = True - except Exception as e: - self.log.exception("Error initializing bundle %s: %s", bundle.name, e) continue - try: - bundle_state = self.get_bundle_state(bundle.name) - except Exception: - self.log.exception("Error fetching state for bundle %s", bundle.name) - continue - if bundle_state is None: - self.log.warning("Bundle model not found for %s", bundle.name) - continue - elapsed_time_since_refresh = (now - (bundle_state.last_refreshed or utc_epoch())).total_seconds() - if bundle.supports_versioning: - # we will also check the version of the bundle to see if another DAG processor has seen - # a new version - pre_refresh_version = self._bundle_versions.get(bundle.name) - # Use `is None` (not falsy) so an empty-string version is treated as a valid cached value. - if pre_refresh_version is None: - pre_refresh_version, _ = unpack_bundle_version(bundle.get_current_version(), bundle) - current_version_matches_db = pre_refresh_version == bundle_state.version - else: - # With no versioning, it always "matches" - current_version_matches_db = True - - previously_seen = bundle.name in self._bundle_versions - if self.should_skip_refresh( - bundle=bundle, - elapsed_time_since_refresh=elapsed_time_since_refresh, - current_version_matches_db=current_version_matches_db, - previously_seen=previously_seen, - ): - self.log.debug("Not time to refresh bundle %s", bundle.name) - continue + if self.api_client is not None and catalog is None: + continue + # TODO: AIP-66 handle errors in the case of incomplete cloning? And test this. + # What if the cloning/refreshing took too long(longer than the dag processor timeout) + if not bundle.is_initialized: + self._invalidate_bundle_parse_source(bundle.name) + try: + bundle.initialize() + any_refreshed = True + except Exception as e: + self.log.exception("Error initializing bundle %s: %s", bundle.name, e) + continue + try: + bundle_state = ( + catalog.get(bundle.name) + if catalog is not None + else self.get_bundle_state(bundle.name) + ) + except httpx2.HTTPError: + self.log.warning("Unable to read metadata for bundle %s; retrying later", bundle.name) + continue + except Exception: + self.log.exception("Error fetching state for bundle %s", bundle.name) + continue + if bundle_state is None: + self.log.warning("Bundle model not found for %s", bundle.name) + continue + elapsed_time_since_refresh = ( + now - (bundle_state.last_refreshed or utc_epoch()) + ).total_seconds() + if bundle.supports_versioning: + # we will also check the version of the bundle to see if another DAG processor has seen + # a new version + pre_refresh_version = self._bundle_versions.get(bundle.name) + # Use `is None` (not falsy) so an empty-string version is treated as a valid cached value. + if pre_refresh_version is None: + pre_refresh_version, _ = unpack_bundle_version(bundle.get_current_version(), bundle) + current_version_matches_db = pre_refresh_version == bundle_state.version + else: + # With no versioning, it always "matches" + current_version_matches_db = True + + previously_seen = bundle.name in self._bundle_versions + if ( + bundle.name not in self._bundles_waiting_for_refresh + and bundle.name in self._bundle_parse_sources + and self.should_skip_refresh( + bundle=bundle, + elapsed_time_since_refresh=elapsed_time_since_refresh, + current_version_matches_db=current_version_matches_db, + previously_seen=previously_seen, + ) + ): + self.log.debug("Not time to refresh bundle %s", bundle.name) + continue - self.log.info("Refreshing bundle %s", bundle.name) + if bundle.name in active_bundles: + # Forced refreshes wait too: refreshing under an import would discard its result. + self._bundles_waiting_for_refresh.add(bundle.name) + continue - try: - bundle.refresh() - any_refreshed = True - except Exception: - self.log.exception("Error refreshing bundle %s", bundle.name) - continue + self._bundles_waiting_for_refresh.discard(bundle.name) + self.log.info("Refreshing bundle %s", bundle.name) + + self._invalidate_bundle_parse_source(bundle.name) + try: + bundle.refresh() + any_refreshed = True + except Exception: + self.log.exception("Error refreshing bundle %s", bundle.name) + continue - self._force_refresh_bundles.discard(bundle.name) + self._force_refresh_bundles.discard(bundle.name) - if bundle.supports_versioning: - # We can short-circuit the rest of this if (1) bundle was seen before by - # this dag processor and (2) the version of the bundle did not change - # after refreshing it - version_after_refresh, version_data_after_refresh = unpack_bundle_version( - bundle.get_current_version(), bundle + if bundle.supports_versioning: + version_after_refresh, version_data_after_refresh = unpack_bundle_version( + bundle.get_current_version(), bundle + ) + else: + version_after_refresh = None + version_data_after_refresh = None + + self._bundle_parse_sources[bundle.name] = DagParseSource( + bundle_version=version_after_refresh, + version_data=deepcopy(version_data_after_refresh), + refresh_generation=self._bundle_refresh_generations[bundle.name], ) - if previously_seen and pre_refresh_version == version_after_refresh: - self.log.debug( - "Bundle %s version not changed after refresh: %s", - bundle.name, - version_after_refresh, + + if bundle.supports_versioning: + # We can short-circuit the rest of this if (1) bundle was seen before by + # this dag processor and (2) the version of the bundle did not change + # after refreshing it + if ( + self.api_client is None + and previously_seen + and pre_refresh_version == version_after_refresh + ): + self.log.debug( + "Bundle %s version not changed after refresh: %s", + bundle.name, + version_after_refresh, + ) + try: + self.update_bundle_state(bundle.name, last_refreshed=now, version=None) + except Exception: + self.log.exception("Error persisting state for bundle %s", bundle.name) + continue + + self.log.info( + "Version changed for %s, new version: %s", bundle.name, version_after_refresh ) - try: - self.update_bundle_state(bundle.name, last_refreshed=now, version=None) - except Exception: - self.log.exception("Error persisting state for bundle %s", bundle.name) + try: + found_files = self._find_files_in_bundle(bundle) + except Exception: + # Keep the bundle's known files and Dags, and leave its version unadvanced so the + # next refresh lists it again. + self.log.exception("Error listing Dag definitions in bundle %s", bundle.name) + if self.api_client is not None: + self._bundle_parse_sources.pop(bundle.name, None) continue - self.log.info("Version changed for %s, new version: %s", bundle.name, version_after_refresh) - else: - version_after_refresh = None - version_data_after_refresh = None - - try: - found_files = self._find_files_in_bundle(bundle) - except Exception: - # Keep the bundle's known files and Dags, and leave its version unadvanced so the - # next refresh lists it again. - self.log.exception("Error listing Dag definitions in bundle %s", bundle.name) - continue + if self.api_client is not None: + self._dispatch_sequence += 1 + try: + inventory = DagBundleInventoryBody( + attempt_id=uuid7(), + dispatch_sequence=self._dispatch_sequence, + expected_revision=bundle_state.revision, + version=version_after_refresh, + files=sorted(self._get_observed_filelocs(found_files)), + failed_sources=sorted( + {loc for file in found_files for loc in file.failed_sources} + ), + ) + except ValidationError: + # Handled like a failed listing, so the API never sees a partial inventory. + self.log.exception("Bundle %s has a listing the API cannot accept", bundle.name) + self._discard_unusable_bundle_source(bundle.name) + continue + self._pending_inventories[bundle.name] = ( + inventory, + found_files, + self._bundle_parse_sources.pop(bundle.name), + ) + self._publish_bundle_inventory(bundle.name, known_files) + continue - # Persistence failure must not skip file scanning (bundle is already refreshed locally). - # _bundle_versions is only advanced on success to stay consistent with the DB. - try: - self.update_bundle_state(bundle.name, last_refreshed=now, version=version_after_refresh) - except Exception: - self.log.exception("Error persisting state for bundle %s", bundle.name) - else: - self._bundle_versions[bundle.name] = version_after_refresh - self._bundle_version_data[bundle.name] = version_data_after_refresh + # Persistence failure must not skip file scanning (bundle is already refreshed locally). + # _bundle_versions is only advanced on success to stay consistent with the DB. + try: + self.update_bundle_state(bundle.name, last_refreshed=now, version=version_after_refresh) + except Exception: + self.log.exception("Error persisting state for bundle %s", bundle.name) + else: + self._bundle_versions[bundle.name] = version_after_refresh - known_files[bundle.name] = found_files + known_files[bundle.name] = found_files - self.deactivate_deleted_dags(bundle_name=bundle.name, present=found_files) - self.clear_orphaned_import_errors( - bundle_name=bundle.name, - observed_filelocs=self._get_observed_filelocs(found_files), - ) + self.deactivate_deleted_dags(bundle_name=bundle.name, present=found_files) + self.clear_orphaned_import_errors( + bundle_name=bundle.name, + observed_filelocs=self._get_observed_filelocs(found_files), + failed_sources={loc for file in found_files for loc in file.failed_sources}, + ) if any_refreshed: # Bundle-to-team assignments can only change on bundle refresh, so clear the cache. - self._bundle_name_to_team_name = {} + if self.api_client is None: + self._bundle_name_to_team_name = {} self.handle_removed_files(known_files=known_files) self._resort_file_queue() self._add_new_files_to_queue(known_files=known_files) @@ -1012,6 +1481,7 @@ def _find_files_in_bundle(self, bundle: BaseDagBundle) -> set[DagFileInfo]: self.log.info("Searching for Dag definitions in %s at %s", bundle.name, bundle.path) registry = get_importer_registry(bundle.name) definition_locs: defaultdict[Path, set[str]] = defaultdict(set) + failed_sources: defaultdict[Path, set[str]] = defaultdict(set) for _, item in registry.list_dag_definitions(bundle, safe_mode=self.dag_discovery_safe_mode): if isinstance(item, DagImportError): # Importers report a source either absolutely or relative to the bundle. @@ -1030,6 +1500,8 @@ def _find_files_in_bundle(self, bundle: BaseDagBundle) -> set[DagFileInfo]: ) continue definition_locs[path.relative_to(bundle.path)].add(rel_fileloc) + if isinstance(item, DagImportError): + failed_sources[path.relative_to(bundle.path)].add(rel_fileloc) self.log.info( "Found %s files for bundle %s (dag_discovery_safe_mode=%s)", len(definition_locs), @@ -1042,6 +1514,7 @@ def _find_files_in_bundle(self, bundle: BaseDagBundle) -> set[DagFileInfo]: bundle_name=bundle.name, bundle_path=bundle.path, definition_locs=frozenset(locs), + failed_sources=frozenset(failed_sources[rel_path]), ) for rel_path, locs in definition_locs.items() } @@ -1054,9 +1527,21 @@ def _get_observed_filelocs(present: set[DagFileInfo]) -> set[str]: def deactivate_deleted_dags(self, bundle_name: str, present: set[DagFileInfo]) -> None: """Deactivate DAGs that come from files that are no longer present in bundle.""" observed_filelocs = self._get_observed_filelocs(present) + failed_sources = {loc for file in present for loc in file.failed_sources} with create_session() as session: try: with with_db_lock_timeout(session=session, lock_timeout=30): + if failed_sources: + observed_filelocs.update( + loc + for loc in session.scalars( + select(DagModel.relative_fileloc).where( + DagModel.bundle_name == bundle_name, + DagModel.relative_fileloc.is_not(None), + ) + ) + if loc is not None and is_source_discovery_failed(loc, failed_sources) + ) any_deactivated = DagModel.deactivate_deleted_dags( bundle_name=bundle_name, rel_filelocs=observed_filelocs, @@ -1087,7 +1572,12 @@ def print_stats(self, known_files: dict[str, set[DagFileInfo]]): @provide_session def clear_orphaned_import_errors( - self, bundle_name: str, observed_filelocs: set[str], *, session: Session = NEW_SESSION + self, + bundle_name: str, + observed_filelocs: set[str], + *, + failed_sources: Collection[str] = (), + session: Session = NEW_SESSION, ): """ Clear import errors for files that no longer exist. @@ -1102,7 +1592,9 @@ def clear_orphaned_import_errors( .options(load_only(ParseImportError.filename)) ) for error in errors: - if error.filename not in observed_filelocs: + if error.filename not in observed_filelocs and not is_source_discovery_failed( + error.filename, failed_sources + ): session.delete(error) except Exception: self.log.exception("Error removing old import errors") @@ -1248,7 +1740,14 @@ def handle_removed_files(self, known_files: dict[str, set[DagFileInfo]]): 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) + + 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]): @@ -1257,6 +1756,9 @@ def remove_orphaned_file_stats(self, present: set[DagFileInfo]): stats_to_remove = {file for file in self._file_stats if file.presence_key not in present_keys} for file in stats_to_remove: del self._file_stats[file] + for file in list(self._pending_publications): + if file.presence_key not in present_keys: + del self._pending_publications[file] def terminate_orphan_processes(self, present: set[DagFileInfo]): """Stop processors that are working on deleted files.""" @@ -1265,6 +1767,8 @@ def terminate_orphan_processes(self, present: set[DagFileInfo]): 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: @@ -1316,6 +1820,19 @@ def handle_parsing_result( if is_callback_only: self.log.debug("Detected callback-only processing for %s", file) + if ( + proc.parsing_result is not None + and proc.parse_source.refresh_generation + != self._bundle_refresh_generations.get(file.bundle_name, 0) + ): + self.log.info("Discarding parse of %s after its bundle refreshed; requeuing", file.rel_path) + stats.incr( + "dag_processing.results_discarded_on_refresh", + tags={"bundle_name": normalize_name_for_stats(file.bundle_name)}, + ) + self._add_files_to_queue([file], mode="front") + return + run_duration = time.monotonic() - proc.start_time finish_time = timezone.utcnow() team_name = self._get_team_name(file.bundle_name) @@ -1332,10 +1849,23 @@ def handle_parsing_result( if proc.parsing_result is not None: try: + if self.api_client is not None: + pending = PendingDagPublication( + body=self._build_parse_result(file, proc, run_duration), + stat=next_stat, + refresh_generation=proc.parse_source.refresh_generation, + ) + if len(pending.body.model_dump_json().encode()) > MAX_PUBLICATION_BYTES: + pending = self._build_rejection_report( + pending, + f"it exceeds the {MAX_PUBLICATION_BYTES // (1024 * 1024)} MiB publication limit", + ) + self._pending_publications[file] = pending + return self.persist_parsing_result( bundle_name=file.bundle_name, - bundle_version=self._bundle_versions[file.bundle_name], - version_data=self._bundle_version_data.get(file.bundle_name), + bundle_version=proc.parse_source.bundle_version, + version_data=proc.parse_source.version_data, parsing_result=proc.parsing_result, run_duration=run_duration, relative_fileloc=str(file.rel_path), @@ -1349,19 +1879,163 @@ def handle_parsing_result( str(file.rel_path), file.bundle_name, ) - current_stat = self._file_stats[file] - self._file_stats[file] = DagFileStat( - num_dags=current_stat.num_dags, - import_errors=current_stat.import_errors, - last_finish_time=finish_time, - last_duration=run_duration, - run_count=current_stat.run_count + 1, - last_num_of_db_queries=current_stat.last_num_of_db_queries, - ) + self._record_failed_publication(file, next_stat) + 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: + self._file_stats[file] = attrs.evolve( + self._file_stats[file], + last_attempt_time=timezone.utcnow(), + last_duration=stat.last_duration, + run_count=stat.run_count, + ) + + def _publish_pending_results(self, *, deadline: float | None = None) -> None: + """Publish queued results until the deadline, stopping early while the API is unavailable.""" + if self.api_client is None: + return + from airflow.dag_processing.api_client import get_error_reason + + if deadline is None: + deadline = time.monotonic() + _API_SEND_BUDGET_SECONDS + files = deque(self._pending_publications) + while files and time.monotonic() < deadline: + file = files.popleft() + pending = self._pending_publications[file] + if pending.refresh_generation != self._bundle_refresh_generations.get(file.bundle_name, 0): + del self._pending_publications[file] + self._add_files_to_queue([file], mode="front") + continue + if time.monotonic() < pending.next_attempt_time: + continue + pending.attempts += 1 + try: + self.api_client.publish_parse_result(pending.body) + except Exception as error: + status_error = error if isinstance(error, httpx2.HTTPStatusError) else None + reason = get_error_reason(status_error) if status_error is not None else None + if reason == "source_changed": + # Another processor's inventory moved the revision; refreshing fetches it. + self.log.info( + "Bundle %s changed before %s was published; reparsing", file.bundle_name, file + ) + del self._pending_publications[file] + self.request_bundle_refresh(file.bundle_name) + self._add_files_to_queue([file], mode="front") + return + retryable = isinstance(error, httpx2.RequestError) or ( + status_error is not None and status_error.response.status_code >= 500 + ) + if retryable and pending.attempts < conf.getint("workers", "execution_api_retries"): + pending.next_attempt_time = time.monotonic() + min( + conf.getfloat("workers", "execution_api_retry_wait_max"), + max( + conf.getfloat("workers", "execution_api_retry_wait_min"), + 2 ** (pending.attempts - 1), + ), + ) + self.log.warning("Unable to publish %s; retrying on a later loop", file) + del self._pending_publications[file] + self._pending_publications[file] = pending + return + if ( + status_error is not None + and not pending.is_rejection_report + and _is_rejected_content(status_error.response.status_code, reason) + ): + self.log.warning("The API server rejected the parsing result for %s", file) + status_code = status_error.response.status_code + detail = f"HTTP {status_code}: {reason}" if reason else f"HTTP {status_code}" + if reason == "invalid_dag": + message = status_error.response.json()["detail"].get("message") + if isinstance(message, str): + detail += f": {message}" + self._pending_publications[file] = self._build_rejection_report( + pending, f"the API server rejected it ({detail})" + ) + files.append(file) + continue + self.log.exception("Failed to publish parsing result for %s", file) + self._record_failed_publication(file, pending.stat) + 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] + "..." + import_errors = { + pending.body.relative_fileloc: f"The parsing result was not published: {reason}.", + } + return PendingDagPublication( + body=DagParseResultBody( + attempt_id=uuid7(), + dispatch_sequence=self._dispatch_sequence, + bundle_name=pending.body.bundle_name, + relative_fileloc=pending.body.relative_fileloc, + bundle_version=pending.body.bundle_version, + bundle_revision=pending.body.bundle_revision, + parse_duration=pending.body.parse_duration, + serialized_dags=[], + source_codes={}, + import_errors=import_errors, + ), + stat=attrs.evolve(pending.stat, num_dags=0, import_errors=len(import_errors)), + refresh_generation=pending.refresh_generation, + is_rejection_report=True, + ) + + def _build_parse_result( + self, file: DagFileInfo, proc: BaseDagFileProcessorProcess, run_duration: float + ) -> DagParseResultBody: + result = proc.parsing_result + if result is None or self.api_client is None: + raise ValueError("API publication requires a client and completed parse result") + published_dag_ids = {dag.dag_id for dag in result.serialized_dags} + warnings = [ + ParseWarning( + **{ + key: warning[key] if isinstance(warning, dict) else getattr(warning, key) + for key in ("dag_id", "warning_type", "message") + } + ) + for warning in result.warnings or [] + ] + # The API accepts warnings only for published Dags; a Dag that failed to serialize has an import error. + warnings = [warning for warning in warnings if warning.dag_id in published_dag_ids] + source_codes = {} + for dag in result.serialized_dags: + source = result.dag_source_codes.get(dag.dag_id) + source_codes[dag.dag_id] = ParseSourceCode( + source_code=source.source_code if source else None, + language=source.language if source else "python", + ) + return DagParseResultBody( + attempt_id=proc.id, + dispatch_sequence=proc.dispatch_sequence, + bundle_name=file.bundle_name, + relative_fileloc=str(file.rel_path), + bundle_version=proc.parse_source.bundle_version, + bundle_revision=proc.parse_source.bundle_revision, + version_data=proc.parse_source.version_data, + parse_duration=run_duration, + serialized_dags=[dag.data for dag in result.serialized_dags], + import_errors=result.import_errors or {}, + parsed_definitions=result.parsed_definitions, + warnings=warnings, + source_codes=source_codes, + ) + def persist_parsing_result( self, *, @@ -1416,6 +2090,7 @@ def _collect_results(self): for file in finished: processor = self._processors.pop(file) + self._finish_callback_attempt(processor) processor.close() def _get_log_dir(self) -> str: @@ -1467,23 +2142,34 @@ def _get_logger_for_dag_file(self, dag_file: DagFileInfo): @functools.cached_property def client(self) -> Client: + if self.api_client is not None: + return self.api_client + from airflow.sdk.api.client import Client + self._api_server = InProcessExecutionAPI() client = Client(base_url=None, token="", dry_run=True, transport=self._api_server.transport) # Mypy is wrong -- the setter accepts a string on the property setter! `URLType = URL | str` client.base_url = "http://in-process.invalid./" 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, @@ -1496,7 +2182,7 @@ def _create_process(self, dag_file: DagFileInfo) -> BaseDagFileProcessorProcess: ) 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, @@ -1509,17 +2195,69 @@ def _create_process(self, dag_file: DagFileInfo) -> BaseDagFileProcessorProcess: client=self.client, ) + def _get_dispatch_source(self, file: DagFileInfo) -> DagParseSource | None: + if file in self._callback_to_execute: + # Callback runs publish no parse result, so they need no refreshed source; they run on the + # bundle's current tree rather than waiting for a refresh that may keep failing. + return DagParseSource(bundle_version=file.bundle_version) + if ( + file.bundle_name in self._bundles_waiting_for_refresh + or file.bundle_name in self._force_refresh_bundles + ): + 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 self._has_requested_work(file): + requested.append(file) + else: + regular.append(file) + if requested: + self._add_files_to_queue(list(reversed(requested)), mode="front") + if regular: + self._add_files_to_queue(regular, mode="back") + def _start_new_processes(self): """Start more processors if we have enough slots and files to process.""" + self._readmit_files_with_source() bundle_to_team = self._get_team_names({file.bundle_name for file in self._file_queue}) - while self._parallelism > len(self._processors) and self._file_queue: + # 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: + 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) + if source is None: + self._files_awaiting_source[file] = None + continue + source = deepcopy(source) processor = self._create_process(file) + 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( @@ -1534,6 +2272,8 @@ def _start_new_processes(self): 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]]): """ @@ -1545,6 +2285,8 @@ def _add_new_files_to_queue(self, known_files: dict[str, set[DagFileInfo]]): tracked_presence_keys = {file.presence_key for file in self._file_queue} tracked_presence_keys.update(file.presence_key for file in self._file_stats) tracked_presence_keys.update(file.presence_key for file in self._processors) + tracked_presence_keys.update(file.presence_key for file in self._pending_publications) + tracked_presence_keys.update(file.presence_key for file in self._files_awaiting_source) for files in known_files.values(): for file in files: if file.presence_key not in tracked_presence_keys: @@ -1583,7 +2325,7 @@ def _sort_by_mtime(self, files: Iterable[DagFileInfo]): modified_datetime = datetime.fromtimestamp(modified_timestamp, tz=timezone.utc) files_with_mtime[file] = modified_timestamp stat = file_stats_by_presence_key.get(file.presence_key) - last_time = stat.last_finish_time if stat else None + last_time = (stat.last_attempt_time or stat.last_finish_time) if stat else None if not last_time: continue if modified_datetime > last_time: @@ -1610,7 +2352,7 @@ def processed_recently(self, now, file): ), None, ) - last_time = stat.last_finish_time if stat else None + last_time = (stat.last_attempt_time or stat.last_finish_time) if stat else None if not last_time: return False elapsed_ss = (now - last_time).total_seconds() @@ -1636,6 +2378,8 @@ def prepare_file_queue(self, known_files: dict[str, set[DagFileInfo]]): # If the file path is already being processed, or if a file was # processed recently, wait until the next batch in_progress_keys = {file.presence_key for file in self._processors} + in_progress_keys.update(file.presence_key for file in self._pending_publications) + in_progress_keys.update(file.presence_key for file in self._files_awaiting_source) file_stats_by_presence_key = {file.presence_key: stat for file, stat in self._file_stats.items()} now = timezone.utcnow() @@ -1647,7 +2391,7 @@ def prepare_file_queue(self, known_files: dict[str, set[DagFileInfo]]): for file in bundle_files: files.append(file) stat = file_stats_by_presence_key.get(file.presence_key) - last_time = stat.last_finish_time if stat else None + last_time = (stat.last_attempt_time or stat.last_finish_time) if stat else None if last_time and (now - last_time).total_seconds() < self._file_process_interval: recently_processed.add(file) @@ -1745,10 +2489,13 @@ def _kill_timed_out_processors(self): 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( @@ -1788,6 +2535,22 @@ def max_runs_reached(self): return False if self._num_run < self.max_runs: return False + if self._pending_publications: + 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] + ) or any(not proc.had_callbacks for proc in self._processors.values()): + return False return all(stat.run_count >= self.max_runs for stat in self._file_stats.values()) def terminate(self): @@ -1810,10 +2573,15 @@ def terminate(self): 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 e6f3f422a0abf..95f6cb2d0824f 100644 --- a/airflow-core/src/airflow/dag_processing/processor.py +++ b/airflow-core/src/airflow/dag_processing/processor.py @@ -21,10 +21,12 @@ import importlib import logging import os +import sys import traceback from collections.abc import Callable, Sequence from pathlib import Path from typing import TYPE_CHECKING, Annotated, BinaryIO, ClassVar, Literal, cast +from uuid import UUID import attrs from pydantic import BaseModel, Field, TypeAdapter @@ -38,6 +40,7 @@ ) 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 @@ -87,6 +90,7 @@ from structlog.typing import FilteringBoundLogger + from airflow.dag_processing.api_client import DagParseContext from airflow.sdk.api.client import Client from airflow.sdk.bases.operator import BaseOperator from airflow.sdk.definitions.context import Context @@ -423,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 ), @@ -570,6 +578,16 @@ def _execute_email_callbacks(dagbag: DagBag, request: EmailRequest, log: Filteri ) +@attrs.frozen +class DagParseSource: + """Bundle metadata captured for one import of the local source tree.""" + + bundle_version: str | None = None + version_data: dict | None = None + refresh_generation: int = 0 + bundle_revision: UUID | None = None + + @attrs.define(kw_only=True) class BaseDagFileProcessorProcess(WatchedSubprocess, LoggingMixin): """ @@ -586,6 +604,8 @@ class BaseDagFileProcessorProcess(WatchedSubprocess, LoggingMixin): ) """The file's parse log, which the process closes. Without one, the output is discarded.""" parsing_result: DagFileParsingResult | None = None + parse_source: DagParseSource = attrs.field(factory=DagParseSource) + dispatch_sequence: int = 0 decoder: ClassVar[TypeAdapter[ToManager]] = TypeAdapter[ToManager](ToManager) had_callbacks: bool = False # Track if this process was started with callbacks to prevent stale DAG detection false positives @@ -598,6 +618,7 @@ class BaseDagFileProcessorProcess(WatchedSubprocess, LoggingMixin): bundle_name: str dag_file_rel_path: str + _api_parse_context: DagParseContext | None = attrs.field(init=False, default=None) def _get_target_loggers(self) -> tuple[FilteringBoundLogger, ...]: base = super()._get_target_loggers() @@ -628,10 +649,15 @@ def _handle_parsing_result( self.parsing_result = msg return None, {} + def _handle_previous_successful_dag_run( + self, msg: GetPrevSuccessfulDagRun, log: FilteringBoundLogger, req_id: int + ) -> RequestResult: + # The run is looked up by task instance, and a parse process is not one, so the answer is always empty. + return PrevSuccessfulDagRunResult(), {"exclude_unset": True} + _client_request_types: ClassVar[tuple[type[BaseModel], ...]] = ( DeleteVariable, GetConnection, - GetPrevSuccessfulDagRun, GetPreviousDagRun, GetPreviousTI, GetTICount, @@ -652,7 +678,12 @@ def _handle_parsing_result( "dict[type[BaseModel], RequestHandler[BaseDagFileProcessorProcess]]", WatchedSubprocess._get_shared_request_handlers(*_client_request_types, MaskSecret), ), - **dict([register_request_method(DagFileParsingResult, _handle_parsing_result)]), + **dict( + [ + register_request_method(DagFileParsingResult, _handle_parsing_result), + register_request_method(GetPrevSuccessfulDagRun, _handle_previous_successful_dag_run), + ] + ), } def _handle_request(self, msg, log: FilteringBoundLogger, req_id: int) -> None: @@ -665,7 +696,25 @@ def _handle_request(self, msg, log: FilteringBoundLogger, req_id: int) -> None: ), ) return - super()._handle_request(msg, log, req_id) + # The HTTP client pulls in the Execution API versions. A DagProcessorAPIClient can only exist once + # its module is loaded, so the database-backed mode never has to import it. + api_client_module = sys.modules.get("airflow.dag_processing.api_client") + if api_client_module is None or not isinstance(self.client, api_client_module.DagProcessorAPIClient): + super()._handle_request(msg, log, req_id) + return + from airflow.api_fastapi.execution_api.datamodels.job import DagParseTokenBody + from airflow.dag_processing.api_client import DagParseContext + + if self._api_parse_context is None: + self._api_parse_context = DagParseContext( + request=DagParseTokenBody( + attempt_id=self.id, + bundle_name=self.bundle_name, + relative_fileloc=self.dag_file_rel_path, + ) + ) + with self.client.use_parse(self._api_parse_context): + super()._handle_request(msg, log, req_id) def _reject_request(self, msg, log: FilteringBoundLogger, req_id: int) -> None: log.error("Unhandled request", msg=msg) diff --git a/airflow-core/src/airflow/jobs/dag_processor_job_runner.py b/airflow-core/src/airflow/jobs/dag_processor_job_runner.py index 3b25d32d4a52c..3af40ae3f91be 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 @@ def __init__( ): 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/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index 40e1fa976ffeb..c24d6df58d2cf 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -71,6 +71,7 @@ ) from airflow.configuration import conf from airflow.dag_processing.bundles.base import BundleUsageTrackingManager +from airflow.dag_processing.maintenance import cleanup_processor_metadata from airflow.exceptions import DagNotFound, TaskNotFound from airflow.executors import workloads from airflow.executors.executor_loader import ExecutorLoader @@ -2011,6 +2012,9 @@ def _run_scheduler_loop(self) -> None: timers.call_regular_interval(60.0, self._update_dag_run_state_for_paused_dags) timers.call_regular_interval(5.0, self._finalize_draining_dags) + timers.call_regular_interval( + conf.getfloat("scheduler", "parsing_cleanup_interval"), cleanup_processor_metadata, non_fatal=True + ) timers.call_regular_interval( conf.getfloat("scheduler", "task_queued_timeout_check_interval"), diff --git a/airflow-core/src/airflow/migrations/versions/0148_3_5_0_add_dag_parse_checkpoint.py b/airflow-core/src/airflow/migrations/versions/0148_3_5_0_add_dag_parse_checkpoint.py new file mode 100644 index 0000000000000..39df787eb794a --- /dev/null +++ b/airflow-core/src/airflow/migrations/versions/0148_3_5_0_add_dag_parse_checkpoint.py @@ -0,0 +1,58 @@ +# 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. + +""" +Add checkpoints for Dag result publication. + +Revision ID: e941ab8243b7 +Revises: c71065a1b14f +Create Date: 2026-10-05 13:00:00.000000 +""" + +from __future__ import annotations + +import sqlalchemy as sa +from alembic import op + +from airflow.utils.sqlalchemy import UtcDateTime + +revision = "e941ab8243b7" +down_revision = "c71065a1b14f" +branch_labels = None +depends_on = None +airflow_version = "3.5.0" + + +def upgrade(): + """Store the latest accepted publication for each Job and source.""" + op.create_table( + "dag_parse_checkpoint", + sa.Column("job_id", sa.Integer(), nullable=False), + sa.Column("source_key", sa.String(64), nullable=False), + sa.Column("attempt_id", sa.Uuid(), nullable=False), + sa.Column("dispatch_sequence", sa.BigInteger(), nullable=False), + sa.Column("payload_hash", sa.String(64), nullable=False), + sa.Column("accepted_at", UtcDateTime(), nullable=False), + sa.Column("bundle_revision", sa.Uuid(), nullable=True), + sa.ForeignKeyConstraint(["job_id"], ["job.id"], ondelete="CASCADE"), + sa.PrimaryKeyConstraint("job_id", "source_key"), + ) + + +def downgrade(): + """Remove publication checkpoints.""" + op.drop_table("dag_parse_checkpoint") diff --git a/airflow-core/src/airflow/migrations/versions/0149_3_5_0_add_processor_reconciliation.py b/airflow-core/src/airflow/migrations/versions/0149_3_5_0_add_processor_reconciliation.py new file mode 100644 index 0000000000000..58ab377cd5e2c --- /dev/null +++ b/airflow-core/src/airflow/migrations/versions/0149_3_5_0_add_processor_reconciliation.py @@ -0,0 +1,77 @@ +# 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. + +""" +Add source revisions for Dag processor reconciliation. + +Revision ID: 687eb253de24 +Revises: e941ab8243b7 +Create Date: 2026-10-05 18:00:00.000000 +""" + +from __future__ import annotations + +import sqlalchemy as sa +from alembic import op + +from airflow.migrations.utils import disable_sqlite_fkeys + +revision = "687eb253de24" +down_revision = "e941ab8243b7" +branch_labels = None +depends_on = None +airflow_version = "3.5.0" + + +def upgrade(): + """Fence publications against accepted discovery snapshots.""" + # SQLite rebuilds these tables, and dropping the old copy would cascade to deadline and team rows. + with disable_sqlite_fkeys(op): + op.create_index("idx_dag_bundle_stale", "dag", ["bundle_name", "is_stale"]) + with op.batch_alter_table("dag_bundle") as batch: + batch.add_column(sa.Column("parse_revision", sa.Uuid(), nullable=True)) + batch.add_column(sa.Column("inventory_hash", sa.String(64), nullable=True)) + for table in ("callback", "dag_priority_parsing_request"): + with op.batch_alter_table(table) as batch: + batch.add_column(sa.Column("processor_job_id", sa.Integer(), nullable=True)) + batch.add_column(sa.Column("processor_claim_id", sa.Uuid(), nullable=True)) + batch.create_index(f"idx_{table}_processor_job_id", ["processor_job_id"]) + batch.create_foreign_key( + f"{table}_processor_job_id_fkey", "job", ["processor_job_id"], ["id"], ondelete="SET NULL" + ) + + +def downgrade(): + """Remove discovery revision metadata.""" + with disable_sqlite_fkeys(op): + if op.get_bind().dialect.name == "mysql": + # MySQL lets idx_dag_bundle_stale back the bundle_name foreign key and refuses to drop it + # while the key exists; recreating the key restores the index MySQL keeps for it. + op.drop_constraint("dag_bundle_name_fkey", "dag", type_="foreignkey") + op.drop_index("idx_dag_bundle_stale", table_name="dag") + op.create_foreign_key("dag_bundle_name_fkey", "dag", "dag_bundle", ["bundle_name"], ["name"]) + else: + op.drop_index("idx_dag_bundle_stale", table_name="dag") + for table in ("dag_priority_parsing_request", "callback"): + with op.batch_alter_table(table) as batch: + batch.drop_constraint(f"{table}_processor_job_id_fkey", type_="foreignkey") + batch.drop_index(f"idx_{table}_processor_job_id") + batch.drop_column("processor_claim_id") + batch.drop_column("processor_job_id") + with op.batch_alter_table("dag_bundle") as batch: + batch.drop_column("inventory_hash") + batch.drop_column("parse_revision") diff --git a/airflow-core/src/airflow/models/__init__.py b/airflow-core/src/airflow/models/__init__.py index ac16af219134c..1a17024b799ca 100644 --- a/airflow-core/src/airflow/models/__init__.py +++ b/airflow-core/src/airflow/models/__init__.py @@ -64,6 +64,7 @@ def import_all_models(): import airflow.models.backfill import airflow.models.connection_test import airflow.models.dag_favorite + import airflow.models.dag_parse_checkpoint import airflow.models.dag_version import airflow.models.dagbag import airflow.models.dagbundle diff --git a/airflow-core/src/airflow/models/callback.py b/airflow-core/src/airflow/models/callback.py index 6295b7dfa7b71..78309784de163 100644 --- a/airflow-core/src/airflow/models/callback.py +++ b/airflow-core/src/airflow/models/callback.py @@ -133,6 +133,10 @@ class Callback(Base, BaseWorkload): # State of the Callback of type: CallbackState. Can be null for instances of DagProcessorCallback. state: Mapped[str | None] = mapped_column(String(10)) + processor_job_id: Mapped[int | None] = mapped_column( + ForeignKey("job.id", ondelete="SET NULL"), nullable=True, index=True + ) + processor_claim_id: Mapped[UUID | None] = mapped_column(Uuid(), nullable=True) # Return value of the callback if successful, otherwise exception details output: Mapped[str | None] = mapped_column(Text, nullable=True) diff --git a/airflow-core/src/airflow/models/dag.py b/airflow-core/src/airflow/models/dag.py index a073155438d20..e2e26ffadd94e 100644 --- a/airflow-core/src/airflow/models/dag.py +++ b/airflow-core/src/airflow/models/dag.py @@ -414,6 +414,7 @@ class DagModel(Base): __table_args__ = ( Index("idx_next_dagrun_create_after", next_dagrun_create_after, unique=False), Index("idx_dag_is_draining", is_draining, unique=False), + Index("idx_dag_bundle_stale", bundle_name, is_stale), CheckConstraint("NOT (is_paused AND is_draining)", name="dag_pause_state_valid"), ) diff --git a/airflow-core/src/airflow/models/dag_parse_checkpoint.py b/airflow-core/src/airflow/models/dag_parse_checkpoint.py new file mode 100644 index 0000000000000..6ede3b1f40924 --- /dev/null +++ b/airflow-core/src/airflow/models/dag_parse_checkpoint.py @@ -0,0 +1,41 @@ +# 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 datetime import datetime +from uuid import UUID + +from sqlalchemy import BigInteger, ForeignKey, String, Uuid +from sqlalchemy.orm import Mapped, mapped_column + +from airflow.models.base import Base +from airflow.utils.sqlalchemy import UtcDateTime + + +class DagParseCheckpoint(Base): + """Latest publication per Job and source; removed when its Job is purged.""" + + __tablename__ = "dag_parse_checkpoint" + + job_id: Mapped[int] = mapped_column(ForeignKey("job.id", ondelete="CASCADE"), primary_key=True) + source_key: Mapped[str] = mapped_column(String(64), primary_key=True) + attempt_id: Mapped[UUID] = mapped_column(Uuid, nullable=False) + dispatch_sequence: Mapped[int] = mapped_column(BigInteger, nullable=False) + payload_hash: Mapped[str] = mapped_column(String(64), nullable=False) + accepted_at: Mapped[datetime] = mapped_column(UtcDateTime, nullable=False) + bundle_revision: Mapped[UUID | None] = mapped_column(Uuid, nullable=True) diff --git a/airflow-core/src/airflow/models/dagbag.py b/airflow-core/src/airflow/models/dagbag.py index 1665699bec197..b64ed3097d4e1 100644 --- a/airflow-core/src/airflow/models/dagbag.py +++ b/airflow-core/src/airflow/models/dagbag.py @@ -27,7 +27,7 @@ from uuid import UUID from cachetools import LRUCache, TTLCache -from sqlalchemy import String, select +from sqlalchemy import ForeignKey, String, Uuid, select from sqlalchemy.orm import Mapped, joinedload, mapped_column from airflow._shared.observability.metrics import stats @@ -303,9 +303,7 @@ class DagPriorityParsingRequest(Base): # Adding a unique constraint to fileloc results in the creation of an index and we have a limitation # on the size of the string we can use in the index for MySQL DB. We also have to keep the fileloc # size consistent with other tables. This is a workaround to enforce the unique constraint. - id: Mapped[str] = mapped_column( - String(32), primary_key=True, default=generate_md5_hash, onupdate=generate_md5_hash - ) + id: Mapped[str] = mapped_column(String(32), primary_key=True, default=generate_md5_hash) bundle_name: Mapped[str] = mapped_column(StringID(), nullable=False) # The location of the file containing the DAG object @@ -313,6 +311,10 @@ class DagPriorityParsingRequest(Base): # packaged DAG, it will point to the subpath of the DAG within the # associated zip. relative_fileloc: Mapped[str] = mapped_column(String(2000), nullable=False) + processor_job_id: Mapped[int | None] = mapped_column( + ForeignKey("job.id", ondelete="SET NULL"), nullable=True, index=True + ) + processor_claim_id: Mapped[UUID | None] = mapped_column(Uuid(), nullable=True) def __init__(self, bundle_name: str, relative_fileloc: str) -> None: super().__init__() diff --git a/airflow-core/src/airflow/models/dagbundle.py b/airflow-core/src/airflow/models/dagbundle.py index 437c6e6b775d8..47b56d4152d55 100644 --- a/airflow-core/src/airflow/models/dagbundle.py +++ b/airflow-core/src/airflow/models/dagbundle.py @@ -19,6 +19,7 @@ from datetime import datetime from string import Formatter from typing import TYPE_CHECKING +from uuid import UUID import sqlalchemy as sa from sqlalchemy import Boolean, String, Text, select @@ -68,6 +69,8 @@ class DagBundleModel(Base, LoggingMixin): active: Mapped[bool | None] = mapped_column(Boolean, default=True, nullable=True) version: Mapped[str | None] = mapped_column(String(200), nullable=True) last_refreshed: Mapped[datetime | None] = mapped_column(UtcDateTime, nullable=True) + parse_revision: Mapped[UUID | None] = mapped_column(sa.Uuid(), nullable=True) + inventory_hash: Mapped[str | None] = mapped_column(String(64), nullable=True) signed_url_template: Mapped[str | None] = mapped_column(Text, nullable=True) template_params: Mapped[dict | None] = mapped_column(sa.JSON(), nullable=True) teams = relationship("Team", secondary=dag_bundle_team_association_table, back_populates="dag_bundles") diff --git a/airflow-core/src/airflow/models/dagcode.py b/airflow-core/src/airflow/models/dagcode.py index 12c02ed00f1b7..05d0af1bb3de7 100644 --- a/airflow-core/src/airflow/models/dagcode.py +++ b/airflow-core/src/airflow/models/dagcode.py @@ -85,7 +85,7 @@ def __init__( ): self.dag_version = dag_version self.fileloc = full_filepath - self.source_code = source_code or DagCode.code(self.dag_version.dag_id) + self.source_code = source_code if source_code is not None else DagCode.code(self.dag_version.dag_id) self.source_code_hash = self.dag_source_hash(self.source_code) self.dag_id = dag_version.dag_id self.language = language diff --git a/airflow-core/src/airflow/models/dagwarning.py b/airflow-core/src/airflow/models/dagwarning.py index 4d1eeac7db59d..25224943925bd 100644 --- a/airflow-core/src/airflow/models/dagwarning.py +++ b/airflow-core/src/airflow/models/dagwarning.py @@ -19,21 +19,15 @@ from datetime import datetime from enum import Enum -from typing import TYPE_CHECKING, Annotated +from typing import Annotated from pydantic import StringConstraints, TypeAdapter -from sqlalchemy import ForeignKeyConstraint, Index, String, Text, delete, select, true +from sqlalchemy import ForeignKeyConstraint, Index, String, Text from sqlalchemy.orm import Mapped, mapped_column, relationship from airflow._shared.timezones import timezone from airflow.models.base import Base, StringID -from airflow.models.dag import DagModel -from airflow.utils.retries import retry_db_transaction -from airflow.utils.session import NEW_SESSION, provide_session -from airflow.utils.sqlalchemy import UtcDateTime, get_dialect_name - -if TYPE_CHECKING: - from sqlalchemy.orm import Session +from airflow.utils.sqlalchemy import UtcDateTime WARNING_TYPE_MAX_LENGTH = 50 @@ -77,24 +71,6 @@ def __eq__(self, other) -> bool: def __hash__(self) -> int: return hash((self.dag_id, self.warning_type)) - @classmethod - @provide_session - @retry_db_transaction - def purge_inactive_dag_warnings(cls, *, session: Session = NEW_SESSION) -> None: - """ - Deactivate DagWarning records for inactive dags. - - :return: None - """ - if get_dialect_name(session) == "sqlite": - dag_ids_stmt = select(DagModel.dag_id).where(DagModel.is_stale == true()) - query = delete(cls).where(cls.dag_id.in_(dag_ids_stmt.scalar_subquery())) - else: - query = delete(cls).where(cls.dag_id == DagModel.dag_id, DagModel.is_stale == true()) - - session.execute(query.execution_options(synchronize_session=False)) - session.commit() - class DagWarningType(str, Enum): """ diff --git a/airflow-core/src/airflow/serialization/definitions/dag.py b/airflow-core/src/airflow/serialization/definitions/dag.py index c529bf95adea9..aa2621bc1e947 100644 --- a/airflow-core/src/airflow/serialization/definitions/dag.py +++ b/airflow-core/src/airflow/serialization/definitions/dag.py @@ -205,6 +205,7 @@ def bulk_write_to_db( dags: Collection[DAG | LazyDeserializedDAG], parse_duration: float | None = None, *, + enforce_bundle_ownership: bool = False, session: Session = NEW_SESSION, ) -> None: """ @@ -224,6 +225,7 @@ def bulk_write_to_db( bundle_name=bundle_name, bundle_version=bundle_version, dags={d.dag_id: LazyDeserializedDAG.from_dag(d) for d in dags}, + enforce_bundle_ownership=enforce_bundle_ownership, ) orm_dags = dag_op.add_dags(session=session) diff --git a/airflow-core/src/airflow/utils/cli.py b/airflow-core/src/airflow/utils/cli.py index d5af192b35016..b2ade41fda90d 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 wrapper(*args, **kwargs): _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/src/airflow/utils/db.py b/airflow-core/src/airflow/utils/db.py index 9d4adfaf82eb5..d99e7cb1ea8be 100644 --- a/airflow-core/src/airflow/utils/db.py +++ b/airflow-core/src/airflow/utils/db.py @@ -118,7 +118,7 @@ class MappedClassProtocol(Protocol): "3.2.0": "1d6611b6ab7c", "3.3.0": "d2f4e1b3c5a7", "3.4.0": "a5d7b9c13e40", - "3.5.0": "c71065a1b14f", + "3.5.0": "687eb253de24", } # Prefix used to identify tables holding data moved during migration. 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 0000000000000..6701c7f4a26c9 --- /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 53bc94b2e9b2e..977cb2ec5c569 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,12 +16,16 @@ # 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 httpx +import httpx2 import jwt import pytest import uvicorn @@ -32,20 +36,32 @@ 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, ) @@ -210,22 +226,22 @@ def test_rotation_stop_and_completion_over_http( } -class LostAcknowledgmentTransport(httpx.BaseTransport): +class LostAcknowledgmentTransport(httpx2.BaseTransport): def __init__(self): - self.transport = httpx.HTTPTransport() + self.transport = httpx2.HTTPTransport() self.lost = set() 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 ): self.lost.add(request.url.path) response.read() response.close() - raise httpx.ReadError("Lost acknowledgment after commit", request=request) + raise httpx2.ReadError("Lost acknowledgment after commit", request=request) return response def close(self): @@ -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(httpx.HTTPStatusError) as error: + 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", } @@ -269,7 +307,7 @@ def test_expired_credentials_fail_until_provisioning_resumes(clock, api_url, pro time_machine.shift(700) clock.return_value = 700 - with pytest.raises(httpx.HTTPStatusError) as error: + with pytest.raises(httpx2.HTTPStatusError) as error: processor.heartbeat() assert error.value.response.status_code == 403 @@ -301,7 +339,7 @@ def test_changed_session_signals_restart_but_allows_drain_and_completion( assert session.get(Job, job_id).state == JobState.SUCCESS.value -class CompletionFailureTransport(httpx.HTTPTransport): +class CompletionFailureTransport(httpx2.HTTPTransport): def __init__(self, *, commit): super().__init__() self.commit = commit @@ -313,7 +351,7 @@ def handle_request(self, request): response = super().handle_request(request) response.read() response.close() - raise httpx.ReadError("Completion acknowledgment unavailable", request=request) + raise httpx2.ReadError("Completion acknowledgment unavailable", request=request) return super().handle_request(request) @@ -329,7 +367,7 @@ def test_completion_recovery_after_token_expiry( base_url=api_url, token_file=token_file, hostname="processor-1", transport=transport ) as processor: job_id = processor.register_job() - with pytest.raises(httpx.ReadError): + with pytest.raises(httpx2.ReadError): processor.complete_job(TerminalJobState.SUCCESS) time_machine.shift(350) clock.return_value = 350 @@ -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/api_fastapi/execution_api/test_token_scope_boundaries.py b/airflow-core/tests/unit/api_fastapi/execution_api/test_token_scope_boundaries.py index 48b19c2369585..f97bd40bb968b 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/test_token_scope_boundaries.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/test_token_scope_boundaries.py @@ -76,6 +76,11 @@ "POST /jobs/{job_id}/heartbeat": {"dag_processor"}, "POST /jobs/{job_id}/complete": {"dag_processor"}, "POST /jobs/{job_id}/parse-token": {"dag_processor"}, + "POST /jobs/{job_id}/parse-results": {"dag_processor"}, + "GET /jobs/{job_id}/bundles": {"dag_processor"}, + "POST /jobs/{job_id}/bundles/{bundle_name:path}/inventory": {"dag_processor"}, + "POST /jobs/{job_id}/requested-work/{kind}/claim": {"dag_processor"}, + "POST /jobs/{job_id}/requested-work/{kind}/{work_id}/ack": {"dag_processor"}, } # Routes that check the Job of a dag_processor token themselves instead of requiring it to be open. @@ -86,6 +91,11 @@ "POST /jobs/{job_id}/heartbeat", "POST /jobs/{job_id}/complete", "POST /jobs/{job_id}/parse-token", + "POST /jobs/{job_id}/parse-results", + "GET /jobs/{job_id}/bundles", + "POST /jobs/{job_id}/bundles/{bundle_name:path}/inventory", + "POST /jobs/{job_id}/requested-work/{kind}/claim", + "POST /jobs/{job_id}/requested-work/{kind}/{work_id}/ack", } # Every message a Dag file parsing process can send its supervisor, mapped to the Execution API diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_dag_parsing.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_dag_parsing.py new file mode 100644 index 0000000000000..b9ff0bb9c43ad --- /dev/null +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_dag_parsing.py @@ -0,0 +1,1243 @@ +# 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 + +import asyncio +import json +import zipfile +from concurrent.futures import ThreadPoolExecutor +from copy import deepcopy +from threading import Barrier, Event +from unittest import mock +from uuid import UUID + +import pytest +from fastapi import APIRouter, FastAPI, Request +from sqlalchemy import delete, event, select +from sqlalchemy.exc import OperationalError +from sqlalchemy.orm import Session +from uuid6 import uuid7 + +from airflow import settings +from airflow._shared.timezones import timezone +from airflow.api_fastapi.execution_api.datamodels.token import DagProcessorClaims, DagProcessorToken +from airflow.api_fastapi.execution_api.routes.dag_parsing import DagParsingRoute, get_bundles +from airflow.api_fastapi.execution_api.security import require_auth +from airflow.callbacks.callback_requests import DagCallbackRequest +from airflow.dag_processing.bundles.local import LocalDagBundle +from airflow.dag_processing.collection import DagModelOperation, update_dag_parsing_results_in_db +from airflow.dag_processing.manager import DagFileProcessorManager +from airflow.jobs.job import Job, JobState +from airflow.models.asset import DagScheduleAssetReference +from airflow.models.callback import DagProcessorCallback +from airflow.models.dag import DagModel +from airflow.models.dag_parse_checkpoint import DagParseCheckpoint +from airflow.models.dag_version import DagVersion +from airflow.models.dagbag import DagPriorityParsingRequest +from airflow.models.dagbundle import DagBundleModel +from airflow.models.dagcode import DagCode +from airflow.models.errors import ParseImportError +from airflow.models.serialized_dag import SerializedDagModel +from airflow.models.team import Team +from airflow.sdk import DAG, Asset +from airflow.serialization.serialized_objects import LazyDeserializedDAG +from airflow.utils.session import create_session + +from tests_common.test_utils.asserts import assert_queries_count +from tests_common.test_utils.config import conf_vars +from tests_common.test_utils.db import ( + clear_db_assets, + clear_db_callbacks, + clear_db_dag_bundles, + clear_db_dags, + clear_db_import_errors, + clear_db_jobs, +) + +pytestmark = pytest.mark.db_test +SESSION_ID = UUID("00000000-0000-0000-0000-0000000000aa") + + +@pytest.fixture(autouse=True) +def clean_db(time_machine, session): + time_machine.move_to("2026-10-05T12:00:00Z", tick=False) + clear_db_jobs() + clear_db_callbacks() + session.execute(delete(DagPriorityParsingRequest)) + session.commit() + clear_db_dags() + clear_db_import_errors() + clear_db_dag_bundles() + yield + clear_db_jobs() + clear_db_callbacks() + session.execute(delete(DagPriorityParsingRequest)) + session.commit() + clear_db_dags() + clear_db_import_errors() + clear_db_dag_bundles() + + +@pytest.fixture +def job(session, exec_app): + job = Job(job_type="DagProcessorJob", state=JobState.RUNNING) + job.session_id = SESSION_ID + job.registration_id = uuid7() + job.bundle_names = ["bundle"] + session.add_all([job, DagBundleModel(name="bundle"), DagBundleModel(name="other")]) + session.commit() + + async def authenticate(request: Request): + return DagProcessorToken( + id=SESSION_ID, + claims=DagProcessorClaims(exp=1, job_id=job.id, dag_bundles=frozenset({"bundle"})), + ) + + exec_app.dependency_overrides[require_auth] = authenticate + return job + + +@pytest.fixture +def body(): + dag = DAG("published", schedule=None) + dag.fileloc = "/worker-only/dags/a.py" + dag.relative_fileloc = "a.py" + return { + "attempt_id": str(uuid7()), + "dispatch_sequence": 1, + "bundle_name": "bundle", + "relative_fileloc": "a.py", + "bundle_version": "v1", + "version_data": {"ref": "v1"}, + "parse_duration": 0.5, + "serialized_dags": [LazyDeserializedDAG.from_dag(dag).data], + "source_codes": {dag.dag_id: {"source_code": "# captured on worker", "language": "python"}}, + } + + +@mock.patch.object( + ParseImportError, "full_file_path", autospec=True, side_effect=AssertionError("Server constructed bundle") +) +@mock.patch("airflow.listeners.listener.get_listener_manager", autospec=True) +def test_import_error_listeners_get_bundle_relative_paths( + get_listener_manager, full_file_path, client, job, body +): + hook = get_listener_manager.return_value.hook + body.update(serialized_dags=[], source_codes={}, import_errors={"a.py": "broken"}) + assert client.post(f"/execution/jobs/{job.id}/parse-results", json=body).status_code == 200 + body.update(attempt_id=str(uuid7()), dispatch_sequence=2, import_errors={"a.py": "still broken"}) + assert client.post(f"/execution/jobs/{job.id}/parse-results", json=body).status_code == 200 + + hook.on_new_dag_import_error.assert_called_once_with(filename="a.py", stacktrace="broken") + hook.on_existing_dag_import_error.assert_called_once_with(filename="a.py", stacktrace="still broken") + full_file_path.assert_not_called() + + +@pytest.mark.parametrize("source", ["# captured on worker", "", None]) +@mock.patch.object(DagCode, "code", autospec=True, side_effect=AssertionError("Server read source")) +def test_publishes_worker_source_without_filesystem_access(mock_code, client, session, job, body, source): + body["source_codes"]["published"]["source_code"] = source + response = client.post(f"/execution/jobs/{job.id}/parse-results", json=body) + assert response.status_code == 200, response.json() + assert session.get(DagModel, "published").relative_fileloc == "a.py" + assert session.scalar(select(SerializedDagModel)).dag_id == "published" + assert session.scalar(select(DagCode)).source_code == ( + source if source is not None else "Source unavailable" + ) + assert session.scalar(select(DagParseCheckpoint)).attempt_id == UUID(body["attempt_id"]) + mock_code.assert_not_called() + + +def test_replay_returns_receipt_without_rewriting_metadata(client, session, job, body, time_machine): + first = client.post(f"/execution/jobs/{job.id}/parse-results", json=body) + assert first.status_code == 200, first.json() + parsed_at = session.get(DagModel, "published").last_parsed_time + time_machine.shift(10) + replay = client.post(f"/execution/jobs/{job.id}/parse-results", json=body) + assert replay.json() == first.json() + session.expire_all() + assert session.get(DagModel, "published").last_parsed_time == parsed_at + assert len(session.scalars(select(DagParseCheckpoint)).all()) == 1 + + +@pytest.mark.parametrize("first", ["parse", "inventory"]) +def test_parse_and_inventory_checkpoints_have_separate_namespaces(client, session, exec_app, job, first): + session.add_all([DagBundleModel(name="inventory"), DagBundleModel(name="a.py")]) + job.bundle_names = ["inventory", "a.py"] + session.commit() + + async def authenticate(request: Request): + return DagProcessorToken( + id=job.session_id, + claims=DagProcessorClaims(exp=1, job_id=job.id, dag_bundles=frozenset(job.bundle_names)), + ) + + exec_app.dependency_overrides[require_auth] = authenticate + publications = { + "parse": ( + f"/execution/jobs/{job.id}/parse-results", + { + "attempt_id": str(uuid7()), + "dispatch_sequence": 2, + "bundle_name": "inventory", + "relative_fileloc": "a.py", + "parse_duration": 0.5, + "serialized_dags": [], + "source_codes": {}, + "import_errors": {"a.py": "Import failed"}, + }, + ), + "inventory": ( + f"/execution/jobs/{job.id}/bundles/a.py/inventory", + { + "attempt_id": str(uuid7()), + "dispatch_sequence": 2, + "expected_revision": None, + "files": [], + }, + ), + } + publications[first][1]["dispatch_sequence"] = 3 + second = "inventory" if first == "parse" else "parse" + receipts = {} + for kind in (first, second, first, second): + url, payload = publications[kind] + response = client.post(url, json=payload) + assert response.status_code == 200, response.json() + if kind in receipts: + assert response.json() == receipts[kind] + else: + receipts[kind] = response.json() + + checkpoints = session.scalars(select(DagParseCheckpoint)).all() + assert len(checkpoints) == 2 + assert {checkpoint.attempt_id for checkpoint in checkpoints} == { + UUID(payload["attempt_id"]) for _, payload in publications.values() + } + error = session.scalar(select(ParseImportError)) + assert (error.bundle_name, error.filename) == ("inventory", "a.py") + + +@pytest.mark.parametrize("paused", [True, False]) +def test_publication_preserves_pause_on_creation_and_existing_pause_state(client, session, job, body, paused): + body["serialized_dags"][0]["dag"]["is_paused_upon_creation"] = paused + assert client.post(f"/execution/jobs/{job.id}/parse-results", json=body).status_code == 200 + assert session.get(DagModel, "published").is_paused is paused + body.update(attempt_id=str(uuid7()), dispatch_sequence=2) + body["serialized_dags"][0]["dag"]["is_paused_upon_creation"] = not paused + assert client.post(f"/execution/jobs/{job.id}/parse-results", json=body).status_code == 200 + session.expire_all() + assert session.get(DagModel, "published").is_paused is paused + + +@pytest.mark.parametrize( + ("change", "reason"), [("payload", "publication_conflict"), ("attempt", "publication_superseded")] +) +def test_rejects_conflicting_or_superseded_publication(client, session, job, body, change, reason): + assert client.post(f"/execution/jobs/{job.id}/parse-results", json=body).status_code == 200 + if change == "payload": + body["parse_duration"] = 99 + else: + body["attempt_id"] = str(uuid7()) + response = client.post(f"/execution/jobs/{job.id}/parse-results", json=body) + assert response.status_code == 409 + assert response.json()["detail"]["reason"] == reason + assert session.get(DagModel, "published").last_parse_duration == 0.5 + + +@conf_vars({("core", "min_serialized_dag_update_interval"): "0"}) +def test_newer_result_supersedes_old_attempt(client, session, job, body): + original = deepcopy(body) + assert client.post(f"/execution/jobs/{job.id}/parse-results", json=body).status_code == 200 + body.update(attempt_id=str(uuid7()), dispatch_sequence=2, parse_duration=2.0) + assert client.post(f"/execution/jobs/{job.id}/parse-results", json=body).status_code == 200 + assert client.post(f"/execution/jobs/{job.id}/parse-results", json=original).status_code == 409 + session.expire_all() + assert session.get(DagModel, "published").last_parse_duration == 2.0 + assert len(session.scalars(select(DagParseCheckpoint)).all()) == 1 + + +@pytest.mark.parametrize("same_bundle", [True, False]) +def test_allows_file_relocation_but_not_cross_bundle_takeover(client, session, job, body, same_bundle): + session.add( + DagModel( + dag_id="published", bundle_name="bundle" if same_bundle else "other", relative_fileloc="old.py" + ) + ) + session.commit() + response = client.post(f"/execution/jobs/{job.id}/parse-results", json=body) + assert response.status_code == (200 if same_bundle else 409), response.json() + session.expire_all() + assert session.get(DagModel, "published").relative_fileloc == ("a.py" if same_bundle else "old.py") + + +@pytest.mark.parametrize( + "failure", ["bundle", "job", "closed", "replaced", "inactive", "deleted", "missing_bundle"] +) +def test_rejects_unauthorized_publication(client, session, job, body, failure): + job_id = job.id + if failure == "bundle": + body["bundle_name"] = "other" + elif failure == "job": + job_id += 1000 + elif failure == "closed": + job.end_date = job.start_date + elif failure == "replaced": + job.session_id = None + elif failure == "deleted": + session.delete(job) + elif failure == "missing_bundle": + session.delete(session.get(DagBundleModel, "bundle")) + else: + session.get(DagBundleModel, "bundle").active = False + session.commit() + response = client.post(f"/execution/jobs/{job_id}/parse-results", json=body) + assert response.status_code == {"job": 404, "inactive": 409, "missing_bundle": 409}.get(failure, 403) + assert session.get(DagModel, "published") is None + assert session.scalar(select(DagParseCheckpoint)) is None + + +@pytest.mark.parametrize( + "failure", + [ + "schema", + "version", + "source", + "path", + "error_path", + "warning", + "duplicate", + "dag_id", + "relative_fileloc", + "fileloc", + "long_path", + ], +) +def test_invalid_result_is_rejected_without_writes(client, session, job, body, failure): + dag = body["serialized_dags"][0] + if failure == "schema": + dag["dag"]["tasks"] = "invalid" + elif failure == "version": + dag["__version"] = 999 + elif failure == "source": + body["source_codes"] = {} + elif failure == "path": + dag["dag"]["relative_fileloc"] = "../other.py" + elif failure == "error_path": + body["import_errors"] = {"other.py": "error"} + elif failure == "warning": + body["warnings"] = [{"dag_id": "other", "warning_type": "duplicate dag id", "message": "error"}] + elif failure == "duplicate": + body["serialized_dags"].append(deepcopy(dag)) + elif failure == "long_path": + body["parsed_definitions"] = ["a.py/" + "x" * 2000] + else: + dag["dag"].pop(failure) + response = client.post(f"/execution/jobs/{job.id}/parse-results", json=body) + assert response.status_code == 422, response.json() + assert session.get(DagModel, "published") is None + assert session.scalar(select(DagParseCheckpoint)) is None + + +@mock.patch("airflow.api_fastapi.execution_api.routes.dag_parsing.MAX_PUBLICATION_BYTES", 10) +@pytest.mark.parametrize("chunked", [False, True], ids=["declared", "chunked"]) +@pytest.mark.parametrize("path", ["parse-results", "bundles/bundle/inventory"]) +def test_rejects_a_body_over_the_limit_before_reading_it(client, session, job, body, path, chunked): + content = json.dumps(body).encode() + response = client.post( + f"/execution/jobs/{job.id}/{path}", + content=iter([content]) if chunked else content, + headers={"Content-Type": "application/json"}, + ) + assert response.status_code == 413, response.json() + assert session.scalar(select(DagParseCheckpoint)) is None + + +def test_client_disconnect_during_the_body_is_a_bad_request(): + app = FastAPI() + router = APIRouter(route_class=DagParsingRoute) + + @router.post("/echo") + def echo(body: dict[str, str]) -> dict[str, str]: + return body + + app.include_router(router) + messages = iter( + [{"type": "http.request", "body": b'{"key":', "more_body": True}, {"type": "http.disconnect"}] + ) + sent = [] + + async def receive(): + return next(messages) + + async def send(message): + sent.append(message) + + scope = { + "type": "http", + "method": "POST", + "path": "/echo", + "root_path": "", + "query_string": b"", + "headers": [(b"content-type", b"application/json")], + "http_version": "1.1", + "scheme": "http", + "server": ("testserver", 80), + "client": ("testclient", 1), + } + + asyncio.run(app(scope, receive, send)) + + assert sent[0]["status"] == 400 + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("dag_id", "d" * 251), + ("dag_id", "nested/dag"), + ("tags", ["t" * 101]), + ("fileloc", "/" + "f" * 2000), + ], + ids=["long_dag_id", "dag_id_with_slash", "long_tag", "long_fileloc"], +) +def test_rejects_dag_metadata_the_database_cannot_hold(client, session, job, body, field, value): + body["serialized_dags"][0]["dag"][field] = value + if field == "dag_id": + body["source_codes"] = {value: body["source_codes"].pop("published")} + + response = client.post(f"/execution/jobs/{job.id}/parse-results", json=body) + + assert response.status_code == 422, response.json() + assert session.scalar(select(DagModel)) is None + assert session.scalar(select(DagParseCheckpoint)) is None + + +def test_rejects_non_json_version_metadata(client, session, job, body): + body["version_data"] = {"invalid": float("nan")} + response = client.post(f"/execution/jobs/{job.id}/parse-results", content=json.dumps(body)) + assert response.status_code == 422 + assert session.scalar(select(DagParseCheckpoint)) is None + + +@mock.patch( + "airflow.dag_processing.collection._update_import_errors", + autospec=True, + side_effect=RuntimeError("write failed"), +) +def test_diagnostic_failure_rolls_back_metadata_and_receipt(mock_update, client, session, job, body): + with pytest.raises(RuntimeError, match="write failed"): + client.post(f"/execution/jobs/{job.id}/parse-results", json=body) + assert session.get(DagModel, "published") is None + assert session.scalar(select(SerializedDagModel)) is None + assert session.scalar(select(DagParseCheckpoint)) is None + + +def test_empty_result_clears_file_errors_and_has_replay_receipt(client, session, job, body): + body.update(serialized_dags=[], source_codes={}, import_errors={"a.py": "broken"}) + assert client.post(f"/execution/jobs/{job.id}/parse-results", json=body).status_code == 200 + assert session.scalar(select(ParseImportError)) is not None + body.update(attempt_id=str(uuid7()), dispatch_sequence=2, import_errors={}) + result = client.post(f"/execution/jobs/{job.id}/parse-results", json=body) + assert result.status_code == 200, result.json() + assert session.scalar(select(ParseImportError)) is None + assert client.post(f"/execution/jobs/{job.id}/parse-results", json=body).json() == result.json() + + +def test_container_result_keeps_member_errors_and_source(client, session, job, body): + body["relative_fileloc"] = "bundle.zip" + body["serialized_dags"][0]["dag"]["relative_fileloc"] = "bundle.zip/a.py" + body["parsed_definitions"] = ["bundle.zip/a.py", "bundle.zip/b.py"] + body["import_errors"] = {"bundle.zip/b.py": "bad member"} + response = client.post(f"/execution/jobs/{job.id}/parse-results", json=body) + assert response.status_code == 200, response.json() + assert session.scalar(select(ParseImportError)).filename == "bundle.zip/b.py" + + +@mock.patch( + "airflow.api_fastapi.execution_api.routes.dag_parsing.DagSerialization.validate_serialized_dag", + autospec=True, +) +@mock.patch( + "airflow.api_fastapi.execution_api.routes.dag_parsing._reject_other_teams_plugin_classes", + autospec=True, + return_value=[], +) +def test_plugin_authorization_precedes_deserialization(reject, deserialize, client, session, job, body): + response = client.post(f"/execution/jobs/{job.id}/parse-results", json=body) + assert response.status_code == 403 + deserialize.assert_not_called() + assert session.scalar(select(DagParseCheckpoint)) is None + + +@pytest.mark.parametrize( + ("stage", "error"), + [ + ("_build_duplicate_dag_id_warnings", ValueError("failed")), + ("_update_dag_warnings", ValueError("failed")), + ("SerializedDagModel.write_dag", ValueError("failed")), + ("SerializedDAG.bulk_write_to_db", OperationalError(None, None, RuntimeError("failed"))), + ], +) +def test_persistence_errors_do_not_leave_accepted_partial_results( + stage, error, client, session, job, body, mocker +): + write = mocker.patch(f"airflow.dag_processing.collection.{stage}", autospec=True, side_effect=error) + if isinstance(error, OperationalError): + response = client.post(f"/execution/jobs/{job.id}/parse-results", json=body) + assert response.status_code == 500 + else: + with pytest.raises(type(error), match="failed"): + client.post(f"/execution/jobs/{job.id}/parse-results", json=body) + write.assert_called_once() + assert session.get(DagModel, "published") is None + assert session.scalar(select(DagParseCheckpoint)) is None + + +def test_other_job_publication_is_not_overwritten_by_receipt_recovery(client, session, exec_app, job, body): + first = client.post(f"/execution/jobs/{job.id}/parse-results", json=body) + assert first.status_code == 200 + other = Job(job_type="DagProcessorJob", state=JobState.RUNNING) + other.session_id = uuid7() + other.registration_id = uuid7() + other.bundle_names = ["bundle"] + session.add(other) + session.commit() + identities = {job.id: SESSION_ID, other.id: other.session_id} + + async def authenticate(request: Request): + job_id = int(request.path_params["job_id"]) + return DagProcessorToken( + id=identities[job_id], + claims=DagProcessorClaims(exp=1, job_id=job_id, dag_bundles=frozenset({"bundle"})), + ) + + exec_app.dependency_overrides[require_auth] = authenticate + newer = {**body, "attempt_id": str(uuid7()), "parse_duration": 3.0} + assert client.post(f"/execution/jobs/{other.id}/parse-results", json=newer).status_code == 200 + assert client.post(f"/execution/jobs/{job.id}/parse-results", json=body).json() == first.json() + session.expire_all() + assert session.get(DagModel, "published").last_parse_duration == 3.0 + assert len(session.scalars(select(DagParseCheckpoint)).all()) == 2 + + +def test_job_cleanup_removes_its_checkpoint(client, session, job, body): + assert client.post(f"/execution/jobs/{job.id}/parse-results", json=body).status_code == 200 + session.delete(job) + session.commit() + assert session.scalar(select(DagParseCheckpoint)) is None + assert session.get(DagModel, "published") is not None + + +@pytest.mark.backend("postgres") +def test_concurrent_replay_commits_one_publication(client, session, job, body): + ready = Barrier(2) + url = f"/execution/jobs/{job.id}/parse-results" + + def publish(): + ready.wait(timeout=10) + return client.post(url, json=body) + + with ThreadPoolExecutor(max_workers=2) as executor: + futures = [executor.submit(publish) for _ in range(2)] + responses = [future.result(timeout=30) for future in futures] + assert [response.status_code for response in responses] == [200, 200] + assert responses[0].json() == responses[1].json() + assert len(session.scalars(select(DagParseCheckpoint)).all()) == 1 + + +@pytest.mark.backend("postgres") +def test_publication_locks_do_not_block_foreign_key_inserts(client, session, job, body): + statements = [] + + def capture(conn, cursor, statement, parameters, context, executemany): + statements.append(statement) + + event.listen(settings.engine, "before_cursor_execute", capture) + try: + assert client.post(f"/execution/jobs/{job.id}/parse-results", json=body).status_code == 200 + inventory_start = len(statements) + inventory = { + "attempt_id": str(uuid7()), + "dispatch_sequence": 2, + "expected_revision": None, + "version": "v1", + "files": [], + } + response = client.post(f"/execution/jobs/{job.id}/bundles/bundle/inventory", json=inventory) + assert response.status_code == 200, response.json() + finally: + event.remove(settings.engine, "before_cursor_execute", capture) + + locks = [statement for statement in statements if " FOR " in statement] + assert any("FROM dag " in lock and "FOR NO KEY UPDATE" in lock for lock in locks) + assert any("FROM dag_bundle " in lock and "FOR NO KEY UPDATE" in lock for lock in locks) + assert all("FROM job " in lock for lock in locks if "FOR UPDATE" in lock) + inventory_dag_locks = [ + statement + for statement in statements[inventory_start:] + if " FOR " in statement and "FROM dag " in statement + ] + assert inventory_dag_locks + assert all("dag.dag_id IN" in lock for lock in inventory_dag_locks) + + +@conf_vars({("core", "auth_manager"): "airflow.providers.fab.auth_manager.fab_auth_manager.FabAuthManager"}) +@pytest.mark.parametrize("fail", [False, True]) +@mock.patch("airflow.dag_processing.collection._update_import_errors", autospec=True) +def test_fab_permissions_share_publication_transaction(update_errors, client, session, job, body, fail): + security = pytest.importorskip("airflow.providers.fab.www.security_appless") + dag_id = f"published_{uuid7().hex}" + body["serialized_dags"][0]["dag"]["dag_id"] = dag_id + body["source_codes"] = {dag_id: body["source_codes"].pop("published")} + if fail: + update_errors.side_effect = RuntimeError("failure after Dag permissions") + with pytest.raises(RuntimeError, match="failure after Dag permissions"): + client.post(f"/execution/jobs/{job.id}/parse-results", json=body) + else: + response = client.post(f"/execution/jobs/{job.id}/parse-results", json=body) + assert response.status_code == 200, response.json() + session.expire_all() + assert (session.get(DagModel, dag_id) is None) == fail + assert (session.scalar(select(SerializedDagModel)) is None) == fail + assert (session.scalar(select(DagParseCheckpoint)) is None) == fail + manager = security.ApplessAirflowSecurityManager(session=session) + assert (manager.get_resource(f"DAG:{dag_id}") is None) == fail + + +@conf_vars({("core", "auth_manager"): "airflow.providers.fab.auth_manager.fab_auth_manager.FabAuthManager"}) +@pytest.mark.parametrize("existing", [False, True]) +@pytest.mark.parametrize("legacy_exception", [False, True]) +@pytest.mark.parametrize( + ("failure", "resource_actions", "message"), + [ + ("role", {"DAGs": {"can_read"}}, "role does not exist"), + ("resource", {"missing_resource": {"can_read"}}, "invalid resource name"), + ("permission", {"DAGs": {"missing_permission"}}, "invalid permissions"), + ], +) +def test_invalid_fab_access_control_rejects_and_rolls_back_publication( + client, session, job, body, existing, legacy_exception, failure, resource_actions, message, monkeypatch +): + security = pytest.importorskip("airflow.providers.fab.www.security_appless") + if legacy_exception: + override = pytest.importorskip("airflow.providers.fab.auth_manager.security_manager.override") + monkeypatch.setattr(override, "DagAccessControlValidationError", override.FabException, raising=False) + dag_id = f"published_{uuid7().hex}" + role_name = f"role_{uuid7().hex}" + manager = security.ApplessAirflowSecurityManager(session=session) + role = manager.add_role(role_name) if failure != "role" else None + body["serialized_dags"][0]["dag"]["dag_id"] = dag_id + body["source_codes"] = {dag_id: body["source_codes"].pop("published")} + original_attempt = UUID(body["attempt_id"]) + if existing: + assert client.post(f"/execution/jobs/{job.id}/parse-results", json=body).status_code == 200 + dag = DAG(dag_id, schedule=None, access_control={role_name: resource_actions}) + body["serialized_dags"][0]["dag"]["access_control"] = LazyDeserializedDAG.from_dag(dag).data["dag"][ + "access_control" + ] + body.update(attempt_id=str(uuid7()), dispatch_sequence=2, parse_duration=2.0) + + response = client.post(f"/execution/jobs/{job.id}/parse-results", json=body) + + assert response.status_code == 422 + assert response.json()["detail"]["reason"] == "invalid_dag" + assert message in response.json()["detail"]["message"] + assert dag_id in response.json()["detail"]["message"] + session.expire_all() + stored_dag = session.get(DagModel, dag_id) + checkpoint = session.scalar(select(DagParseCheckpoint)) + if existing: + assert stored_dag.last_parse_duration == 0.5 + assert not stored_dag.is_stale + assert checkpoint.attempt_id == original_attempt + else: + assert stored_dag is None + assert checkpoint is None + assert len(session.scalars(select(SerializedDagModel)).all()) == int(existing) + assert len(session.scalars(select(DagVersion)).all()) == int(existing) + assert len(session.scalars(select(DagCode)).all()) == int(existing) + assert session.scalar(select(ParseImportError)) is None + assert (manager.get_resource(f"DAG:{dag_id}") is not None) == existing + if role is not None: + session.delete(role) + session.commit() + + +@conf_vars({("core", "auth_manager"): "airflow.providers.fab.auth_manager.fab_auth_manager.FabAuthManager"}) +@mock.patch.object(Session, "commit", autospec=True) +def test_fab_backend_exception_is_not_reported_as_invalid_dag(commit, client, session, job, body): + override = pytest.importorskip("airflow.providers.fab.auth_manager.security_manager.override") + dag_id = f"published_{uuid7().hex}" + body["serialized_dags"][0]["dag"]["dag_id"] = dag_id + body["source_codes"] = {dag_id: body["source_codes"].pop("published")} + commit.side_effect = OperationalError(None, None, RuntimeError("temporary database failure")) + + with pytest.raises(override.FabException) as error: + client.post(f"/execution/jobs/{job.id}/parse-results", json=body) + + assert isinstance(error.value.__cause__, OperationalError) + commit.assert_called_once() + assert session.get(DagModel, dag_id) is None + assert session.scalar(select(DagParseCheckpoint)) is None + assert session.scalar(select(ParseImportError)) is None + + +@conf_vars({("core", "auth_manager"): "airflow.providers.fab.auth_manager.fab_auth_manager.FabAuthManager"}) +def test_transient_permission_failure_remains_server_error(client, session, job, body, mocker): + fab = pytest.importorskip("airflow.providers.fab.auth_manager.fab_auth_manager") + assert client.post(f"/execution/jobs/{job.id}/parse-results", json=body).status_code == 200 + original_attempt = UUID(body["attempt_id"]) + body.update(attempt_id=str(uuid7()), dispatch_sequence=2, parse_duration=2.0) + sync_permissions = mocker.patch.object( + fab.FabAuthManager, + "sync_dag_perms", + autospec=True, + side_effect=OperationalError(None, None, RuntimeError("temporary database failure")), + ) + + response = client.post(f"/execution/jobs/{job.id}/parse-results", json=body) + + assert response.status_code == 500 + sync_permissions.assert_called_once() + session.expire_all() + stored_dag = session.get(DagModel, "published") + assert stored_dag.last_parse_duration == 0.5 + assert not stored_dag.is_stale + assert session.scalar(select(DagParseCheckpoint)).attempt_id == original_attempt + assert session.scalar(select(ParseImportError)) is None + + +@pytest.mark.parametrize("ineligible", [None, "active_bundle", "active_dag", "file_location", "version"]) +def test_recovers_only_orphaned_legacy_dags(client, session, job, body, ineligible): + previous_bundle = session.get(DagBundleModel, "other") + previous_bundle.active = ineligible == "active_bundle" + session.add( + DagModel( + dag_id="published", + bundle_name="other", + relative_fileloc="old.py" if ineligible == "file_location" else None, + fileloc="/old-host/dags/a.py", + is_stale=ineligible != "active_dag", + ) + ) + session.flush() + if ineligible == "version": + session.add(DagVersion(dag_id="published", bundle_name="other")) + session.commit() + response = client.post(f"/execution/jobs/{job.id}/parse-results", json=body) + assert response.status_code == (409 if ineligible else 200), response.json() + session.expire_all() + dag = session.get(DagModel, "published") + assert dag.bundle_name == ("other" if ineligible else "bundle") + if ineligible is None: + assert dag.relative_fileloc == "a.py" + assert not dag.is_stale + + +@conf_vars({("core", "multi_team"): "True"}) +@pytest.mark.parametrize("same_team", [True, False]) +def test_legacy_dag_recovery_stays_within_a_team(client, session, job, body, same_team): + publishing_team, other_team = Team(name=f"team_{uuid7().hex}"), Team(name=f"team_{uuid7().hex}") + session.get(DagBundleModel, "bundle").teams = [publishing_team] + previous_bundle = session.get(DagBundleModel, "other") + previous_bundle.active = False + previous_bundle.teams = [publishing_team if same_team else other_team] + session.add( + DagModel( + dag_id="published", bundle_name="other", relative_fileloc=None, fileloc="/old/a.py", is_stale=True + ) + ) + session.commit() + try: + response = client.post(f"/execution/jobs/{job.id}/parse-results", json=body) + + assert response.status_code == (200 if same_team else 409), response.json() + session.expire_all() + assert session.get(DagModel, "published").bundle_name == ("bundle" if same_team else "other") + finally: + session.rollback() + for name in ("bundle", "other"): + session.get(DagBundleModel, name).teams = [] + session.flush() + session.execute(delete(Team).where(Team.name.in_([publishing_team.name, other_team.name]))) + session.commit() + + +@conf_vars({("core", "auth_manager"): "airflow.providers.fab.auth_manager.fab_auth_manager.FabAuthManager"}) +def test_failed_publication_rolls_back_permissions_when_dag_metadata_is_unchanged( + client, session, job, body, mocker +): + security = pytest.importorskip("airflow.providers.fab.www.security_appless") + assert client.post(f"/execution/jobs/{job.id}/parse-results", json=body).status_code == 200 + original_attempt = UUID(body["attempt_id"]) + manager = security.ApplessAirflowSecurityManager(session=session) + role_name = f"publication_{uuid7().hex}" + role = manager.add_role(role_name) + dag = DAG("published", schedule=None, access_control={role_name: {"can_read"}}) + body["serialized_dags"][0]["dag"]["access_control"] = LazyDeserializedDAG.from_dag(dag).data["dag"][ + "access_control" + ] + body.update(attempt_id=str(uuid7()), dispatch_sequence=2) + mocker.patch( + "airflow.dag_processing.collection._update_import_errors", + autospec=True, + side_effect=RuntimeError("failure after permissions"), + ) + with pytest.raises(RuntimeError, match="failure after permissions"): + client.post(f"/execution/jobs/{job.id}/parse-results", json=body) + session.expire_all() + assert not role.permissions + assert session.scalar(select(DagParseCheckpoint)).attempt_id == original_attempt + session.delete(role) + session.commit() + + +@pytest.mark.backend("postgres") +@pytest.mark.parametrize("race_stage", ["before_persistence", "missing_row"]) +def test_concurrent_creation_cannot_overwrite_another_bundle( + client, session, exec_app, job, body, mocker, race_stage +): + other = Job(job_type="DagProcessorJob", state=JobState.RUNNING) + other.session_id = uuid7() + other.registration_id = uuid7() + other.bundle_names = ["other"] + session.add(other) + session.commit() + identities = {job.id: (job.session_id, "bundle"), other.id: (other.session_id, "other")} + + async def authenticate(request: Request): + job_id = int(request.path_params["job_id"]) + session_id, bundle = identities[job_id] + return DagProcessorToken( + id=session_id, + claims=DagProcessorClaims(exp=1, job_id=job_id, dag_bundles=frozenset({bundle})), + ) + + exec_app.dependency_overrides[require_auth] = authenticate + checked_ownership = Event() + second_published = Event() + + def wait_for_competitor(): + checked_ownership.set() + assert second_published.wait(timeout=30) + + if race_stage == "before_persistence": + + def persist(**kwargs): + if kwargs["bundle_name"] == "bundle": + wait_for_competitor() + return update_dag_parsing_results_in_db(**kwargs) + + mocker.patch( + "airflow.api_fastapi.execution_api.routes.dag_parsing.update_dag_parsing_results_in_db", + autospec=True, + side_effect=persist, + ) + else: + find_orm_dags = DagModelOperation.find_orm_dags + + def find(op, *, session): + dags = find_orm_dags(op, session=session) + if op.bundle_name == "bundle" and not dags: + wait_for_competitor() + return dags + + mocker.patch.object(DagModelOperation, "find_orm_dags", autospec=True, side_effect=find) + + with ThreadPoolExecutor(max_workers=1) as executor: + first = executor.submit(client.post, f"/execution/jobs/{job.id}/parse-results", json=body) + try: + assert checked_ownership.wait(timeout=30) + second = client.post( + f"/execution/jobs/{other.id}/parse-results", json={**body, "bundle_name": "other"} + ) + assert second.status_code == 200, second.json() + finally: + second_published.set() + first_response = first.result(timeout=30) + assert first_response.status_code == 409, first_response.json() + session.expire_all() + assert session.get(DagModel, "published").bundle_name == "other" + assert [row.job_id for row in session.scalars(select(DagParseCheckpoint))] == [other.id] + + +def test_complete_inventory_reconciles_only_its_bundle_and_replays(client, session, job, body): + assert client.post(f"/execution/jobs/{job.id}/parse-results", json=body).status_code == 200 + session.add(DagModel(dag_id="other_bundle", bundle_name="other", relative_fileloc="a.py", is_stale=False)) + session.commit() + inventory = { + "attempt_id": str(uuid7()), + "dispatch_sequence": 2, + "expected_revision": None, + "version": "v2", + "files": [], + } + url = f"/execution/jobs/{job.id}/bundles/bundle/inventory" + response = client.post(url, json=inventory) + assert response.status_code == 200, response.json() + assert client.post(url, json=inventory).json() == response.json() + session.expire_all() + assert session.get(DagModel, "published").is_stale + assert not session.get(DagModel, "other_bundle").is_stale + body.update(attempt_id=str(uuid7()), dispatch_sequence=3) + assert client.post(f"/execution/jobs/{job.id}/parse-results", json=body).status_code == 409 + assert client.post(url.replace("bundle/inventory", "other/inventory"), json=inventory).status_code == 403 + + +@pytest.fixture +def clean_assets(): + clear_db_assets() + yield + clear_db_assets() + + +@pytest.mark.usefixtures("clean_assets") +def test_inventory_removes_asset_references_of_removed_dags(client, session, job, body): + dag = DAG("published", schedule=[Asset(name="key", uri="s3://bucket/key")]) + dag.fileloc = "/worker-only/dags/a.py" + dag.relative_fileloc = "a.py" + body["serialized_dags"] = [LazyDeserializedDAG.from_dag(dag).data] + assert client.post(f"/execution/jobs/{job.id}/parse-results", json=body).status_code == 200 + assert session.scalars(select(DagScheduleAssetReference.dag_id)).all() == ["published"] + inventory = { + "attempt_id": str(uuid7()), + "dispatch_sequence": 2, + "expected_revision": None, + "version": "v2", + "files": [], + } + + response = client.post(f"/execution/jobs/{job.id}/bundles/bundle/inventory", json=inventory) + + assert response.status_code == 200, response.json() + session.expire_all() + assert session.get(DagModel, "published").is_stale + assert session.scalar(select(DagScheduleAssetReference)) is None + + +@pytest.mark.parametrize("api_mode", [False, True], ids=["direct", "api"]) +@pytest.mark.parametrize("recovery", ["member_present", "member_removed", "archive_removed"]) +def test_unreadable_archive_preserves_members(client, session, job, body, tmp_path, api_mode, recovery): + body["relative_fileloc"] = "archive.zip" + body["serialized_dags"][0]["dag"]["relative_fileloc"] = "archive.zip/dag.py" + assert client.post(f"/execution/jobs/{job.id}/parse-results", json=body).status_code == 200 + session.add_all( + [ + DagModel(dag_id="deleted", bundle_name="bundle", relative_fileloc="gone.py", is_stale=False), + DagModel( + dag_id="similar_prefix", + bundle_name="bundle", + relative_fileloc="archive.zip.other/dag.py", + is_stale=False, + ), + ParseImportError( + filename="archive.zip/broken.py", + bundle_name="bundle", + timestamp=timezone.utcnow(), + stacktrace="previous error", + ), + ParseImportError( + filename=None, bundle_name="bundle", timestamp=timezone.utcnow(), stacktrace="legacy error" + ), + ] + ) + session.commit() + (tmp_path / "archive.zip").write_bytes(b"unreadable archive") + bundle = LocalDagBundle(name="bundle", path=tmp_path) + manager = DagFileProcessorManager(max_runs=1) + discovered = manager._find_files_in_bundle(bundle) + failed_sources = {loc for file in discovered for loc in file.failed_sources} + assert failed_sources == {"archive.zip"} + files = manager._get_observed_filelocs(discovered) + url = f"/execution/jobs/{job.id}/bundles/bundle/inventory" + inventory = dict( + attempt_id=str(uuid7()), + dispatch_sequence=2, + expected_revision=None, + version="v1", + files=sorted(files), + failed_sources=sorted(failed_sources), + ) + if api_mode: + response = client.post(url, json=inventory) + assert response.status_code == 200, response.json() + assert client.post(url, json=inventory).json() == response.json() + assert client.post(url, json={**inventory, "failed_sources": []}).status_code == 409 + else: + manager.deactivate_deleted_dags("bundle", discovered) + manager.clear_orphaned_import_errors("bundle", files, failed_sources=failed_sources) + session.expire_all() + assert not session.get(DagModel, "published").is_stale + assert session.get(DagModel, "deleted").is_stale + assert session.get(DagModel, "similar_prefix").is_stale + assert session.scalars(select(ParseImportError.filename)).all() == ["archive.zip/broken.py"] + + if recovery == "archive_removed": + (tmp_path / "archive.zip").unlink() + else: + with zipfile.ZipFile(tmp_path / "archive.zip", "w") as archive: + archive.writestr( + "dag.py" if recovery == "member_present" else "replacement.py", + "from airflow.sdk import DAG\ndag = DAG('published')\n", + ) + recovered = manager._find_files_in_bundle(bundle) + recovered_files = manager._get_observed_filelocs(recovered) + assert not {loc for file in recovered for loc in file.failed_sources} + if api_mode: + inventory.update( + attempt_id=str(uuid7()), + dispatch_sequence=3, + expected_revision=response.json()["revision"], + files=sorted(recovered_files), + failed_sources=[], + ) + complete = client.post(url, json=inventory) + assert complete.status_code == 200, complete.json() + assert complete.json()["revision"] != response.json()["revision"] + else: + manager.deactivate_deleted_dags("bundle", recovered) + manager.clear_orphaned_import_errors("bundle", recovered_files) + session.expire_all() + assert session.get(DagModel, "published").is_stale is (recovery != "member_present") + assert not session.scalars(select(ParseImportError)).all() + + +def test_inventory_revision_fences_late_refreshes(client, session, job): + url = f"/execution/jobs/{job.id}/bundles/bundle/inventory" + inventory = { + "attempt_id": str(uuid7()), + "dispatch_sequence": 1, + "expected_revision": None, + "version": "v1", + "files": ["a.py"], + } + first = client.post(url, json=inventory) + assert first.status_code == 200, first.json() + retry = {**inventory, "attempt_id": str(uuid7()), "dispatch_sequence": 2, "version": "v0"} + assert client.post(url, json=retry).status_code == 409 + retry.update(expected_revision=first.json()["revision"], version="v1") + current = client.post(url, json=retry) + assert current.status_code == 200 + assert current.json()["revision"] == first.json()["revision"] + session.expire_all() + assert session.get(DagBundleModel, "bundle").version == "v1" + + +@pytest.mark.parametrize("bundle_count", [2, 8]) +def test_catalog_loads_teams_with_bounded_queries(session, job, bundle_count): + bundles = [DagBundleModel(name=f"bundle_{index}") for index in range(bundle_count)] + team = Team(name="catalog_team") + bundles[0].teams = [team] + session.add_all(bundles) + session.commit() + token = DagProcessorToken( + id=SESSION_ID, + claims=DagProcessorClaims(exp=1, job_id=job.id, dag_bundles=frozenset(b.name for b in bundles)), + ) + try: + with create_session() as fresh, assert_queries_count(4, session=fresh): + result = get_bundles(job.id, session=fresh, token=token) + assert [item.name for item in result] == [bundle.name for bundle in bundles] + assert [item.team_name for item in result] == [team.name, *([None] * (bundle_count - 1))] + finally: + session.delete(team) + session.commit() + + +@pytest.mark.parametrize("field", ["files", "failed_sources"]) +@pytest.mark.parametrize("files", [["../a.py"], ["a.py", "a.py"], ["/a.py"]]) +def test_invalid_inventory_cannot_remove_dags(client, session, job, body, files, field): + assert client.post(f"/execution/jobs/{job.id}/parse-results", json=body).status_code == 200 + response = client.post( + f"/execution/jobs/{job.id}/bundles/bundle/inventory", + json={ + "attempt_id": str(uuid7()), + "dispatch_sequence": 2, + "expected_revision": None, + "files": [], + field: files, + }, + ) + assert response.status_code == 422 + session.expire_all() + assert not session.get(DagModel, "published").is_stale + + +@pytest.mark.parametrize("rejected", [False, True]) +def test_empty_publication_deactivates_only_after_acceptance(client, session, job, body, rejected): + url = f"/execution/jobs/{job.id}/parse-results" + assert client.post(url, json=body).status_code == 200 + body.update(attempt_id=str(uuid7()), dispatch_sequence=2, serialized_dags=[], source_codes={}) + if rejected: + body["bundle_revision"] = str(uuid7()) + response = client.post(url, json=body) + assert response.status_code == (409 if rejected else 200) + session.expire_all() + assert session.get(DagModel, "published").is_stale is not rejected + + +@pytest.mark.parametrize("kind", ["callbacks", "priority"]) +def test_requested_work_claim_survives_response_loss_and_ack_replay(client, session, job, kind): + if kind == "callbacks": + work = DagProcessorCallback( + priority_weight=5, + callback=DagCallbackRequest( + filepath="a.py", + bundle_name="bundle", + bundle_version="v1", + dag_id="published", + run_id="run", + ), + ) + else: + work = DagPriorityParsingRequest(bundle_name="bundle", relative_fileloc="a.py") + session.add(work) + session.commit() + base = f"/execution/jobs/{job.id}/requested-work/{kind}" + claim = {"claim_id": str(uuid7()), "bundle_names": ["bundle"], "limit": 1} + first = client.post(f"{base}/claim", json=claim) + assert first.status_code == 200, first.json() + assert len(first.json()) == 1 + assert client.post(f"{base}/claim", json=claim).json() == first.json() + assert client.post(f"{base}/claim", json={**claim, "claim_id": str(uuid7())}).json() == [] + ack = {"claim_id": claim["claim_id"], "state": "success"} + key = work.id + url = f"{base}/{key}/ack" + assert client.post(url, json={**ack, "claim_id": str(uuid7())}).status_code == 409 + assert client.post(url, json=ack).status_code == 204 + assert client.post(url, json=ack).status_code == 204 + session.expire_all() + model = DagProcessorCallback if kind == "callbacks" else DagPriorityParsingRequest + assert session.get(model, key) is None + + +@pytest.mark.parametrize("retired", [False, True]) +def test_priority_claim_preserves_live_owner_and_recovers_retired_owner(client, session, job, retired): + owner = Job(job_type="DagProcessorJob", state=JobState.RUNNING) + session.add(owner) + session.flush() + work = DagPriorityParsingRequest(bundle_name="bundle", relative_fileloc="recover.py") + work.processor_job_id = owner.id + work.processor_claim_id = uuid7() + if retired: + owner.end_date = timezone.utcnow() + session.add(work) + session.commit() + response = client.post( + f"/execution/jobs/{job.id}/requested-work/priority/claim", + json={ + "claim_id": str(uuid7()), + "bundle_names": ["bundle"], + "limit": 1, + }, + ) + assert response.status_code == 200, response.json() + assert len(response.json()) == int(retired) + session.delete(work) + session.commit() + + +@pytest.mark.backend("postgres", "mysql") +@pytest.mark.parametrize("kind", ["callbacks", "priority"]) +def test_competing_jobs_claim_requested_work_once(client, exec_app, session, job, kind): + second = Job(job_type="DagProcessorJob", state=JobState.RUNNING) + second.session_id = uuid7() + second.registration_id = uuid7() + second.bundle_names = ["bundle"] + session.add(second) + if kind == "callbacks": + work = DagProcessorCallback( + priority_weight=1, + callback=DagCallbackRequest( + filepath="a.py", bundle_name="bundle", bundle_version=None, dag_id="dag", run_id="run" + ), + ) + else: + work = DagPriorityParsingRequest(bundle_name="bundle", relative_fileloc="a.py") + session.add(work) + session.commit() + identities = {job.id: job.session_id, second.id: second.session_id} + + async def authenticate(request: Request): + job_id = int(request.path_params["job_id"]) + return DagProcessorToken( + id=identities[job_id], + claims=DagProcessorClaims(exp=1, job_id=job_id, dag_bundles=frozenset({"bundle"})), + ) + + exec_app.dependency_overrides[require_auth] = authenticate + ready = Barrier(2) + + def claim(job_id): + ready.wait(timeout=10) + return client.post( + f"/execution/jobs/{job_id}/requested-work/{kind}/claim", + json={"claim_id": str(uuid7()), "bundle_names": ["bundle"]}, + ) + + with ThreadPoolExecutor(max_workers=2) as pool: + futures = [pool.submit(claim, job_id) for job_id in identities] + responses = [future.result(timeout=30) for future in futures] + assert [response.status_code for response in responses] == [200, 200] + assert sorted(len(response.json()) for response in responses) == [0, 1] + + +def test_catalog_and_inventory_support_granted_bundle_with_slash(client, exec_app, session, job): + session.add(DagBundleModel(name="team/bundle")) + session.commit() + + async def authenticate(request: Request): + return DagProcessorToken( + id=SESSION_ID, + claims=DagProcessorClaims(exp=1, job_id=job.id, dag_bundles=frozenset({"team/bundle"})), + ) + + exec_app.dependency_overrides[require_auth] = authenticate + response = client.get(f"/execution/jobs/{job.id}/bundles") + assert response.status_code == 200 + assert [bundle["name"] for bundle in response.json()] == ["team/bundle"] + response = client.post( + f"/execution/jobs/{job.id}/bundles/team%2Fbundle/inventory", + json={"attempt_id": str(uuid7()), "dispatch_sequence": 1, "expected_revision": None, "files": []}, + ) + assert response.status_code == 200, response.json() + + +@pytest.mark.parametrize("kind", ["callbacks", "priority"]) +@pytest.mark.parametrize("retired", [False, True]) +def test_direct_mode_recovers_retired_api_claims_only(session, job, tmp_path, kind, retired): + if kind == "callbacks": + work = DagProcessorCallback( + priority_weight=1, + callback=DagCallbackRequest( + filepath="a.py", bundle_name="bundle", bundle_version=None, dag_id="dag", run_id="run" + ), + ) + else: + work = DagPriorityParsingRequest(bundle_name="bundle", relative_fileloc="a.py") + work.processor_job_id = job.id + work.processor_claim_id = uuid7() + if retired: + job.end_date = timezone.utcnow() + session.add(work) + session.commit() + manager = DagFileProcessorManager(max_runs=1) + manager._dag_bundles = [LocalDagBundle(name="bundle", path=tmp_path)] + result = manager.fetch_callbacks() if kind == "callbacks" else manager.claim_priority_files() + assert len(result) == int(retired) diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2027_02_28/test_jobs.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2027_02_28/test_jobs.py index a03eb96771c90..4fcefa6f4e17b 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2027_02_28/test_jobs.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2027_02_28/test_jobs.py @@ -46,6 +46,10 @@ def authenticate(): [ MISSING_JOB_HEARTBEAT_URL, "/execution/jobs/0/parse-token", + "/execution/jobs/0/parse-results", + "/execution/jobs/0/bundles/bundle/inventory", + "/execution/jobs/0/requested-work/priority/claim", + "/execution/jobs/0/requested-work/priority/id/ack", ], ) def test_old_version_returns_404(self, client, path): @@ -62,6 +66,10 @@ def test_head_version_routes_to_endpoint(self, client): assert response.status_code == 404 assert response.json()["detail"]["reason"] == "not_found" + def test_old_version_has_no_bundle_catalog(self, client): + client.headers["Airflow-API-Version"] = "2026-10-30" + assert client.get("/execution/jobs/0/bundles").status_code == 404 + def test_head_version_routes_to_parse_token_exchange(self, client): response = client.post( "/execution/jobs/0/parse-token", 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 77de6f928d33e..b0a69dd9378f5 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 @@ 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 @@ def test_memray_traces_the_job_and_not_the_parent_process(self, mock_runner, moc 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/bundles/test_dag_bundle_manager.py b/airflow-core/tests/unit/dag_processing/bundles/test_dag_bundle_manager.py index 8938d2d33b0fc..f8599034e9c85 100644 --- a/airflow-core/tests/unit/dag_processing/bundles/test_dag_bundle_manager.py +++ b/airflow-core/tests/unit/dag_processing/bundles/test_dag_bundle_manager.py @@ -277,6 +277,63 @@ def clear_db(): clear_db_dag_bundles() +@pytest.mark.db_test +@pytest.mark.parametrize("existing", [False, True]) +@conf_vars({("core", "load_examples"): "False"}) +@patch.object(BasicBundle, "__init__", autospec=True, side_effect=AssertionError("constructed too early")) +def test_metadata_sync_does_not_construct_bundles_or_clear_urls(constructor, clear_db, session, existing): + if existing: + bundle = DagBundleModel(name="my-test-bundle") + bundle.signed_url_template = "saved" + bundle.template_params = {"x": 1} + session.add(bundle) + session.commit() + with conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps(BASIC_BUNDLE_CONFIG)}): + DagBundlesManager().sync_bundles_to_db(include_bundle_urls=False) + session.expire_all() + bundle = session.get(DagBundleModel, "my-test-bundle") + assert bundle.active + if existing: + assert (bundle.signed_url_template, bundle.template_params) == ("saved", {"x": 1}) + constructor.assert_not_called() + + +@pytest.mark.parametrize("operation", ["get", "list", "sync"]) +@pytest.mark.db_test +@conf_vars({("core", "load_examples"): "False"}) +def test_bundle_construction_and_url_resolution_use_the_bundle_context(clear_db, operation): + active = [] + seen = [] + + @contextlib.contextmanager + def bundle_context(name): + active.append(name) + try: + yield + finally: + active.pop() + + def read_secret(*args, **kwargs): + seen.append(active[-1]) + + config = [*BASIC_BUNDLE_CONFIG, {"name": "excluded", "classpath": "missing.Bundle", "kwargs": {}}] + with ( + conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps(config)}), + patch.object(BasicBundle, "__init__", autospec=True, side_effect=read_secret), + patch.object(BasicBundle, "view_url_template", autospec=True, side_effect=read_secret), + ): + manager = DagBundlesManager(bundle_names=["my-test-bundle"], bundle_context=bundle_context) + if operation == "get": + manager.get_bundle("my-test-bundle") + elif operation == "list": + list(manager.get_all_dag_bundles()) + else: + manager.sync_bundles_to_db(deactivate_missing=False) + assert seen + assert set(seen) == {"my-test-bundle"} + assert active == [] + + @pytest.mark.db_test @conf_vars({("core", "LOAD_EXAMPLES"): "False"}) def test_sync_bundles_to_db(clear_db, session): @@ -1780,7 +1837,7 @@ class TestSkippedRowLifecycle: A legacy 2.x row with an unconfigured bundle and a fileloc outside every configured bundle is intentionally left untouched by the repair. The - inactive-bundle branch of ``deactivate_stale_dags`` then marks it stale, + scheduler's ``cleanup_processor_metadata`` then marks it stale, and a later successful parse from a now-configured bundle must restore it end-to-end. Each unit-level test exercises one stage of this lifecycle; this test pins the whole sequence so a future refactor can't silently @@ -1791,7 +1848,7 @@ def test_skipped_row_is_recoverable_after_operator_fix( self, clear_dags_and_bundles, session, tmp_path ) -> None: from airflow.dag_processing.collection import update_dag_parsing_results_in_db - from airflow.dag_processing.manager import DagFileProcessorManager + from airflow.dag_processing.maintenance import cleanup_processor_metadata from airflow.sdk import DAG from airflow.serialization.serialized_objects import LazyDeserializedDAG @@ -1846,10 +1903,8 @@ def test_skipped_row_is_recoverable_after_operator_fix( assert refreshed.relative_fileloc is None assert refreshed.is_stale is False - # Stale-Dag scan deactivates the row via the inactive-bundle - # branch, which runs before the NULL relative_fileloc guard. - dfp_manager = DagFileProcessorManager(max_runs=1, processor_timeout=10 * 60) - dfp_manager.deactivate_stale_dags(last_parsed={}) + # Scheduler maintenance deactivates the row because its bundle is inactive. + cleanup_processor_metadata() session.expire_all() stale_row = session.get(DagModel, "legacy_dag") assert stale_row.is_stale is True @@ -1868,7 +1923,7 @@ def test_skipped_row_is_recoverable_after_operator_fix( {}, None, set(), - session, + session=session, ) session.commit() diff --git a/airflow-core/tests/unit/dag_processing/test_api_client.py b/airflow-core/tests/unit/dag_processing/test_api_client.py index 631eecd4d585c..9e685a38463af 100644 --- a/airflow-core/tests/unit/dag_processing/test_api_client.py +++ b/airflow-core/tests/unit/dag_processing/test_api_client.py @@ -21,11 +21,16 @@ from contextvars import Context from unittest import mock -import httpx +import httpx2 import jwt import pytest from tenacity import wait_none +from uuid6 import uuid7 +from airflow.api_fastapi.execution_api.datamodels.dag_parsing import ( + DagBundleInventoryBody, + DagParseResultBody, +) from airflow.api_fastapi.execution_api.datamodels.job import DagParseTokenBody, JobState, TerminalJobState from airflow.dag_processing import api_client from airflow.dag_processing.api_client import ( @@ -50,7 +55,7 @@ def make_job_token(**claims): def make_registration_response(*, job_id=1, token=None, expires_in=300, **kwargs): - return httpx.Response( + return httpx2.Response( 201, json={"job_id": job_id, "token": token or make_job_token(job_id=job_id), "expires_in": expires_in}, **kwargs, @@ -94,7 +99,7 @@ def handle(request): base_url="http://api/execution/", token_file=token_file, hostname="processor-1", - transport=httpx.MockTransport(handle), + transport=httpx2.MockTransport(handle), **kwargs, ) clients.append(client) @@ -108,8 +113,8 @@ def handle(request): def test_registration_and_runtime_use_distinct_credentials(make_client): client, requests = make_client( make_registration_response(headers={"Refreshed-API-Token": "wrong-token"}), - httpx.Response(200, json={"state": "restarting"}, headers={"Refreshed-API-Token": "wrong-token"}), - httpx.Response(200, json={"state": "running"}), + httpx2.Response(200, json={"state": "restarting"}, headers={"Refreshed-API-Token": "wrong-token"}), + httpx2.Response(200, json={"state": "running"}), unixname="airflow", bundle_names=["bundle-a"], ) @@ -133,17 +138,17 @@ def test_registration_and_runtime_use_distinct_credentials(make_client): def test_unread_non_json_error_body_is_a_plain_status_error(make_client): client, _ = make_client( - httpx.Response( - 409, headers={"content-type": "text/html"}, stream=httpx.ByteStream(b"proxy error") + httpx2.Response( + 409, headers={"content-type": "text/html"}, stream=httpx2.ByteStream(b"proxy error") ) ) - with pytest.raises(httpx.HTTPStatusError): + with pytest.raises(httpx2.HTTPStatusError): client.register_job(retry=False) def test_registration_identity_survives_response_loss_but_not_a_process_restart(make_client): - first, requests = make_client(httpx.ReadError("Lost acknowledgment"), make_registration_response()) + first, requests = make_client(httpx2.ReadError("Lost acknowledgment"), make_registration_response()) second, _ = make_client() first.register_job() @@ -160,16 +165,16 @@ def test_registration_reloads_a_rotated_credential_after_authentication_failure( ): def rotate(request): token_file.write_text("session-2") - return httpx.Response(status) + return httpx2.Response(status) client, requests = make_client( - rotate, make_registration_response() if succeeds else httpx.Response(status) + rotate, make_registration_response() if succeeds else httpx2.Response(status) ) if succeeds: client.register_job() else: - with pytest.raises(httpx.HTTPStatusError): + with pytest.raises(httpx2.HTTPStatusError): client.register_job() assert [request.headers["Authorization"] for request in requests] == [ @@ -181,10 +186,10 @@ def rotate(request): @pytest.mark.parametrize("status", [401, 403, 404, 409]) def test_registration_errors_do_not_start_another_job(make_client, status): - client, requests = make_client(httpx.Response(status)) + client, requests = make_client(httpx2.Response(status)) registration_id = client.registration_id - with pytest.raises(httpx.HTTPStatusError) as error: + with pytest.raises(httpx2.HTTPStatusError) as error: client.register_job() assert error.value.response.status_code == status @@ -195,7 +200,7 @@ def test_registration_errors_do_not_start_another_job(make_client, status): def test_registration_retry_after_a_crashed_job_keeps_its_identity(make_client): client, requests = make_client( - httpx.Response(409, json={"detail": {"reason": "job_running"}}), + httpx2.Response(409, json={"detail": {"reason": "job_running"}}), make_registration_response(), ) registration_id = client.registration_id @@ -212,9 +217,9 @@ def test_registration_retry_after_a_crashed_job_keeps_its_identity(make_client): def test_rotation_is_detected_within_the_cache_interval(make_client, token_file, clock): client, requests = make_client( make_registration_response(), - httpx.Response(204), + httpx2.Response(204), make_registration_response(token=make_job_token(iat=30, exp=330)), - httpx.Response(204), + httpx2.Response(204), ) client.register_job() token_file.write_text("session-2") @@ -235,7 +240,7 @@ def test_rotation_is_detected_within_the_cache_interval(make_client, token_file, def test_job_token_renewal_uses_monotonic_time(make_client, clock, time_machine): client, requests = make_client( - make_registration_response(), httpx.Response(204), make_registration_response(), httpx.Response(204) + make_registration_response(), httpx2.Response(204), make_registration_response(), httpx2.Response(204) ) client.register_job() time_machine.move_to("2030-01-01", tick=False) @@ -251,10 +256,10 @@ def test_job_token_renewal_uses_monotonic_time(make_client, clock, time_machine) @pytest.mark.parametrize("status", [401, 403, 404, 409]) def test_runtime_errors_do_not_renew_credentials(make_client, status): - client, requests = make_client(make_registration_response(), httpx.Response(status)) + client, requests = make_client(make_registration_response(), httpx2.Response(status)) client.register_job() - with pytest.raises(httpx.HTTPStatusError) as error: + with pytest.raises(httpx2.HTTPStatusError) as error: client.get("resource") assert error.value.response.status_code == status @@ -262,14 +267,14 @@ def test_runtime_errors_do_not_renew_credentials(make_client, status): @pytest.mark.parametrize("expired", [False, True]) -@pytest.mark.parametrize("failure", [httpx.ReadError("Unavailable"), httpx.Response(503)]) +@pytest.mark.parametrize("failure", [httpx2.ReadError("Unavailable"), httpx2.Response(503)]) def test_heartbeat_does_not_retry_transport_failures(make_client, clock, expired, failure): client, requests = make_client(make_registration_response(), failure) client.register_job() if expired: clock.return_value = 300 - with pytest.raises(httpx.HTTPError): + with pytest.raises(httpx2.HTTPError): client.heartbeat() assert len(requests) == 2 @@ -334,7 +339,7 @@ def test_runtime_operations_require_registration(make_client, operation): def test_completion_recovers_a_lost_acknowledgment_with_the_same_token_and_outcome(make_client): client, requests = make_client( - make_registration_response(), httpx.ReadError("Lost acknowledgment"), httpx.Response(204) + make_registration_response(), httpx2.ReadError("Lost acknowledgment"), httpx2.Response(204) ) client.register_job() @@ -352,11 +357,11 @@ def test_completion_retries_remain_possible_after_transport_retries_are_exhauste ): client, requests = make_client( make_registration_response(), - *[httpx.ReadError("Lost acknowledgment") for _ in range(API_RETRIES)], - httpx.Response(204), + *[httpx2.ReadError("Lost acknowledgment") for _ in range(API_RETRIES)], + httpx2.Response(204), ) client.register_job() - with pytest.raises(httpx.ReadError): + with pytest.raises(httpx2.ReadError): client.complete_job(TerminalJobState.FAILED) token_file.unlink() clock.return_value = 290 @@ -370,12 +375,12 @@ def test_completion_retries_remain_possible_after_transport_retries_are_exhauste @pytest.mark.parametrize("status", [204, 401, 403, 404]) def test_completion_stops_normal_work_even_when_acknowledgment_is_uncertain(make_client, status): - client, requests = make_client(make_registration_response(), httpx.Response(status)) + client, requests = make_client(make_registration_response(), httpx2.Response(status)) client.register_job() if status == 204: client.complete_job(TerminalJobState.SUCCESS) else: - with pytest.raises(httpx.HTTPStatusError): + with pytest.raises(httpx2.HTTPStatusError): client.complete_job(TerminalJobState.SUCCESS) with pytest.raises(RuntimeError, match="only completion retries"): @@ -395,17 +400,17 @@ def test_invalid_reload_interval_is_rejected(make_client, interval): @pytest.mark.parametrize("operation", ["runtime", "heartbeat"]) @pytest.mark.parametrize( - "failure", [httpx.ReadTimeout("Unavailable"), httpx.Response(503), httpx.Response(403)] + "failure", [httpx2.ReadTimeout("Unavailable"), httpx2.Response(503), httpx2.Response(403)] ) def test_failed_early_renewal_keeps_using_the_valid_token(make_client, clock, operation, failure): response_body = {"state": "running"} if operation == "heartbeat" else {"key": "key", "value": "value"} client, requests = make_client( make_registration_response(), failure, - httpx.Response(200, json=response_body), - httpx.Response(200, json=response_body), + httpx2.Response(200, json=response_body), + httpx2.Response(200, json=response_body), make_registration_response(), - httpx.Response(204), + httpx2.Response(204), ) client.register_job() clock.return_value = 240 @@ -429,7 +434,7 @@ def test_failed_early_renewal_keeps_using_the_valid_token(make_client, clock, op @pytest.mark.parametrize("missing", [False, True]) def test_unreadable_rotated_file_does_not_discard_a_valid_token(make_client, token_file, clock, missing): client, requests = make_client( - make_registration_response(), httpx.Response(200, json={"state": "running"}) + make_registration_response(), httpx2.Response(200, json={"state": "running"}) ) client.register_job() token_file.unlink() if missing else token_file.write_text("") @@ -442,10 +447,10 @@ def test_unreadable_rotated_file_does_not_discard_a_valid_token(make_client, tok def test_renewal_that_outlasts_the_token_requires_synchronous_recovery(make_client, clock): def time_out(request): clock.return_value = 300 - raise httpx.ReadTimeout("Outlasted the token") + raise httpx2.ReadTimeout("Outlasted the token") client, requests = make_client( - make_registration_response(), time_out, make_registration_response(), httpx.Response(204) + make_registration_response(), time_out, make_registration_response(), httpx2.Response(204) ) client.register_job() clock.return_value = 299 @@ -459,10 +464,10 @@ def time_out(request): def test_retired_registration_signals_restart_without_blocking_valid_requests(make_client, clock, complete): client, requests = make_client( make_registration_response(), - httpx.Response(409, json={"detail": {"reason": "registration_retired"}}), - httpx.Response(204), - httpx.Response(204), - httpx.Response(204), + httpx2.Response(409, json={"detail": {"reason": "registration_retired"}}), + httpx2.Response(204), + httpx2.Response(204), + httpx2.Response(204), ) client.register_job() clock.return_value = 240 @@ -484,8 +489,8 @@ def test_retired_registration_signals_restart_without_blocking_valid_requests(ma "payload", [{"detail": "conflict"}, ["unexpected"], {"detail": {"reason": "registration_conflict"}}] ) def test_other_conflicts_are_not_mistaken_for_retirement(make_client, payload): - client, _ = make_client(httpx.Response(409, json=payload)) - with pytest.raises(httpx.HTTPStatusError): + client, _ = make_client(httpx2.Response(409, json=payload)) + with pytest.raises(httpx2.HTTPStatusError): client.register_job() assert not client.restart_required @@ -493,7 +498,7 @@ def test_other_conflicts_are_not_mistaken_for_retirement(make_client, payload): def test_closed_job_signals_restart(make_client): client, _ = make_client( make_registration_response(), - httpx.Response(403, json={"detail": {"reason": "job_closed", "message": "replaced"}}), + httpx2.Response(403, json={"detail": {"reason": "job_closed", "message": "replaced"}}), ) client.register_job() @@ -507,10 +512,10 @@ def test_closed_job_signals_restart(make_client): "payload", [{"detail": "Invalid auth token"}, {"detail": {"reason": "bundle_not_granted"}}] ) def test_other_authorization_failures_are_not_mistaken_for_a_closed_job(make_client, payload): - client, _ = make_client(make_registration_response(), httpx.Response(403, json=payload)) + client, _ = make_client(make_registration_response(), httpx2.Response(403, json=payload)) client.register_job() - with pytest.raises(httpx.HTTPStatusError): + with pytest.raises(httpx2.HTTPStatusError): client.get("resource") assert not client.restart_required @@ -519,9 +524,9 @@ def test_other_authorization_failures_are_not_mistaken_for_a_closed_job(make_cli def test_sdk_helpers_keep_bundle_context_for_reads_and_writes(make_client): client, requests = make_client( make_registration_response(), - httpx.Response(200, json={"key": "key", "value": "value"}), - httpx.Response(204), - httpx.Response( + httpx2.Response(200, json={"key": "key", "value": "value"}), + httpx2.Response(204), + httpx2.Response( 200, json={ "conn_id": "connection", @@ -549,14 +554,14 @@ def test_sdk_helpers_keep_bundle_context_for_reads_and_writes(make_client): def test_bundle_context_preserves_headers_and_restores_context_on_error(make_client): client, requests = make_client( make_registration_response(), - httpx.Response(204), - httpx.ReadError("Subprocess failure"), - *[httpx.Response(204) for _ in range(3)], + httpx2.Response(204), + httpx2.ReadError("Subprocess failure"), + *[httpx2.Response(204) for _ in range(3)], ) client.register_job() with client.use_bundle("bundle-a"): client.put("resource", content=b"data", headers={"content-type": "text/plain", "custom": "value"}) - with pytest.raises(httpx.ReadError), client.use_bundle("bundle-b"): + with pytest.raises(httpx2.ReadError), client.use_bundle("bundle-b"): client.request( "PUT", "resource", @@ -590,18 +595,18 @@ def test_empty_bundle_context_is_rejected(make_client): @pytest.mark.parametrize("lost_acknowledgment", [False, True]) def test_completion_renews_an_expired_token_if_the_job_is_still_open(make_client, clock, lost_acknowledgment): failures = ( - [httpx.ReadError("Lost acknowledgment") for _ in range(API_RETRIES)] if lost_acknowledgment else [] + [httpx2.ReadError("Lost acknowledgment") for _ in range(API_RETRIES)] if lost_acknowledgment else [] ) renewed_token = make_job_token(iat=300, exp=600) client, requests = make_client( make_registration_response(), *failures, make_registration_response(token=renewed_token), - httpx.Response(204), + httpx2.Response(204), ) client.register_job() if lost_acknowledgment: - with pytest.raises(httpx.ReadError): + with pytest.raises(httpx2.ReadError): client.complete_job(TerminalJobState.SUCCESS) clock.return_value = 300 @@ -616,7 +621,8 @@ def test_completion_renews_an_expired_token_if_the_job_is_still_open(make_client def test_completion_does_not_claim_success_when_expired_registration_is_retired(make_client, clock): client, requests = make_client( - make_registration_response(), httpx.Response(409, json={"detail": {"reason": "registration_retired"}}) + make_registration_response(), + httpx2.Response(409, json={"detail": {"reason": "registration_retired"}}), ) client.register_job() clock.return_value = 300 @@ -634,19 +640,19 @@ def test_completion_does_not_claim_success_when_expired_registration_is_retired( def test_completion_recovers_if_the_token_expires_during_the_request(make_client, clock, status, succeeds): def expire(request): clock.return_value += 300 - return httpx.Response(status) + return httpx2.Response(status) client, requests = make_client( make_registration_response(), expire, make_registration_response(), - httpx.Response(204) if succeeds else expire, + httpx2.Response(204) if succeeds else expire, ) client.register_job() if succeeds: client.complete_job(TerminalJobState.SUCCESS) else: - with pytest.raises(httpx.HTTPStatusError): + with pytest.raises(httpx2.HTTPStatusError): client.complete_job(TerminalJobState.SUCCESS) assert requests[0].content == requests[2].content @@ -674,7 +680,7 @@ def test_control_contracts_do_not_require_new_sdk_models(monkeypatch, token_file def test_secrets_comms_answers_lookups_for_the_selected_bundle(make_client): client, requests = make_client( - make_registration_response(), httpx.Response(200, json={"key": "my_key", "value": "my_value"}) + make_registration_response(), httpx2.Response(200, json={"key": "my_key", "value": "my_value"}) ) client.register_job() comms = DagProcessorSecretsComms(client) @@ -725,7 +731,7 @@ def make_parse_token_response(*, expires_in=60, **claims): }, SECRET, ) - return httpx.Response(200, json={"token": token, "expires_in": expires_in}) + return httpx2.Response(200, json={"token": token, "expires_in": expires_in}) def test_parse_requests_exchange_once_and_restore_the_manager_credential(make_client): @@ -733,9 +739,9 @@ def test_parse_requests_exchange_once_and_restore_the_manager_credential(make_cl client, requests = make_client( make_registration_response(), exchanged, - httpx.Response(200, json={"key": "key", "value": "value"}), - httpx.Response(200, json={"key": "key", "value": "value"}), - httpx.Response(200, json={"state": "running"}), + httpx2.Response(200, json={"key": "key", "value": "value"}), + httpx2.Response(200, json={"key": "key", "value": "value"}), + httpx2.Response(200, json={"state": "running"}), ) client.register_job() context = make_parse_context() @@ -760,18 +766,75 @@ def test_parse_requests_exchange_once_and_restore_the_manager_credential(make_cl def test_parse_exchange_recovers_a_lost_response_without_changing_attempt(make_client): client, requests = make_client( make_registration_response(), - httpx.ReadError("exchange acknowledgment lost"), + httpx2.ReadError("exchange acknowledgment lost"), make_parse_token_response(), - httpx.Response(200), + httpx2.Response(200), ) client.register_job() with client.use_parse(make_parse_context()): + with pytest.raises(httpx2.ReadError): + client.get("variables/key") client.get("variables/key") assert requests[1].content == requests[2].content assert len(requests) == 4 +@pytest.mark.parametrize( + ("request_kwargs", "expected_timeout"), + [ + pytest.param({}, httpx2.Timeout(5), id="client-default"), + pytest.param({"timeout": 2.5}, httpx2.Timeout(2.5), id="explicit-timeout"), + pytest.param({"timeout": None}, httpx2.Timeout(None), id="disabled-timeout"), + pytest.param( + {"timeout": httpx2.Timeout(4, connect=2)}, + httpx2.Timeout(4, connect=2), + id="per-phase-timeout", + ), + ], +) +def test_parse_requests_make_one_attempt_with_the_request_timeout( + make_client, request_kwargs, expected_timeout +): + client, requests = make_client( + make_registration_response(), make_parse_token_response(), httpx2.Response(503), timeout=5 + ) + client.register_job() + + with client.use_parse(make_parse_context()), pytest.raises(httpx2.HTTPStatusError): + client.get("variables/key", **request_kwargs) + + assert len(requests) == 3 + assert requests[1].extensions["timeout"] == expected_timeout.as_dict() + assert requests[2].extensions["timeout"] == expected_timeout.as_dict() + + +@pytest.mark.parametrize(("expired", "tick"), [("parse", 60), ("job", 300)]) +def test_required_token_renewals_use_the_runtime_request_timeout(make_client, clock, expired, tick): + client, requests = make_client( + make_registration_response(), + make_parse_token_response(), + httpx2.Response(200), + *([make_registration_response()] if expired == "job" else []), + make_parse_token_response(), + httpx2.Response(200), + timeout=5, + ) + client.register_job() + + with client.use_parse(make_parse_context()): + client.get("variables/key") + clock.return_value = tick + client.get("variables/key", timeout=2.5) + + assert [request.url.path for request in requests[3:]] == [ + *(["/execution/jobs"] if expired == "job" else []), + "/execution/jobs/1/parse-token", + "/execution/variables/key", + ] + assert all(request.extensions["timeout"] == httpx2.Timeout(2.5).as_dict() for request in requests[3:]) + + def test_interleaved_parses_keep_distinct_file_credentials(make_client): first_response = make_parse_token_response() second_response = make_parse_token_response( @@ -782,10 +845,10 @@ def test_interleaved_parses_keep_distinct_file_credentials(make_client): client, requests = make_client( make_registration_response(), first_response, - httpx.Response(200), + httpx2.Response(200), second_response, - httpx.Response(200), - httpx.Response(200), + httpx2.Response(200), + httpx2.Response(200), ) client.register_job() second = make_parse_context( @@ -812,10 +875,10 @@ def test_parse_credential_renews_with_the_same_file_and_attempt(make_client, clo client, requests = make_client( make_registration_response(), make_parse_token_response(), - httpx.Response(200), - httpx.Response(200), + httpx2.Response(200), + httpx2.Response(200), make_parse_token_response(iat=renew_at, exp=renew_at + 60), - httpx.Response(200), + httpx2.Response(200), ) client.register_job() with client.use_parse(make_parse_context()): @@ -832,7 +895,7 @@ def test_parse_credential_renews_with_the_same_file_and_attempt(make_client, clo @pytest.mark.parametrize( "failure", - [httpx.ReadError("unavailable"), httpx.Response(503), httpx.Response(200, json={"token": "invalid"})], + [httpx2.ReadError("unavailable"), httpx2.Response(503), httpx2.Response(200, json={"token": "invalid"})], ) @pytest.mark.parametrize("recovers_at_expiry", [False, True]) def test_failed_early_parse_renewal_keeps_the_valid_token_until_expiry( @@ -842,12 +905,12 @@ def test_failed_early_parse_renewal_keeps_the_valid_token_until_expiry( client, requests = make_client( make_registration_response(), token, - httpx.Response(200), + httpx2.Response(200), failure, - httpx.Response(200), - httpx.Response(200), - make_parse_token_response(iat=60, exp=120) if recovers_at_expiry else httpx.Response(503), - httpx.Response(200), + httpx2.Response(200), + httpx2.Response(200), + make_parse_token_response(iat=60, exp=120) if recovers_at_expiry else httpx2.Response(503), + httpx2.Response(200), ) client.register_job() with client.use_parse(make_parse_context()): @@ -858,7 +921,7 @@ def test_failed_early_parse_renewal_keeps_the_valid_token_until_expiry( if recovers_at_expiry: client.request("GET", "variables/key", retry=False) else: - with pytest.raises(httpx.HTTPStatusError): + with pytest.raises(httpx2.HTTPStatusError): client.request("GET", "variables/key", retry=False) assert [request.url.path for request in requests[:7]] == [ @@ -881,15 +944,15 @@ def test_failed_early_parse_renewal_keeps_the_valid_token_until_expiry( def test_parse_renewal_that_outlasts_the_token_requires_synchronous_recovery(make_client, clock): def expire(request): clock.return_value = 60 - return httpx.Response(503) + return httpx2.Response(503) client, requests = make_client( make_registration_response(), make_parse_token_response(), - httpx.Response(200), + httpx2.Response(200), expire, make_parse_token_response(iat=60, exp=120), - httpx.Response(200), + httpx2.Response(200), ) client.register_job() with client.use_parse(make_parse_context()): @@ -906,15 +969,15 @@ def test_early_parse_renewal_does_not_hide_authorization_errors(make_client, clo client, requests = make_client( make_registration_response(), make_parse_token_response(), - httpx.Response(200), - httpx.Response(403, json={"detail": {"reason": reason}}), + httpx2.Response(200), + httpx2.Response(403, json={"detail": {"reason": reason}}), ) client.register_job() with client.use_parse(make_parse_context()): client.get("variables/key") clock.return_value = 48 with pytest.raises( - DagProcessorRegistrationRetired if reason == "job_closed" else httpx.HTTPStatusError + DagProcessorRegistrationRetired if reason == "job_closed" else httpx2.HTTPStatusError ): client.get("variables/key") @@ -940,8 +1003,8 @@ def test_invalid_exchanged_lifetimes_never_reach_a_runtime_endpoint(make_client, def test_exchange_failure_restores_context_and_signals_a_closed_job(make_client): client, requests = make_client( make_registration_response(), - httpx.Response(403, json={"detail": {"reason": "job_closed"}}), - httpx.Response(200, json={"state": "restarting"}), + httpx2.Response(403, json={"detail": {"reason": "job_closed"}}), + httpx2.Response(200, json={"state": "restarting"}), ) client.register_job() with client.use_parse(make_parse_context()), pytest.raises(DagProcessorRegistrationRetired): @@ -950,3 +1013,69 @@ def test_exchange_failure_restores_context_and_signals_a_closed_job(make_client) assert client.restart_required assert client.heartbeat() == JobState.RESTARTING assert requests[-1].headers["Authorization"] == f"Bearer {make_job_token()}" + + +@pytest.mark.parametrize("expired", [False, True]) +@pytest.mark.parametrize("failure", [httpx2.ReadTimeout("Lost response"), httpx2.Response(503)]) +def test_publication_attempt_does_not_sleep_through_heartbeats(make_client, clock, expired, failure): + body = DagParseResultBody( + attempt_id=uuid7(), + dispatch_sequence=1, + bundle_name="bundle-a", + relative_fileloc="dag.py", + parse_duration=0.1, + serialized_dags=[], + source_codes={}, + ) + responses = [make_registration_response()] + if expired: + responses.append(make_registration_response()) + responses.extend( + [ + failure, + httpx2.Response(200, json={"state": "running"}), + httpx2.Response( + 200, json={"attempt_id": str(body.attempt_id), "accepted_at": "2026-10-05T12:00:00Z"} + ), + ] + ) + client, requests = make_client(*responses) + client.register_job() + if expired: + clock.return_value = 300 + with pytest.raises(httpx2.HTTPError): + client.publish_parse_result(body) + assert len(requests) == (3 if expired else 2) + publication_timeout = requests[-1].extensions["timeout"] + assert publication_timeout["connect"] <= 1 + assert publication_timeout["pool"] <= 1 + assert publication_timeout["read"] == client.timeout.read > 1 + if expired: + assert all(value <= 1 for value in requests[1].extensions["timeout"].values()) + assert client.heartbeat() == JobState.RUNNING + assert client.publish_parse_result(body).attempt_id == body.attempt_id + assert requests[-1].content == requests[-3].content + + +def test_inventory_waits_for_a_slow_write_but_connects_quickly(make_client): + client, requests = make_client( + make_registration_response(), + httpx2.Response( + 200, + json={ + "attempt_id": str(uuid7()), + "accepted_at": "2026-10-05T12:00:00Z", + "revision": str(uuid7()), + }, + ), + ) + client.register_job() + + client.publish_inventory( + "bundle-a", + DagBundleInventoryBody(attempt_id=uuid7(), dispatch_sequence=1, expected_revision=None, files=[]), + ) + + timeout = requests[-1].extensions["timeout"] + assert timeout["connect"] <= 1 + assert timeout["read"] == client.timeout.read > 1 diff --git a/airflow-core/tests/unit/dag_processing/test_collection.py b/airflow-core/tests/unit/dag_processing/test_collection.py index 6e9a61c8756ca..cc9cd8c1d5ccc 100644 --- a/airflow-core/tests/unit/dag_processing/test_collection.py +++ b/airflow-core/tests/unit/dag_processing/test_collection.py @@ -727,7 +727,7 @@ class CustomAuthManager(FabAuthManager): with patch.object(FabAuthManager, "sync_dag_perms", autospec=True) as mock_sync_dag_perms: assert ( airflow.dag_processing.collection._serialize_dag_capturing_errors( - dag, "testing", session, None + dag, "testing", session=session, bundle_version=None ) == [] ) @@ -742,7 +742,9 @@ def test_serialize_without_auth_manager_sync_hook( dag = DAG(dag_id="test") assert ( - airflow.dag_processing.collection._serialize_dag_capturing_errors(dag, "testing", session, None) + airflow.dag_processing.collection._serialize_dag_capturing_errors( + dag, "testing", session=session, bundle_version=None + ) == [] ) @@ -766,7 +768,9 @@ def _sync_to_db(): time_machine.shift(20) import_errors = {} - update_dag_parsing_results_in_db("testing", None, [dag], import_errors, None, set(), session) + update_dag_parsing_results_in_db( + "testing", None, [dag], import_errors, None, set(), session=session + ) assert not import_errors _sync_to_db() @@ -816,9 +820,9 @@ def test_sync_to_db_is_retried( # Test that 3 attempts were made to run 'DAG.bulk_write_to_db' successfully mock_bulk_write_to_db.assert_has_calls( [ - mock.call("testing", None, mock.ANY, None, session=mock.ANY), - mock.call("testing", None, mock.ANY, None, session=mock.ANY), - mock.call("testing", None, mock.ANY, None, session=mock.ANY), + mock.call("testing", None, mock.ANY, None, enforce_bundle_ownership=False, session=mock.ANY), + mock.call("testing", None, mock.ANY, None, enforce_bundle_ownership=False, session=mock.ANY), + mock.call("testing", None, mock.ANY, None, enforce_bundle_ownership=False, session=mock.ANY), ] ) # Assert that rollback is called twice (i.e. whenever OperationalError occurs) @@ -1023,7 +1027,9 @@ def test_parse_time_written_to_db_on_sync(self, testing_dag_bundle, session): parse_duration = 1.25 dag = DAG(dag_id="test") - update_dag_parsing_results_in_db("testing", None, [dag], dict(), parse_duration, set(), session) + update_dag_parsing_results_in_db( + "testing", None, [dag], dict(), parse_duration, set(), session=session + ) dag_model: DagModel = session.get(DagModel, (dag.dag_id,)) assert dag_model.last_parse_duration == parse_duration @@ -1047,7 +1053,7 @@ def test_timetable_asset_gated_written_to_db_on_sync(self, testing_dag_bundle, s {}, None, set(), - session, + session=session, ) assert session.get(DagModel, gated_dag.dag_id).timetable_asset_gated is True @@ -1072,7 +1078,14 @@ def test_serialized_dag_errors_are_import_errors( import_errors = {} update_dag_parsing_results_in_db( - "testing", None, [dag], import_errors, None, set(), session, files_parsed={("testing", "abc.py")} + "testing", + None, + [dag], + import_errors, + None, + set(), + session=session, + files_parsed={("testing", "abc.py")}, ) assert "SerializationError" in caplog.text @@ -1130,7 +1143,7 @@ def test_import_error_persist_for_invalid_access_control_role( # the DAG processor should raise an import error when processing the DAG above. import_errors = {} # run the DAG parsing. - update_dag_parsing_results_in_db("testing", None, [dag], import_errors, None, set(), session) + update_dag_parsing_results_in_db("testing", None, [dag], import_errors, None, set(), session=session) # expect to get an error with "role does not exist" message. err = import_errors.get(("testing", dag.relative_fileloc)) assert "AirflowException" in err @@ -1160,7 +1173,7 @@ def test_import_error_persist_for_invalid_access_control_role( # run the update again. Even though the DAG is not updated, the processor should raise import error since the access control is not fixed. time_machine.move_to(tz.datetime(2020, 1, 5, 0, 0, 5), tick=False) - update_dag_parsing_results_in_db("testing", None, [dag], dict(), None, set(), session) + update_dag_parsing_results_in_db("testing", None, [dag], dict(), None, set(), session=session) dag_model: DagModel = session.get(DagModel, (dag.dag_id,)) # the DAG should contain an import error. @@ -1190,7 +1203,7 @@ def test_import_error_persist_for_invalid_access_control_role( # run the update again, but the incorrect access control configuration is removed. time_machine.move_to(tz.datetime(2020, 1, 5, 0, 0, 10), tick=False) dag.access_control = None - update_dag_parsing_results_in_db("testing", None, [dag], dict(), None, set(), session) + update_dag_parsing_results_in_db("testing", None, [dag], dict(), None, set(), session=session) dag_model: DagModel = session.get(DagModel, (dag.dag_id,)) # the import error should be cleared. @@ -1538,13 +1551,13 @@ def test_dagmodel_properties(self, attrs, expected, session, time_machine, testi def test_existing_dag_is_paused_upon_creation(self, testing_dag_bundle, session, dag_maker): with dag_maker("dag_paused", schedule=None) as dag: ... - update_dag_parsing_results_in_db("testing", None, [dag], {}, 0.1, set(), session) + update_dag_parsing_results_in_db("testing", None, [dag], {}, 0.1, set(), session=session) orm_dag = session.get(DagModel, ("dag_paused",)) assert orm_dag.is_paused is False with dag_maker("dag_paused", schedule=None, is_paused_upon_creation=True) as dag: ... - update_dag_parsing_results_in_db("testing", None, [dag], {}, 0.1, set(), session) + update_dag_parsing_results_in_db("testing", None, [dag], {}, 0.1, set(), session=session) # Since the dag existed before, it should not follow the pause flag upon creation orm_dag = session.get(DagModel, ("dag_paused",)) assert orm_dag.is_paused is False @@ -1552,7 +1565,7 @@ def test_existing_dag_is_paused_upon_creation(self, testing_dag_bundle, session, def test_bundle_name_and_version_are_stored(self, testing_dag_bundle, session, dag_maker): with dag_maker("mydag", schedule=None) as dag: ... - update_dag_parsing_results_in_db("testing", "1.0", [dag], {}, 0.1, set(), session) + update_dag_parsing_results_in_db("testing", "1.0", [dag], {}, 0.1, set(), session=session) orm_dag = session.get(DagModel, "mydag") assert orm_dag.bundle_name == "testing" assert orm_dag.bundle_version == "1.0" @@ -1560,7 +1573,7 @@ def test_bundle_name_and_version_are_stored(self, testing_dag_bundle, session, d def test_max_active_tasks_explicit_value_is_used(self, testing_dag_bundle, session, dag_maker): with dag_maker("dag_max_tasks", schedule=None, max_active_tasks=5) as dag: ... - update_dag_parsing_results_in_db("testing", None, [dag], {}, 0.1, set(), session) + update_dag_parsing_results_in_db("testing", None, [dag], {}, 0.1, set(), session=session) orm_dag = session.get(DagModel, "dag_max_tasks") assert orm_dag.max_active_tasks == 5 @@ -1569,14 +1582,14 @@ def test_max_active_tasks_defaults_from_conf_when_none(self, testing_dag_bundle, with conf_vars({("core", "max_active_tasks_per_dag"): "7"}): with dag_maker("dag_max_tasks_default", schedule=None) as dag: ... - update_dag_parsing_results_in_db("testing", None, [dag], {}, 0.1, set(), session) + update_dag_parsing_results_in_db("testing", None, [dag], {}, 0.1, set(), session=session) orm_dag = session.get(DagModel, "dag_max_tasks_default") assert orm_dag.max_active_tasks == 7 def test_max_active_runs_explicit_value_is_used(self, testing_dag_bundle, session, dag_maker): with dag_maker("dag_max_runs", schedule=None, max_active_runs=3) as dag: ... - update_dag_parsing_results_in_db("testing", None, [dag], {}, 0.1, set(), session) + update_dag_parsing_results_in_db("testing", None, [dag], {}, 0.1, set(), session=session) orm_dag = session.get(DagModel, "dag_max_runs") assert orm_dag.max_active_runs == 3 @@ -1594,7 +1607,7 @@ def test_config_driven_field_equal_to_schema_default_not_overridden_by_conf( with conf_vars({("core", cfg_key): "1"}): with dag_maker(f"dag_{field}_schema_default", schedule=None, **{field: schema_default}) as dag: ... - update_dag_parsing_results_in_db("testing", None, [dag], {}, 0.1, set(), session) + update_dag_parsing_results_in_db("testing", None, [dag], {}, 0.1, set(), session=session) orm_dag = session.get(DagModel, f"dag_{field}_schema_default") assert getattr(orm_dag, field) == schema_default @@ -1602,7 +1615,7 @@ def test_max_active_runs_defaults_from_conf_when_none(self, testing_dag_bundle, with conf_vars({("core", "max_active_runs_per_dag"): "4"}): with dag_maker("dag_max_runs_default", schedule=None) as dag: ... - update_dag_parsing_results_in_db("testing", None, [dag], {}, 0.1, set(), session) + update_dag_parsing_results_in_db("testing", None, [dag], {}, 0.1, set(), session=session) orm_dag = session.get(DagModel, "dag_max_runs_default") assert orm_dag.max_active_runs == 4 @@ -1611,7 +1624,7 @@ def test_max_consecutive_failed_dag_runs_explicit_value_is_used( ): with dag_maker("dag_max_failed_runs", schedule=None, max_consecutive_failed_dag_runs=2) as dag: ... - update_dag_parsing_results_in_db("testing", None, [dag], {}, 0.1, set(), session) + update_dag_parsing_results_in_db("testing", None, [dag], {}, 0.1, set(), session=session) orm_dag = session.get(DagModel, "dag_max_failed_runs") assert orm_dag.max_consecutive_failed_dag_runs == 2 @@ -1621,7 +1634,7 @@ def test_max_consecutive_failed_dag_runs_defaults_from_conf_when_none( with conf_vars({("core", "max_consecutive_failed_dag_runs_per_dag"): "6"}): with dag_maker("dag_max_failed_runs_default", schedule=None) as dag: ... - update_dag_parsing_results_in_db("testing", None, [dag], {}, 0.1, set(), session) + update_dag_parsing_results_in_db("testing", None, [dag], {}, 0.1, set(), session=session) orm_dag = session.get(DagModel, "dag_max_failed_runs_default") assert orm_dag.max_consecutive_failed_dag_runs == 6 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 0000000000000..ffd118fdb31e7 --- /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_maintenance.py b/airflow-core/tests/unit/dag_processing/test_maintenance.py new file mode 100644 index 0000000000000..8e42e2054bfb2 --- /dev/null +++ b/airflow-core/tests/unit/dag_processing/test_maintenance.py @@ -0,0 +1,111 @@ +# 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 datetime import timedelta +from unittest import mock + +import pytest +from sqlalchemy import select +from uuid6 import uuid7 + +from airflow._shared.timezones import timezone +from airflow.dag_processing.maintenance import cleanup_processor_metadata +from airflow.jobs.job import Job, JobState +from airflow.models.dag import DagModel +from airflow.models.dagbundle import DagBundleModel +from airflow.models.dagwarning import DagWarning + +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 + +pytestmark = pytest.mark.db_test + + +@pytest.fixture(autouse=True) +def clean_db(): + clear_db_jobs() + clear_db_dags() + clear_db_dag_bundles() + yield + clear_db_jobs() + clear_db_dags() + clear_db_dag_bundles() + + +@mock.patch("airflow.dag_processing.maintenance.log", autospec=True) +def test_cleanup_is_bounded_and_retains_live_bundles(log, session, time_machine): + time_machine.move_to("2026-10-05T12:00:00Z", tick=False) + inactive = DagBundleModel(name="removed") + inactive.active = False + session.add_all([inactive, DagBundleModel(name="live")]) + session.flush() + session.add_all( + [DagModel(dag_id=f"removed_{i}", bundle_name="removed", is_stale=False) for i in range(3)] + + [DagModel(dag_id="current", bundle_name="live", is_stale=False)] + ) + session.commit() + cleanup_processor_metadata(batch_size=2) + session.expire_all() + assert len(session.scalars(select(DagModel).where(DagModel.is_stale.is_(True))).all()) == 2 + assert [call.args[1:] for call in log.info.call_args_list] == [ + ("removed_0", "removed"), + ("removed_1", "removed"), + ] + cleanup_processor_metadata(batch_size=2) + session.expire_all() + assert len(session.scalars(select(DagModel).where(DagModel.is_stale.is_(True))).all()) == 3 + assert not session.get(DagModel, "current").is_stale + + +@conf_vars( + { + ("dag_processor", "health_check_threshold"): "30", + ("dag_processor", "job_heartbeat_timeout"): "120", + } +) +def test_cleanup_retires_only_expired_api_processor_jobs(session, time_machine): + time_machine.move_to("2026-10-05T12:00:00Z", tick=False) + jobs = [Job(job_type="DagProcessorJob", state=JobState.RUNNING) for _ in range(4)] + for job in jobs[:3]: + job.session_id = uuid7() + jobs[0].latest_heartbeat = jobs[3].latest_heartbeat = timezone.utcnow() - timedelta(seconds=180) + jobs[1].latest_heartbeat = timezone.utcnow() - timedelta(seconds=60) + session.add_all(jobs) + session.commit() + job_ids = [job.id for job in jobs] + cleanup_processor_metadata() + jobs = [session.get(Job, job_id) for job_id in job_ids] + assert jobs[0].end_date is not None + assert jobs[0].state == JobState.FAILED + assert [job.end_date for job in jobs[1:]] == [None, None, None] + + +def test_cleanup_warnings_is_bounded_and_preserves_active_dags(session): + session.add(DagBundleModel(name="live")) + session.flush() + session.add_all([DagModel(dag_id=f"dag_{i}", bundle_name="live", is_stale=i < 3) for i in range(4)]) + session.flush() + session.add_all( + [DagWarning(dag_id=f"dag_{i}", warning_type="python:test", message="warning") for i in range(4)] + ) + session.commit() + cleanup_processor_metadata(batch_size=2) + assert len(session.scalars(select(DagWarning)).all()) == 2 + cleanup_processor_metadata(batch_size=2) + assert [warning.dag_id for warning in session.scalars(select(DagWarning))] == ["dag_3"] diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index 2a92835edba39..e51b2753a0fad 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py @@ -23,20 +23,25 @@ import os import random import re +import runpy import selectors import shutil import signal +import sys import textwrap import time import zipfile from collections import Counter, OrderedDict, defaultdict, namedtuple -from contextlib import contextmanager +from contextlib import contextmanager, nullcontext from datetime import datetime, timedelta from pathlib import Path from socket import socket, socketpair from unittest import mock from unittest.mock import MagicMock +from uuid import UUID +import httpx2 +import jwt import msgspec import pytest import structlog @@ -46,10 +51,23 @@ from uuid6 import uuid7 from airflow._shared.timezones import timezone +from airflow.api_fastapi.execution_api.datamodels.dag_parsing import ( + DagBundleInventoryBody, + DagBundleInventoryResponse, + DagBundleStateResponse, + DagParseResultBody, + ParseSourceCode, + 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 -from airflow.dag_processing.bundles.manager import DagBundlesManager +from airflow.dag_processing.bundles.manager import DagBundlesManager, _load_bundle_config_snapshot from airflow.dag_processing.collection import update_dag_parsing_results_in_db from airflow.dag_processing.dagbag import DagBag from airflow.dag_processing.lang_sdk_processor import LangSDKDagFileProcessorProcess @@ -58,11 +76,15 @@ DagFileInfo, DagFileProcessorManager, DagFileStat, + DeferredCallback, + PendingDagPublication, + PriorityParseClaim, ) from airflow.dag_processing.processor import ( DagFileParseRequest, DagFileParsingResult, DagFileProcessorProcess, + DagParseSource, _parse_file, ) from airflow.models import DagModel, DbCallbackRequest @@ -73,7 +95,8 @@ from airflow.models.serialized_dag import SerializedDagModel from airflow.models.team import Team from airflow.providers.standard.operators.empty import EmptyOperator -from airflow.sdk import DAG as SdkDAG, BaseOperator +from airflow.sdk import DAG as SdkDAG, BaseOperator, Connection, SecretCache +from airflow.sdk.execution_time import task_runner from airflow.sdk.importers import DagDefinition, DagImporterRegistry, DagImportError, DagSourceCode from airflow.serialization.serialized_objects import DagSerialization, LazyDeserializedDAG from airflow.utils.net import get_hostname @@ -457,10 +480,11 @@ def test_find_files_in_bundle_parses_zip_once_and_observes_its_members(self, tmp "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" @@ -468,8 +492,9 @@ def test_find_files_in_bundle_makes_discovery_errors_bundle_relative( (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 found_files == { @@ -555,14 +580,18 @@ def test_sync_bundles_deactivates_missing_when_owning_all_bundles(self): manager = DagFileProcessorManager(max_runs=1) with mock.patch("airflow.dag_processing.manager.DagBundlesManager") as mock_bundles_manager: manager.sync_bundles() - mock_bundles_manager.return_value.sync_bundles_to_db.assert_called_once_with(deactivate_missing=True) + mock_bundles_manager.return_value.sync_bundles_to_db.assert_called_once_with( + deactivate_missing=True, include_bundle_urls=True + ) def test_sync_bundles_does_not_deactivate_missing_when_filtered(self): """A processor started with ``--bundle-name`` owns a subset and must not deactivate others.""" manager = DagFileProcessorManager(max_runs=1, bundle_names_to_parse=["only-mine"]) with mock.patch("airflow.dag_processing.manager.DagBundlesManager") as mock_bundles_manager: manager.sync_bundles() - mock_bundles_manager.return_value.sync_bundles_to_db.assert_called_once_with(deactivate_missing=False) + mock_bundles_manager.return_value.sync_bundles_to_db.assert_called_once_with( + deactivate_missing=False, include_bundle_urls=True + ) @pytest.mark.parametrize( "safe_mode", @@ -694,6 +723,7 @@ def test_start_new_processes_with_same_filepath(self, configure_testing_dag_bund ) file_3 = DagFileInfo(bundle_name="testing", rel_path=Path("file_3.py"), bundle_path=TEST_DAGS_FOLDER) manager._file_queue = OrderedDict.fromkeys([file_1, file_2, file_3]) + manager._bundle_parse_sources["testing"] = DagParseSource() # Mock that only one processor exists. This processor runs with 'file_1' manager._processors[file_1] = MagicMock() @@ -1416,7 +1446,7 @@ def test_scan_stale_dags_deactivates_zip_packaged_dags(self, session, test_zip_p assert serialized_dag_count == 1 @pytest.mark.usefixtures("testing_dag_bundle") - def test_deactivate_stale_dags_marks_dags_in_inactive_bundles(self, session): + def test_deactivate_stale_dags_leaves_inactive_bundles_to_scheduler(self, session): """Dags whose bundle is no longer active should be marked stale even without a parse signal.""" session.add(DagBundleModel(name="gone-bundle")) session.flush() @@ -1453,7 +1483,7 @@ def test_deactivate_stale_dags_marks_dags_in_inactive_bundles(self, session): ) ).all() ) - assert is_stale_by_dag == {"dag_in_inactive_bundle": True, "dag_in_active_bundle": False} + assert is_stale_by_dag == {"dag_in_inactive_bundle": False, "dag_in_active_bundle": False} @pytest.mark.usefixtures("testing_dag_bundle") def test_deactivate_stale_dags_matches_dags_nested_in_a_container(self, session): @@ -1479,7 +1509,7 @@ def test_deactivate_stale_dags_matches_dags_nested_in_a_container(self, session) assert session.scalar(select(DagModel.is_stale).where(DagModel.dag_id == "nested_dag")) @pytest.mark.usefixtures("testing_dag_bundle") - def test_deactivate_stale_dags_marks_dags_with_null_bundle_name(self, session): + def test_deactivate_stale_dags_leaves_null_bundle_to_scheduler(self, session): """Dags carried over from Airflow 2.x keep a NULL bundle_name and must still be deactivated. Their files were removed during the upgrade, so nothing will ever parse them and fill the @@ -1516,7 +1546,7 @@ def execute_with_null_bundle_name(statement, *args, **kwargs): with mock.patch.object(session, "execute", side_effect=execute_with_null_bundle_name): manager.deactivate_stale_dags(last_parsed={}, session=session) - assert session.scalar(select(DagModel.is_stale).where(DagModel.dag_id == "legacy_dag")) + assert not session.scalar(select(DagModel.is_stale).where(DagModel.dag_id == "legacy_dag")) @pytest.mark.usefixtures("testing_dag_bundle") def test_deactivate_stale_dags_tolerates_null_relative_fileloc(self, session): @@ -1588,9 +1618,6 @@ def test_deactivate_stale_dags_handles_lock_timeout(self, mock_is_lock_not_avail session.add(DagBundleModel(name="gone-bundle")) session.flush() - session.execute( - DagBundleModel.__table__.update().where(DagBundleModel.name == "gone-bundle").values(active=False) - ) session.add( DagModel( dag_id="dag_in_inactive_bundle", @@ -1615,7 +1642,12 @@ def mock_execute(*args, **kwargs): mock_is_lock_not_available.return_value = True with mock.patch.object(session, "execute", side_effect=mock_execute): - manager.deactivate_stale_dags(last_parsed={}) + manager.deactivate_stale_dags( + last_parsed={ + DagFileInfo(rel_path=Path("some_file.py"), bundle_name="gone-bundle"): timezone.utcnow() + + timedelta(hours=1) + } + ) assert "Lock not available when deactivating stale DAGs" in caplog.text @@ -2066,6 +2098,411 @@ def test_terminate_normalizes_file_path_stats_tag(self): ) processor.kill.assert_called_once_with(signal.SIGTERM, escalation_delay=5.0) + @mock.patch.object(DagFileProcessorManager, "persist_parsing_result", autospec=True) + @mock.patch.object(DagFileProcessorManager, "_create_process", autospec=True) + def test_publication_uses_a_copy_of_the_source_metadata_at_dispatch(self, create, persist, session): + manager = DagFileProcessorManager(max_runs=1) + data = {"files": {"dag.py": "original"}} + manager._bundle_parse_sources["testing"] = DagParseSource(bundle_version="v1", version_data=data) + file = DagFileInfo(bundle_name="testing", rel_path=Path("dag.py"), bundle_path=TEST_DAGS_FOLDER) + manager._file_queue[file] = None + processor, _ = self.mock_processor() + processor.had_callbacks = False + processor.parsing_result = DagFileParsingResult(fileloc="dag.py", serialized_dags=[]) + create.return_value = processor + + manager._start_new_processes() + dispatched = next(iter(manager._processors)) + data["files"]["dag.py"] = "changed" + manager._bundle_versions["testing"] = "v2" + manager.handle_parsing_result(dispatched, processor, session=session) + + assert dispatched == file + persist.assert_called_once_with( + manager, + bundle_name="testing", + bundle_version="v1", + version_data={"files": {"dag.py": "original"}}, + parsing_result=processor.parsing_result, + run_duration=mock.ANY, + relative_fileloc="dag.py", + session=session, + ) + + @pytest.mark.parametrize("publication_fails", [False, True]) + @mock.patch.object(DagFileProcessorManager, "persist_parsing_result", autospec=True) + @mock.patch.object(DagFileProcessorManager, "_create_process", autospec=True) + def test_api_publication_uses_dispatch_identity_without_local_persistence( + self, create, persist, session, publication_fails + ): + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + manager = DagFileProcessorManager(max_runs=1, api_client=client) + manager._bundle_parse_sources["testing"] = DagParseSource(bundle_version="v1") + file = _get_file_infos(["dag.py"])[0] + manager._file_queue[file] = None + proc, read_end = self.mock_processor() + proc.had_callbacks = False + proc.parsing_result = DagFileParsingResult(fileloc="dag.py", serialized_dags=[]) + create.return_value = proc + if publication_fails: + client.publish_parse_result.side_effect = RuntimeError("publication failed") + try: + manager._start_new_processes() + manager.handle_parsing_result(file, proc, session=session) + client.publish_parse_result.assert_not_called() + manager._publish_pending_results() + body = client.publish_parse_result.call_args.args[0] + assert body.attempt_id == proc.id + assert body.dispatch_sequence == proc.dispatch_sequence == 1 + assert body.bundle_version == "v1" + assert body.relative_fileloc == "dag.py" + assert body.serialized_dags == [] + assert manager._file_stats[file].run_count == 1 + persist.assert_not_called() + finally: + read_end.close() + proc.close() + + @pytest.mark.parametrize("source_available", [True, False]) + def test_api_publication_carries_diagnostics_and_explicit_source_availability(self, source_available): + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + manager = DagFileProcessorManager(max_runs=1, api_client=client) + file = _get_file_infos(["dag.py"])[0] + dag = SdkDAG("published", schedule=None) + dag.fileloc = "/worker/dag.py" + dag.relative_fileloc = "dag.py" + proc, read_end = self.mock_processor() + proc.dispatch_sequence = 1 + warning = {"dag_id": dag.dag_id, "warning_type": "python:deprecated", "message": "warning"} + unserialized_warning = {"dag_id": "failed", "warning_type": "python:deprecated", "message": "warning"} + proc.parsing_result = DagFileParsingResult( + fileloc=dag.fileloc, + serialized_dags=[LazyDeserializedDAG.from_dag(dag)], + warnings=[warning, unserialized_warning], + parsed_definitions=["dag.py"], + import_errors={"dag.py": "import error"}, + dag_source_codes={dag.dag_id: DagSourceCode(source_code="captured", language="python")} + if source_available + else {}, + ) + try: + body = manager._build_parse_result(file, proc, 0.5) + assert [published.model_dump() for published in body.warnings] == [warning] + assert body.import_errors == {"dag.py": "import error"} + assert body.parsed_definitions == ["dag.py"] + assert body.source_codes[dag.dag_id].source_code == ("captured" if source_available else None) + finally: + read_end.close() + proc.close() + + @pytest.mark.usefixtures("testing_dag_bundle") + @conf_vars({("workers", "execution_api_retries"): "1"}) + @pytest.mark.parametrize("failure", [413, 422, 403, 500, "timeout"]) + def test_rejected_publication_preserves_stale_detection(self, session, time_machine, failure): + time_machine.move_to("2026-10-05T12:00:00Z", tick=False) + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + request = httpx2.Request("POST", "http://api/jobs/1/parse-results") + client.publish_parse_result.side_effect = ( + httpx2.ReadTimeout("Lost response", request=request) + if failure == "timeout" + else httpx2.HTTPStatusError("Rejected", request=request, response=httpx2.Response(failure)) + ) + manager = DagFileProcessorManager(max_runs=1, api_client=client) + file = _get_file_infos(["dag.py"])[0] + previous = timezone.utcnow() - timedelta(seconds=manager.stale_dag_threshold + 1) + manager._file_stats[file] = DagFileStat(num_dags=1, last_finish_time=previous) + dag = SdkDAG("published", schedule=None) + dag.fileloc = "/worker/dag.py" + dag.relative_fileloc = "dag.py" + session.add( + DagModel( + dag_id=dag.dag_id, + bundle_name="testing", + relative_fileloc="dag.py", + fileloc=dag.fileloc, + last_parsed_time=previous, + is_stale=False, + ) + ) + session.commit() + proc, read_end = self.mock_processor() + proc.dispatch_sequence = 1 + proc.parsing_result = DagFileParsingResult( + fileloc=dag.fileloc, serialized_dags=[LazyDeserializedDAG.from_dag(dag)] + ) + try: + manager.handle_parsing_result(file, proc) + manager._publish_pending_results() + manager._scan_stale_dags() + session.expire_all() + assert not session.get(DagModel, dag.dag_id).is_stale + assert manager._file_stats[file].last_finish_time == previous + assert manager._file_stats[file].num_dags == 1 + assert manager.processed_recently(timezone.utcnow(), file) + assert not manager._pending_publications + finally: + read_end.close() + proc.close() + + def test_publication_retries_leave_the_loop_responsive(self, mocker): + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + client.publish_parse_result.side_effect = [httpx2.ReadTimeout("Lost acknowledgment"), None, None] + heartbeat_observations = [] + with selectors.DefaultSelector() as selector: + manager = DagFileProcessorManager( + max_runs=1, + api_client=client, + selector=selector, + heartbeat=lambda: heartbeat_observations.append(client.publish_parse_result.call_count), + ) + for method in ( + "_refresh_dag_bundles", + "_queue_requested_files_for_parsing", + "_scan_stale_dags", + "_cleanup_stale_bundle_versions", + "print_stats", + ): + mocker.patch.object(DagFileProcessorManager, method, autospec=True) + mocker.patch.object(DagFileProcessorManager, "fetch_callbacks", autospec=True, return_value=[]) + for file in _get_file_infos(["a.py", "b.py"]): + proc, read_end = self.mock_processor() + proc.dispatch_sequence = 1 + proc.parsing_result = DagFileParsingResult(fileloc=str(file.rel_path), serialized_dags=[]) + try: + manager.handle_parsing_result(file, proc) + finally: + read_end.close() + proc.close() + assert not manager.max_runs_reached() + manager._run_parsing_loop() + calls = client.publish_parse_result.call_args_list + assert [call.args[0].relative_fileloc for call in calls] == ["a.py", "b.py", "a.py"] + assert calls[0].args[0] is calls[2].args[0] + assert {0, 1}.issubset(heartbeat_observations) + assert not manager._pending_publications + assert all(stat.run_count == 1 for stat in manager._file_stats.values()) + + @pytest.mark.parametrize("change", ["refresh", "delete"]) + def test_pending_publication_is_discarded_when_source_changes(self, change): + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + manager = DagFileProcessorManager(max_runs=1, api_client=client) + file = _get_file_infos(["dag.py"])[0] + proc, read_end = self.mock_processor() + proc.dispatch_sequence = 1 + proc.parsing_result = DagFileParsingResult(fileloc="dag.py", serialized_dags=[]) + try: + manager.handle_parsing_result(file, proc) + if change == "refresh": + manager._invalidate_bundle_parse_source("testing") + else: + manager.remove_orphaned_file_stats(present=set()) + manager._publish_pending_results() + client.publish_parse_result.assert_not_called() + assert not manager._pending_publications + assert list(manager._file_queue) == ([file] if change == "refresh" else []) + finally: + read_end.close() + proc.close() + + @mock.patch.object(DagFileProcessorManager, "_create_process", autospec=True) + def test_pending_publications_bound_admission_and_prevent_reimports(self, create): + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + manager = DagFileProcessorManager(max_runs=1, api_client=client, parallelism=1) + manager._file_parsing_sort_mode = "alphabetical" + file, other = _get_file_infos(["dag.py", "other.py"]) + proc, read_end = self.mock_processor() + proc.dispatch_sequence = 1 + proc.parsing_result = DagFileParsingResult(fileloc="dag.py", serialized_dags=[]) + manager._bundle_parse_sources["testing"] = DagParseSource() + try: + manager.handle_parsing_result(file, proc) + manager.prepare_file_queue(known_files={"testing": {file, other}}) + assert list(manager._file_queue) == [other] + manager._start_new_processes() + create.assert_not_called() + assert list(manager._file_queue) == [other] + manager._parallelism = 2 + manager._file_queue = OrderedDict.fromkeys([file]) + manager._start_new_processes() + create.assert_not_called() + finally: + read_end.close() + proc.close() + + @pytest.mark.parametrize("current_version", [None, "v1", "v2"]) + @mock.patch.object(DagFileProcessorManager, "persist_parsing_result", autospec=True) + @mock.patch.object(DagFileProcessorManager, "_create_process", autospec=True) + def test_forced_refresh_waits_for_in_flight_imports(self, create, persist, current_version): + manager = DagFileProcessorManager(max_runs=1) + bundle = self._make_refresh_bundle( + supports_versioning=current_version is not None, current_version=current_version + ) + manager._bundle_versions[bundle.name] = "v1" + file = DagFileInfo(bundle_name=bundle.name, rel_path=Path("dag.py"), bundle_path=bundle.path) + manager._file_stats[file] = DagFileStat() + manager._file_queue = OrderedDict.fromkeys([file]) + manager._bundle_parse_sources = {bundle.name: DagParseSource()} + processor = self.mock_processor()[0] + processor.had_callbacks = False + processor.parsing_result = DagFileParsingResult(fileloc="dag.py", serialized_dags=[]) + create.return_value = processor + manager._start_new_processes() + + self._refresh_with_mocked_state( + manager, bundle, BundleState(last_refreshed=None, version="v1"), force=True + ) + bundle.refresh.assert_not_called() + assert manager._bundles_waiting_for_refresh == {bundle.name} + + manager._collect_results() + persist.assert_called_once() + assert not manager._file_queue + + self._refresh_with_mocked_state( + manager, bundle, BundleState(last_refreshed=None, version="v1"), force=True + ) + bundle.refresh.assert_called_once_with() + + @pytest.mark.parametrize("current_version", [None, "v1", "v2"]) + @mock.patch("airflow.dag_processing.manager.stats.incr", autospec=True) + @mock.patch.object(DagFileProcessorManager, "persist_parsing_result", autospec=True) + @mock.patch.object(DagFileProcessorManager, "_create_process", autospec=True) + def test_reinitialized_bundle_discards_only_affected_results_without_advancing_freshness( + self, create, persist, incr, current_version + ): + manager = DagFileProcessorManager(max_runs=1) + manager._num_run = 1 + manager._parallelism = 2 + bundle = self._make_refresh_bundle( + supports_versioning=current_version is not None, current_version=current_version + ) + manager._bundle_versions[bundle.name] = "v1" + file = DagFileInfo(bundle_name=bundle.name, rel_path=Path("dag.py"), bundle_path=bundle.path) + other = DagFileInfo(bundle_name="other", rel_path=Path("other.py"), bundle_path=bundle.path) + original_stat = DagFileStat(num_dags=1, run_count=1, last_finish_time=DEFAULT_DATE) + manager._file_stats[file] = original_stat + manager._file_queue = OrderedDict.fromkeys([file, other]) + manager._bundle_parse_sources = {name: DagParseSource() for name in (bundle.name, "other")} + processors = [self.mock_processor()[0] for _ in range(2)] + for processor in processors: + processor.had_callbacks = False + processor.parsing_result = DagFileParsingResult(fileloc="dag.py", serialized_dags=[]) + create.side_effect = processors + manager._start_new_processes() + bundle.initialize.side_effect = OSError("Source may have changed") + bundle.is_initialized = False + + self._refresh_with_mocked_state( + manager, bundle, BundleState(last_refreshed=None, version="v1"), force=True + ) + manager._collect_results() + + assert manager._file_stats[file] is original_stat + assert list(manager._file_queue) == [file] + assert manager._processors == {} + assert not manager.max_runs_reached() + persist.assert_called_once() + assert persist.call_args.kwargs["bundle_name"] == "other" + incr.assert_any_call("dag_processing.results_discarded_on_refresh", tags={"bundle_name": bundle.name}) + assert bundle.name not in manager._bundle_parse_sources + + @mock.patch.object(DagFileProcessorManager, "_create_process", autospec=True) + def test_unavailable_bundle_does_not_block_ready_bundles_or_pinned_callbacks(self, create): + manager = DagFileProcessorManager(max_runs=1) + manager._parallelism = 2 + blocked = DagFileInfo(bundle_name="blocked", rel_path=Path("dag.py")) + healthy = DagFileInfo(bundle_name="healthy", rel_path=Path("dag.py")) + pinned = DagFileInfo(bundle_name="blocked", rel_path=Path("dag.py"), bundle_version="v1") + manager._file_queue = OrderedDict.fromkeys([blocked, healthy, pinned]) + manager._bundle_parse_sources["healthy"] = DagParseSource() + manager._callback_to_execute[pinned] = [mock.create_autospec(DagCallbackRequest, instance=True)] + create.side_effect = [self.mock_processor()[0] for _ in range(2)] + + manager._start_new_processes() + + assert list(manager._processors) == [healthy, pinned] + assert not manager._file_queue + assert list(manager._files_awaiting_source) == [blocked] + assert manager._processors[pinned].parse_source.bundle_version == "v1" + + @mock.patch.object(DagFileProcessorManager, "_create_process", autospec=True) + def test_callbacks_run_while_their_bundle_has_no_parse_source(self, create): + manager = DagFileProcessorManager(max_runs=1) + file = DagFileInfo(bundle_name="unrefreshable", rel_path=Path("dag.py")) + manager._callback_to_execute[file] = [mock.create_autospec(DagCallbackRequest, instance=True)] + manager._file_queue[file] = None + create.return_value = self.mock_processor()[0] + + manager._start_new_processes() + + assert manager._processors[file].parse_source == DagParseSource() + assert not manager._files_awaiting_source + + 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] + + @conf_vars({("dag_processor", "file_parsing_sort_mode"): "alphabetical"}) + def test_files_awaiting_a_bundle_source_do_not_stop_other_bundles_refilling(self): + manager = DagFileProcessorManager(max_runs=1) + blocked = DagFileInfo(bundle_name="blocked", rel_path=Path("dag.py")) + removed = DagFileInfo(bundle_name="blocked", rel_path=Path("removed.py")) + healthy = DagFileInfo(bundle_name="healthy", rel_path=Path("dag.py")) + manager._files_awaiting_source = dict.fromkeys([blocked, removed]) + manager._num_run = 1 + + assert not manager.max_runs_reached() + + manager.prepare_file_queue(known_files={"blocked": {blocked}, "healthy": {healthy}}) + assert list(manager._file_queue) == [healthy] + + manager._add_new_files_to_queue({"blocked": {blocked}}) + assert list(manager._file_queue) == [healthy] + + manager.purge_removed_files_from_queue(present={blocked, healthy}) + assert list(manager._files_awaiting_source) == [blocked] + + @mock.patch.object(DagFileProcessorManager, "persist_parsing_result", autospec=True) + def test_refresh_does_not_requeue_completed_callbacks(self, persist): + manager = DagFileProcessorManager(max_runs=1) + file = DagFileInfo(bundle_name="testing", rel_path=Path("dag.py")) + processor, _ = self.mock_processor() + processor.had_callbacks = True + processor.parsing_result = None + manager._bundle_refresh_generations["testing"] = 1 + + manager.handle_parsing_result(file, processor) + + persist.assert_not_called() + assert not manager._file_queue + assert manager._file_stats[file].last_finish_time is None + def test_handle_parsing_result_provides_its_own_session_when_caller_omits(self): """``handle_parsing_result`` is wrapped in ``@provide_session`` so subclasses overriding it can run without a caller-supplied session.""" manager = DagFileProcessorManager(max_runs=1) @@ -2103,13 +2540,18 @@ def test_handle_parsing_result_throttles_retry_when_first_persist_fails(self, se assert manager._file_stats[file].num_dags == 0 assert manager._file_stats[file].import_errors == 0 assert manager._file_stats[file].run_count == 1 - assert manager._file_stats[file].last_finish_time is not None + assert manager._file_stats[file].last_finish_time is None + assert manager._file_stats[file].last_attempt_time is not None assert manager._file_stats[file].last_duration is not None assert manager.processed_recently(timezone.utcnow(), file) is True def test_handle_parsing_result_updates_stats_after_successful_persist(self, session): manager = DagFileProcessorManager(max_runs=1) - file = DagFileInfo(bundle_name="testing", rel_path=Path("abc.txt"), bundle_path=TEST_DAGS_FOLDER) + file = DagFileInfo( + bundle_name="testing", + rel_path=Path("abc.txt"), + bundle_path=TEST_DAGS_FOLDER, + ) original_stat = DagFileStat( num_dags=1, import_errors=0, @@ -2130,7 +2572,7 @@ def test_handle_parsing_result_updates_stats_after_successful_persist(self, sess mock_persist.assert_called_once_with( bundle_name="testing", - bundle_version="v1", + bundle_version=None, version_data=None, parsing_result=processor.parsing_result, run_duration=mock.ANY, @@ -2181,7 +2623,8 @@ def test_collect_results_processes_remaining_files_when_one_persist_fails(self, assert manager._file_stats[file_a].num_dags == stat_a_before.num_dags assert manager._file_stats[file_a].import_errors == stat_a_before.import_errors assert manager._file_stats[file_a].run_count == stat_a_before.run_count + 1 - assert manager._file_stats[file_a].last_finish_time is not None + assert manager._file_stats[file_a].last_finish_time == stat_a_before.last_finish_time + assert manager._file_stats[file_a].last_attempt_time is not None assert manager._file_stats[file_b] is not stat_b_before assert manager._file_stats[file_b].run_count == 2 assert len(manager._processors) == 0 @@ -2977,8 +3420,9 @@ def test_callback_queue(self, mock_get_logger, configure_testing_dag_bundle): ] with mock.patch.object( - DagFileProcessorProcess, "start", side_effect=lambda *args, **kwargs: self.mock_processor() + 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 == [ @@ -3033,6 +3477,26 @@ def test_prepare_callback_bundle_initializes_versioned_bundle(self, mock_bundle_ assert manager.prepare_callback_bundle(request) is bundle bundle.initialize.assert_called_once() + @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"] + @mock.patch("airflow.dag_processing.manager.DagBundlesManager") def test_prepare_callback_bundle_forwards_version_data(self, mock_bundle_manager): manager = DagFileProcessorManager(max_runs=1) @@ -3886,29 +4350,6 @@ def test_prepare_bundles_can_be_overridden_without_sync(self, tmp_path, configur sync_mock.assert_not_called() assert [b.name for b in manager._dag_bundles] == ["testing"] - def test_purge_inactive_dag_warnings_delegates_to_dagwarning(self): - """Default `purge_inactive_dag_warnings` calls `DagWarning.purge_inactive_dag_warnings`.""" - manager = DagFileProcessorManager(max_runs=1) - with mock.patch( - "airflow.dag_processing.manager.DagWarning.purge_inactive_dag_warnings" - ) as purge_mock: - manager.purge_inactive_dag_warnings() - purge_mock.assert_called_once_with() - - def test_run_parsing_loop_uses_overridable_purge(self, tmp_path, configure_testing_dag_bundle): - """`_run_parsing_loop` calls the overridable `purge_inactive_dag_warnings` seam.""" - with configure_testing_dag_bundle(tmp_path): - manager = DagFileProcessorManager(max_runs=1) - with ( - mock.patch.object(manager, "purge_inactive_dag_warnings") as purge_mock, - mock.patch( - "airflow.dag_processing.manager.DagWarning.purge_inactive_dag_warnings" - ) as direct_mock, - ): - manager.run() - purge_mock.assert_called() - direct_mock.assert_not_called() - @mock.patch("airflow.dag_processing.manager.stats.gauge") def test_stats_total_parse_time(self, statsd_gauge_mock, tmp_path, configure_testing_dag_bundle): key = "dag_processing.total_parse_time" @@ -4047,7 +4488,7 @@ def _make_refresh_bundle(self, *, supports_versioning=False, current_version=Non bundle.get_current_version.return_value = current_version return bundle - def _refresh_with_mocked_state(self, manager, bundle, initial_state): + def _refresh_with_mocked_state(self, manager, bundle, initial_state, *, force=False): """Run _refresh_dag_bundles with get/update_bundle_state mocked out. Returns the two MagicMock objects for post-call assertions. MagicMock retains its @@ -4056,10 +4497,28 @@ def _refresh_with_mocked_state(self, manager, bundle, initial_state): them normally after this method returns. """ manager._dag_bundles = [bundle] - manager._force_refresh_bundles = set() + manager._force_refresh_bundles = {bundle.name} if force else set() mock_get = mock.patch.object(manager, "get_bundle_state", return_value=initial_state) mock_update = mock.patch.object(manager, "update_bundle_state") + mock_catalog = ( + mock.patch.object( + manager.api_client, + "get_bundles", + return_value=[ + DagBundleStateResponse( + name=bundle.name, + last_refreshed=initial_state.last_refreshed, + version=initial_state.version, + revision=initial_state.revision, + team_name=None, + ) + ], + ) + if manager.api_client is not None + else nullcontext() + ) with ( + mock_catalog, mock_get as patched_get, mock_update as patched_update, mock.patch.object(manager, "_find_files_in_bundle", return_value=set()), @@ -4072,6 +4531,158 @@ def _refresh_with_mocked_state(self, manager, bundle, initial_state): manager._refresh_dag_bundles({}) return patched_get, patched_update + @pytest.mark.parametrize("failure", ["state", "discovery"]) + @mock.patch.object(DagFileProcessorManager, "_create_process", autospec=True) + @mock.patch.object(DagFileProcessorManager, "handle_removed_files", autospec=True) + @mock.patch.object(DagFileProcessorManager, "clear_orphaned_import_errors", autospec=True) + @mock.patch.object(DagFileProcessorManager, "deactivate_deleted_dags", autospec=True) + @mock.patch.object(DagFileProcessorManager, "_find_files_in_bundle", autospec=True) + @mock.patch.object(DagFileProcessorManager, "update_bundle_state", autospec=True) + @mock.patch.object(DagFileProcessorManager, "get_bundle_state", autospec=True) + def test_dispatch_uses_refreshed_source_even_when_metadata_update_fails( + self, get_state, update_state, find_files, deactivate, clear, removed, create, failure, tmp_path + ): + manager = DagFileProcessorManager(max_runs=1) + bundle = self._make_refresh_bundle( + supports_versioning=True, + current_version=BundleVersion(version="v2", data={"files": {"dag.py": "new"}}), + ) + bundle.path = tmp_path + (tmp_path / "dag.py").touch() + manager._dag_bundles = [bundle] + manager._bundle_versions[bundle.name] = "v1" + file = DagFileInfo(bundle_name=bundle.name, rel_path=Path("dag.py"), bundle_path=bundle.path) + manager._file_queue[file] = None + get_state.return_value = BundleState(last_refreshed=None, version="v1") + find_files.return_value = {file} + if failure == "state": + update_state.side_effect = OSError("Database unavailable") + else: + find_files.side_effect = OSError("Listing failed") + + manager._refresh_dag_bundles({bundle.name: {file}}) + manager._start_new_processes() + + assert manager._bundle_versions[bundle.name] == "v1" + assert create.return_value.parse_source == DagParseSource( + bundle_version="v2", version_data={"files": {"dag.py": "new"}}, refresh_generation=1 + ) + + @pytest.mark.parametrize("failure", ["invalid_path", "too_many_files"]) + @mock.patch.object(DagFileProcessorManager, "_get_observed_filelocs", autospec=True) + def test_listing_the_api_cannot_accept_releases_priority_capacity(self, observed, failure): + observed.return_value = ( + {"bad\\dag.py"} if failure == "invalid_path" else {f"{index}.py" for index in range(100001)} + ) + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + manager = DagFileProcessorManager( + max_runs=1, api_client=client, parallelism=1, bundle_refresh_check_interval=0 + ) + bundle = self._make_refresh_bundle() + 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" + + self._refresh_with_mocked_state( + manager, bundle, BundleState(last_refreshed=None, version=None), force=True + ) + + bundle.refresh.assert_called_once() + client.publish_inventory.assert_not_called() + assert bundle.name not in manager._bundle_parse_sources + assert bundle.name not in manager._pending_inventories + assert manager._bundle_versions[bundle.name] == "previous" + assert not manager._force_refresh_bundles + assert not manager._bundles_waiting_for_refresh + 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() + self._refresh_with_mocked_state(manager, bundle, BundleState(last_refreshed=None, version=None)) + assert bundle.name in manager._bundle_parse_sources + assert manager._bundle_versions[bundle.name] is None + assert not manager._pending_inventories + + @mock.patch.object(DagFileProcessorManager, "_create_process", autospec=True) + def test_failed_refresh_blocks_dispatch_until_the_source_recovers(self, create): + manager = DagFileProcessorManager(max_runs=1, bundle_refresh_check_interval=0) + bundle = self._make_refresh_bundle() + bundle.refresh_interval = 300 + manager._bundle_versions[bundle.name] = None + file = DagFileInfo(bundle_name=bundle.name, rel_path=Path("dag.py")) + manager._file_queue[file] = None + bundle.refresh.side_effect = [OSError("Partial refresh"), None] + state = BundleState(last_refreshed=timezone.utcnow(), version=None) + + self._refresh_with_mocked_state(manager, bundle, state) + manager._start_new_processes() + create.assert_not_called() + assert not manager._file_queue + assert list(manager._files_awaiting_source) == [file] + + self._refresh_with_mocked_state(manager, bundle, state) + manager._start_new_processes() + create.assert_called_once() + assert create.return_value.parse_source.refresh_generation == 2 + assert not manager._file_queue + assert not manager._files_awaiting_source + + @pytest.mark.parametrize("api_mode", [False, True]) + @mock.patch.object(DagFileProcessorManager, "_create_process", autospec=True) + @mock.patch.object(DagFileProcessorManager, "persist_parsing_result", autospec=True) + def test_periodic_refresh_drains_imports_before_admitting_more(self, persist, create, api_mode): + client = mock.create_autospec(DagProcessorAPIClient, instance=True) if api_mode else None + if client is not None: + client.publish_inventory.return_value.revision = uuid7() + manager = DagFileProcessorManager(max_runs=1, bundle_refresh_check_interval=300, api_client=client) + bundle = self._make_refresh_bundle(supports_versioning=True, current_version="v2") + bundle.refresh_interval = 300 + manager._bundle_versions[bundle.name] = "v1" + manager._bundle_parse_sources[bundle.name] = DagParseSource(bundle_version="v1") + active = DagFileInfo(bundle_name=bundle.name, rel_path=Path("active.py")) + queued = DagFileInfo(bundle_name=bundle.name, rel_path=Path("queued.py")) + processor, _ = self.mock_processor() + processor.had_callbacks = False + processor.dispatch_sequence = 1 + processor.parse_source = DagParseSource(bundle_version="v1") + processor.parsing_result = DagFileParsingResult(fileloc="active.py", serialized_dags=[]) + manager._processors[active] = processor + manager._file_queue[queued] = None + + self._refresh_with_mocked_state(manager, bundle, BundleState(last_refreshed=None, version="v1")) + manager._start_new_processes() + bundle.refresh.assert_not_called() + create.assert_not_called() + assert manager._bundles_waiting_for_refresh == {bundle.name} + + manager._collect_results() + if api_mode: + persist.assert_not_called() + self._refresh_with_mocked_state(manager, bundle, BundleState(last_refreshed=None, version="v1")) + bundle.refresh.assert_not_called() + manager._publish_pending_results() + assert client.publish_parse_result.call_args.args[0].bundle_version == "v1" + else: + persist.assert_called_once() + assert persist.call_args.kwargs["bundle_version"] == "v1" + self._refresh_with_mocked_state( + manager, bundle, BundleState(last_refreshed=timezone.utcnow(), version="v1") + ) + manager._start_new_processes() + + bundle.refresh.assert_called_once_with() + create.assert_called_once_with(manager, queued) + assert create.return_value.parse_source.bundle_version == "v2" + assert not manager._bundles_waiting_for_refresh + def test_refresh_dag_bundles_non_versioned_calls_update_bundle_state(self): """Non-versioned bundle: update_bundle_state called with version=None.""" manager = DagFileProcessorManager(max_runs=1) @@ -4221,6 +4832,34 @@ def test_refresh_dag_bundles_initialize_non_airflow_exception_skips_bundle(self) healthy_bundle.refresh.assert_called_once() + 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"] + def test_refresh_dag_bundles_update_bundle_state_failure_still_scans_files(self): """A failure in update_bundle_state() logs but does not skip file scanning. @@ -4392,7 +5031,7 @@ def test_unpack_bundle_version_with_string_non_versioned_bundle(self): @pytest.mark.usefixtures("testing_dag_bundle") def test_bundle_version_data_stored_after_refresh(self, session): - """Test that version_data from BundleVersion is stored in _bundle_version_data.""" + """Keep the local source metadata independently of its persisted bundle state.""" from airflow.dag_processing.bundles.base import BundleVersion manager = DagFileProcessorManager(max_runs=1) @@ -4408,7 +5047,7 @@ def test_bundle_version_data_stored_after_refresh(self, session): self._refresh_with_mocked_state(manager, bundle, BundleState(last_refreshed=None, version="oldhash")) assert manager._bundle_versions["mock_bundle"] == "newhash" - assert manager._bundle_version_data["mock_bundle"] == test_data + assert manager._bundle_parse_sources["mock_bundle"].version_data == test_data # --- statement budget --- # @@ -4438,14 +5077,14 @@ def _measure_persistence_call( errors: dict = {} update_dag_parsing_results_in_db( - "testing", None, dags, errors, 0.1, set(), session, files_parsed=files_parsed + "testing", None, dags, errors, 0.1, set(), session=session, files_parsed=files_parsed ) session.commit() assert not errors, f"fixture Dags must serialize cleanly: {errors}" with _count_statements(session) as counts: update_dag_parsing_results_in_db( - "testing", None, counted, errors, 0.1, set(), session, files_parsed=files_parsed + "testing", None, counted, errors, 0.1, set(), session=session, files_parsed=files_parsed ) session.flush() return counts @@ -4478,14 +5117,14 @@ def _measure_import_error_call(session, rel_path: str) -> Counter[tuple[str, str files_parsed = {("testing", rel_path)} recorded = {("testing", rel_path): "boom"} update_dag_parsing_results_in_db( - "testing", None, [], recorded, 0.1, set(), session, files_parsed=files_parsed + "testing", None, [], recorded, 0.1, set(), session=session, files_parsed=files_parsed ) session.commit() again = {("testing", rel_path): "boom again"} with _count_statements(session) as counts: update_dag_parsing_results_in_db( - "testing", None, [], again, 0.1, set(), session, files_parsed=files_parsed + "testing", None, [], again, 0.1, set(), session=session, files_parsed=files_parsed ) session.flush() return counts @@ -4648,6 +5287,7 @@ def test_log_file_processing_stats_omits_team_name_when_not_multi_team( @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"), @@ -4933,3 +5573,1590 @@ def test_normalized_file_path_for_stats_does_not_warn(caplog): assert result == "dags_test_test_dag.py" assert caplog.entries == [] + + +API_CLIENT_BUNDLES = [ + { + "name": name, + "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", + "kwargs": {"path": "/"}, + } + for name in ("bundle_a", "bundle_b") +] + + +def _read_bundle_connection() -> None: + Connection.get("bundle_conn") + + +class TestDagFileProcessorManagerWithAPIClient: + """With an API client, bundle code resolves its connections through it, for the bundle it belongs to.""" + + @pytest.fixture + def api_requests(self) -> list[httpx2.Request]: + return [] + + @pytest.fixture + def api_client(self, tmp_path, api_requests): + token_file = tmp_path / "processor.jwt" + token_file.write_text("session-token") + job_token = jwt.encode({"iat": 0, "exp": 300, "scope": "dag_processor", "job_id": 1}, "secret") + + def handle(request: httpx2.Request) -> httpx2.Response: + api_requests.append(request) + if request.url.path.endswith("/jobs"): + return httpx2.Response(201, json={"job_id": 1, "token": job_token, "expires_in": 300}) + if request.url.path.endswith("/bundles"): + return httpx2.Response( + 200, + json=[ + { + "name": bundle["name"], + "version": None, + "last_refreshed": None, + "revision": None, + "team_name": None, + } + for bundle in API_CLIENT_BUNDLES + ], + ) + return httpx2.Response( + 200, + json={ + "conn_id": "bundle_conn", + "conn_type": "http", + "host": None, + "schema": None, + "login": None, + "password": None, + "port": None, + "extra": None, + }, + ) + + with DagProcessorAPIClient( + base_url="http://api/execution/", + token_file=token_file, + hostname="processor-1", + transport=httpx2.MockTransport(handle), + ) as client: + client.register_job() + yield client + + @pytest.fixture + def make_manager(self, api_client): + managers = [] + + def make(**kwargs) -> DagFileProcessorManager: + manager = DagFileProcessorManager(max_runs=1, api_client=api_client, **kwargs) + manager.prepare_api_secrets_context(api_client) + managers.append(manager) + return manager + + yield make + for manager in managers: + manager.after_run() + + @pytest.fixture(autouse=True) + def configured_bundles(self): + with conf_vars( + { + ("dag_processor", "dag_bundle_config_list"): json.dumps(API_CLIENT_BUNDLES), + ("core", "load_examples"): "False", + } + ): + _load_bundle_config_snapshot.cache_clear() + yield + _load_bundle_config_snapshot.cache_clear() + + @staticmethod + def _get_connection_bundles(api_requests: list[httpx2.Request]) -> list[str]: + return [ + request.headers["Airflow-Dag-Bundle"] + for request in api_requests + if request.url.path.startswith("/execution/connections/") + ] + + def test_uses_the_api_client_instead_of_the_in_process_api(self, make_manager, api_client): + manager = make_manager() + + assert manager.client is api_client + assert manager._api_server is None + + @conf_vars({("secrets", "use_cache"): "True"}) + def test_api_mode_does_not_share_sdk_secret_cache_between_bundles(self, make_manager, monkeypatch): + monkeypatch.setattr(SecretCache, "_cache", {"inherited": "secret"}) + manager = make_manager() + manager.heartbeat = mock.Mock(spec=lambda: None, side_effect=SystemExit) + with pytest.raises(SystemExit): + manager._run_parsing_loop() + assert SecretCache._cache is None + + @mock.patch("airflow.dag_processing.manager.get_importer_registry", autospec=True) + def test_importer_warmup_selects_its_bundle(self, registry, make_manager, api_requests): + manager = make_manager(bundle_names_to_parse=["bundle_b"]) + manager._dag_bundles = manager.get_all_bundles() + registry.return_value.warm_importers.side_effect = _read_bundle_connection + manager.warm_importers() + assert self._get_connection_bundles(api_requests) == ["bundle_b"] + + @mock.patch.object( + DagFileProcessorManager, + "prepare_process_context", + autospec=True, + side_effect=ValueError("startup failed"), + ) + def test_startup_failure_clears_the_api_secrets_context(self, prepare, api_client): + manager = DagFileProcessorManager(max_runs=1, api_client=api_client) + with pytest.raises(ValueError, match="startup failed"): + manager.run() + assert not hasattr(task_runner, "SUPERVISOR_COMMS") + + @mock.patch("airflow.dag_processing.manager.gc", autospec=True) + def test_run_resolves_secrets_through_the_api_client_instead_of_the_database(self, _, api_client): + manager = DagFileProcessorManager(max_runs=1, api_client=api_client) + with ( + mock.patch.object(manager, "prepare_server_process_context") as server_context, + mock.patch.object(manager, "prepare_process_context"), + mock.patch.object(manager, "register_exit_signals"), + mock.patch.object(manager, "prepare_bundles"), + mock.patch.object(manager, "_symlink_latest_log_directory"), + mock.patch.object(manager, "warm_importers"), + ): + manager.before_run() + server_context.assert_not_called() + assert isinstance(task_runner.SUPERVISOR_COMMS, DagProcessorSecretsComms) + + manager.after_run() + + assert not hasattr(task_runner, "SUPERVISOR_COMMS") + + @pytest.mark.parametrize( + ("bundle_names_to_parse", "expected"), + [(None, ["bundle_a", "bundle_b"]), (["bundle_b"], ["bundle_b"])], + ) + @mock.patch("airflow.dag_processing.bundles.local.LocalDagBundle.__init__", autospec=True) + def test_creates_each_bundle_for_its_bundle( + self, mock_bundle, make_manager, api_requests, bundle_names_to_parse, expected + ): + def create(bundle, *, name, **kwargs): + _read_bundle_connection() + bundle.name = name + + mock_bundle.side_effect = create + + bundles = make_manager(bundle_names_to_parse=bundle_names_to_parse).get_all_bundles() + + assert [bundle.name for bundle in bundles] == expected + assert self._get_connection_bundles(api_requests) == expected + + def test_refreshes_each_bundle_for_its_bundle(self, make_manager, api_requests): + manager = make_manager() + bundle = MagicMock(spec=BaseDagBundle) + bundle.name = "bundle_b" + bundle.is_initialized = False + + def initialize(): + _read_bundle_connection() + raise OSError("clone failed") + + bundle.initialize.side_effect = initialize + manager._dag_bundles = [bundle] + + manager._refresh_dag_bundles(known_files={}) + + assert self._get_connection_bundles(api_requests) == ["bundle_b"] + + @mock.patch.object(DagBundlesManager, "get_bundle", autospec=True) + def test_prepares_a_callback_bundle_for_its_bundle(self, mock_get_bundle, make_manager, api_requests): + def create(_, name, version=None, version_data=None): + _read_bundle_connection() + bundle = MagicMock(spec=BaseDagBundle) + bundle.supports_versioning = False + return bundle + + mock_get_bundle.side_effect = create + request = DagCallbackRequest( + filepath="dag.py", + dag_id="dag_id", + run_id="run_id", + bundle_name="bundle_b", + bundle_version="v1", + context_from_server=None, + is_failure_callback=False, + ) + + assert make_manager().prepare_callback_bundle(request) is not None + assert self._get_connection_bundles(api_requests) == ["bundle_b"] + + @mock.patch("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"}) +def test_refresh_reads_one_fresh_catalog_per_pass(find_files, tmp_path): + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + manager = DagFileProcessorManager(max_runs=1, api_client=client, bundle_refresh_check_interval=0) + states = [] + known_files = {} + for name in ("a", "b", "c"): + bundle = _make_bundle(tmp_path) + bundle.name = name + bundle.is_initialized = True + bundle.supports_versioning = False + bundle.refresh_interval = 300 + manager._dag_bundles.append(bundle) + manager._bundle_versions[name] = None + manager._bundle_parse_sources[name] = DagParseSource() + known_files[name] = set() + states.append( + DagBundleStateResponse( + name=name, version=None, last_refreshed=timezone.utcnow(), revision=uuid7(), team_name=None + ) + ) + client.get_bundles.return_value = states + client.publish_inventory.return_value.revision = uuid7() + manager._refresh_dag_bundles(known_files) + client.get_bundles.assert_called_once_with() + find_files.assert_not_called() + + states[0] = states[0].model_copy( + update={"last_refreshed": None, "team_name": "new_team", "revision": uuid7()} + ) + states.pop() + manager._refresh_dag_bundles(known_files) + assert client.get_bundles.call_count == 2 + find_files.assert_called_once_with(manager, manager._dag_bundles[0]) + assert client.publish_inventory.call_args.args[1].expected_revision == states[0].revision + assert manager._bundle_name_to_team_name == {"a": "new_team", "b": None} + manager._dag_bundles[2].refresh.assert_not_called() + + file = DagFileInfo(bundle_name="b", rel_path=Path("a.py")) + known_files["b"] = {file} + pending_body = DagBundleInventoryBody( + attempt_id=uuid7(), dispatch_sequence=3, expected_revision=states[0].revision, files=[] + ) + manager._pending_inventories["a"] = (pending_body, set(), DagParseSource()) + client.get_bundles.side_effect = httpx2.ConnectError("API unavailable") + manager._refresh_dag_bundles(known_files) + assert client.get_bundles.call_count == 3 + assert known_files["b"] == {file} + assert client.publish_inventory.call_args.args[1] is pending_body + assert find_files.call_count == 1 + + +@pytest.mark.parametrize("catalog_failed", [False, True]) +@pytest.mark.parametrize("has_pending_inventory", [False, True]) +@mock.patch("airflow.dag_processing.manager.time", autospec=True) +@mock.patch.object(DagFileProcessorManager, "_find_files_in_bundle", autospec=True) +def test_catalog_omission_releases_priority_capacity_and_allows_recovery( + find_files, clock, catalog_failed, has_pending_inventory, tmp_path +): + clock.monotonic.return_value = 100 + (tmp_path / "a.py").touch() + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + manager = DagFileProcessorManager(max_runs=1, api_client=client, bundle_refresh_check_interval=10) + bundles = {} + files = {} + claims = {} + for name in ("missing", "ready"): + bundle = _make_bundle(tmp_path) + bundle.name = name + bundle.is_initialized = True + bundle.supports_versioning = False + bundle.refresh_interval = 300 + bundles[name] = bundle + file = DagFileInfo(bundle_name=name, rel_path=Path("a.py"), bundle_path=tmp_path) + files[name] = {file} + manager._bundle_versions[name] = None + manager._bundle_parse_sources[name] = DagParseSource() + 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() + ready_source = manager._bundle_parse_sources["ready"] + missing_source = manager._bundle_parse_sources["missing"] + pending_body = DagBundleInventoryBody( + attempt_id=uuid7(), dispatch_sequence=1, expected_revision=None, files=["a.py"] + ) + if has_pending_inventory: + manager._pending_inventories["missing"] = (pending_body, files["missing"], missing_source) + manager._bundle_parse_sources.pop("missing") + manager._force_refresh_bundles.add("missing") + manager._bundles_waiting_for_refresh.add("missing") + ready_state = DagBundleStateResponse( + name="ready", version=None, last_refreshed=timezone.utcnow(), revision=None, team_name=None + ) + client.get_bundles.return_value = [ready_state] + if catalog_failed: + client.get_bundles.side_effect = httpx2.ConnectError("Catalog unavailable") + client.publish_inventory.side_effect = _make_status_error(503) + + manager._refresh_dag_bundles(files) + + client.get_bundles.assert_called_once_with() + find_files.assert_not_called() + assert files == previous_files + assert manager._bundle_versions == previous_versions + assert manager._bundle_parse_sources["ready"] is ready_source + 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: + client.publish_inventory.assert_called_once_with("missing", pending_body) + assert manager._pending_inventories["missing"][0] is pending_body + else: + client.publish_inventory.assert_not_called() + assert manager._bundle_parse_sources["missing"] is missing_source + return + + client.publish_inventory.assert_not_called() + assert 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"})] + client.publish_inventory.side_effect = None + revision = uuid7() + client.publish_inventory.return_value.revision = revision + find_files.return_value = files["missing"] + manager._refresh_dag_bundles(files) + + bundles["missing"].refresh.assert_called_once() + bundles["ready"].refresh.assert_not_called() + assert manager._bundle_parse_sources["missing"].bundle_revision == revision + assert files == previous_files + assert manager._bundle_versions == previous_versions + + +@pytest.mark.parametrize("status_code", [409, 503]) +def test_api_inventory_failure_preserves_known_files(status_code): + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + manager = DagFileProcessorManager(max_runs=1, api_client=client) + body = DagBundleInventoryBody(attempt_id=uuid7(), dispatch_sequence=1, expected_revision=None, files=[]) + manager._pending_inventories["testing"] = (body, set(), DagParseSource()) + known_files = {"testing": set(_get_file_infos(["previous.py"]))} + previous = known_files.copy() + response = httpx2.Response(status_code, request=httpx2.Request("POST", "http://api/inventory")) + client.publish_inventory.side_effect = httpx2.HTTPStatusError( + "Unavailable", request=response.request, response=response + ) + manager._publish_bundle_inventory("testing", known_files) + assert known_files == previous + assert not manager._bundle_parse_sources + if status_code == 409: + assert not manager._pending_inventories + assert "testing" in manager._force_refresh_bundles + else: + assert manager._pending_inventories["testing"][0] is body + revision = uuid7() + client.publish_inventory.side_effect = None + client.publish_inventory.return_value = DagBundleInventoryResponse( + attempt_id=body.attempt_id, accepted_at=timezone.utcnow(), revision=revision + ) + manager._publish_bundle_inventory("testing", known_files) + assert client.publish_inventory.call_args.args[1] is body + assert known_files == {"testing": set()} + assert manager._bundle_parse_sources["testing"].bundle_revision == revision + + +@pytest.mark.parametrize("rejection", ["local_limit", "server_limit"]) +@mock.patch("airflow.dag_processing.manager.MAX_PUBLICATION_BYTES", 1024) +@mock.patch("airflow.dag_processing.manager.time", autospec=True) +@mock.patch.object(DagFileProcessorManager, "_create_process", autospec=True) +def test_oversized_inventory_preserves_previous_inventory_and_other_bundle_progress( + create, clock, rejection, tmp_path +): + clock.monotonic.return_value = 100 + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + client.get_bundles.return_value = [ + DagBundleStateResponse(name=name, version=None, last_refreshed=None, revision=None, team_name=None) + for name in ("oversized", "ready") + ] + manager = DagFileProcessorManager(max_runs=1, api_client=client, bundle_refresh_check_interval=10) + old_file, ready_file = [ + DagFileInfo(bundle_name=name, rel_path=Path("a.py"), bundle_path=tmp_path) + for name in ("oversized", "ready") + ] + (tmp_path / "a.py").touch() + known_files = {"oversized": {old_file}} + manager._bundle_versions["oversized"] = "previous" + manager._force_refresh_bundles.add("oversized") + for name, files in ( + ("oversized", [f"{index}/{'é' * 80}.py" for index in range(6)] if rejection == "local_limit" else []), + ("ready", ["a.py"]), + ): + bundle = _make_bundle(tmp_path) + bundle.name = name + manager._dag_bundles.append(bundle) + body = DagBundleInventoryBody( + attempt_id=uuid7(), dispatch_sequence=1, expected_revision=None, version="new", files=files + ) + if name == "oversized" and rejection == "local_limit": + assert len(body.model_dump_json()) < 1024 < len(body.model_dump_json().encode()) + manager._pending_inventories[name] = ( + body, + {ready_file} if name == "ready" else set(), + DagParseSource(bundle_version="new"), + ) + + def publish(bundle_name, body): + if bundle_name == "oversized": + assert rejection == "server_limit", "The oversized inventory must be rejected before sending" + raise _make_status_error(413) + return DagBundleInventoryResponse( + attempt_id=body.attempt_id, accepted_at=timezone.utcnow(), revision=uuid7() + ) + + client.publish_inventory.side_effect = publish + manager._refresh_dag_bundles(known_files) + manager._start_new_processes() + + assert known_files == {"oversized": {old_file}, "ready": {ready_file}} + assert manager._bundle_versions == {"oversized": "previous", "ready": "new"} + assert not manager._pending_inventories + assert not manager._force_refresh_bundles + assert set(manager._bundle_parse_sources) == {"ready"} + assert list(manager._files_awaiting_source) == [old_file] + create.assert_called_once_with(manager, ready_file) + calls = client.get_bundles.call_count + manager._refresh_dag_bundles(known_files) + assert client.get_bundles.call_count == calls + + +@pytest.mark.parametrize("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) + manager = DagFileProcessorManager(max_runs=1, api_client=client) + body = DagBundleInventoryBody( + attempt_id=uuid7(), dispatch_sequence=1, expected_revision=None, files=["é.py"] + ) + pending = (body, set(), DagParseSource()) + manager._pending_inventories["testing"] = pending + client.publish_inventory.return_value = DagBundleInventoryResponse( + attempt_id=body.attempt_id, accepted_at=timezone.utcnow(), revision=uuid7() + ) + size = len(body.model_dump_json().encode()) + with mock.patch("airflow.dag_processing.manager.MAX_PUBLICATION_BYTES", size): + manager._publish_bundle_inventory("testing", {}) + client.publish_inventory.assert_called_once() + manager._pending_inventories["testing"] = pending + with mock.patch("airflow.dag_processing.manager.MAX_PUBLICATION_BYTES", size - 1): + manager._publish_bundle_inventory("testing", {}) + client.publish_inventory.assert_called_once() + assert not manager._pending_inventories + assert not manager._bundle_parse_sources + manager._pending_inventories["testing"] = pending + client.publish_inventory.side_effect = _make_status_error(auth_status) + with pytest.raises(httpx2.HTTPStatusError): + manager._publish_bundle_inventory("testing", {}) + + +@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 = manager._priority_claims[work.id].attempt_id + response = httpx2.Response( + 409, + json={"detail": {"reason": "source_changed"}}, + request=httpx2.Request("POST", "http://api/parse-results"), + ) + client.publish_parse_result.side_effect = httpx2.HTTPStatusError( + "Rejected", request=response.request, response=response + ) + + manager._publish_pending_results() + + assert not manager._pending_publications + assert manager._force_refresh_bundles == {"testing"} + assert list(manager._file_queue) == [file] + 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"}) +def test_api_team_lookup_failure_is_retried_without_stopping_the_processor(): + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + client.get_bundles.side_effect = [ + httpx2.ConnectError("unavailable"), + [ + DagBundleStateResponse( + name="testing", version=None, last_refreshed=None, revision=None, team_name="team_a" + ) + ], + ] + manager = DagFileProcessorManager(max_runs=1, api_client=client) + + assert manager._get_team_names({"testing"}) == {"testing": None} + assert manager._get_team_names({"testing"}) == {"testing": "team_a"} + + +@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: + dag = SdkDAG(f"dag_in_{file.rel_path.stem}", schedule=None) + dag.fileloc = f"/worker/{file.rel_path}" + dag.relative_fileloc = str(file.rel_path) + return PendingDagPublication( + body=DagParseResultBody( + attempt_id=attempt_id or uuid7(), + dispatch_sequence=1, + bundle_name=file.bundle_name, + relative_fileloc=str(file.rel_path), + parse_duration=0.1, + serialized_dags=[LazyDeserializedDAG.from_dag(dag).data], + source_codes={dag.dag_id: ParseSourceCode(source_code="# source", language="python")}, + import_errors=import_errors or {}, + ), + stat=DagFileStat(num_dags=1, run_count=1, last_finish_time=timezone.utcnow()), + refresh_generation=0, + ) + + +def _make_status_error( + status_code: int, reason: str | None = None, *, message: str | None = None +) -> httpx2.HTTPStatusError: + response = httpx2.Response( + status_code, + json={"detail": {"reason": reason, "message": message}} if reason else None, + request=httpx2.Request("POST", "http://api/parse-results"), + ) + return httpx2.HTTPStatusError("Rejected", request=response.request, response=response) + + +def test_publication_sends_every_ready_result_within_the_loop_budget(): + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + manager = DagFileProcessorManager(max_runs=1, api_client=client) + files = _get_file_infos(["a.py", "b.py"]) + for file in files: + manager._pending_publications[file] = _make_pending_publication(file) + + manager._publish_pending_results(deadline=time.monotonic()) + client.publish_parse_result.assert_not_called() + + manager._publish_pending_results() + sent = [call.args[0].relative_fileloc for call in client.publish_parse_result.call_args_list] + assert sent == ["a.py", "b.py"] + assert not manager._pending_publications + assert all(manager._file_stats[file].num_dags == 1 for file in files) + + +@conf_vars({("workers", "execution_api_retries"): "1"}) +@pytest.mark.parametrize( + ("error", "message"), + [ + (_make_status_error(409, "dag_owned_by_another_bundle"), "HTTP 409: dag_owned_by_another_bundle"), + (_make_status_error(403, "plugin_not_granted"), "HTTP 403: plugin_not_granted"), + (_make_status_error(413), "HTTP 413"), + (_make_status_error(422), "HTTP 422"), + ( + _make_status_error(422, "invalid_dag", message="Unknown role: missing_role"), + "HTTP 422: invalid_dag: Unknown role: missing_role", + ), + ], +) +def test_result_the_api_refuses_is_replaced_by_an_import_error(error, message): + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + client.publish_parse_result.side_effect = [error, None] + manager = DagFileProcessorManager(max_runs=1, api_client=client) + manager._dispatch_sequence = 7 + file = _get_file_infos(["a.py"])[0] + work = _claim_priority_parse(manager, file) + original = _make_pending_publication( + file, + attempt_id=manager._priority_claims[work.id].attempt_id, + import_errors={"a.py/member.py": "member failed"}, + ) + manager._pending_publications[file] = original + + manager._publish_pending_results() + + report = client.publish_parse_result.call_args_list[1].args[0] + assert report.attempt_id != original.body.attempt_id + assert report.dispatch_sequence == 8 + assert report.serialized_dags == [] + assert report.source_codes == {} + assert report.import_errors == { + "a.py": f"The parsing result was not published: the API server rejected it ({message}).", + } + assert not manager._pending_publications + assert manager._file_stats[file].num_dags == 0 + assert manager._file_stats[file].import_errors == 1 + assert manager._pending_work_acks == [("priority", work, True)] + + +@conf_vars({("workers", "execution_api_retries"): "1"}) +def test_server_failure_after_publication_retries_preserves_previous_result(): + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + client.publish_parse_result.side_effect = [_make_status_error(500), None] + manager = DagFileProcessorManager(max_runs=1, api_client=client) + file = _get_file_infos(["a.py"])[0] + 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=manager._priority_claims[work.id].attempt_id) + manager._pending_publications[file] = pending + + manager._publish_pending_results() + + client.publish_parse_result.assert_called_once_with(pending.body) + assert not manager._pending_publications + assert manager._file_stats[file].num_dags == previous.num_dags + assert manager._file_stats[file].last_finish_time == previous.last_finish_time + assert manager._file_stats[file].last_attempt_time is not None + assert manager._pending_work_acks == [("priority", work, True)] + + +@mock.patch("airflow.dag_processing.manager.MAX_PUBLICATION_BYTES", 32 * 1024) +@pytest.mark.parametrize("oversized_field", ["import_errors", "version_data", "parsed_definitions"]) +def test_rejection_report_discards_oversized_parse_metadata(oversized_field): + limit = 32 * 1024 + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + client.publish_parse_result.side_effect = [_make_status_error(413), None] + manager = DagFileProcessorManager(max_runs=1, api_client=client) + file = _get_file_infos(["a.py"])[0] + original = _make_pending_publication(file) + original.body.bundle_version = "v1" + original.body.bundle_revision = uuid7() + oversized = { + "import_errors": {"a.py/member.py": "é" * limit}, + "version_data": {"metadata": "é" * limit}, + "parsed_definitions": ["a.py/member.py"] * limit, + } + setattr(original.body, oversized_field, oversized[oversized_field]) + assert len(original.body.model_dump_json().encode()) > limit + manager._pending_publications[file] = original + + manager._publish_pending_results() + + report = client.publish_parse_result.call_args_list[1].args[0] + assert len(report.model_dump_json().encode()) < limit + assert DagParseResultBody.model_validate_json(report.model_dump_json()) == report + assert report.bundle_version == original.body.bundle_version + assert report.bundle_revision == original.body.bundle_revision + assert report.version_data is None + assert report.parsed_definitions == [] + assert list(report.import_errors) == ["a.py"] + assert manager._file_stats[file].import_errors == 1 + assert not manager._pending_publications + + +@mock.patch("airflow.dag_processing.manager.MAX_PUBLICATION_BYTES", 32 * 1024) +def test_rejection_report_bounds_server_validation_message(): + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + message = "Unknown role: " + "é" * (32 * 1024) + client.publish_parse_result.side_effect = [_make_status_error(422, "invalid_dag", message=message), None] + manager = DagFileProcessorManager(max_runs=1, api_client=client) + file = _get_file_infos(["a.py"])[0] + manager._pending_publications[file] = _make_pending_publication(file) + + manager._publish_pending_results() + + report = client.publish_parse_result.call_args_list[1].args[0] + assert "Unknown role: é" in report.import_errors["a.py"] + assert len(report.model_dump_json().encode()) < 32 * 1024 + + +def test_direct_manager_does_not_import_api_client(monkeypatch, mocker): + monkeypatch.setitem(sys.modules, "airflow.dag_processing.api_client", None) + namespace = runpy.run_path(sys.modules[DagFileProcessorManager.__module__].__file__) + manager_class = namespace["DagFileProcessorManager"] + bundle_manager = mocker.patch.object(manager_class, "_create_bundle_manager", autospec=True) + bundle_manager.return_value.get_bundle.side_effect = ValueError("missing bundle") + manager = manager_class(max_runs=1) + + manager._publish_pending_results() + manager.after_run() + request = DagCallbackRequest( + filepath="a.py", + dag_id="dag1", + run_id="run1", + is_failure_callback=False, + bundle_name="testing", + bundle_version="v1", + ) + assert manager.prepare_callback_bundle(request) is None + + +@pytest.mark.parametrize("reason", ["bundle_inactive", "publication_superseded", None]) +def test_refusal_unrelated_to_content_is_not_reported_as_an_import_error(reason): + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + client.publish_parse_result.side_effect = _make_status_error(409, reason) + manager = DagFileProcessorManager(max_runs=1, api_client=client) + file = _get_file_infos(["a.py"])[0] + manager._pending_publications[file] = _make_pending_publication(file) + + manager._publish_pending_results() + + client.publish_parse_result.assert_called_once() + assert not manager._pending_publications + + +def test_refused_import_error_report_is_not_replaced_again(): + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + client.publish_parse_result.side_effect = _make_status_error(413) + manager = DagFileProcessorManager(max_runs=1, api_client=client) + file = _get_file_infos(["a.py"])[0] + manager._pending_publications[file] = _make_pending_publication(file) + + manager._publish_pending_results() + + assert client.publish_parse_result.call_count == 2 + assert not manager._pending_publications + assert manager._file_stats[file].num_dags == 0 + + +@mock.patch("airflow.dag_processing.manager.MAX_PUBLICATION_BYTES", 1024 * 1024) +def test_result_over_the_publication_limit_is_reported_without_sending_it(): + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + manager = DagFileProcessorManager(max_runs=1, api_client=client) + file = _get_file_infos(["a.py"])[0] + 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 = manager._priority_claims[work.id].attempt_id + proc.start_time = time.monotonic() + proc.had_callbacks = False + proc.dispatch_sequence = 1 + proc.parse_source = DagParseSource() + proc.parsing_result = DagFileParsingResult( + fileloc="a.py", serialized_dags=[LazyDeserializedDAG.from_dag(dag)] + ) + + manager.handle_parsing_result(file, proc) + + pending = manager._pending_publications[file] + assert pending.is_rejection_report + assert pending.body.serialized_dags == [] + assert pending.body.import_errors == { + "a.py": "The parsing result was not published: it exceeds the 1 MiB publication limit." + } + assert manager._pending_work_acks == [("priority", work, True)] + client.publish_parse_result.assert_not_called() + + +def test_acknowledgments_are_sent_within_the_loop_budget(): + 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") + 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() + + 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/dag_processing/test_processor.py b/airflow-core/tests/unit/dag_processing/test_processor.py index 629cb6b1c49b1..d00fe91987e34 100644 --- a/airflow-core/tests/unit/dag_processing/test_processor.py +++ b/airflow-core/tests/unit/dag_processing/test_processor.py @@ -18,6 +18,7 @@ from __future__ import annotations import inspect +import json import logging import os import pathlib @@ -29,8 +30,10 @@ from collections.abc import Callable, Iterable from socket import socket, socketpair from typing import TYPE_CHECKING, Any, BinaryIO -from unittest.mock import MagicMock, PropertyMock, patch +from unittest.mock import ANY, MagicMock, PropertyMock, patch +import httpx2 +import jwt import pytest import structlog from pydantic import TypeAdapter @@ -51,6 +54,7 @@ EmailRequest, TaskCallbackRequest, ) +from airflow.dag_processing.api_client import DagProcessorAPIClient from airflow.dag_processing.dagbag import DagBag from airflow.dag_processing.manager import DagFileProcessorManager, process_parse_results from airflow.dag_processing.processor import ( @@ -85,6 +89,7 @@ TICount, ToSupervisor, ToTask, + VariableResult, XComResult, XComSequenceSliceResult, ) @@ -2414,7 +2419,11 @@ def test_registered_message_types(self): @pytest.mark.parametrize( "message_type", - sorted(set(typing.get_args(typing.get_args(ToManager)[0])) - {DagFileParsingResult}, key=str), + sorted( + set(typing.get_args(typing.get_args(ToManager)[0])) + - {DagFileParsingResult, comms.GetPrevSuccessfulDagRun}, + key=str, + ), ids=lambda message_type: message_type.__name__, ) def test_reuses_shared_request_handlers(self, message_type): @@ -2480,15 +2489,92 @@ def test_dispatch_parsing_result(self, send_msg, proc): send_msg.assert_called_once_with(proc, None, request_id=42, error=None) @patch.object(DagFileProcessorProcess, "send_msg", autospec=True) - def test_previous_successful_run_uses_process_id(self, send_msg, proc): - proc.client.task_instances.get_previous_successful_dagrun.return_value = ( - comms.PrevSuccessfulDagRunResult() - ) + def test_previous_successful_run_is_answered_locally(self, send_msg, proc): proc._handle_request( comms.GetPrevSuccessfulDagRun(ti_id=uuid.uuid4()), structlog.get_logger(), req_id=42 ) - proc.client.task_instances.get_previous_successful_dagrun.assert_called_once_with(proc.id) + assert not proc.client.mock_calls + send_msg.assert_called_once_with( + proc, comms.PrevSuccessfulDagRunResult(), request_id=42, error=None, exclude_unset=True + ) + + @patch.object(supervisor.WatchedSubprocess, "_handle_request", autospec=True) + def test_default_mode_requests_do_not_load_the_api_client(self, handle_request, proc, monkeypatch): + monkeypatch.setitem(sys.modules, "airflow.dag_processing.api_client", None) + msg = GetVariable(key="my_key") + + proc._handle_request(msg, structlog.get_logger(), req_id=42) + + handle_request.assert_called_once_with(proc, msg, ANY, 42) + + @patch.object(DagFileProcessorProcess, "send_msg", autospec=True) + def test_api_client_requests_are_made_for_the_files_bundle(self, send_msg, proc, tmp_path): + token_file = tmp_path / "processor.jwt" + token_file.write_text("session-token") + job_token = jwt.encode({"iat": 0, "exp": 300, "scope": "dag_processor", "job_id": 1}, "secret") + parse_token = jwt.encode( + { + "iat": 0, + "exp": 300, + "sub": str(proc.id), + "scope": "dag_parse", + "job_id": 1, + "dag_bundles": [proc.bundle_name], + "relative_fileloc": proc.dag_file_rel_path, + }, + "secret", + ) + requests = [] + + def handle(request): + requests.append(request) + if request.url.path.endswith("/jobs"): + return httpx2.Response(201, json={"job_id": 1, "token": job_token, "expires_in": 300}) + if request.url.path.endswith("/parse-token"): + return httpx2.Response(200, json={"token": parse_token, "expires_in": 300}) + return httpx2.Response(200, json={"key": "my_key", "value": "my_value"}) + + with DagProcessorAPIClient( + base_url="http://api/execution/", + token_file=token_file, + hostname="processor-1", + transport=httpx2.MockTransport(handle), + ) as client: + client.register_job() + proc.client = client + proc._handle_request(GetVariable(key="my_key"), structlog.get_logger(), req_id=42) + proc._handle_request(GetVariable(key="my_key"), structlog.get_logger(), req_id=43) + + assert requests[-1].url.path == "/execution/variables/my_key" + assert requests[-1].headers["Airflow-Dag-Bundle"] == "mybundle" + assert requests[-1].headers["Authorization"] == f"Bearer {parse_token}" + assert json.loads(requests[1].content) == { + "attempt_id": str(proc.id), + "bundle_name": proc.bundle_name, + "relative_fileloc": proc.dag_file_rel_path, + } + assert sum(request.url.path.endswith("/parse-token") for request in requests) == 1 + assert send_msg.call_count == 2 + send_msg.assert_any_call( + proc, + VariableResult(key="my_key", value="my_value"), + request_id=42, + error=None, + exclude_unset=True, + ) + + @patch.object(DagFileProcessorProcess, "send_msg", autospec=True) + def test_local_parse_messages_do_not_require_a_token_exchange(self, send_msg, proc, tmp_path): + with DagProcessorAPIClient( + base_url="http://unused/", token_file=tmp_path / "missing-token", hostname="processor-1" + ) as client: + proc.client = client + proc._handle_request( + comms.GetPrevSuccessfulDagRun(ti_id=uuid.uuid4()), structlog.get_logger(), req_id=42 + ) + + assert client.job_id is None send_msg.assert_called_once_with( proc, comms.PrevSuccessfulDagRunResult(), request_id=42, error=None, exclude_unset=True ) diff --git a/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 0000000000000..f55ed3ca5242c --- /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/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index 7118d08edc61f..b7eb013f76e6c 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -59,6 +59,7 @@ from airflow.callbacks.database_callback_sink import DatabaseCallbackSink from airflow.dag_processing.collection import AssetModelOperation, DagModelOperation from airflow.dag_processing.dagbag import DagBag, sync_bag_to_db +from airflow.dag_processing.maintenance import cleanup_processor_metadata from airflow.exceptions import AirflowException, DeserializationError, TaskNotFound from airflow.executors.base_executor import BaseExecutor from airflow.executors.executor_constants import MOCK_EXECUTOR @@ -161,6 +162,7 @@ PartitionedAssetTimetable as CorePartitionedAssetTimetable, PartitionedAtRuntime, ) +from airflow.utils.event_scheduler import EventScheduler from airflow.utils.session import NEW_SESSION, create_session, provide_session from airflow.utils.sqlalchemy import CommitProhibitorGuard, with_row_locks from airflow.utils.state import CallbackState, DagRunState, DagSchedulingState, State, TaskInstanceState @@ -4256,6 +4258,19 @@ def _create_dagruns(): assert total_enqueued == 31 session.rollback() + @conf_vars({("scheduler", "parsing_cleanup_interval"): "123"}) + @mock.patch.object(EventScheduler, "call_regular_interval", autospec=True) + def test_dag_processor_maintenance_runs_on_the_parsing_cleanup_interval(self, call_regular_interval): + scheduler_job = Job() + scheduler_job.heartbeat = mock.MagicMock() + self.job_runner = SchedulerJobRunner( + job=scheduler_job, num_runs=1, executors=[MockExecutor(do_update=False)] + ) + + self.job_runner._run_scheduler_loop() + + call_regular_interval.assert_any_call(mock.ANY, 123.0, cleanup_processor_metadata, non_fatal=True) + def test_adopt_or_reset_orphaned_tasks(self, dag_maker, session): with dag_maker("test_execute_helper_reset_orphaned_tasks", session=session): op1 = EmptyOperator(task_id="op1") diff --git a/airflow-core/tests/unit/models/test_dagwarning.py b/airflow-core/tests/unit/models/test_dagwarning.py index 0f07d7a91924b..b62d25fd48d8a 100644 --- a/airflow-core/tests/unit/models/test_dagwarning.py +++ b/airflow-core/tests/unit/models/test_dagwarning.py @@ -17,18 +17,10 @@ from __future__ import annotations -from unittest import mock -from unittest.mock import MagicMock - import pytest -from sqlalchemy import select -from sqlalchemy.exc import OperationalError -from airflow.models import DagModel from airflow.models.dagwarning import DagWarning, DagWarningType, get_warning_type_value -from tests_common.test_utils.db import clear_db_dags - class TestGetWarningTypeValue: @pytest.mark.parametrize( @@ -61,49 +53,3 @@ def test_invalid_type_is_rejected(self, warning_type): def test_dag_warning_stores_importer_type(self): assert DagWarning("dag_1", "yaml:schema_violation", "message").warning_type == "yaml:schema_violation" - - -@pytest.mark.db_test -class TestDagWarning: - def setup_method(self): - clear_db_dags() - - def test_purge_inactive_dag_warnings(self, session, testing_dag_bundle): - """ - Test that the purge_inactive_dag_warnings method deletes inactive dag warnings - """ - dags = [ - DagModel(dag_id="dag_1", bundle_name="testing", is_stale=True), - DagModel(dag_id="dag_2", bundle_name="testing", is_stale=False), - ] - session.add_all(dags) - session.commit() - - dag_warnings = [ - DagWarning("dag_1", "non-existent pool", "non-existent pool"), - DagWarning("dag_2", "non-existent pool", "non-existent pool"), - ] - session.add_all(dag_warnings) - session.commit() - - DagWarning.purge_inactive_dag_warnings(session=session) - - remaining_dag_warnings = session.scalars(select(DagWarning)).all() - assert len(remaining_dag_warnings) == 1 - assert remaining_dag_warnings[0].dag_id == "dag_2" - - @mock.patch("airflow.models.dagwarning.delete") - def test_retry_purge_inactive_dag_warnings(self, delete_mock): - """ - Test that the purge_inactive_dag_warnings method calls the delete method twice - if the query throws an operationalError on the first call and works on the second attempt - """ - self.session_mock = MagicMock() - - self.session_mock.execute.side_effect = [OperationalError(None, None, "database timeout"), None] - - DagWarning.purge_inactive_dag_warnings(session=self.session_mock) - - # Assert that the delete method was called twice - assert delete_mock.call_count == 2 - assert self.session_mock.execute.call_count == 2 diff --git a/airflow-core/tests/unit/utils/test_db_cleanup.py b/airflow-core/tests/unit/utils/test_db_cleanup.py index 57eeb9c76fd6c..24badfd966e5f 100644 --- a/airflow-core/tests/unit/utils/test_db_cleanup.py +++ b/airflow-core/tests/unit/utils/test_db_cleanup.py @@ -1223,6 +1223,7 @@ def test_no_models_missing(self): "asset_partition_dag_run", "asset_watcher", # cascade from trigger "dag_favorite", # cascade from dag + "dag_parse_checkpoint", # cascade from job "deadline_alert", # cascade from serialized_dag, which cascades from dag_version "dynamic_region", # cascade from dag_run; archived with it, never aged out on its own "hitl_detail", # cascade from task_instance diff --git a/chart/templates/_helpers.yaml b/chart/templates/_helpers.yaml index 9a5238fdc8c91..fb84e9c7d3b3c 100644 --- a/chart/templates/_helpers.yaml +++ b/chart/templates/_helpers.yaml @@ -838,8 +838,13 @@ server_tls_key_file = /etc/pgbouncer/server.key - sh - -c - | + {{- if (get (.Values.config.dag_processor | default dict) "execution_api_token_file") }} + CONNECTION_CHECK_MAX_COUNT=0 AIRFLOW__LOGGING__LOGGING_LEVEL=ERROR exec /entrypoint \ + airflow dag-processor --check-health + {{- else }} CONNECTION_CHECK_MAX_COUNT=0 AIRFLOW__LOGGING__LOGGING_LEVEL=ERROR exec /entrypoint \ airflow jobs check --local --job-type DagProcessorJob + {{- end }} {{- end }} {{- define "registry_docker_config" }} diff --git a/chart/templates/dag-processor/dag-processor-deployment.yaml b/chart/templates/dag-processor/dag-processor-deployment.yaml index ca8453c2bfe83..28e3d45f213ef 100644 --- a/chart/templates/dag-processor/dag-processor-deployment.yaml +++ b/chart/templates/dag-processor/dag-processor-deployment.yaml @@ -21,6 +21,9 @@ ## Airflow Dag Processor Deployment ################################# {{- if .Values.dagProcessor.enabled }} +{{- if and (get (.Values.config.dag_processor | default dict) "execution_api_token_file") (gt (int .Values.dagProcessor.replicas) 1) }} +{{- fail "`dagProcessor.replicas` must not exceed 1 when `config.dag_processor.execution_api_token_file` is set, because the pods would share one session token. Deploy a separate release, with its own token, for each additional processor." }} +{{- end }} {{- $nodeSelector := or .Values.dagProcessor.nodeSelector .Values.nodeSelector }} {{- $affinity := or .Values.dagProcessor.affinity .Values.affinity }} {{- $tolerations := or .Values.dagProcessor.tolerations .Values.tolerations }} @@ -58,7 +61,11 @@ spec: tier: airflow component: dag-processor release: {{ .Release.Name }} - {{- if .Values.dagProcessor.strategy }} + {{- if (get (.Values.config.dag_processor | default dict) "execution_api_token_file") }} + {{- /* Pods share one session token, and a new pod can register its Job only after the old one stops. */}} + strategy: + type: Recreate + {{- else if .Values.dagProcessor.strategy }} strategy: {{- toYaml .Values.dagProcessor.strategy | nindent 4 }} {{- end }} template: @@ -189,6 +196,18 @@ spec: {{- include "custom_airflow_environment" . | indent 10 }} {{- include "standard_airflow_environment" . | indent 10 }} {{- include "container_extra_envs" (list . .Values.dagProcessor.env) | indent 10 }} + {{- if (get (.Values.config.dag_processor | default dict) "execution_api_token_file") }} + readinessProbe: + timeoutSeconds: {{ .Values.dagProcessor.livenessProbe.timeoutSeconds }} + periodSeconds: {{ .Values.dagProcessor.livenessProbe.periodSeconds }} + exec: + command: + - sh + - -c + - >- + CONNECTION_CHECK_MAX_COUNT=0 AIRFLOW__LOGGING__LOGGING_LEVEL=ERROR exec /entrypoint + airflow dag-processor --check-ready + {{- end }} livenessProbe: initialDelaySeconds: {{ .Values.dagProcessor.livenessProbe.initialDelaySeconds }} timeoutSeconds: {{ .Values.dagProcessor.livenessProbe.timeoutSeconds }} diff --git a/chart/tests/helm_tests/airflow_core/test_dag_processor.py b/chart/tests/helm_tests/airflow_core/test_dag_processor.py index ec6dcaab47513..7fdff7fb35a7d 100644 --- a/chart/tests/helm_tests/airflow_core/test_dag_processor.py +++ b/chart/tests/helm_tests/airflow_core/test_dag_processor.py @@ -16,8 +16,13 @@ # under the License. from __future__ import annotations +from configparser import ConfigParser +from pathlib import Path +from subprocess import CalledProcessError + import jmespath import pytest +import yaml from chart_utils.helm_template_generator import render_chart from chart_utils.log_groomer import LogGroomerTestBase @@ -434,6 +439,108 @@ def test_livenessprobe_command(self): in jmespath.search("spec.template.spec.containers[0].livenessProbe.exec.command", docs[0])[-1] ) + def test_api_mode_uses_local_liveness_probe(self): + docs = render_chart( + values={"config": {"dag_processor": {"execution_api_token_file": "/run/processor.jwt"}}}, + show_only=["templates/dag-processor/dag-processor-deployment.yaml"], + ) + command = jmespath.search("spec.template.spec.containers[0].livenessProbe.exec.command", docs[0])[-1] + assert "airflow dag-processor --check-health" in command + readiness = docs[0]["spec"]["template"]["spec"]["containers"][0]["readinessProbe"]["exec"]["command"] + assert "airflow dag-processor --check-ready" in readiness[-1] + assert "airflow jobs check" not in command + + def test_api_mode_replaces_pods_with_recreate(self): + docs = render_chart( + values={ + "config": {"dag_processor": {"execution_api_token_file": "/run/processor.jwt"}}, + "dagProcessor": { + "strategy": {"rollingUpdate": {"maxSurge": "100%", "maxUnavailable": "50%"}} + }, + }, + show_only=["templates/dag-processor/dag-processor-deployment.yaml"], + ) + + assert jmespath.search("spec.strategy", docs[0]) == {"type": "Recreate"} + + @pytest.mark.parametrize(("replicas", "rejected"), [(0, False), (1, False), (2, True)]) + def test_api_mode_rejects_replicas_sharing_a_session_token(self, replicas, rejected): + values = { + "config": {"dag_processor": {"execution_api_token_file": "/run/processor.jwt"}}, + "dagProcessor": {"replicas": replicas}, + } + show_only = ["templates/dag-processor/dag-processor-deployment.yaml"] + if rejected: + with pytest.raises(CalledProcessError, match="must not exceed 1"): + render_chart(values=values, show_only=show_only) + else: + docs = render_chart(values=values, show_only=show_only) + assert jmespath.search("spec.replicas", docs[0]) == replicas + + def test_isolated_api_deployment_example(self): + examples = Path(__file__).resolve().parents[4] / "airflow-core/docs/security/dag-processor-api" + values = yaml.safe_load((examples / "helm-values.yaml").read_text()) + docs = render_chart(values=values, namespace="dag-processors") + workloads = [doc for doc in docs if doc["kind"] in {"Deployment", "StatefulSet", "Job", "CronJob"}] + assert len(workloads) == 1 + deployment = workloads[0] + assert deployment["metadata"]["labels"]["component"] == "dag-processor" + assert deployment["spec"]["strategy"] == {"type": "Recreate"} + pod = deployment["spec"]["template"]["spec"] + assert not pod.get("initContainers") + for container in pod["containers"]: + assert not container.get("envFrom") + assert all("valueFrom" not in env for env in container.get("env", [])) + processor = next(container for container in pod["containers"] if container["name"] == "dag-processor") + assert {env["name"]: env["value"] for env in processor["env"]} == { + "AIRFLOW_HOME": "/opt/airflow", + "CONNECTION_CHECK_MAX_COUNT": "0", + } + assert "--check-health" in processor["livenessProbe"]["exec"]["command"][-1] + assert "--check-ready" in processor["readinessProbe"]["exec"]["command"][-1] + assert {v["secret"]["secretName"] for v in pod["volumes"] if "secret" in v} == { + "dag-processor-session" + } + token_mount = next(m for m in processor["volumeMounts"] if m["name"] == "processor-token") + assert token_mount["readOnly"] + assert "subPath" not in token_mount + account = next( + doc + for doc in docs + if doc["kind"] == "ServiceAccount" and doc["metadata"]["name"] == pod["serviceAccountName"] + ) + assert account["automountServiceAccountToken"] is False + configmap = next( + doc for doc in docs if doc["kind"] == "ConfigMap" and "airflow.cfg" in doc.get("data", {}) + ) + config = ConfigParser(interpolation=None) + config.read_string(configmap["data"]["airflow.cfg"]) + for section, key in ( + ("database", "sql_alchemy_conn"), + ("core", "fernet_key"), + ("api_auth", "jwt_secret"), + ("api_auth", "jwt_private_key_path"), + ): + assert not config.has_option(section, key) + assert config["dag_processor"]["execution_api_token_file"] == "/run/dag-processor/processor.jwt" + policy = yaml.safe_load((examples / "network-policy.yaml").read_text())["spec"] + labels = deployment["spec"]["template"]["metadata"]["labels"] + assert policy["podSelector"]["matchLabels"].items() <= labels.items() + assert set(policy["policyTypes"]) == {"Ingress", "Egress"} + assert not policy["ingress"] + assert len(policy["egress"]) == 2 + assert policy["egress"][0] == { + "to": [ + { + "namespaceSelector": {"matchLabels": {"kubernetes.io/metadata.name": "airflow"}}, + "podSelector": {"matchLabels": {"component": "api-server", "release": "airflow"}}, + } + ], + "ports": [{"protocol": "TCP", "port": 8080}], + } + assert policy["egress"][1]["to"][0]["podSelector"]["matchLabels"] == {"k8s-app": "kube-dns"} + assert {port["port"] for port in policy["egress"][1]["ports"]} == {53} + @pytest.mark.parametrize( ("log_values", "expected_volume"), [ diff --git a/chart/values.schema.json b/chart/values.schema.json index 49251ffa3d2ff..6e13f38afe191 100644 --- a/chart/values.schema.json +++ b/chart/values.schema.json @@ -4639,7 +4639,7 @@ } }, "replicas": { - "description": "Number of dag processors to run.", + "description": "Number of dag processors to run. At most 1 when ``config.dag_processor.execution_api_token_file`` is set, since the pods would share one session token.", "type": "integer", "default": 1 }, @@ -4679,7 +4679,7 @@ ] }, "strategy": { - "description": "Specifies the strategy used to replace old Pods by new ones when deployed as a Deployment.", + "description": "Specifies the strategy used to replace old Pods by new ones when deployed as a Deployment. Ignored when ``config.dag_processor.execution_api_token_file`` is set; those pods share one session token and are replaced with the Recreate strategy.", "type": [ "null", "object" diff --git a/chart/values.yaml b/chart/values.yaml index 3edbbf55c58fd..e8976363680b0 100644 --- a/chart/values.yaml +++ b/chart/values.yaml @@ -2185,7 +2185,8 @@ dagProcessor: # classpath: "airflow.dag_processing.bundles.local.LocalDagBundle" # kwargs: {} - # Number of Airflow dag processors in the Deployment + # Number of Airflow dag processors in the Deployment. At most 1 when + # config.dag_processor.execution_api_token_file is set, since the pods would share one session token. replicas: 1 # Max number of old ReplicaSets to retain @@ -2197,7 +2198,9 @@ dagProcessor: # Args to use when running Airflow dag processor (templated). args: ["bash", "-c", "exec airflow dag-processor"] - # Update Strategy for dag processors + # Update Strategy for dag processors. Ignored when + # config.dag_processor.execution_api_token_file is set: those pods share one session + # token, so they are always replaced with the Recreate strategy. strategy: rollingUpdate: maxSurge: "100%" diff --git a/docker-tests/tests/docker_tests/test_dag_processor_api_deployment.py b/docker-tests/tests/docker_tests/test_dag_processor_api_deployment.py new file mode 100644 index 0000000000000..ba679c3774340 --- /dev/null +++ b/docker-tests/tests/docker_tests/test_dag_processor_api_deployment.py @@ -0,0 +1,65 @@ +# 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 + +import pytest +from python_on_whales import DockerClient + +from docker_tests.constants import AIRFLOW_ROOT_PATH + + +@pytest.mark.parametrize("uid", [50000, 54321]) +def test_isolated_processor_compose_example(tmp_path, monkeypatch, uid): + monkeypatch.setenv("AIRFLOW_IMAGE_NAME", "processor-api:test") + monkeypatch.setenv("DAG_PROCESSOR_API_NETWORK", "test-api-only") + monkeypatch.setenv("AIRFLOW_UID", str(uid)) + credentials = { + "AIRFLOW__DATABASE__SQL_ALCHEMY_CONN": "postgresql://must-not-reach-the-processor", + "AIRFLOW__CORE__FERNET_KEY": "must-not-reach-the-processor", + "AIRFLOW__API_AUTH__JWT_SECRET": "must-not-reach-the-processor", + "AIRFLOW__API_AUTH__JWT_PRIVATE_KEY_PATH": "/server/private.key", + } + for name, value in credentials.items(): + monkeypatch.setenv(name, value) + (tmp_path / ".env").write_text("\n".join(f"{name}={value}" for name, value in credentials.items())) + example = AIRFLOW_ROOT_PATH / "airflow-core/docs/security/dag-processor-api/docker-compose.yaml" + rendered = DockerClient(compose_files=[example], compose_project_directory=tmp_path).compose.config( + return_json=True + ) + assert set(rendered["services"]) == {"dag-processor"} + processor = rendered["services"]["dag-processor"] + assert processor["command"] == ["dag-processor"] + assert processor["user"] == f"{uid}:0" + assert processor["environment"]["USER"] == "airflow" + assert not processor.get("depends_on") + assert not processor.get("env_file") + assert not credentials.keys() & processor["environment"].keys() + assert processor["environment"]["CONNECTION_CHECK_MAX_COUNT"] == "0" + assert ( + processor["environment"]["AIRFLOW__DAG_PROCESSOR__EXECUTION_API_TOKEN_FILE"] + == "/run/dag-processor/processor.jwt" + ) + assert processor["environment"]["AIRFLOW_CONFIG"] == "/tmp/airflow.cfg" + assert processor["healthcheck"]["test"] == ["CMD", "airflow", "dag-processor", "--check-health"] + assert set(processor["networks"]) == {"processor-api"} + assert rendered["networks"]["processor-api"] == {"name": "test-api-only", "external": True} + mounts = {mount["target"]: mount for mount in processor["volumes"]} + assert set(mounts) == {"/run/dag-processor", "/opt/airflow/dags", "/opt/airflow/logs"} + assert mounts["/run/dag-processor"]["read_only"] + assert mounts["/opt/airflow/dags"]["read_only"] + assert processor["read_only"] + assert processor["cap_drop"] == ["ALL"] diff --git a/generated/known_sdk_imports_in_core.txt b/generated/known_sdk_imports_in_core.txt index 91d224dbcce3f..9c761d53a59e0 100644 --- a/generated/known_sdk_imports_in_core.txt +++ b/generated/known_sdk_imports_in_core.txt @@ -7,7 +7,7 @@ airflow-core/src/airflow/dag_processing/api_client.py::3 airflow-core/src/airflow/dag_processing/dagbag.py::3 airflow-core/src/airflow/dag_processing/importer_routing.py::1 airflow-core/src/airflow/dag_processing/lang_sdk_processor.py::6 -airflow-core/src/airflow/dag_processing/manager.py::5 +airflow-core/src/airflow/dag_processing/manager.py::7 airflow-core/src/airflow/dag_processing/processor.py::14 airflow-core/src/airflow/exceptions.py::1 airflow-core/src/airflow/executors/base_executor.py::3 diff --git a/providers/fab/src/airflow/providers/fab/auth_manager/security_manager/override.py b/providers/fab/src/airflow/providers/fab/auth_manager/security_manager/override.py index bd138d43914d4..8f193335136bd 100644 --- a/providers/fab/src/airflow/providers/fab/auth_manager/security_manager/override.py +++ b/providers/fab/src/airflow/providers/fab/auth_manager/security_manager/override.py @@ -163,6 +163,10 @@ class FabException(Exception): """Custom exception for FAB security manager.""" +class DagAccessControlValidationError(FabException, ValueError): + """The Dag's access control names an invalid role, resource, or permission.""" + + class AzureTenantResolutionError(FabException): """Raised when an Azure AD tenant identifier cannot be resolved to a canonical GUID.""" @@ -1129,7 +1133,7 @@ def _get_or_create_dag_permission(action_name: str, dag_resource_name: str) -> P for rolename, resource_actions_raw in access_control.items(): role = self.find_role(rolename) if not role: - raise FabException( + raise DagAccessControlValidationError( f"The access_control mapping for DAG '{dag_id}' includes a role named " f"'{rolename}', but that role does not exist" ) @@ -1143,7 +1147,7 @@ def _get_or_create_dag_permission(action_name: str, dag_resource_name: str) -> P for resource_name, actions in resource_actions.items(): if resource_name not in self.RESOURCE_DETAILS_MAP: - raise FabException( + raise DagAccessControlValidationError( f"The access_control map for DAG '{dag_id}' includes the following invalid " f"resource name: '{resource_name}'; " f"The set of valid resource names is: {self.RESOURCE_DETAILS_MAP.keys()}" @@ -1155,7 +1159,7 @@ def _get_or_create_dag_permission(action_name: str, dag_resource_name: str) -> P invalid_actions = set(actions) - self.RESOURCE_DETAILS_MAP[resource_name]["actions"] if invalid_actions: - raise FabException( + raise DagAccessControlValidationError( f"The access_control map for DAG '{dag_resource_name}' includes " f"the following invalid permissions: {invalid_actions}; " f"The set of valid permissions is: {self.RESOURCE_DETAILS_MAP[resource_name]['actions']}" diff --git a/providers/fab/tests/unit/fab/auth_manager/test_security.py b/providers/fab/tests/unit/fab/auth_manager/test_security.py index c73b94b96695f..a9b0bee0eae38 100644 --- a/providers/fab/tests/unit/fab/auth_manager/test_security.py +++ b/providers/fab/tests/unit/fab/auth_manager/test_security.py @@ -55,6 +55,7 @@ class _TestBase(DeclarativeBase): from airflow.providers.fab.auth_manager.fab_auth_manager import FabAuthManager from airflow.providers.fab.auth_manager.models.anonymous_user import AnonymousUser from airflow.providers.fab.auth_manager.security_manager.override import ( + DagAccessControlValidationError, FabAirflowSecurityManagerOverride, FabException, ) @@ -866,7 +867,7 @@ def test_has_all_dag_access(app, security_manager): def test_access_control_with_non_existent_role(security_manager): - with pytest.raises(FabException) as ctx: + with pytest.raises(DagAccessControlValidationError) as ctx: security_manager._sync_dag_view_permissions( dag_id="access-control-test", access_control={ @@ -874,10 +875,12 @@ def test_access_control_with_non_existent_role(security_manager): }, ) assert "role does not exist" in str(ctx.value) + assert isinstance(ctx.value, FabException) + assert isinstance(ctx.value, ValueError) def test_access_control_with_non_allowed_resource(security_manager): - with pytest.raises(FabException) as ctx: + with pytest.raises(DagAccessControlValidationError) as ctx: security_manager._sync_dag_view_permissions( dag_id="access-control-test", access_control={ @@ -908,11 +911,14 @@ def test_all_dag_access_doesnt_give_non_dag_access(app, security_manager): ) -def test_access_control_with_invalid_permission(app, security_manager): - invalid_actions = [ +@pytest.mark.parametrize( + "action", + [ "can_varimport", # a real action, but not a member of DAG_ACTIONS "can_eat_pudding", # clearly not a real action - ] + ], +) +def test_access_control_with_invalid_permission(app, security_manager, action): username = "LaUser" rolename = "team-a" with create_user_scope( @@ -920,18 +926,17 @@ def test_access_control_with_invalid_permission(app, security_manager): username=username, role_name=rolename, ): - for action in invalid_actions: - with pytest.raises(FabException) as ctx: - security_manager._sync_dag_view_permissions( - "access_control_test", - access_control={rolename: {action}}, - ) - assert "invalid permissions" in str(ctx.value) - if hasattr(permissions, "resource_name"): - assert "invalid permission" in str(ctx.value) - else: - # Test with old airflow running new FAB - assert "invalid resource name" in str(ctx.value) + with pytest.raises(DagAccessControlValidationError) as ctx: + security_manager._sync_dag_view_permissions( + "access_control_test", + access_control={rolename: {action}}, + ) + assert "invalid permissions" in str(ctx.value) + if hasattr(permissions, "resource_name"): + assert "invalid permission" in str(ctx.value) + else: + # Test with old airflow running new FAB + assert "invalid resource name" in str(ctx.value) def test_access_control_is_set_on_init( 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 4782f69b7d714..ec76646de07a7 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" @@ -143,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" diff --git a/task-sdk/src/airflow/sdk/api/client.py b/task-sdk/src/airflow/sdk/api/client.py index 42afaace779d8..a29617c82e502 100644 --- a/task-sdk/src/airflow/sdk/api/client.py +++ b/task-sdk/src/airflow/sdk/api/client.py @@ -1432,7 +1432,8 @@ def from_response(cls, response: httpx2.Response) -> ServerResponseError | None: if not (400 <= response.status_code < 600): return None - if response.headers.get("content-type") != "application/json": + media_type = response.headers.get("content-type", "").partition(";")[0].strip().lower() + if media_type != "application/json": return None detail: list[RemoteValidationError] | dict[str, Any] | None = None diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py index 02d51e15b57e5..8489c9f98a48e 100644 --- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py +++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py @@ -109,6 +109,70 @@ class ConnectionTestState(str, Enum): FAILED = "failed" +class Version(RootModel[str]): + root: Annotated[str, Field(max_length=200, title="Version")] + + +class DagBundleInventoryBody(BaseModel): + """ + A discovery snapshot with explicit failures, conditional on the previously observed revision. + """ + + model_config = ConfigDict( + extra="forbid", + ) + attempt_id: Annotated[UUID, Field(title="Attempt Id")] + dispatch_sequence: Annotated[int, Field(ge=1, le=9007199254740991, title="Dispatch Sequence")] + expected_revision: Annotated[UUID | None, Field(title="Expected Revision")] + version: Annotated[Version | None, Field(title="Version")] = None + files: Annotated[list[str], Field(max_length=100000, title="Files")] + failed_sources: Annotated[list[str] | None, Field(max_length=100000, title="Failed Sources")] = None + + +class DagBundleInventoryResponse(BaseModel): + """ + The source revision to attach to imports dispatched from this inventory. + """ + + model_config = ConfigDict( + extra="forbid", + ) + attempt_id: Annotated[UUID, Field(title="Attempt Id")] + accepted_at: Annotated[AwareDatetime, Field(title="Accepted At")] + revision: Annotated[UUID, Field(title="Revision")] + + +class DagBundleStateResponse(BaseModel): + """ + Server-owned metadata for an authorized bundle. + """ + + model_config = ConfigDict( + extra="forbid", + ) + name: Annotated[str, Field(title="Name")] + version: Annotated[str | None, Field(title="Version")] + last_refreshed: Annotated[AwareDatetime | None, Field(title="Last Refreshed")] + revision: Annotated[UUID | None, Field(title="Revision")] + team_name: Annotated[str | None, Field(title="Team Name")] + + +class BundleVersion(RootModel[str]): + root: Annotated[str, Field(max_length=200, title="Bundle Version")] + + +class DagParseResultResponse(BaseModel): + """ + Receipt returned unchanged when the latest accepted publication is replayed. + """ + + model_config = ConfigDict( + extra="forbid", + ) + attempt_id: Annotated[UUID, Field(title="Attempt Id")] + accepted_at: Annotated[AwareDatetime, Field(title="Accepted At")] + + class DagParseTokenBody(BaseModel): """ Exchange a Job credential for access on behalf of one file-parsing attempt. @@ -221,6 +285,21 @@ class DagRunType(str, Enum): ASSET_MATERIALIZATION = "asset_materialization" +class DagWarningType(str, Enum): + """ + Enum for DAG warning types. + + This is the set of allowable values for the ``warning_type`` field + in the DagWarning model. + """ + + ASSET_CONFLICT = "asset conflict" + DUPLICATE_DAG_ID = "duplicate dag id" + NON_EXISTENT_POOL = "non-existent pool" + RUNTIME_VARYING_VALUE = "runtime varying value" + TASK_GROUP_CYCLE = "task group cycle" + + class HITLUser(BaseModel): """ Schema for a Human-in-the-loop users. @@ -344,6 +423,37 @@ class LoopContext(BaseModel): terminal_is_mapped: Annotated[bool, Field(title="Terminal Is Mapped")] +class ParseSourceCode(BaseModel): + """ + Source captured by the importer; null explicitly means unavailable. + """ + + model_config = ConfigDict( + extra="forbid", + ) + source_code: Annotated[str | None, Field(title="Source Code")] + language: Annotated[str, Field(max_length=64, min_length=1, title="Language")] + + +class WarningType(RootModel[str]): + root: Annotated[ + str, Field(max_length=50, pattern="^[a-z][a-z0-9_]*:[a-z0-9_.\\-]+$", title="Warning Type") + ] + + +class ParseWarning(BaseModel): + """ + A warning attached to a Dag returned by this import. + """ + + model_config = ConfigDict( + extra="forbid", + ) + dag_id: Annotated[str, Field(max_length=250, min_length=1, title="Dag Id")] + warning_type: Annotated[DagWarningType | WarningType, Field(title="Warning Type")] + message: Annotated[str, Field(title="Message")] + + class PrevSuccessfulDagRunResponse(BaseModel): """ Schema for response with previous successful DagRun information for Task Template Context. @@ -374,6 +484,51 @@ class PreviousTIResponse(BaseModel): duration: Annotated[float | None, Field(title="Duration")] = None +class State(str, Enum): + SUCCESS = "success" + FAILED = "failed" + + +class ProcessorWorkAckBody(BaseModel): + """ + Acknowledge delivery, retaining the claim identity across retries. + """ + + model_config = ConfigDict( + extra="forbid", + ) + claim_id: Annotated[UUID, Field(title="Claim Id")] + state: Annotated[State, Field(title="State")] + + +class ProcessorWorkClaimBody(BaseModel): + """ + A retryable bounded claim for requested work in ready bundles. + """ + + model_config = ConfigDict( + extra="forbid", + ) + claim_id: Annotated[UUID, Field(title="Claim Id")] + bundle_names: Annotated[list[str], Field(max_length=1000, min_length=1, title="Bundle Names")] + limit: Annotated[int | None, Field(ge=1, le=100, title="Limit")] = 20 + + +class ProcessorWorkItem(BaseModel): + """ + A claimed callback or priority parse request. + """ + + model_config = ConfigDict( + extra="forbid", + ) + id: Annotated[str, Field(title="Id")] + claim_id: Annotated[UUID, Field(title="Claim Id")] + bundle_name: Annotated[str, Field(title="Bundle Name")] + relative_fileloc: Annotated[str, Field(title="Relative Fileloc")] + callback: Annotated[str | None, Field(title="Callback")] = None + + class TIAwaitingInputStatePayload(BaseModel): """ Schema for parking a TaskInstance in an awaiting_input state (Human-in-the-loop, no trigger). @@ -810,6 +965,29 @@ class ConnectionTestResultBody(BaseModel): result_message: Annotated[ResultMessage | None, Field(title="Result Message")] = None +class DagParseResultBody(BaseModel): + """ + One completed file or container import, including an empty result. + """ + + model_config = ConfigDict( + extra="forbid", + ) + attempt_id: Annotated[UUID, Field(title="Attempt Id")] + bundle_name: Annotated[str, Field(max_length=250, min_length=1, title="Bundle Name")] + relative_fileloc: Annotated[str, Field(max_length=2000, min_length=1, title="Relative Fileloc")] + dispatch_sequence: Annotated[int, Field(ge=1, le=9007199254740991, title="Dispatch Sequence")] + bundle_revision: Annotated[UUID | None, Field(title="Bundle Revision")] = None + bundle_version: Annotated[BundleVersion | None, Field(title="Bundle Version")] = None + version_data: Annotated[dict[str, Any] | None, Field(title="Version Data")] = None + parse_duration: Annotated[float, Field(ge=0.0, title="Parse Duration")] + serialized_dags: Annotated[list[dict[str, Any]], Field(title="Serialized Dags")] + import_errors: Annotated[dict[str, str] | None, Field(title="Import Errors")] = None + warnings: Annotated[list[ParseWarning] | None, Field(title="Warnings")] = None + parsed_definitions: Annotated[list[str] | None, Field(title="Parsed Definitions")] = None + source_codes: Annotated[dict[str, ParseSourceCode], Field(title="Source Codes")] + + class HITLDetailRequest(BaseModel): """ Schema for the request part of a Human-in-the-loop detail for a specific task instance. diff --git a/task-sdk/src/airflow/sdk/execution_time/cache.py b/task-sdk/src/airflow/sdk/execution_time/cache.py index 0c220fd795cf4..1d1910bf93f6f 100644 --- a/task-sdk/src/airflow/sdk/execution_time/cache.py +++ b/task-sdk/src/airflow/sdk/execution_time/cache.py @@ -97,7 +97,7 @@ def init(cls): @classmethod def reset(cls): - """Use for test purposes only.""" + """Detach this process from the shared cache, including entries inherited across a fork.""" cls._cache = None @classmethod diff --git a/task-sdk/tests/task_sdk/api/test_client.py b/task-sdk/tests/task_sdk/api/test_client.py index 62fde3373e007..9425c81e8b67f 100644 --- a/task-sdk/tests/task_sdk/api/test_client.py +++ b/task-sdk/tests/task_sdk/api/test_client.py @@ -199,6 +199,19 @@ def test_error_parsing_other_json(self): assert err.value.args == ("Not found",) assert err.value.detail is None + @pytest.mark.parametrize("content_type", ["application/json; charset=utf-8", "Application/JSON"]) + def test_error_parsing_json_media_type_variants(self, content_type): + responses = [ + httpx2.Response( + 409, content=b'{"detail": {"reason": "job_running"}}', headers={"content-type": content_type} + ) + ] + client = make_client_w_responses(responses) + + with pytest.raises(ServerResponseError) as err: + client.get("http://error") + assert err.value.detail == {"reason": "job_running"} + def test_server_response_error_pickling(self): responses = [httpx2.Response(404, json={"detail": {"message": "Invalid input"}})] client = make_client_w_responses(responses)