diff --git a/airflow-core/newsfragments/68705.bugfix.rst b/airflow-core/newsfragments/68705.bugfix.rst new file mode 100644 index 0000000000000..5d48af708e134 --- /dev/null +++ b/airflow-core/newsfragments/68705.bugfix.rst @@ -0,0 +1 @@ +A failed backfill creation no longer leaves an orphaned backfill record that blocks all future backfills for the dag. If creating the backfill runs fails, the backfill is now rolled back so it can be retried. diff --git a/airflow-core/src/airflow/models/backfill.py b/airflow-core/src/airflow/models/backfill.py index 13b9e6d004f6a..393276036ea3c 100644 --- a/airflow-core/src/airflow/models/backfill.py +++ b/airflow-core/src/airflow/models/backfill.py @@ -36,18 +36,17 @@ Integer, String, UniqueConstraint, - delete, func, select, ) -from sqlalchemy.exc import IntegrityError, OperationalError +from sqlalchemy.exc import IntegrityError from sqlalchemy.orm import Mapped, mapped_column, relationship, validates from airflow._shared.timezones import timezone from airflow.exceptions import AirflowException, DagNotFound, DagRunTypeNotAllowed from airflow.models.base import Base, StringID from airflow.utils.session import create_session -from airflow.utils.sqlalchemy import UtcDateTime, is_lock_not_available_error, with_row_locks +from airflow.utils.sqlalchemy import UtcDateTime, with_row_locks from airflow.utils.state import DagRunState from airflow.utils.types import DagRunTriggeredByType, DagRunType @@ -699,54 +698,32 @@ def _create_backfill( triggering_user_name=triggering_user_name, ) session.add(backfill) - # Commit immediately so the backfill is visible to concurrent requests - # checking num_active backfills, preventing duplicate active backfills - # for the same dag. - session.commit() + # Flush (not commit) so backfill.id is populated while keeping the whole + # creation in one transaction: any failure below rolls back the Backfill + # row together with its runs, leaving no orphan behind. + session.flush() session.scalars(select(DagModel).where(DagModel.dag_id == dag_id)).one() first_info = dagrun_info_list[0] - try: - if first_info.partition_key: - _create_runs_partitioned( - backfill=backfill, - dag=dag, - dagrun_info_list=dagrun_info_list, - session=session, - ) - else: - _create_runs_non_partitioned( - backfill=backfill, - dag=dag, - dagrun_info_list=dagrun_info_list, - run_on_latest_version=run_on_latest_version, - session=session, - ) - except OperationalError as e: - if is_lock_not_available_error(e): - # Lock error: clean up the orphan so the user can retry. The - # helper is best-effort; if it fails the original error still - # surfaces and the route returns 503. - _cleanup_partial_backfill(backfill, session) - raise + if first_info.partition_key: + _create_runs_partitioned( + backfill=backfill, + dag=dag, + dagrun_info_list=dagrun_info_list, + session=session, + ) + else: + _create_runs_non_partitioned( + backfill=backfill, + dag=dag, + dagrun_info_list=dagrun_info_list, + run_on_latest_version=run_on_latest_version, + session=session, + ) return backfill -def _cleanup_partial_backfill(backfill: Backfill, session: Session) -> None: - """Best-effort removal of a partially-created backfill after a lock error.""" - from airflow.models.dagrun import DagRun - - try: - session.rollback() - session.execute(delete(BackfillDagRun).where(BackfillDagRun.backfill_id == backfill.id)) - session.execute(delete(DagRun).where(DagRun.backfill_id == backfill.id)) - session.delete(backfill) - session.commit() - except Exception: - session.rollback() - - def _create_runs_partitioned( *, backfill: Backfill, diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_backfills.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_backfills.py index fd9d0d5c9c390..7969812303435 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_backfills.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_backfills.py @@ -24,7 +24,7 @@ import pendulum import pytest from sqlalchemy import and_, func, select -from sqlalchemy.exc import OperationalError, ProgrammingError +from sqlalchemy.exc import OperationalError from airflow._shared.timezones import timezone from airflow.dag_processing.dagbag import DagBag @@ -35,6 +35,7 @@ NoBackfillRunsToCreate, ReprocessBehavior, _create_backfill, + _create_backfill_dag_run_non_partitioned, ) from airflow.models.dag import DAG from airflow.models.dagbundle import DagBundleModel @@ -43,7 +44,6 @@ from airflow.sdk import CronPartitionTimetable from airflow.utils.session import provide_session from airflow.utils.state import DagRunState -from airflow.utils.types import DagRunType from tests_common.test_utils.asserts import assert_queries_count from tests_common.test_utils.db import ( @@ -585,164 +585,58 @@ def test_create_backfill_database_locked(self, session, dag_maker, test_client): assert response.status_code == 503 assert "database is locked" in response.json()["detail"].lower() - def test_create_backfill_cleans_up_orphan_on_lock_error(self, session, dag_maker): - """The partial Backfill row is removed when the cleanup runs after a lock error.""" - from airflow.models.backfill import Backfill, _cleanup_partial_backfill + def test_create_backfill_lock_error_rolls_back_partial_state(self, session, dag_maker, test_client): + """A lock error partway through run creation rolls back the whole backfill. - with dag_maker(session=session, dag_id="TEST_DAG_CLEANUP", schedule="0 * * * *") as dag: - EmptyOperator(task_id="mytask") - session.commit() - - bf = Backfill( - dag_id=dag.dag_id, - from_date=pendulum.parse("2024-01-01"), - to_date=pendulum.parse("2024-02-01"), - max_active_runs=5, - dag_run_conf={}, - reprocess_behavior="none", - dag_model=dag, - triggering_user_name="test", - ) - session.add(bf) - session.commit() - bf_id = bf.id - - assert session.scalar(select(func.count()).select_from(Backfill).where(Backfill.id == bf_id)) == 1 - - _cleanup_partial_backfill(bf, session) - - assert session.scalar(select(func.count()).select_from(Backfill).where(Backfill.id == bf_id)) == 0 - - def test_create_backfill_cleans_up_after_failed_transaction(self, session, dag_maker): - """The cleanup works when the session is in a deactivated state. - - Mirrors the real flow inside ``_create_backfill`` after an - ``OperationalError``: SQLAlchemy deactivates the session until - an explicit ``rollback()``. The cleanup must call it before any - further operation; without that, ``session.execute()`` raises - ``InvalidRequestError`` and the cleanup is a silent no-op. + Creation is a single transaction, so the Backfill row, the DagRuns (and + their TaskInstances), and the BackfillDagRun rows created before the + failure must all be gone, and the route must return 503. """ - from sqlalchemy import text - - from airflow.models.backfill import Backfill, _cleanup_partial_backfill - - with dag_maker(session=session, dag_id="TEST_DAG_CLEANUP_DEACT", schedule="0 * * * *") as dag: - EmptyOperator(task_id="mytask") - session.commit() - - bf = Backfill( - dag_id=dag.dag_id, - from_date=pendulum.parse("2024-01-01"), - to_date=pendulum.parse("2024-02-01"), - max_active_runs=5, - dag_run_conf={}, - reprocess_behavior="none", - dag_model=dag, - triggering_user_name="test", - ) - session.add(bf) - session.commit() - bf_id = bf.id - - # Force the session into a deactivated state (same shape as after - # a failed flush). SQLite raises OperationalError; Postgres/MySQL raise ProgrammingError. - with pytest.raises((OperationalError, ProgrammingError)): - session.execute(text("INVALID SQL STATEMENT")) - - _cleanup_partial_backfill(bf, session) - - assert session.scalar(select(func.count()).select_from(Backfill).where(Backfill.id == bf_id)) == 0 - - def test_create_backfill_cleanup_removes_partial_dag_runs(self, session, dag_maker): - """Cleanup removes partial DagRuns, TIs, and BackfillDagRun rows alongside the Backfill.""" - from airflow.models.backfill import _cleanup_partial_backfill - - with dag_maker(session=session, dag_id="TEST_DAG_CLEANUP_DR", schedule="0 * * * *") as dag: + with dag_maker(session=session, dag_id="TEST_DAG_LOCK_ROLLBACK", schedule="0 0 * * *") as dag: EmptyOperator(task_id="mytask") + session.scalars(select(DagModel)).all() session.commit() - bf = Backfill( - dag_id=dag.dag_id, - from_date=pendulum.parse("2024-01-01"), - to_date=pendulum.parse("2024-02-01"), - max_active_runs=5, - dag_run_conf={}, - reprocess_behavior="none", - dag_model=dag, - triggering_user_name="test", - ) - session.add(bf) - session.commit() + data = { + "dag_id": dag.dag_id, + "from_date": to_iso(pendulum.parse("2024-01-01")), + "to_date": to_iso(pendulum.parse("2024-01-05")), + "max_active_runs": 5, + "run_backwards": False, + "dag_run_conf": {}, + } - dr1 = dag_maker.create_dagrun( - logical_date=pendulum.parse("2024-01-01"), - run_type=DagRunType.BACKFILL_JOB, - backfill_id=bf.id, - state=DagRunState.QUEUED, - ) - session.add( - BackfillDagRun( - backfill_id=bf.id, - dag_run_id=dr1.id, - logical_date=pendulum.parse("2024-01-01"), - sort_ordinal=1, - ) - ) + call_count = 0 - dr2 = dag_maker.create_dagrun( - logical_date=pendulum.parse("2024-01-02"), - run_type=DagRunType.BACKFILL_JOB, - backfill_id=bf.id, - state=DagRunState.QUEUED, - ) - session.add( - BackfillDagRun( - backfill_id=bf.id, - dag_run_id=dr2.id, - logical_date=pendulum.parse("2024-01-02"), - sort_ordinal=2, - ) - ) - session.commit() + def fail_on_third_run(*args, **kwargs): + nonlocal call_count + call_count += 1 + if call_count == 3: + raise OperationalError("statement", "params", sqlite3.OperationalError("database is locked")) + return _create_backfill_dag_run_non_partitioned(*args, **kwargs) - bf_id = bf.id - dr1_id = dr1.id - dr2_id = dr2.id - run_ids = [dr1.run_id, dr2.run_id] + with mock.patch( + "airflow.models.backfill._create_backfill_dag_run_non_partitioned", + side_effect=fail_on_third_run, + ): + response = test_client.post("/backfills", json=data) - assert session.scalar(select(func.count()).select_from(Backfill).where(Backfill.id == bf_id)) == 1 - assert session.scalar(select(func.count()).select_from(DagRun).where(DagRun.id == dr1_id)) == 1 - assert session.scalar(select(func.count()).select_from(DagRun).where(DagRun.id == dr2_id)) == 1 - assert ( - session.scalar( - select(func.count()).select_from(BackfillDagRun).where(BackfillDagRun.backfill_id == bf_id) - ) - == 2 - ) + assert response.status_code == 503 + assert call_count == 3 assert ( - session.scalar( - select(func.count()).select_from(TaskInstance).where(TaskInstance.run_id.in_(run_ids)) - ) - >= 2 + session.scalar(select(func.count()).select_from(Backfill).where(Backfill.dag_id == dag.dag_id)) + == 0 ) - - _cleanup_partial_backfill(bf, session) - - assert session.scalar(select(func.count()).select_from(Backfill).where(Backfill.id == bf_id)) == 0 - assert session.scalar(select(func.count()).select_from(DagRun).where(DagRun.id == dr1_id)) == 0 - assert session.scalar(select(func.count()).select_from(DagRun).where(DagRun.id == dr2_id)) == 0 assert ( - session.scalar( - select(func.count()).select_from(BackfillDagRun).where(BackfillDagRun.backfill_id == bf_id) - ) - == 0 + session.scalar(select(func.count()).select_from(DagRun).where(DagRun.dag_id == dag.dag_id)) == 0 ) assert ( session.scalar( - select(func.count()).select_from(TaskInstance).where(TaskInstance.run_id.in_(run_ids)) + select(func.count()).select_from(TaskInstance).where(TaskInstance.dag_id == dag.dag_id) ) == 0 ) + assert session.scalar(select(func.count()).select_from(BackfillDagRun)) == 0 @pytest.mark.parametrize( "run_backwards", diff --git a/airflow-core/tests/unit/models/test_backfill.py b/airflow-core/tests/unit/models/test_backfill.py index ea621be84d6c1..0f23273ca9cf5 100644 --- a/airflow-core/tests/unit/models/test_backfill.py +++ b/airflow-core/tests/unit/models/test_backfill.py @@ -108,6 +108,49 @@ def test_reverse_and_depends_on_past_fails(dep_on_past, dag_maker, session): assert b is not None +def test_create_backfill_no_orphan_on_run_creation_failure(dag_maker, session): + """A failure while creating backfill runs must not leave an orphaned Backfill row. + + The Backfill row is created first; if run creation then fails, the row must be + rolled back, otherwise the ``num_active > 0`` check blocks all future backfills + for the dag. + """ + with dag_maker(schedule="@daily") as dag: + PythonOperator(task_id="hi", python_callable=print) + session.commit() + + with mock.patch( + "airflow.models.backfill._create_backfill_dag_run_non_partitioned", + side_effect=RuntimeError("boom"), + ): + with pytest.raises(RuntimeError, match="boom"): + _create_backfill( + dag_id=dag.dag_id, + from_date=pendulum.parse("2021-01-01"), + to_date=pendulum.parse("2021-01-05"), + max_active_runs=2, + reverse=False, + triggering_user_name="pytest", + dag_run_conf={}, + ) + + assert ( + session.scalar(select(func.count()).select_from(Backfill).where(Backfill.dag_id == dag.dag_id)) == 0 + ) + + # A subsequent backfill must not be blocked by a leftover row. + b = _create_backfill( + dag_id=dag.dag_id, + from_date=pendulum.parse("2021-01-01"), + to_date=pendulum.parse("2021-01-05"), + max_active_runs=2, + reverse=False, + triggering_user_name="pytest", + dag_run_conf={}, + ) + assert b is not None + + @pytest.mark.parametrize("reverse", [True, False]) @pytest.mark.parametrize("existing", [["2021-01-02", "2021-01-03"], []]) def test_create_backfill_simple(reverse, existing, dag_maker, session):