This is an automated email from the ASF dual-hosted git repository. ephraimbuddy pushed a commit to branch codex/dag-processor-work-delivery in repository https://gitbox.apache.org/repos/asf/airflow.git
commit 78aaedc405aa6a8cc76372f77cd36535b9bda005 Author: Ephraim Anierobi <[email protected]> AuthorDate: Thu Oct 15 10:46:29 2026 +0100 Retain Dag processor work until delivery is acknowledged Expose bounded, replayable claims for callback and priority parsing requests through the Execution API. Keep ownership until acknowledgment or processor retirement, and let the scheduler clean stale metadata in bounded batches. --- airflow-core/docs/migrations-ref.rst | 4 +- .../execution_api/datamodels/dag_parsing.py | 27 +++- .../execution_api/routes/dag_parsing.py | 124 ++++++++++++++- .../api_fastapi/execution_api/routes/jobs.py | 3 +- .../execution_api/versions/v2027_02_28.py | 4 +- .../src/airflow/config_templates/config.yml | 13 +- .../src/airflow/dag_processing/maintenance.py | 91 +++++++++++ airflow-core/src/airflow/dag_processing/manager.py | 50 +++--- .../src/airflow/jobs/scheduler_job_runner.py | 4 + .../0150_3_5_0_add_processor_work_claims.py | 61 ++++++++ airflow-core/src/airflow/models/callback.py | 4 + airflow-core/src/airflow/models/dagbag.py | 10 +- airflow-core/src/airflow/models/dagwarning.py | 30 +--- airflow-core/src/airflow/utils/db.py | 2 +- .../execution_api/test_token_scope_boundaries.py | 4 + .../versions/head/test_dag_parsing.py | 168 +++++++++++++++++++++ .../versions/v2027_02_28/test_jobs.py | 2 + .../tests/unit/dag_processing/test_maintenance.py | 111 ++++++++++++++ .../tests/unit/dag_processing/test_manager.py | 41 ++--- airflow-core/tests/unit/jobs/test_scheduler_job.py | 15 ++ airflow-core/tests/unit/models/test_dagwarning.py | 54 ------- .../src/airflow/sdk/api/datamodels/_generated.py | 45 ++++++ 22 files changed, 714 insertions(+), 153 deletions(-) diff --git a/airflow-core/docs/migrations-ref.rst b/airflow-core/docs/migrations-ref.rst index ec17da20497..a0b88dc5588 100644 --- a/airflow-core/docs/migrations-ref.rst +++ b/airflow-core/docs/migrations-ref.rst @@ -39,7 +39,9 @@ Here's the list of all the Database Migrations that are executed via when you ru +-------------------------+------------------+-------------------+--------------------------------------------------------------+ | Revision ID | Revises ID | Airflow Version | Description | +=========================+==================+===================+==============================================================+ -| ``687eb253de24`` (head) | ``e941ab8243b7`` | ``3.5.0`` | Add source revisions for Dag processor reconciliation. | +| ``d8c6f3a129b0`` (head) | ``687eb253de24`` | ``3.5.0`` | Add durable Dag processor work claims. | ++-------------------------+------------------+-------------------+--------------------------------------------------------------+ +| ``687eb253de24`` | ``e941ab8243b7`` | ``3.5.0`` | Add source revisions for Dag processor reconciliation. | +-------------------------+------------------+-------------------+--------------------------------------------------------------+ | ``e941ab8243b7`` | ``c71065a1b14f`` | ``3.5.0`` | Add checkpoints for Dag result publication. | +-------------------------+------------------+-------------------+--------------------------------------------------------------+ 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 index 74d479a26fa..cbd003f4dc8 100644 --- 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 @@ -18,7 +18,7 @@ from __future__ import annotations from datetime import datetime -from typing import Any +from typing import Any, Literal from uuid import UUID from pydantic import Field, model_validator @@ -143,3 +143,28 @@ 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/dag_parsing.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/dag_parsing.py index f799f777a7c..d0c3fe3ad38 100644 --- 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 @@ -20,7 +20,8 @@ from __future__ import annotations import hashlib import json from datetime import timedelta -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Literal, cast +from uuid import UUID from cadwyn import VersionedAPIRouter from fastapi import HTTPException, Request, Security, status @@ -39,6 +40,9 @@ from airflow.api_fastapi.execution_api.datamodels.dag_parsing import ( DagBundleStateResponse, DagParseResultBody, DagParseResultResponse, + ProcessorWorkAckBody, + ProcessorWorkClaimBody, + ProcessorWorkItem, ) from airflow.api_fastapi.execution_api.security import CurrentExecutionToken, ExecutionAPIRoute, require_auth from airflow.configuration import conf @@ -53,13 +57,17 @@ from airflow.dag_processing.collection import ( 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 @@ -399,3 +407,117 @@ def publish_inventory( accepted_at=checkpoint.accepted_at, revision=bundle.parse_revision, ) + + [email protected]( + "/{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 + + [email protected]( + "/{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 f0b5f4f82ee..018e5faaab4 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 b6db627b5bc..a02678ac362 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 @@ from cadwyn import VersionChange, endpoint class AddDagProcessorJobEndpoints(VersionChange): - """Add Dag processor lifecycle, discovery inventory, and publication endpoints.""" + """Add Dag processor lifecycle, discovery inventory, publication, and requested-work endpoints.""" description = __doc__ @@ -33,4 +33,6 @@ class AddDagProcessorJobEndpoints(VersionChange): 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/config_templates/config.yml b/airflow-core/src/airflow/config_templates/config.yml index 873b519576f..5eb6a37879e 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: ~ @@ -3434,6 +3436,15 @@ dag_processor: 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. + 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/maintenance.py b/airflow-core/src/airflow/dag_processing/maintenance.py new file mode 100644 index 00000000000..c5ab9877634 --- /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 8c81af7429d..19aee0476ed 100644 --- a/airflow-core/src/airflow/dag_processing/manager.py +++ b/airflow-core/src/airflow/dag_processing/manager.py @@ -39,7 +39,7 @@ from typing import TYPE_CHECKING, Any, Literal, NamedTuple, cast import attrs import structlog -from sqlalchemy import select, update +from sqlalchemy import or_, select, update from sqlalchemy.exc import OperationalError from sqlalchemy.orm import load_only from tabulate import tabulate @@ -65,6 +65,7 @@ from airflow.dag_processing.processor import ( 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 @@ -487,23 +488,7 @@ class DagFileProcessorManager(LoggingMixin): 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 @@ -608,7 +593,6 @@ class DagFileProcessorManager(LoggingMixin): 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 @@ -701,9 +685,16 @@ class DagFileProcessorManager(LoggingMixin): """Fetch priority parsing requests from the metadata database.""" 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())) + query = 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)) + ), + ), ) + requests = session.scalars(with_row_locks(query, session=session, skip_locked=True)) for request in requests: bundle = bundles[request.bundle_name] files.append( @@ -748,7 +739,15 @@ class DagFileProcessorManager(LoggingMixin): 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, @@ -877,15 +876,6 @@ class DagFileProcessorManager(LoggingMixin): 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. - - 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 diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index 40e1fa976ff..c24d6df58d2 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.callbacks.callback_requests import ( ) 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 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin): 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/0150_3_5_0_add_processor_work_claims.py b/airflow-core/src/airflow/migrations/versions/0150_3_5_0_add_processor_work_claims.py new file mode 100644 index 00000000000..d938a9ed9ca --- /dev/null +++ b/airflow-core/src/airflow/migrations/versions/0150_3_5_0_add_processor_work_claims.py @@ -0,0 +1,61 @@ +# 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 durable Dag processor work claims. + +Revision ID: d8c6f3a129b0 +Revises: 687eb253de24 +Create Date: 2026-10-15 09:00:00.000000 +""" + +from __future__ import annotations + +import sqlalchemy as sa +from alembic import op + +from airflow.migrations.utils import disable_sqlite_fkeys + +revision = "d8c6f3a129b0" +down_revision = "687eb253de24" +branch_labels = None +depends_on = None +airflow_version = "3.5.0" + + +def upgrade(): + """Retain callback and priority requests until their processor acknowledges delivery.""" + with disable_sqlite_fkeys(op): + 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 processor ownership from callback and priority requests.""" + with disable_sqlite_fkeys(op): + 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") diff --git a/airflow-core/src/airflow/models/callback.py b/airflow-core/src/airflow/models/callback.py index 6295b7dfa7b..78309784de1 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/dagbag.py b/airflow-core/src/airflow/models/dagbag.py index 1665699bec1..b64ed3097d4 100644 --- a/airflow-core/src/airflow/models/dagbag.py +++ b/airflow-core/src/airflow/models/dagbag.py @@ -27,7 +27,7 @@ from typing import TYPE_CHECKING, Any, NamedTuple 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/dagwarning.py b/airflow-core/src/airflow/models/dagwarning.py index 4d1eeac7db5..25224943925 100644 --- a/airflow-core/src/airflow/models/dagwarning.py +++ b/airflow-core/src/airflow/models/dagwarning.py @@ -19,21 +19,15 @@ from __future__ import annotations 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 @@ class DagWarning(Base): 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/utils/db.py b/airflow-core/src/airflow/utils/db.py index d99e7cb1ea8..12f40b7f76a 100644 --- a/airflow-core/src/airflow/utils/db.py +++ b/airflow-core/src/airflow/utils/db.py @@ -118,7 +118,7 @@ _REVISION_HEADS_MAP: dict[str, str] = { "3.2.0": "1d6611b6ab7c", "3.3.0": "d2f4e1b3c5a7", "3.4.0": "a5d7b9c13e40", - "3.5.0": "687eb253de24", + "3.5.0": "d8c6f3a129b0", } # Prefix used to identify tables holding data moved during migration. 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 7e84d4d6745..f97bd40bb96 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 @@ -79,6 +79,8 @@ NON_DEFAULT_TOKEN_POLICY: dict[str, set[str]] = { "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. @@ -92,6 +94,8 @@ DAG_PROCESSOR_LIFECYCLE_ROUTES = { "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 index 9cb187c5e2e..79ae13752b0 100644 --- 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 @@ -37,12 +37,17 @@ 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 @@ -56,6 +61,7 @@ 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, @@ -70,11 +76,17 @@ SESSION_ID = UUID("00000000-0000-0000-0000-0000000000aa") 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() @@ -1076,3 +1088,159 @@ def test_catalog_and_inventory_support_granted_bundle_with_slash(client, exec_ap json={"attempt_id": str(uuid7()), "dispatch_sequence": 1, "expected_revision": None, "files": []}, ) assert response.status_code == 200, response.json() + + [email protected]("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 + + [email protected]("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() + + [email protected]("postgres", "mysql") [email protected]("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] + + [email protected]("kind", ["callbacks", "priority"]) [email protected]("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) + + [email protected]("postgres", "mysql") +def test_direct_priority_claim_skips_request_locked_by_api(session, job, tmp_path): + work = DagPriorityParsingRequest(bundle_name="bundle", relative_fileloc="a.py") + session.add(work) + session.commit() + key, job_id = work.id, job.id + manager = DagFileProcessorManager(max_runs=1) + manager._dag_bundles = [LocalDagBundle(name="bundle", path=tmp_path)] + + with create_session() as claim_session, ThreadPoolExecutor(max_workers=1) as pool: + claimed = claim_session.scalar( + select(DagPriorityParsingRequest).where(DagPriorityParsingRequest.id == key).with_for_update() + ) + claimed.processor_job_id = job_id + claimed.processor_claim_id = uuid7() + claim_session.flush() + direct_claim = pool.submit(manager.claim_priority_files) + try: + assert direct_claim.result(timeout=10) == [] + finally: + claim_session.commit() + + session.expire_all() + assert session.get(DagPriorityParsingRequest, key).processor_job_id == job_id 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 0b781ccb6b0..4fcefa6f4e1 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 @@ -48,6 +48,8 @@ class TestDagProcessorJobEndpointsVersioning: "/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): 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 00000000000..8e42e2054bf --- /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 + + [email protected](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() + + [email protected]("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 a9dbbea950f..9418be4f43d 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py @@ -1419,7 +1419,7 @@ class TestDagFileProcessorManager: 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() @@ -1456,7 +1456,7 @@ class TestDagFileProcessorManager: ) ).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): @@ -1482,7 +1482,7 @@ class TestDagFileProcessorManager: 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 @@ -1519,7 +1519,7 @@ class TestDagFileProcessorManager: 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): @@ -1591,9 +1591,6 @@ class TestDagFileProcessorManager: 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", @@ -1618,7 +1615,12 @@ class TestDagFileProcessorManager: 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 @@ -4223,29 +4225,6 @@ class TestDagFileProcessorManager: 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" diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index 7118d08edc6..b7eb013f76e 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.callback_requests import ( 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 @@ from airflow.timetables.simple import ( 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 @@ class TestSchedulerJob: 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 0f07d7a9192..b62d25fd48d 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 @@ class TestGetWarningTypeValue: def test_dag_warning_stores_importer_type(self): assert DagWarning("dag_1", "yaml:schema_violation", "message").warning_type == "yaml:schema_violation" - - [email protected]_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/task-sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py index 7e5fdb7fa13..8489c9f98a4 100644 --- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py +++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py @@ -484,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).
