kaxil commented on code in PR #70972: URL: https://github.com/apache/airflow/pull/70972#discussion_r3705982822
########## airflow-core/src/airflow/migrations/versions/0128_3_4_0_add_asset_event_id_to_adrq.py: ########## @@ -0,0 +1,196 @@ +# +# 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. + +""" +Reference the asset event from asset_dag_run_queue (consume-by-reference). + +Table ``asset_dag_run_queue`` gets a new column ``asset_event_id``, and the +primary key is made ``(target_dag_id, asset_event_id)``, so the scheduler +consumes queued asset events by reference instead of by a ``created_at`` time +window. ``asset_id`` is kept as a denormalized column. + +Existing rows are coalesced (one per ``(asset_id, target_dag_id)``) and carry no +event reference. Rather than dropping them (which would silently skip pending +asset-triggered Dag runs), the pre-migration scheduler's own consumption window is +replayed per dag to rebuild per-event rows: + + triggered_date = + MAX(asset_dag_run_queue.created_at) (per dag) + floor = + MAX(dag_run.run_after) for asset-triggered runs of the dag with + run_after < triggered_date (per dag) + contributing events = + the queued asset's events with + floor < asset_event.timestamp <= triggered_date + +The expansion is staged in a side table, the queue is cleared, the (now empty) +table is reshaped, and the staged rows are inserted back. + +Revision ID: b2f1a9c7d4e0 +Revises: 7a98f1b7dbd3 +Create Date: 2026-08-03 12:00:00.000000 +""" + +from __future__ import annotations + +from datetime import datetime, timezone +from textwrap import dedent + +import sqlalchemy as sa +from alembic import context, op + +# revision identifiers, used by Alembic. +revision = "b2f1a9c7d4e0" +down_revision = "7a98f1b7dbd3" +branch_labels = None +depends_on = None +airflow_version = "3.4.0" + +_STAGING = "_adrq_migration_staging" + +_STAGE_SQL = f""" +CREATE TABLE {_STAGING} AS +SELECT DISTINCT + adrq.target_dag_id AS target_dag_id, + adrq.asset_id AS asset_id, + ae.id AS asset_event_id +FROM asset_dag_run_queue adrq +JOIN ( + SELECT t.target_dag_id, + t.triggered_date, + ( + SELECT MAX(dr.run_after) + FROM dag_run dr + WHERE dr.dag_id = t.target_dag_id + AND dr.run_type = 'asset_triggered' + AND dr.run_after < t.triggered_date + ) AS floor_date + FROM ( + SELECT target_dag_id, MAX(created_at) AS triggered_date + FROM asset_dag_run_queue + GROUP BY target_dag_id + ) t +) td ON td.target_dag_id = adrq.target_dag_id +JOIN asset_event ae + ON ae.asset_id = adrq.asset_id + AND ae.timestamp <= td.triggered_date + AND ae.timestamp > COALESCE(td.floor_date, :floor_min) +""" + + +def upgrade(): + """Directly reference the asset event from asset_dag_run_queue.""" + # 1. Add the reference column, nullable for now; it is made NOT NULL after the rebuild. + op.add_column("asset_dag_run_queue", sa.Column("asset_event_id", sa.Integer(), nullable=True)) + + # 2. Stage the faithful per-event expansion, then clear the coalesced rows. + if context.is_offline_mode(): + print( + dedent(""" + ------------ + -- WARNING: asset_dag_run_queue cannot be rebuilt in offline mode; + -- any pending (unprocessed) queued asset events will be dropped. + ------------ + """) + ) + op.execute("DELETE FROM asset_dag_run_queue") + else: + conn = op.get_bind() + floor_min = datetime(1970, 1, 1, tzinfo=timezone.utc) + conn.execute(sa.text(_STAGE_SQL), {"floor_min": floor_min}) + conn.execute(sa.text("DELETE FROM asset_dag_run_queue")) + + # 3. Make the reference NOT NULL, move the PK, add the FKs. + with op.batch_alter_table("asset_dag_run_queue") as batch_op: + batch_op.alter_column("asset_event_id", existing_type=sa.Integer(), nullable=False) + batch_op.drop_constraint("adrq_asset_fkey", type_="foreignkey") + batch_op.drop_constraint("assetdagrunqueue_pkey", type_="primary") + batch_op.create_primary_key("assetdagrunqueue_pkey", ["target_dag_id", "asset_event_id"]) Review Comment: With target_dag_id now leading the PK, `idx_asset_dag_run_queue_target_dag_id` is redundant and could be dropped in this migration. Conversely asset_id loses its index on postgres (the old PK led with it, and postgres doesn't auto-index FKs); the /assets/{assetId}/queuedEvents endpoints filter on it, though the table is usually small enough that this may not matter. ########## airflow-core/src/airflow/models/asset.py: ########## @@ -749,21 +749,28 @@ def __repr__(self): class AssetDagRunQueue(Base): """Model for storing asset events that need processing.""" - asset_id: Mapped[int] = mapped_column(Integer, primary_key=True, nullable=False) target_dag_id: Mapped[str] = mapped_column(StringID(), primary_key=True, nullable=False) + asset_event_id: Mapped[int] = mapped_column(Integer, primary_key=True, nullable=False) Review Comment: Now that ADRQ can hold multiple rows per (asset_id, target_dag_id), the LEFT JOIN in `get_asset_triggered_next_run_info` (models/dag.py) fans out and inflates both counts: a dag scheduled on 2 assets with 3 queued events for one of them reports "3 of 4 assets updated" (checked with a quick sqlite repro). `total` and `ready` need distinct counting on the asset, e.g. `count(distinct DagScheduleAssetReference.asset_id)`. The single-uri `case((func.count() == 1, ...))` hack breaks the same way. Multiple rows per pair are routine here: an AND-condition dag accumulates one row per event while it waits, and the migration expansion can create many at once. ########## airflow-core/src/airflow/migrations/versions/0128_3_4_0_add_asset_event_id_to_adrq.py: ########## @@ -0,0 +1,196 @@ +# +# 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. + +""" +Reference the asset event from asset_dag_run_queue (consume-by-reference). + +Table ``asset_dag_run_queue`` gets a new column ``asset_event_id``, and the +primary key is made ``(target_dag_id, asset_event_id)``, so the scheduler +consumes queued asset events by reference instead of by a ``created_at`` time +window. ``asset_id`` is kept as a denormalized column. + +Existing rows are coalesced (one per ``(asset_id, target_dag_id)``) and carry no +event reference. Rather than dropping them (which would silently skip pending +asset-triggered Dag runs), the pre-migration scheduler's own consumption window is +replayed per dag to rebuild per-event rows: + + triggered_date = + MAX(asset_dag_run_queue.created_at) (per dag) + floor = + MAX(dag_run.run_after) for asset-triggered runs of the dag with + run_after < triggered_date (per dag) + contributing events = + the queued asset's events with + floor < asset_event.timestamp <= triggered_date + +The expansion is staged in a side table, the queue is cleared, the (now empty) +table is reshaped, and the staged rows are inserted back. + +Revision ID: b2f1a9c7d4e0 +Revises: 7a98f1b7dbd3 +Create Date: 2026-08-03 12:00:00.000000 +""" + +from __future__ import annotations + +from datetime import datetime, timezone +from textwrap import dedent + +import sqlalchemy as sa +from alembic import context, op + +# revision identifiers, used by Alembic. +revision = "b2f1a9c7d4e0" +down_revision = "7a98f1b7dbd3" +branch_labels = None +depends_on = None +airflow_version = "3.4.0" + +_STAGING = "_adrq_migration_staging" + +_STAGE_SQL = f""" +CREATE TABLE {_STAGING} AS +SELECT DISTINCT + adrq.target_dag_id AS target_dag_id, + adrq.asset_id AS asset_id, + ae.id AS asset_event_id +FROM asset_dag_run_queue adrq +JOIN ( + SELECT t.target_dag_id, + t.triggered_date, + ( + SELECT MAX(dr.run_after) + FROM dag_run dr + WHERE dr.dag_id = t.target_dag_id + AND dr.run_type = 'asset_triggered' + AND dr.run_after < t.triggered_date + ) AS floor_date + FROM ( + SELECT target_dag_id, MAX(created_at) AS triggered_date + FROM asset_dag_run_queue + GROUP BY target_dag_id + ) t +) td ON td.target_dag_id = adrq.target_dag_id +JOIN asset_event ae + ON ae.asset_id = adrq.asset_id + AND ae.timestamp <= td.triggered_date + AND ae.timestamp > COALESCE(td.floor_date, :floor_min) Review Comment: This floor only replays the catchup=True window. The pre-migration scheduler also floored catchup=False dags at `min(DagScheduleAssetReference.created_at)`, so a catchup=False dag whose first asset-triggered run is pending at upgrade time gets its whole pre-subscription backlog expanded here, and the consume-by-reference path will consume every staged row in that first run (the #39456 behaviour the floor existed for). I think the dsar floor can go into this COALESCE unconditionally: catchup=True dags don't need backlog rows staged at all, since the catchup branch in `_create_dag_runs_asset_triggered` re-derives unconsumed backlog events from the schedule references without a queue row. ########## airflow-core/src/airflow/assets/manager.py: ########## @@ -827,71 +779,33 @@ def _get_or_create_apdr( return apdr @classmethod - def _queue_dagruns_nonpartitioned_slow_path( + def _queue_dagruns_nonpartitioned( cls, asset_id: int, dags_to_queue: set[DagModel], event: AssetEvent, session: Session ) -> None: - def _queue_dagrun_if_needed(dag: DagModel) -> str | None: - item = AssetDagRunQueue(target_dag_id=dag.dag_id, asset_id=asset_id, created_at=event.timestamp) - # Don't error whole transaction when a single RunQueue item conflicts. - # https://docs.sqlalchemy.org/en/20/orm/session_transaction.html#using-savepoint - try: - with session.begin_nested(): - existing = session.get( - AssetDagRunQueue, {"target_dag_id": dag.dag_id, "asset_id": asset_id} - ) - if existing and existing.created_at >= event.timestamp: - cls.logger().debug("Skipping record %s due to newer timestamp", item) - return dag.dag_id # already queued with a newer timestamp - session.merge(item) - except exc.IntegrityError: - cls.logger().debug("Skipping record %s", item, exc_info=True) - return dag.dag_id - - queued_results = (_queue_dagrun_if_needed(dag) for dag in dags_to_queue) - if queued_dag_ids := [r for r in queued_results if r is not None]: - cls.logger().debug("consuming dag ids %s", queued_dag_ids) - - @classmethod - def _queue_dagruns_nonpartitioned_mysql( - cls, asset_id: int, dags_to_queue: set[DagModel], event: AssetEvent, session: Session - ) -> None: - from sqlalchemy import case - from sqlalchemy.dialects.mysql import insert + if not dags_to_queue: + return + values = [ + {"asset_id": asset_id, "target_dag_id": dag.dag_id, "asset_event_id": event.id} + for dag in dags_to_queue + ] - values = [{"target_dag_id": dag.dag_id} for dag in dags_to_queue] - stmt = insert(AssetDagRunQueue).values(asset_id=asset_id, created_at=event.timestamp) + if (dialect_name := get_dialect_name(session)) == "mysql": + from sqlalchemy.dialects.mysql import insert as my_insert - update_stmt = stmt.on_duplicate_key_update( - created_at=case( - (stmt.inserted.created_at >= AssetDagRunQueue.created_at, stmt.inserted.created_at), - else_=AssetDagRunQueue.created_at, - ) - ) - session.execute(update_stmt, values) + session.execute(my_insert(AssetDagRunQueue).prefix_with("IGNORE"), values) Review Comment: INSERT IGNORE downgrades more than duplicate-key errors to warnings on MySQL (FK violations, NOT NULL, truncation are all silently swallowed), while the postgres/sqlite path only ignores the PK conflict. A bad event id, for example, would raise on postgres but insert nothing here without a sound. `on_duplicate_key_update(asset_id=stmt.inserted.asset_id)` is the usual no-op trick that keeps the ignore scoped to duplicates. ########## airflow-core/src/airflow/models/asset.py: ########## @@ -749,21 +749,28 @@ def __repr__(self): class AssetDagRunQueue(Base): """Model for storing asset events that need processing.""" - asset_id: Mapped[int] = mapped_column(Integer, primary_key=True, nullable=False) target_dag_id: Mapped[str] = mapped_column(StringID(), primary_key=True, nullable=False) + asset_event_id: Mapped[int] = mapped_column(Integer, primary_key=True, nullable=False) + asset_id: Mapped[int] = mapped_column(Integer, nullable=False) created_at: Mapped[datetime] = mapped_column(UtcDateTime, default=timezone.utcnow, nullable=False) asset: Mapped[AssetModel] = relationship("AssetModel", viewonly=True) dag_model: Mapped[DagModel] = relationship("DagModel", viewonly=True) __tablename__ = "asset_dag_run_queue" __table_args__ = ( - PrimaryKeyConstraint(asset_id, target_dag_id, name="assetdagrunqueue_pkey"), + PrimaryKeyConstraint(target_dag_id, asset_event_id, name="assetdagrunqueue_pkey"), ForeignKeyConstraint( (asset_id,), ["asset.id"], name="adrq_asset_fkey", ondelete="CASCADE", ), + ForeignKeyConstraint( + (asset_event_id,), + ["asset_event.id"], + name="adrq_asset_event_fkey", Review Comment: asset_event is a `db clean` table (utils/db_cleanup.py), so this CASCADE gives `airflow db clean --tables asset_event` a new side effect: deleting old events now deletes pending queue rows. A queued row lives until the rest of the dag's condition is satisfied, so for an AND-condition dag pairing a frequent asset with a monthly/quarterly one (or one whose upstream producer stalls), the row can outlive a typical 30-90 day retention window. Cleaning past it silently un-readies the asset and the dag misses its trigger when the other asset finally fires. Before, the ADRQ row survived event cleanup and the run still happened. Worth a note in the db clean docs at least. Relatedly, while a condition is unmet, rows now accumulate per event rather than upserting one per asset, and `dags_needing_dagruns` loads all of them (with two joinedloads) every loop. ########## airflow-core/src/airflow/jobs/scheduler_job_runner.py: ########## @@ -2658,53 +2656,31 @@ def _create_dag_runs_asset_triggered( ) continue - triggered_date: DateTime = timezone.coerce_datetime(queued_adrqs[0].created_at) - self.log.debug( - "Creating asset-triggered DagRun for '%s': %d queued assets, triggered_date=%s", - dag.dag_id, - len(queued_adrqs), - triggered_date, - ) - cte = ( - select(func.max(DagRun.run_after).label("previous_dag_run_run_after")) - .where( - DagRun.dag_id == dag.dag_id, - DagRun.run_type == DagRunType.ASSET_TRIGGERED, - DagRun.run_after < triggered_date, - ) - .cte() - ) - - # A first asset-triggered run has no previous run to floor the event window. With - # catchup off, floor it at when the Dag started scheduling on its assets so the - # backlog is skipped; with catchup on, only date.min applies and the backlog replays. - event_window_floor: list[Any] = [cte.c.previous_dag_run_run_after] - if not dag.catchup: - event_window_floor.append( - select(func.min(DagScheduleAssetReference.created_at)) - .where(DagScheduleAssetReference.dag_id == dag.dag_id) - .scalar_subquery() + referenced_event_ids = {adrq.asset_event_id for adrq in queued_adrqs} + event_predicate: ColumnElement[bool] = AssetEvent.id.in_(referenced_event_ids) Review Comment: One behaviour delta worth confirming: the old select required the event's asset to still be in the dag's schedule (the DagScheduleAssetReference/alias join). Consuming by reference means a queued event for an asset that was since removed from the dag's schedule now gets consumed and shows up in the run's consumed_asset_events / triggering_asset_events, where before it was silently dropped along with its ADRQ row. Intended? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
