diff --git a/airflow-core/docs/authoring-and-scheduling/asset-scheduling.rst b/airflow-core/docs/authoring-and-scheduling/asset-scheduling.rst index 1bfe98d715ec2..69eb5070a2d3f 100644 --- a/airflow-core/docs/authoring-and-scheduling/asset-scheduling.rst +++ b/airflow-core/docs/authoring-and-scheduling/asset-scheduling.rst @@ -51,6 +51,21 @@ In addition to scheduling Dags based on time, you can also schedule Dags to run :ref:`asset_definitions` for how to declare assets. +.. note:: + + Asset events are registered asynchronously. When a task succeeds, Airflow durably records its + asset events and the scheduler registers them shortly afterwards rather than inside the + task-completion request, usually within one scheduler loop and longer under heavy load. + Downstream asset-scheduled Dags are therefore triggered after a brief, eventually consistent + delay. The scheduler retries a queued registration up to ``[scheduler] + asset_event_queue_max_attempts`` times, after which the row is parked and reported via the + ``asset.event_queue.failures`` metric. Use the ``[scheduler] asset_event_queue_drain_interval`` + option to tune how often the scheduler processes queued asset events. + + ``dag.test`` runs without a scheduler, so it registers each task's asset events itself once + the task finishes; the events stay visible during the test, though downstream asset-scheduled + Dags are still only triggered by a running scheduler. + Schedule Dags with assets ------------------------- diff --git a/airflow-core/docs/migrations-ref.rst b/airflow-core/docs/migrations-ref.rst index de64f8123f059..2bb682744be6c 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 | +=========================+==================+===================+==============================================================+ -| ``7a98f1b7dbd3`` (head) | ``c4e7a1f9b2d0`` | ``3.4.0`` | Add index on asset_event (asset_id, partition_key). | +| ``a63ae7baae97`` (head) | ``7a98f1b7dbd3`` | ``3.4.0`` | Add asset_event_queue table. | ++-------------------------+------------------+-------------------+--------------------------------------------------------------+ +| ``7a98f1b7dbd3`` | ``c4e7a1f9b2d0`` | ``3.4.0`` | Add index on asset_event (asset_id, partition_key). | +-------------------------+------------------+-------------------+--------------------------------------------------------------+ | ``c4e7a1f9b2d0`` | ``436dc127462c`` | ``3.4.0`` | Add index on asset.uri. | +-------------------------+------------------+-------------------+--------------------------------------------------------------+ diff --git a/airflow-core/newsfragments/66854.significant.rst b/airflow-core/newsfragments/66854.significant.rst new file mode 100644 index 0000000000000..42883e3cd303d --- /dev/null +++ b/airflow-core/newsfragments/66854.significant.rst @@ -0,0 +1,34 @@ +Asset events from successful tasks are registered asynchronously + +When a task succeeds, its asset events (and the resulting scheduling of downstream +asset-triggered Dags) are no longer written synchronously inside the task state-update +request. The execution API now records the events as a durable marker in a new +``asset_event_queue`` table, committed in the same transaction as the task state, and the +scheduler drains that queue and registers the events. This removes the asset registration +work from the task-completion request path, which previously held the ``task_instance`` row +lock while it ran and could pile up requests and exhaust the API server under high fan-out. + +As a result there can be a short delay between a task succeeding and its asset events being +visible (typically one scheduler loop, longer under sustained load). A queued registration is +retried up to ``[scheduler] asset_event_queue_max_attempts`` times; a row that still fails +after that is parked in ``asset_event_queue`` and surfaced via the ``asset.event_queue.failures`` +metric rather than retried forever. ``dag.test`` runs without a scheduler, so it drains the queue +itself after each task, keeping asset events visible during a test run. The drain cadence and +batch size are tunable via ``[scheduler] asset_event_queue_drain_interval``, +``asset_event_queue_batch_size`` and ``asset_event_queue_max_attempts``. + +* Types of change + + * [ ] Dag changes + * [x] Config changes + * [ ] API changes + * [ ] CLI changes + * [x] Behaviour changes + * [ ] Plugin changes + * [ ] Dependency changes + * [ ] Code interface changes + +* Migration rules needed + + * None - a new ``asset_event_queue`` table is created by the schema migration; no user + action is required. diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py index c1bac7960234d..7247fb1aebf87 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py @@ -40,6 +40,7 @@ from sqlalchemy.sql import select from structlog.contextvars import bind_contextvars +from airflow._shared.observability.metrics import stats from airflow._shared.observability.traces import override_ids from airflow._shared.state import TaskScope from airflow._shared.timezones import timezone @@ -77,7 +78,7 @@ ) from airflow.configuration import conf from airflow.exceptions import InvalidPartitionKeyError, TaskNotFound -from airflow.models.asset import AssetActive +from airflow.models.asset import AssetActive, AssetEventQueue from airflow.models.base import ID_LEN from airflow.models.dag import DagModel from airflow.models.dagrun import DagRun as DR @@ -97,6 +98,8 @@ if TYPE_CHECKING: from sqlalchemy.sql.dml import Update + from airflow.api_fastapi.execution_api.datamodels.asset import AssetProfile + router = VersionedAPIRouter() ti_id_router = VersionedAPIRouter( @@ -499,6 +502,30 @@ def ti_update_state( extra=json.dumps({"host_name": hostname}) if hostname else None, ) ) + # Durably record the successful task's asset events in the SAME transaction that commits the + # state, so the marker can never be lost in a crash between this commit and their + # registration. The scheduler drain is the single writer that registers them; the request + # never runs asset registration itself, which is what keeps the task_instance row lock short + # under high fan-out. Only a genuine RUNNING->SUCCESS transition reaches here (a duplicate + # SUCCESS->SUCCESS short-circuits earlier), so a completion is never enqueued twice. + if ( + updated_state == TaskInstanceState.SUCCESS + and isinstance(ti_patch_payload, TISuccessStatePayload) + and (ti_patch_payload.task_outlets or ti_patch_payload.outlet_events) + ): + _enqueue_asset_events( + task_instance_id=task_instance_id, + dag_id=dag_id, + run_id=run_id, + task_id=task_id, + map_index=map_index, + task_outlets=ti_patch_payload.task_outlets, + outlet_events=ti_patch_payload.outlet_events, + session=session, + ) + # Commit the state, log entry and durable queue marker together so they land atomically and + # the row lock is released promptly. + session.commit() except DataError: # Let DataErrorHandler return a 422 (not the opaque 500 below). raise @@ -525,7 +552,9 @@ def ti_update_state( task_id=task_id, map_index=map_index, ) + session.commit() except Exception: + session.rollback() log.warning( "Failed to clear task state on success", dag_id=dag_id, @@ -622,6 +651,58 @@ def _validate_outlet_event_partition_keys(outlet_events: list[dict[str, Any]]) - ) +def _asset_event_payload( + task_outlets: list[AssetProfile], + outlet_events: list[dict[str, Any]], + ti_key: dict[str, Any], +) -> dict[str, Any]: + """Serialize a task's outlets, outlet events, and natural key into the queue row's JSON payload.""" + return { + "task_outlets": [outlet.model_dump(mode="json") for outlet in task_outlets], + "outlet_events": outlet_events, + # The scheduler drain resolves the live task instance by this natural key rather than by the + # surrogate ``ti_id``. Clearing a task reassigns its uuid7 id, and relying on the foreign + # key's ON UPDATE CASCADE to re-point the row only works on Postgres -- SQLite does not + # enforce foreign keys in Airflow's production engine -- so a natural-key lookup is what + # keeps the pending events reachable on every backend. + "ti_key": ti_key, + } + + +def _enqueue_asset_events( + *, + task_instance_id: UUID, + dag_id: str, + run_id: str, + task_id: str, + map_index: int, + task_outlets: list[AssetProfile], + outlet_events: list[dict[str, Any]], + session: SessionDep, +) -> None: + """ + Add a durable queue marker for a completed task's asset events, without committing. + + The caller commits this in the same transaction as the task-state update, so the marker and the + ``SUCCESS`` state land atomically: the events can never be lost in a crash between the state + commit and their registration. The scheduler drain is the single writer that later registers the + events and deletes the row. ``merge`` overwrites any pending row a retried completion left behind + rather than conflicting on the ``ti_id`` primary key. + """ + session.merge( + AssetEventQueue( + ti_id=task_instance_id, + payload=_asset_event_payload( + task_outlets, + outlet_events, + {"dag_id": dag_id, "run_id": run_id, "task_id": task_id, "map_index": map_index}, + ), + attempts=0, + ) + ) + stats.incr("asset.event_queue.enqueued") + + def _create_ti_state_update_query_and_update_state( *, ti_patch_payload: TIStateUpdate, @@ -659,14 +740,6 @@ def _create_ti_state_update_query_and_update_state( # Store retry policy overrides so next_retry_datetime() can read them. # These are cleared when the task enters RUNNING (ti_run). query = query.values(retry_delay_override=retry_delay_override, retry_reason=retry_reason) - elif isinstance(ti_patch_payload, TISuccessStatePayload): - if ti is not None: - TI.register_asset_changes_in_db( - ti, - ti_patch_payload.task_outlets, - ti_patch_payload.outlet_events, - session=session, - ) try: _emit_task_span(ti, state=updated_state) except Exception: diff --git a/airflow-core/src/airflow/assets/manager.py b/airflow-core/src/airflow/assets/manager.py index 1dd26379b9d12..2f01daa681fac 100644 --- a/airflow-core/src/airflow/assets/manager.py +++ b/airflow-core/src/airflow/assets/manager.py @@ -22,7 +22,7 @@ from typing import TYPE_CHECKING import structlog -from sqlalchemy import exc, or_, select +from sqlalchemy import exc, insert, or_, select from sqlalchemy.orm import joinedload from airflow._shared.observability.metrics import stats @@ -41,6 +41,7 @@ DagScheduleAssetUriReference, PartitionedAssetKeyLog, TaskOutletAssetReference, + asset_alias_asset_event_association_table, ) from airflow.models.log import Log from airflow.timetables.base import compute_rollup_fingerprint @@ -359,8 +360,14 @@ def register_asset_change( ).unique() for asset_alias_model in asset_alias_models: - asset_alias_model.asset_events.append(asset_event) - session.add(asset_alias_model) + # Direct INSERT instead of asset_alias_model.asset_events.append(...), which + # would lazy-load the alias's entire asset_events collection just to add one row. + session.execute( + insert(asset_alias_asset_event_association_table).values( + alias_id=asset_alias_model.id, + event_id=asset_event.id, + ) + ) dags_to_queue_from_asset_alias |= { alias_ref.dag diff --git a/airflow-core/src/airflow/config_templates/config.yml b/airflow-core/src/airflow/config_templates/config.yml index b3dc3f133d229..3520c41310830 100644 --- a/airflow-core/src/airflow/config_templates/config.yml +++ b/airflow-core/src/airflow/config_templates/config.yml @@ -2639,6 +2639,31 @@ scheduler: type: integer example: ~ default: "60" + asset_event_queue_drain_interval: + description: | + How often (in seconds) the scheduler drains the asset event queue, registering asset + events that the execution API deferred to the queue (for example when the synchronous + write on task completion hit database lock contention). + version_added: 3.4.0 + type: float + example: ~ + default: "5.0" + asset_event_queue_batch_size: + description: | + The maximum number of pending asset event registrations the scheduler processes in a + single drain pass. + version_added: 3.4.0 + type: integer + example: ~ + default: "100" + asset_event_queue_max_attempts: + description: | + How many times the scheduler retries a failing asset event registration before parking + the row for inspection and emitting the ``asset.event_queue.failures`` metric. + version_added: 3.4.0 + type: integer + example: ~ + default: "5" pool_metrics_interval: description: | How often (in seconds) should pool usage stats be sent to StatsD (if statsd_on is enabled) diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index 9b9e85c07e353..85c324a80ce43 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -77,6 +77,7 @@ AssetAliasModel, AssetDagRunQueue, AssetEvent, + AssetEventQueue, AssetModel, AssetPartitionDagRun, AssetWatcherModel, @@ -85,6 +86,7 @@ PartitionedAssetKeyLog, TaskInletAssetReference, TaskOutletAssetReference, + _register_queued_asset_event, ) from airflow.models.asset_state_store import AssetStateStoreModel from airflow.models.backfill import Backfill, BackfillDagRun @@ -335,6 +337,17 @@ def __init__( self._parallelism = conf.getint("core", "parallelism") self._multi_team = conf.getboolean("core", "multi_team") self._dag_tags_in_metrics = conf.getboolean("metrics", "dag_tags_in_metrics", fallback=False) + # Asset-event queue drain tuning. Read here (not per drain tick) so the drain loop stays free + # of conf lookups, which can hit the secret backend and break HA locks (see note above). + self._asset_event_queue_batch_size = conf.getint( + "scheduler", "asset_event_queue_batch_size", fallback=100 + ) + self._asset_event_queue_max_attempts = conf.getint( + "scheduler", "asset_event_queue_max_attempts", fallback=5 + ) + self._asset_event_queue_drain_interval = conf.getfloat( + "scheduler", "asset_event_queue_drain_interval", fallback=5.0 + ) self._max_partition_dag_runs_per_loop = MAX_PARTITION_DAG_RUNS_PER_LOOP self._dag_id_to_team_name: dict[str, str | None] = {} @@ -1788,6 +1801,11 @@ def _run_scheduler_loop(self) -> None: self._remove_unreferenced_triggers, ) + timers.call_regular_interval( + self._asset_event_queue_drain_interval, + self._drain_asset_event_queue, + ) + if any(x.is_local for x in self.executors): bundle_cleanup_mgr = BundleUsageTrackingManager() check_interval = conf.getint( @@ -3698,6 +3716,66 @@ def _remove_unreferenced_triggers(self, *, session: Session = NEW_SESSION) -> No .execution_options(synchronize_session="fetch") ) + @provide_session + def _drain_asset_event_queue(self, *, session: Session = NEW_SESSION) -> None: + """ + Process pending asset-event registrations enqueued by the execution API. + + The task-success path commits an :class:`~airflow.models.asset.AssetEventQueue` marker + atomically with the task state instead of registering the events inline, so the scheduler + is the sole writer that runs ``register_asset_changes_in_db`` in a live deployment. We claim + a batch of pending rows in a single ``FOR UPDATE SKIP LOCKED`` query (so multiple schedulers + do not collide), register each row's events and delete it in the same transaction so it is + processed exactly once. A row that keeps failing is retried on later passes until it exceeds + ``asset_event_queue_max_attempts``, after which it is left in place (parked) and surfaced via + the ``asset.event_queue.failures`` metric for an operator to inspect. + """ + max_attempts = self._asset_event_queue_max_attempts + rows = session.scalars( + with_row_locks( + select(AssetEventQueue) + .where(AssetEventQueue.attempts < max_attempts) + .order_by(AssetEventQueue.created_at) + .limit(self._asset_event_queue_batch_size), + of=AssetEventQueue, + session=session, + skip_locked=True, + ) + ).all() + + processed = 0 + for row in rows: + ti_id = row.ti_id + try: + # Each row is registered in its own savepoint so one poison row rolls back alone + # instead of discarding the whole batch; the outer transaction commits once at the end. + with session.begin_nested(): + _register_queued_asset_event(row, session=session) + processed += 1 + except Exception: + with session.begin_nested(): + row.attempts += 1 + if row.attempts >= max_attempts: + self.log.error( + "Giving up on asset-event registration for task instance %s after %s " + "attempts; the row is left in asset_event_queue for inspection.", + ti_id, + row.attempts, + ) + stats.incr("asset.event_queue.failures") + else: + self.log.warning( + "Asset-event registration for task instance %s failed; will retry (attempt %s).", + ti_id, + row.attempts, + ) + + session.commit() + if processed: + stats.incr("asset.event_queue.processed", processed) + pending = session.scalar(select(func.count()).select_from(AssetEventQueue)) + stats.gauge("asset.event_queue.pending", pending or 0) + @provide_session def _update_asset_orphanage(self, *, session: Session = NEW_SESSION) -> None: """ diff --git a/airflow-core/src/airflow/migrations/versions/0128_3_4_0_add_asset_event_queue.py b/airflow-core/src/airflow/migrations/versions/0128_3_4_0_add_asset_event_queue.py new file mode 100644 index 0000000000000..7e0f3fd9aaa9a --- /dev/null +++ b/airflow-core/src/airflow/migrations/versions/0128_3_4_0_add_asset_event_queue.py @@ -0,0 +1,71 @@ +# +# 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 asset_event_queue table. + +Revision ID: a63ae7baae97 +Revises: 7a98f1b7dbd3 +Create Date: 2026-07-15 00:00:00.000000 +""" + +from __future__ import annotations + +import sqlalchemy as sa +from alembic import op + +from airflow.utils.sqlalchemy import UtcDateTime + +# revision identifiers, used by Alembic. +revision = "a63ae7baae97" +down_revision = "7a98f1b7dbd3" +branch_labels = None +depends_on = None +airflow_version = "3.4.0" + + +def upgrade(): + """Apply Add asset_event_queue table.""" + op.create_table( + "asset_event_queue", + sa.Column("ti_id", sa.Uuid(), nullable=False), + # Emitted task outlets and outlet events are stored together in a single JSON payload. The + # queue is a durable buffer only ever read back in full by the scheduler drain, so separate + # columns buy nothing and a single write keeps the enqueue on the task-success path cheap. + sa.Column("payload", sa.JSON(), nullable=False), + sa.Column("attempts", sa.Integer(), nullable=False, server_default="0"), + sa.Column("created_at", UtcDateTime(timezone=True), nullable=False), + sa.PrimaryKeyConstraint("ti_id", name="asset_event_queue_pkey"), + sa.ForeignKeyConstraint( + ["ti_id"], + ["task_instance.id"], + name="aeq_ti_fkey", + ondelete="CASCADE", + # Cascade a TI id change (clear/retry reassigns a new uuid7) so pending asset events + # follow the task instance instead of being orphaned and dropped by the scheduler drain. + onupdate="CASCADE", + ), + ) + with op.batch_alter_table("asset_event_queue", schema=None) as batch_op: + batch_op.create_index("idx_asset_event_queue_created_at", ["created_at"], unique=False) + + +def downgrade(): + """Unapply Add asset_event_queue table.""" + with op.batch_alter_table("asset_event_queue", schema=None) as batch_op: + batch_op.drop_index("idx_asset_event_queue_created_at") + op.drop_table("asset_event_queue") diff --git a/airflow-core/src/airflow/models/asset.py b/airflow-core/src/airflow/models/asset.py index 7aafddb3b9490..f7262df6172d3 100644 --- a/airflow-core/src/airflow/models/asset.py +++ b/airflow-core/src/airflow/models/asset.py @@ -20,6 +20,7 @@ from datetime import datetime from typing import TYPE_CHECKING from urllib.parse import urlsplit +from uuid import UUID import sqlalchemy as sa from sqlalchemy import ( @@ -788,6 +789,101 @@ def __repr__(self): return f"{self.__class__.__name__}({', '.join(args)})" +class AssetEventQueue(Base): + """ + Durable marker of asset events a successful task emitted, awaiting registration. + + On task success the execution API commits one of these rows atomically with the task + state instead of registering the asset events inline, so the ``ti_update_state`` request + holds the ``task_instance`` row lock only for the state write plus this insert rather than + for the whole ``register_asset_changes_in_db`` call (which was the source of API-server + lock contention under high fan-out). The scheduler drains this table, resolves the live + task instance by natural key, runs + :meth:`~airflow.models.taskinstance.TaskInstance.register_asset_changes_in_db` to create + the ``AssetEvent`` and ``AssetDagRunQueue`` rows, and deletes the queue row once that write + commits. The in-process runner behind ``dag.test`` has no scheduler, so it drains the row + itself via :func:`register_pending_asset_events` right after the task finishes. + + ``ti_id`` is the primary key: at most one pending registration exists per task + instance, and the row is cascade-deleted if the task instance is removed. + """ + + ti_id: Mapped[UUID] = mapped_column(sa.Uuid(), primary_key=True, nullable=False) + # Both the emitted task outlets and the outlet events live in one JSON payload + # (``{"task_outlets": [...], "outlet_events": [...], "ti_key": {...}}``). The queue is a durable + # buffer only ever read back in full when draining, so a single column keeps the enqueue cheap. + payload: Mapped[dict] = mapped_column(sa.JSON(), nullable=False, default=dict) + attempts: Mapped[int] = mapped_column(Integer, nullable=False, default=0, server_default="0") + created_at: Mapped[datetime] = mapped_column(UtcDateTime, default=timezone.utcnow, nullable=False) + + __tablename__ = "asset_event_queue" + __table_args__ = ( + PrimaryKeyConstraint(ti_id, name="asset_event_queue_pkey"), + ForeignKeyConstraint( + (ti_id,), + ["task_instance.id"], + name="aeq_ti_fkey", + ondelete="CASCADE", + # Referential cleanup only. Correctness on clear does not rely on these cascades: + # SQLite does not enforce foreign keys in Airflow's production engine, so they silently + # no-op there. The drain re-resolves the task instance by natural key instead, which + # survives the id reassignment a clear performs on every backend. + onupdate="CASCADE", + ), + Index("idx_asset_event_queue_created_at", created_at), + ) + + def __repr__(self): + return f"AssetEventQueue(ti_id={self.ti_id!r}, attempts={self.attempts!r})" + + +def _register_queued_asset_event(row: AssetEventQueue, *, session: Session) -> None: + """ + Register the asset events captured in one :class:`AssetEventQueue` row, then delete it. + + Resolves the live task instance by natural key (``dag_id``/``run_id``/``task_id``/``map_index``) + rather than the surrogate ``ti_id``: clearing a task reassigns its id, so a lookup by the + enqueued id would miss the row on any backend that does not cascade the id change. If the task + instance no longer exists there is nothing to register and the row is simply dropped. The caller + owns the surrounding transaction (the scheduler wraps each row in a savepoint; the in-process + runner commits the session). + """ + from airflow.api_fastapi.execution_api.datamodels.asset import AssetProfile + from airflow.models.taskinstance import TaskInstance + + payload = row.payload + ti_key = payload["ti_key"] + ti = session.scalar( + select(TaskInstance).where( + TaskInstance.dag_id == ti_key["dag_id"], + TaskInstance.run_id == ti_key["run_id"], + TaskInstance.task_id == ti_key["task_id"], + TaskInstance.map_index == ti_key["map_index"], + ) + ) + if ti is not None: + task_outlets = [AssetProfile.model_validate(outlet) for outlet in payload["task_outlets"]] + TaskInstance.register_asset_changes_in_db(ti, task_outlets, payload["outlet_events"], session=session) + session.delete(row) + + +def register_pending_asset_events(*, ti_ids: Iterable[UUID], session: Session) -> None: + """ + Register queued asset events for the given task instances in line and commit. + + The execution API records asset events on task success as durable + :class:`AssetEventQueue` markers that the scheduler drains. The in-process task runner + behind ``dag.test`` has no scheduler, so it calls this to register the events itself right + after a task succeeds, the same way it runs the triggerer in process. Best-effort and + unbatched: production contention handling and retry/park bookkeeping stay in the scheduler + drain (``SchedulerJobRunner._drain_asset_event_queue``). + """ + rows = session.scalars(select(AssetEventQueue).where(AssetEventQueue.ti_id.in_(ti_ids))).all() + for row in rows: + _register_queued_asset_event(row, session=session) + session.commit() + + association_table = Table( "dagrun_asset_event", Base.metadata, diff --git a/airflow-core/src/airflow/utils/db.py b/airflow-core/src/airflow/utils/db.py index 1076715ee7180..03713248c6ddd 100644 --- a/airflow-core/src/airflow/utils/db.py +++ b/airflow-core/src/airflow/utils/db.py @@ -117,7 +117,7 @@ class MappedClassProtocol(Protocol): "3.1.8": "509b94a1042d", "3.2.0": "1d6611b6ab7c", "3.3.0": "d2f4e1b3c5a7", - "3.4.0": "7a98f1b7dbd3", + "3.4.0": "a63ae7baae97", } # Prefix used to identify tables holding data moved during migration. diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py index 542ce7eaaf15b..41994cf59c2a9 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py @@ -45,7 +45,7 @@ from airflow.api_fastapi.execution_api.security import require_auth from airflow.exceptions import AirflowSkipException from airflow.models import RenderedTaskInstanceFields, TaskReschedule, Trigger -from airflow.models.asset import AssetActive, AssetAliasModel, AssetEvent, AssetModel +from airflow.models.asset import AssetActive, AssetEvent, AssetEventQueue, AssetModel from airflow.models.dag import DagModel from airflow.models.log import Log from airflow.models.task_state_store import TaskStateStoreModel @@ -103,19 +103,6 @@ def _is_select_for_update(statement) -> bool: return getattr(statement, "is_select", False) and "FOR UPDATE" in str(statement.compile()).upper() -def _create_asset_aliases(session, num: int = 2) -> None: - asset_aliases = [ - AssetAliasModel( - id=i, - name=f"simple{i}", - group="alias", - ) - for i in range(1, 1 + num) - ] - session.add_all(asset_aliases) - session.commit() - - @pytest.fixture def _use_real_jwt_bearer(exec_app): """Remove the mock require_auth override so the real JWT validation runs end-to-end.""" @@ -1228,17 +1215,9 @@ def test_ti_update_state_to_terminal_with_rendered_map_index( assert ti.rendered_map_index == rendered_map_index @pytest.mark.parametrize( - "task_outlets", + "outlet_events", [ - pytest.param([{"name": "my-task", "uri": "s3://bucket/my-task", "type": "Asset"}], id="asset"), - pytest.param([{"name": "my-task", "type": "AssetNameRef"}], id="name-ref"), - pytest.param([{"uri": "s3://bucket/my-task", "type": "AssetUriRef"}], id="uri-ref"), - ], - ) - @pytest.mark.parametrize( - ("outlet_events", "expected_extra"), - [ - pytest.param([], {}, id="default"), + pytest.param([], id="default"), pytest.param( [ { @@ -1250,23 +1229,18 @@ def test_ti_update_state_to_terminal_with_rendered_map_index( "extra": {"foo": 2}, }, ], - {"foo": 1}, id="extra", ), ], ) - def test_ti_update_state_to_success_with_asset_events( - self, client, session, create_task_instance, task_outlets, outlet_events, expected_extra + def test_ti_update_state_to_success_enqueues_asset_events( + self, client, session, create_task_instance, outlet_events ): - asset = AssetModel( - id=1, - name="my-task", - uri="s3://bucket/my-task", - group="asset", - extra={}, - ) - asset_active = AssetActive.for_asset(asset) - session.add_all([asset, asset_active]) + """A successful task with outlets records a durable ``asset_event_queue`` marker in the same + transaction as the state. Registration is the scheduler drain's job, so no ``AssetEvent`` + exists yet -- the request only enqueues the raw outlets/events plus the TI natural key the + drain resolves the live task instance with.""" + task_outlets = [{"name": "my-task", "uri": "s3://bucket/my-task", "type": "Asset"}] ti = create_task_instance( task_id="test_ti_update_state_to_success_with_asset_events", @@ -1289,52 +1263,37 @@ def test_ti_update_state_to_success_with_asset_events( assert response.text == "" session.expire_all() - event = session.scalars(select(AssetEvent)).all() - assert len(event) == 1 - assert event[0].asset == AssetModel(name="my-task", uri="s3://bucket/my-task", extra={}) - assert event[0].extra == expected_extra + # Registration is deferred to the scheduler drain: nothing registered inline. + assert session.scalars(select(AssetEvent)).all() == [] + + rows = session.scalars(select(AssetEventQueue)).all() + assert len(rows) == 1 + assert rows[0].ti_id == ti.id + assert rows[0].attempts == 0 + assert rows[0].payload == { + "task_outlets": task_outlets, + "outlet_events": outlet_events, + "ti_key": { + "dag_id": ti.dag_id, + "run_id": ti.run_id, + "task_id": ti.task_id, + "map_index": ti.map_index, + }, + } - @pytest.mark.parametrize( - ("outlet_events", "expected_extra"), - [ - pytest.param([], None, id="default"), - pytest.param( - [ - { - "dest_asset_key": {"name": "my-task", "uri": "s3://bucket/my-task"}, - "source_alias_name": "simple1", - "extra": {"foo": 1}, - }, - { - "dest_asset_key": {"name": "my-task-2", "uri": "s3://bucket/my-task-2"}, - "extra": {"foo": 2}, - }, - { - "dest_asset_key": {"name": "my-task-2", "uri": "s3://bucket/my-task-2"}, - "source_alias_name": "simple2", - "extra": {"foo": 3}, - }, - ], - {"foo": 1}, - id="extra", - ), - ], - ) - def test_ti_update_state_to_success_with_asset_alias_events( - self, client, session, create_task_instance, outlet_events, expected_extra + def test_ti_update_state_to_success_with_asset_alias_events_enqueues( + self, client, session, create_task_instance ): - asset = AssetModel( - id=1, - name="my-task", - uri="s3://bucket/my-task", - group="asset", - extra={}, - ) - asset_active = AssetActive.for_asset(asset) - session.add_all([asset, asset_active]) - - _create_asset_aliases(session, num=2) - + """Alias outlets are enqueued verbatim; the scheduler drain resolves the alias to concrete + AssetEvents when it registers (see the manager/drain tests), so the request itself creates + only the queue marker and registers nothing inline.""" + outlet_events = [ + { + "dest_asset_key": {"name": "my-task", "uri": "s3://bucket/my-task"}, + "source_alias_name": "simple1", + "extra": {"foo": 1}, + }, + ] ti = create_task_instance( task_id="test_ti_update_state_to_success_with_asset_events", start_date=DEFAULT_START_DATE, @@ -1356,13 +1315,72 @@ def test_ti_update_state_to_success_with_asset_alias_events( assert response.text == "" session.expire_all() - events = session.scalars(select(AssetEvent)).all() - if expected_extra is None: - assert events == [] - else: - assert len(events) == 1 - assert events[0].asset == AssetModel(name="my-task", uri="s3://bucket/my-task", extra={}) - assert events[0].extra == expected_extra + assert session.scalars(select(AssetEvent)).all() == [] + rows = session.scalars(select(AssetEventQueue)).all() + assert len(rows) == 1 + assert rows[0].payload["task_outlets"] == [{"name": "simple1", "uri": None, "type": "AssetAlias"}] + assert rows[0].payload["outlet_events"] == outlet_events + + def test_ti_update_state_forced_failure_skips_asset_queue(self, client, session, create_task_instance): + """When an unexpected error while building the state update forces the TI to FAILED, the + success-only asset path must not run: nothing is registered and nothing is enqueued.""" + ti = create_task_instance( + task_id="test_ti_update_state_forced_failure_skips_asset_queue", + start_date=DEFAULT_START_DATE, + state=State.RUNNING, + ) + session.commit() + + with mock.patch( + "airflow.api_fastapi.execution_api.routes.task_instances." + "_create_ti_state_update_query_and_update_state", + side_effect=Exception("kaboom"), + ): + response = client.patch( + f"/execution/task-instances/{ti.id}/state", + json={ + "state": "success", + "end_date": DEFAULT_END_DATE.isoformat(), + "task_outlets": [{"name": "my-task", "uri": "s3://bucket/my-task", "type": "Asset"}], + "outlet_events": [], + }, + ) + + assert response.status_code == 204 + assert response.text == "" + session.expire_all() + + assert session.get(TaskInstance, ti.id).state == State.FAILED + assert session.scalars(select(AssetEventQueue)).all() == [] + assert session.scalars(select(AssetEvent)).all() == [] + + def test_ti_update_state_success_no_outlets_no_queue(self, client, session, create_task_instance): + """A successful TI with no outlets and no outlet events short-circuits: no registration, no + queue row, no AssetEvent.""" + ti = create_task_instance( + task_id="test_ti_update_state_success_no_outlets_no_queue", + start_date=DEFAULT_START_DATE, + state=State.RUNNING, + ) + session.commit() + + response = client.patch( + f"/execution/task-instances/{ti.id}/state", + json={ + "state": "success", + "end_date": DEFAULT_END_DATE.isoformat(), + "task_outlets": [], + "outlet_events": [], + }, + ) + + assert response.status_code == 204 + assert response.text == "" + session.expire_all() + + assert session.get(TaskInstance, ti.id).state == State.SUCCESS + assert session.scalars(select(AssetEventQueue)).all() == [] + assert session.scalars(select(AssetEvent)).all() == [] @pytest.mark.parametrize( ("partition_key", "expected_status"), diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index dff3adea48b21..a10fd26b79de4 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -63,6 +63,7 @@ AssetAliasModel, AssetDagRunQueue, AssetEvent, + AssetEventQueue, AssetModel, AssetPartitionDagRun, DagScheduleAssetReference, @@ -5712,6 +5713,151 @@ def dict_from_obj(obj): assert created_run.creating_job_id == scheduler_job.id + def _seed_asset_producer_and_queue_row(self, session, dag_maker, *, outlet_events=None): + """Create a producer DAG whose task has an asset outlet, a downstream DAG scheduled on that + asset, a successful producer TI, and a matching ``AssetEventQueue`` row. Returns the TI.""" + asset = Asset(uri="test://drain-asset", name="drain_asset", group="test_group") + + with dag_maker(dag_id="drain-producer", start_date=timezone.utcnow(), session=session) as dag: + BashOperator(task_id="task", bash_command="echo 1", outlets=[asset]) + dr = dag_maker.create_dagrun(run_id="producer-run") + [ti] = dr.get_task_instances(session=session) + ti.state = State.SUCCESS + # Capture the outlet profiles before creating other DAGs staled the serialized reference. + task_outlets = [o.asprofile().model_dump(mode="json") for o in dag.get_task("task").outlets] + session.flush() + + # Downstream DAG scheduled on the asset -> registration should enqueue it (ADRQ row). + with dag_maker(dag_id="drain-consumer", schedule=[asset], session=session): + EmptyOperator(task_id="downstream") + session.add( + AssetEventQueue( + ti_id=ti.id, + payload={ + "task_outlets": task_outlets, + "outlet_events": outlet_events or [], + "ti_key": { + "dag_id": ti.dag_id, + "run_id": ti.run_id, + "task_id": ti.task_id, + "map_index": ti.map_index, + }, + }, + attempts=0, + ) + ) + session.commit() + return ti + + @pytest.mark.need_serialized_dag + def test_drain_asset_event_queue_registers_and_deletes(self, session, dag_maker): + """A pending queue row for a successful producer TI is registered (AssetEvent + downstream + AssetDagRunQueue), the queue row is deleted, and the processed metric is emitted.""" + ti = self._seed_asset_producer_and_queue_row(session, dag_maker) + + scheduler_job = Job() + self.job_runner = SchedulerJobRunner(job=scheduler_job, executors=[self.null_exec]) + + with mock.patch("airflow.jobs.scheduler_job_runner.stats") as mock_stats: + self.job_runner._drain_asset_event_queue(session=session) + + # The queue row was consumed. + assert session.scalars(select(AssetEventQueue)).all() == [] + # An AssetEvent was produced by the (now drained) TI. + event = session.scalar(select(AssetEvent).where(AssetEvent.source_run_id == ti.run_id)) + assert event is not None + # The downstream asset-scheduled DAG was queued. + adrq = session.scalars( + select(AssetDagRunQueue).where(AssetDagRunQueue.target_dag_id == "drain-consumer") + ).all() + assert len(adrq) == 1 + mock_stats.incr.assert_any_call("asset.event_queue.processed", 1) + + @pytest.mark.need_serialized_dag + def test_drain_asset_event_queue_retries_then_parks(self, session, dag_maker): + """A row whose registration keeps failing is retried across drain passes, one attempt per + pass, until it reaches max_attempts, at which point it is parked (no longer claimed) and the + failures metric is emitted exactly once. Further passes leave attempts capped.""" + self._seed_asset_producer_and_queue_row(session, dag_maker) + + scheduler_job = Job() + self.job_runner = SchedulerJobRunner(job=scheduler_job, executors=[self.null_exec]) + + max_attempts = self.job_runner._asset_event_queue_max_attempts + + with mock.patch( + "airflow.models.taskinstance.TaskInstance.register_asset_changes_in_db", + side_effect=RuntimeError("boom"), + ): + # Run one extra pass beyond max_attempts to prove the counter caps and does not overflow. + for cur_attempt in range(max_attempts + 1): + with mock.patch("airflow.jobs.scheduler_job_runner.stats") as mock_stats: + self.job_runner._drain_asset_event_queue(session=session) + session.expire_all() + row = session.scalars(select(AssetEventQueue)).one() + # attempts increases by one per pass until it reaches max_attempts, then stays there + # because the drain only claims rows WHERE attempts < max_attempts. + assert row.attempts == min(cur_attempt + 1, max_attempts) + if cur_attempt + 1 == max_attempts: + # The pass that reaches the cap parks the row and reports it once. + mock_stats.incr.assert_any_call("asset.event_queue.failures") + else: + assert mock.call("asset.event_queue.failures") not in mock_stats.incr.mock_calls + + @pytest.mark.need_serialized_dag + def test_drain_asset_event_queue_missing_ti_purged(self, session, dag_maker): + """A queue row whose task instance can no longer be resolved by natural key is dropped + without error or registration. We point the row's ``ti_key`` at a run that does not exist so + the drain's lookup misses, exercising the purge branch.""" + self._seed_asset_producer_and_queue_row(session, dag_maker) + row = session.scalars(select(AssetEventQueue)).one() + row.payload = {**row.payload, "ti_key": {**row.payload["ti_key"], "run_id": "does-not-exist"}} + session.commit() + + scheduler_job = Job() + self.job_runner = SchedulerJobRunner(job=scheduler_job, executors=[self.null_exec]) + + self.job_runner._drain_asset_event_queue(session=session) + + assert session.scalars(select(AssetEventQueue)).all() == [] + # Nothing was registered because the task instance could not be resolved. + assert session.scalars(select(AssetEvent)).all() == [] + + @pytest.mark.need_serialized_dag + def test_drain_asset_event_queue_empty_noop(self, session, dag_maker): + """With no pending rows the drain is a no-op and creates nothing.""" + scheduler_job = Job() + self.job_runner = SchedulerJobRunner(job=scheduler_job, executors=[self.null_exec]) + + self.job_runner._drain_asset_event_queue(session=session) + + assert session.scalars(select(AssetEventQueue)).all() == [] + assert session.scalars(select(AssetEvent)).all() == [] + + @pytest.mark.need_serialized_dag + def test_clearing_ti_still_registers_via_natural_key(self, session, dag_maker): + """Clearing/retrying a TI reassigns its uuid7 id. The drain resolves the live task instance + by natural key (dag_id, run_id, task_id, map_index), not by the enqueued surrogate id, so the + pending asset events are still registered after a clear. This runs on the default sqlite + engine (no ``PRAGMA foreign_keys=ON``), which does not honor the FK's ON UPDATE CASCADE -- + exactly the production condition under which a cascade-dependent re-point would drop them.""" + ti = self._seed_asset_producer_and_queue_row(session, dag_maker) + old_id = ti.id + + # Clearing reassigns the id; the same run_id keeps the DagRun partition context intact. + ti.prepare_db_for_next_try(session=session) + session.commit() + assert ti.id != old_id + + scheduler_job = Job() + self.job_runner = SchedulerJobRunner(job=scheduler_job, executors=[self.null_exec]) + self.job_runner._drain_asset_event_queue(session=session) + + # The events were registered despite the id change, and the queue row was consumed. + assert session.scalars(select(AssetEventQueue)).all() == [] + event = session.scalar(select(AssetEvent).where(AssetEvent.source_run_id == ti.run_id)) + assert event is not None + @pytest.mark.need_serialized_dag @pytest.mark.parametrize( ("catchup", "expects_old_event"), 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 c0d43641559d6..012e588d6b5e6 100644 --- a/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml +++ b/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml @@ -318,6 +318,30 @@ metrics: legacy_name: "-" name_variables: [] + - name: "asset.event_queue.enqueued" + description: "Number of task-completion asset registrations written as a durable queue marker" + type: "counter" + legacy_name: "-" + name_variables: [] + + - name: "asset.event_queue.processed" + description: "Number of asset event registrations drained from the queue by the scheduler" + type: "counter" + legacy_name: "-" + name_variables: [] + + - name: "asset.event_queue.failures" + description: "Number of asset event registrations parked after exceeding the retry limit" + type: "counter" + legacy_name: "-" + name_variables: [] + + - name: "asset.event_queue.pending" + description: "Number of asset event registrations currently pending in the queue" + type: "gauge" + legacy_name: "-" + name_variables: [] + - name: "deadline_alerts.deadline_created" description: "Number of deadline alerts created for a Dag run" type: "counter" diff --git a/task-sdk-integration-tests/tests/task_sdk_tests/conftest.py b/task-sdk-integration-tests/tests/task_sdk_tests/conftest.py index b2c75db65ef52..f290411a094b7 100644 --- a/task-sdk-integration-tests/tests/task_sdk_tests/conftest.py +++ b/task-sdk-integration-tests/tests/task_sdk_tests/conftest.py @@ -725,10 +725,12 @@ def sdk_client_for_assets(asset_test_setup): @pytest.fixture(scope="session") def asset_test_setup(docker_compose_setup, airflow_ready): """Setup assets for testing by triggering asset_producer_dag.""" + import time + headers = airflow_ready["headers"] auth_token = airflow_ready["auth_token"] - return setup_dag_and_get_client( + setup = setup_dag_and_get_client( dag_id="asset_producer_dag", headers=headers, auth_token=auth_token, @@ -741,3 +743,22 @@ def asset_test_setup(docker_compose_setup, airflow_ready): "alias_name": "test_asset_alias", }, ) + + # A successful task only enqueues its asset events; the scheduler registers them when it + # drains the queue, so there is a short delay between the producer Dag reaching "success" + # and the events being queryable. Wait for the first event to appear before handing the + # setup to the tests, otherwise they race the scheduler drain and see zero events. + sdk_client = setup["sdk_client"] + for attempt in range(30): + response = sdk_client.asset_events.get(name=setup["name"]) + if response.asset_events: + console.print(f"[green]✅ Asset event registered after {attempt + 1} attempt(s)") + break + console.print( + f"[blue]Waiting for asset event registration (attempt {attempt + 1}/30, drained by scheduler)" + ) + time.sleep(2) + else: + raise TimeoutError("Asset events for asset_producer_dag were not registered within timeout period") + + return setup diff --git a/task-sdk/src/airflow/sdk/definitions/dag.py b/task-sdk/src/airflow/sdk/definitions/dag.py index cadfe6978300e..f46b3d859d120 100644 --- a/task-sdk/src/airflow/sdk/definitions/dag.py +++ b/task-sdk/src/airflow/sdk/definitions/dag.py @@ -1571,6 +1571,16 @@ def _run_task( break raise + # On success the execution API only enqueues the task's asset events as a durable marker; + # in a live deployment the scheduler drains that queue. dag.test runs in process with no + # scheduler, so register the events in line here (the same way the triggerer is run in + # process above) rather than leaving them stuck in the queue. + if ti.state == TaskInstanceState.SUCCESS: + from airflow.models.asset import register_pending_asset_events + + with create_session() as session: + register_pending_asset_events(ti_ids=[ti.id], session=session) + log.info("[DAG TEST] end task task_id=%s map_index=%s", ti.task_id, ti.map_index) return taskrun_result