Rename DagRun scheduling state update method

Author: fat-catTW

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

@@ -1740,13 +1740,13 @@ def _update_dag_run_state_for_paused_dags(self, *, session: Session = NEW_SESSIO
                     .group_by(DagRun)
                 )
             )
-            # Team name should be added before listeners are called in update_state()
+            # Team name should be added before listeners are called in schedule_dag_run()
             self._stamp_team_names(paused_runs, session)
             for dag_run in paused_runs:
                 dag = self.scheduler_dag_bag.get_dag_for_run(dag_run=dag_run, session=session)
                 if dag is not None:
                     dag_run.dag = dag
-                    _, callback_to_run = dag_run.update_state(execute_callbacks=False, session=session)
+                    _, callback_to_run = dag_run.schedule_dag_run(execute_callbacks=False, session=session)
                     if callback_to_run:
                         self._send_dag_callbacks_to_processor(dag, callback_to_run)
         except Exception as e:  # should not fail the scheduler
@@ -3000,8 +3000,7 @@ def _schedule_dag_run(
 
         dag_run.scheduled_by_job_id = self.job.id
 
-        # TODO[HA]: Rename update_state -> schedule_dag_run, ?? something else?
-        schedulable_tis, callback_to_run = dag_run.update_state(session=session, execute_callbacks=False)
+        schedulable_tis, callback_to_run = dag_run.schedule_dag_run(session=session, execute_callbacks=False)
 
         if dag_run.state in State.finished_dr_states and dag_run.run_type in (
             DagRunType.SCHEDULED,

airflow-core/src/airflow/models/dagrun.py

@@ -1248,7 +1248,7 @@ def _emit_dagrun_span(self, state: DagRunState):
             span.end()
 
     @provide_session
-    def update_state(
+    def schedule_dag_run(
         self, *, session: Session = NEW_SESSION, execute_callbacks: bool = True
     ) -> tuple[list[TI], DagCallbackRequest | None]:
         """
@@ -1446,6 +1446,12 @@ def recalculate(self) -> _UnfinishedStates:
 
         return schedulable_tis, callback
 
+    @provide_session
+    def update_state(
+        self, *, session: Session = NEW_SESSION, execute_callbacks: bool = True
+    ) -> tuple[list[TI], DagCallbackRequest | None]:
+        return self.schedule_dag_run(session=session, execute_callbacks=execute_callbacks)
+
     @provide_session
     def task_instance_scheduling_decisions(self, *, session: Session = NEW_SESSION) -> TISchedulingDecision:
         tis = self.get_task_instances(session=session, state=State.task_states)
@@ -1562,7 +1568,7 @@ def produce_dag_callback(
     def execute_dag_callbacks(
         self, dag: SDKDAG, success: bool = True, relevant_ti: TI | None = None, reason: str = "success"
     ):
-        """Only needed for `dag.test` where `execute_callbacks=True` is passed to `update_state`."""
+        """Only needed for `dag.test` where `execute_callbacks=True` is passed to `schedule_dag_run`."""
         from airflow.api_fastapi.execution_api.datamodels.taskinstance import (
             TaskInstance as TIDataModel,
             TIRunContext,
@@ -1762,7 +1768,7 @@ def _emit_true_scheduling_delay_stats_for_finished_state(self, finished_tis: lis
         The true scheduling delay stats is defined as the time when the first
         task in DAG starts minus the expected DAG run datetime.
 
-        This helper method is used in ``update_state`` when the state of the
+        This helper method is used in ``schedule_dag_run`` when the state of the
         DAG run is updated to a completed status (either success or failure).
         It finds the first started task within the DAG, calculates the run's
         expected start time based on the logical date and timetable, and gets

airflow-core/tests/unit/models/test_dag.py

@@ -944,7 +944,7 @@ def add_failed_dag_run(dag, id, logical_date):
             )
             ti_op1 = dr.get_task_instance(task_id=op1.task_id, session=session)
             ti_op1.set_state(state=TaskInstanceState.FAILED, session=session)
-            dr.update_state(session=session)
+            dr.schedule_dag_run(session=session)
 
         dag_id = "dag_paused_after_limit"
         dag = DAG(dag_id, schedule=None, is_paused_upon_creation=False, max_consecutive_failed_dag_runs=2)
@@ -969,9 +969,9 @@ def add_failed_dag_run(dag, id, logical_date):
 
     @staticmethod
     def _add_dag_run(scheduler_dag, op1, session, run_id, logical_date, run_after, ti_state, run_state):
-        """Create a dagrun, set the task-instance state, and call update_state.
+        """Create a dagrun, set the task-instance state, and call schedule_dag_run.
 
-        update_state triggers _check_last_n_dagruns_failed only when the run
+        schedule_dag_run triggers _check_last_n_dagruns_failed only when the run
         transitions to FAILED, so it is a no-op for SUCCESS runs.
         """
         dr = scheduler_dag.create_dagrun(
@@ -986,7 +986,7 @@ def _add_dag_run(scheduler_dag, op1, session, run_id, logical_date, run_after, t
         )
         ti = dr.get_task_instance(task_id=op1.task_id, session=session)
         ti.set_state(state=ti_state, session=session)
-        dr.update_state(session=session)
+        dr.schedule_dag_run(session=session)
         return dr
 
     def test_dag_paused_after_limit_orders_by_run_after(self, testing_dag_bundle):
@@ -1014,7 +1014,7 @@ def test_dag_paused_after_limit_orders_by_run_after(self, testing_dag_bundle):
         assert not session.get(DagModel, dag.dag_id).is_paused
 
         # Run 1: oldest by run_after but LATEST logical_date — SUCCESS.
-        # update_state is a no-op here because _check_last_n_dagruns_failed
+        # schedule_dag_run is a no-op here because _check_last_n_dagruns_failed
         # is only invoked on the FAILED branch.
         self._add_dag_run(
             scheduler_dag,