Roll back orphaned backfill when run creation fails

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.

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,

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",

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):