Optimize scheduler backfill completion updates

airflow-core/src/airflow/jobs/scheduler_job_runner.py

@@ -2458,30 +2458,34 @@ def _mark_backfills_complete(self, *, session: Session = NEW_SESSION) -> None:
         self.log.debug("checking for completed backfills.")
         unfinished_states = (DagRunState.RUNNING, DagRunState.QUEUED)
         now = timezone.utcnow()
-        # todo: AIP-78 simplify this function to an update statement
         initializing_cutoff = now - timedelta(minutes=2)
-        query = select(Backfill).where(
-            Backfill.completed_at.is_(None),
-            # Guard: backfill must have at least one association,
-            # otherwise it is still being set up (see #61375).
-            # Allow cleanup of orphaned backfills older than 2 minutes
-            # that failed during initialization and never got any associations.
-            or_(
-                exists(select(BackfillDagRun.id).where(BackfillDagRun.backfill_id == Backfill.id)),
-                Backfill.created_at < initializing_cutoff,
-            ),
-            ~exists(
-                select(DagRun.id).where(
-                    and_(DagRun.backfill_id == Backfill.id, DagRun.state.in_(unfinished_states))
+        result = cast(
+            "CursorResult",
+            session.execute(
+                update(Backfill)
+                .where(
+                    Backfill.completed_at.is_(None),
+                    # Guard: backfill must have at least one association,
+                    # otherwise it is still being set up (see #61375).
+                    # Allow cleanup of orphaned backfills older than 2 minutes
+                    # that failed during initialization and never got any associations.
+                    or_(
+                        exists(select(BackfillDagRun.id).where(BackfillDagRun.backfill_id == Backfill.id)),
+                        Backfill.created_at < initializing_cutoff,
+                    ),
+                    ~exists(
+                        select(DagRun.id).where(
+                            and_(DagRun.backfill_id == Backfill.id, DagRun.state.in_(unfinished_states))
+                        )
+                    ),
                 )
+                .values(completed_at=now)
+                .execution_options(synchronize_session=False)
             ),
         )
-        backfills = list(session.scalars(query))
-        if not backfills:
-            return
-        self.log.info("marking %s backfills as complete", len(backfills))
-        for b in backfills:
-            b.completed_at = now
+        updated_count = max(result.rowcount or 0, 0)
+        if updated_count > 0:
+            self.log.info("marking %s backfills as complete", updated_count)
 
     def _create_dag_runs(self, dag_models: Collection[DagModel], session: Session) -> None:
         """Create a DAG run and update the dag_model to control if/when the next DAGRun should be created."""

airflow-core/tests/unit/jobs/test_scheduler_job.py

@@ -10862,6 +10862,54 @@ def test_mark_backfills_completed(dag_maker, session):
     assert b.completed_at.timestamp() > 0
 
 
+def test_mark_backfills_complete_uses_single_update_statement(dag_maker, session):
+    clear_db_backfills()
+    dag_id = "test_backfill_single_update"
+    with dag_maker(serialized=True, dag_id=dag_id, schedule="@daily"):
+        BashOperator(task_id="hi", bash_command="echo hi")
+    b = Backfill(
+        dag_id=dag_id,
+        from_date=pendulum.parse("2021-01-01"),
+        to_date=pendulum.parse("2021-01-03"),
+        max_active_runs=10,
+        dag_run_conf={},
+        reprocess_behavior=ReprocessBehavior.NONE,
+    )
+    session.add(b)
+    session.commit()
+    backfill_id = b.id
+    dr = DagRun(
+        dag_id=dag_id,
+        run_id="backfill__2021-01-01T00:00:00+00:00",
+        run_type=DagRunType.BACKFILL_JOB,
+        logical_date=pendulum.parse("2021-01-01"),
+        data_interval=(pendulum.parse("2021-01-01"), pendulum.parse("2021-01-02")),
+        run_after=pendulum.parse("2021-01-02"),
+        state=DagRunState.SUCCESS,
+        backfill_id=backfill_id,
+    )
+    session.add(dr)
+    session.flush()
+    session.add(
+        BackfillDagRun(
+            backfill_id=backfill_id,
+            dag_run_id=dr.id,
+            logical_date=pendulum.parse("2021-01-01"),
+            sort_ordinal=1,
+        )
+    )
+    session.commit()
+    session.expunge_all()
+    runner = SchedulerJobRunner(
+        job=Job(job_type=SchedulerJobRunner.job_type), executors=[MockExecutor(do_update=False)]
+    )
+    with assert_queries_count(1, session=session):
+        runner._mark_backfills_complete(session=session)
+    session.expire_all()
+    b = session.get(Backfill, backfill_id)
+    assert b.completed_at is not None
+
+
 def test_mark_backfills_complete_skips_initializing_backfill(dag_maker, session):
     clear_db_backfills()
     dag_id = "test_backfill_race_lifecycle"