diff --git a/airflow-core/src/airflow/assets/manager.py b/airflow-core/src/airflow/assets/manager.py index c3491574e9148..e2ed339237bee 100644 --- a/airflow-core/src/airflow/assets/manager.py +++ b/airflow-core/src/airflow/assets/manager.py @@ -518,13 +518,21 @@ def _queue_dagruns( # mapped) tasks update the same asset, this can fail with a unique # constraint violation. # - # If we support it, use ON CONFLICT to do nothing, otherwise + # If we support it, use ON CONFLICT / ON DUPLICATE KEY to do nothing, otherwise # "fallback" to running this in a nested transaction. This is needed # so that the adding of these rows happens in the same transaction # where `ti.state` is changed. - if get_dialect_name(session) == "postgresql": - return cls._queue_dagruns_nonpartitioned_postgres(asset_id, non_partitioned_dags, session) - return cls._queue_dagruns_nonpartitioned_slow_path(asset_id, non_partitioned_dags, session) + # + # Rows are inserted in sorted order so concurrent fan-outs acquire + # row locks in a consistent order, which prevents deadlocks (a + # set's iteration order differs between processes). + dag_ids = sorted(dag.dag_id for dag in non_partitioned_dags) + dialect = get_dialect_name(session) + if dialect == "postgresql": + return cls._queue_dagruns_nonpartitioned_postgres(asset_id, dag_ids, session) + if dialect == "mysql": + return cls._queue_dagruns_nonpartitioned_mysql(asset_id, dag_ids, session) + return cls._queue_dagruns_nonpartitioned_slow_path(asset_id, dag_ids, session) @classmethod def _queue_partitioned_dags( @@ -790,10 +798,10 @@ def _get_or_create_apdr( @classmethod def _queue_dagruns_nonpartitioned_slow_path( - cls, asset_id: int, dags_to_queue: set[DagModel], session: Session + cls, asset_id: int, dag_ids: list[str], session: Session ) -> None: - def _queue_dagrun_if_needed(dag: DagModel) -> str | None: - item = AssetDagRunQueue(target_dag_id=dag.dag_id, asset_id=asset_id) + for dag_id in dag_ids: + item = AssetDagRunQueue(target_dag_id=dag_id, asset_id=asset_id) # Don't error whole transaction when a single RunQueue item conflicts. # https://docs.sqlalchemy.org/en/20/orm/session_transaction.html#using-savepoint try: @@ -801,19 +809,26 @@ def _queue_dagrun_if_needed(dag: DagModel) -> str | None: session.merge(item) except exc.IntegrityError: cls.logger().debug("Skipping record %s", item, exc_info=True) - return dag.dag_id + cls.logger().debug("consuming dag ids %s", dag_ids) + + @classmethod + def _queue_dagruns_nonpartitioned_mysql(cls, asset_id: int, dag_ids: list[str], session: Session) -> None: + from sqlalchemy.dialects.mysql import insert - 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) + # The ON DUPLICATE KEY UPDATE skips duplicates while still throwing errors for other critical issues. + stmt = insert(AssetDagRunQueue).values( + [{"asset_id": asset_id, "target_dag_id": dag_id} for dag_id in dag_ids] + ) + stmt = stmt.on_duplicate_key_update(target_dag_id=stmt.inserted.target_dag_id) + session.execute(stmt) @classmethod def _queue_dagruns_nonpartitioned_postgres( - cls, asset_id: int, dags_to_queue: set[DagModel], session: Session + cls, asset_id: int, dag_ids: list[str], session: Session ) -> None: from sqlalchemy.dialects.postgresql import insert - values = [{"target_dag_id": dag.dag_id} for dag in dags_to_queue] + values = [{"target_dag_id": dag_id} for dag_id in dag_ids] stmt = insert(AssetDagRunQueue).values(asset_id=asset_id).on_conflict_do_nothing() session.execute(stmt, values) diff --git a/airflow-core/src/airflow/models/trigger.py b/airflow-core/src/airflow/models/trigger.py index 386a6c931fcc6..9455a30d6acfd 100644 --- a/airflow-core/src/airflow/models/trigger.py +++ b/airflow-core/src/airflow/models/trigger.py @@ -37,7 +37,7 @@ from airflow.models.taskinstance import TaskInstance from airflow.serialization.enums import stringify_encoding_keys from airflow.triggers.base import BaseTaskEndEvent -from airflow.utils.retries import run_with_db_retries +from airflow.utils.retries import retry_db_transaction, run_with_db_retries from airflow.utils.session import NEW_SESSION, provide_session from airflow.utils.sqlalchemy import UtcDateTime, get_dialect_name, with_row_locks from airflow.utils.state import TaskInstanceState @@ -270,12 +270,17 @@ def clean_unused(cls, *, session: Session = NEW_SESSION) -> None: @classmethod @provide_session + @retry_db_transaction def submit_event(cls, trigger_id, event: TriggerEvent, *, session: Session = NEW_SESSION) -> None: """ Fire an event. Resume all tasks that were in deferred state. Send an event to all assets associated to the trigger. + + Retried as a whole on transient database errors, e.g. a deadlock + while fanning an asset event out to consumer Dags. Such errors abort + the entire transaction, so the retry must replay it from the start. """ # Resume deferred tasks for task_instance in session.scalars( diff --git a/airflow-core/tests/unit/assets/test_manager.py b/airflow-core/tests/unit/assets/test_manager.py index b788b9ab28699..50a2d166b73c0 100644 --- a/airflow-core/tests/unit/assets/test_manager.py +++ b/airflow-core/tests/unit/assets/test_manager.py @@ -26,6 +26,7 @@ import pytest from sqlalchemy import delete, func, select +from sqlalchemy.dialects import mysql from sqlalchemy.orm import Session from airflow import settings @@ -149,6 +150,91 @@ def test_register_asset_change(self, session, mock_task_instance): ) assert session.scalar(select(func.count()).select_from(AssetDagRunQueue)) == 2 + @pytest.mark.usefixtures("dag_maker", "testing_dag_bundle") + def test_register_asset_change_twice_deduplicates_queue_records(self, session, mock_task_instance): + """Fanning out the same asset again must tolerate the already-queued rows, on every backend.""" + asset_manager = AssetManager() + + asset = Asset(uri="test://asset1", name="test_asset_uri", group="asset") + bundle_name = "testing" + + dag1 = DagModel(dag_id="dag1", is_stale=False, bundle_name=bundle_name) + dag2 = DagModel(dag_id="dag2", is_stale=False, bundle_name=bundle_name) + session.add_all([dag1, dag2]) + + asm = AssetModel(uri="test://asset1/", name="test_asset_uri", group="asset") + session.add(asm) + asm.scheduled_dags = [DagScheduleAssetReference(dag_id=dag.dag_id) for dag in (dag1, dag2)] + session.execute(delete(AssetDagRunQueue)) + session.flush() + + for _ in range(2): + asset_manager.register_asset_change( + task_instance=mock_task_instance, asset=asset, session=session + ) + session.flush() + + assert ( + session.scalar(select(func.count()).select_from(AssetEvent).where(AssetEvent.asset_id == asm.id)) + == 2 + ) + assert session.scalar(select(func.count()).select_from(AssetDagRunQueue)) == 2 + + @pytest.mark.parametrize( + ("dialect", "expected_helper"), + [ + ("postgresql", "_queue_dagruns_nonpartitioned_postgres"), + ("mysql", "_queue_dagruns_nonpartitioned_mysql"), + ("sqlite", "_queue_dagruns_nonpartitioned_slow_path"), + ], + ) + @mock.patch.object(AssetManager, "_queue_dagruns_nonpartitioned_slow_path", autospec=True) + @mock.patch.object(AssetManager, "_queue_dagruns_nonpartitioned_mysql", autospec=True) + @mock.patch.object(AssetManager, "_queue_dagruns_nonpartitioned_postgres", autospec=True) + @mock.patch("airflow.assets.manager.get_dialect_name", autospec=True) + def test_queue_dagruns_dispatches_by_dialect_with_sorted_dag_ids( + self, mock_get_dialect_name, mock_postgres, mock_mysql, mock_slow_path, dialect, expected_helper + ): + mock_get_dialect_name.return_value = dialect + dags_to_queue = {DagModel(dag_id=dag_id) for dag_id in ("dag_b", "dag_a", "dag_c")} + session = mock.Mock(spec=Session) + + AssetManager._queue_dagruns( + asset_id=1, + dags_to_queue=dags_to_queue, + partition_key=None, + partition_date=None, + event=mock.Mock(spec=AssetEvent), + task_instance=None, + session=session, + ) + + helper_mocks = { + "_queue_dagruns_nonpartitioned_postgres": mock_postgres, + "_queue_dagruns_nonpartitioned_mysql": mock_mysql, + "_queue_dagruns_nonpartitioned_slow_path": mock_slow_path, + } + for name, helper_mock in helper_mocks.items(): + if name == expected_helper: + helper_mock.assert_called_once_with(1, ["dag_a", "dag_b", "dag_c"], session) + else: + helper_mock.assert_not_called() + + def test_queue_dagruns_nonpartitioned_mysql_single_multi_row_upsert(self): + session = mock.Mock(spec=Session) + + AssetManager._queue_dagruns_nonpartitioned_mysql(42, ["dag_a", "dag_b"], session) + + session.execute.assert_called_once() + stmt = session.execute.call_args.args[0] + sql = str(stmt.compile(dialect=mysql.dialect(), compile_kwargs={"literal_binds": True})) + # NULL is a compile-time placeholder: the created_at default is evaluated per row at execution. + assert ( + "INSERT INTO asset_dag_run_queue (asset_id, target_dag_id, created_at) " + "VALUES (42, 'dag_a', NULL), (42, 'dag_b', NULL) " + "ON DUPLICATE KEY UPDATE target_dag_id = VALUES(target_dag_id)" + ) in sql + @pytest.mark.usefixtures("clear_assets") def test_register_asset_change_with_alias( self, session, dag_maker, mock_task_instance, testing_dag_bundle diff --git a/airflow-core/tests/unit/models/test_trigger.py b/airflow-core/tests/unit/models/test_trigger.py index 91a6a92e27cb3..471047231128e 100644 --- a/airflow-core/tests/unit/models/test_trigger.py +++ b/airflow-core/tests/unit/models/test_trigger.py @@ -27,6 +27,7 @@ import pytz from cryptography.fernet import Fernet from sqlalchemy import delete, func, select +from sqlalchemy.exc import OperationalError from airflow._shared.timezones import timezone from airflow.jobs.job import Job @@ -256,6 +257,33 @@ def test_submit_event_no_n_plus_one_for_assets(_, session, asset_count, expected Trigger.submit_event(trigger_id, TriggerEvent("payload"), session=session) +@patch("airflow.models.trigger.AssetManager.register_asset_change") +def test_submit_event_retries_on_transient_db_error(mock_register_asset_change, session): + """ + A transient database error during the asset-event fan-out rolls back the whole transaction, + so submit_event must replay it from the start. + """ + trigger = Trigger(classpath="airflow.triggers.testing.SuccessTrigger", kwargs={}) + session.add(trigger) + session.flush() + asset = AssetModel("test") + asset.add_trigger(trigger, "test_asset_watcher") + session.add(asset) + session.commit() + trigger_id = trigger.id + + deadlock = OperationalError( + "INSERT INTO asset_dag_run_queue", + {}, + Exception("(1213, 'Deadlock found when trying to get lock; try restarting transaction')"), + ) + mock_register_asset_change.side_effect = [deadlock, None] + + Trigger.submit_event(trigger_id, TriggerEvent("payload"), session=session) + + assert mock_register_asset_change.call_count == 2 + + def test_submit_failure(session, create_task_instance): """ Tests that failures submitted to a trigger fail their dependent