diff --git a/airflow-core/src/airflow/assets/manager.py b/airflow-core/src/airflow/assets/manager.py index c3491574e9148..1dd26379b9d12 100644 --- a/airflow-core/src/airflow/assets/manager.py +++ b/airflow-core/src/airflow/assets/manager.py @@ -518,12 +518,17 @@ 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 - # "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": + # Where the dialect supports a single-statement "insert, ignore on + # conflict" we use it; it is atomic, avoids the per-row SAVEPOINT churn, + # and holds locks for far less time (which on MySQL/InnoDB also makes the + # concurrent fan-out much less deadlock-prone). Otherwise we "fallback" to + # a nested transaction per row. Either way the rows are added in the same + # transaction where `ti.state` is changed. + dialect_name = get_dialect_name(session) + if dialect_name == "postgresql": return cls._queue_dagruns_nonpartitioned_postgres(asset_id, non_partitioned_dags, session) + if dialect_name == "mysql": + return cls._queue_dagruns_nonpartitioned_mysql(asset_id, non_partitioned_dags, session) return cls._queue_dagruns_nonpartitioned_slow_path(asset_id, non_partitioned_dags, session) @classmethod @@ -817,6 +822,20 @@ def _queue_dagruns_nonpartitioned_postgres( stmt = insert(AssetDagRunQueue).values(asset_id=asset_id).on_conflict_do_nothing() session.execute(stmt, values) + @classmethod + def _queue_dagruns_nonpartitioned_mysql( + cls, asset_id: int, dags_to_queue: set[DagModel], session: Session + ) -> None: + from sqlalchemy.dialects.mysql import insert + + values = [{"target_dag_id": dag.dag_id} for dag in dags_to_queue] + stmt = insert(AssetDagRunQueue).values(asset_id=asset_id) + # MySQL has no "ON CONFLICT DO NOTHING"; a no-op ON DUPLICATE KEY UPDATE turns a + # conflicting (asset_id, target_dag_id) row into a no-op rather than an error, + # matching the Postgres path. + stmt = stmt.on_duplicate_key_update(target_dag_id=stmt.inserted.target_dag_id) + session.execute(stmt, values) + def resolve_asset_manager() -> AssetManager: """Retrieve the asset manager.""" diff --git a/airflow-core/tests/unit/assets/test_manager.py b/airflow-core/tests/unit/assets/test_manager.py index b788b9ab28699..c290d7f5a6328 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 @@ -211,6 +212,46 @@ def test_register_asset_change_no_downstreams(self, session, mock_task_instance) ) assert session.scalar(select(func.count()).select_from(AssetDagRunQueue)) == 0 + @pytest.mark.parametrize( + ("dialect_name", "expected_helper"), + [ + ("postgresql", "_queue_dagruns_nonpartitioned_postgres"), + ("mysql", "_queue_dagruns_nonpartitioned_mysql"), + ("sqlite", "_queue_dagruns_nonpartitioned_slow_path"), + ], + ) + def test_queue_dagruns_routes_by_dialect(self, dialect_name, expected_helper): + """Test that _queue_dagruns routes to the dialect-appropriate queue helper.""" + dag = DagModel(dag_id="dag1") + session = mock.MagicMock(spec=Session) + with ( + mock.patch("airflow.assets.manager.get_dialect_name", return_value=dialect_name), + mock.patch.object(AssetManager, "_queue_partitioned_dags"), + mock.patch.object(AssetManager, expected_helper) as mock_helper, + ): + AssetManager._queue_dagruns( + asset_id=1, + dags_to_queue={dag}, + partition_key=None, + partition_date=None, + event=mock.MagicMock(), + task_instance=None, + session=session, + ) + mock_helper.assert_called_once_with(1, {dag}, session) + + def test_queue_dagruns_nonpartitioned_mysql_builds_upsert(self): + """Test that the MySQL queue path emits an INSERT ... ON DUPLICATE KEY UPDATE.""" + dag = DagModel(dag_id="dag1") + session = mock.MagicMock(spec=Session) + + AssetManager._queue_dagruns_nonpartitioned_mysql(asset_id=1, dags_to_queue={dag}, session=session) + + stmt, values = session.execute.call_args.args + compiled = str(stmt.compile(dialect=mysql.dialect())).upper() + assert "ON DUPLICATE KEY UPDATE" in compiled + assert values == [{"target_dag_id": "dag1"}] + def test_register_asset_change_notifies_asset_listener( self, session, mock_task_instance, testing_dag_bundle, listener_manager ):