diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index c727a93182cf0..5e07bf3cf704e 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -3152,8 +3152,7 @@ def _schedule_dag_run( default=None, ) for task_instance in unfinished_task_instances: - task_instance.state = TaskInstanceState.SKIPPED - session.merge(task_instance) + task_instance.set_state(TaskInstanceState.SKIPPED, session=session) session.flush() self.log.info("Run %s of %s has timed-out", dag_run.run_id, dag_run.dag_id) diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index f0c92158e12f8..08d9588b8bca3 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -4549,6 +4549,70 @@ def test_dagrun_timeout_fails_run(self, dag_maker): session.rollback() session.close() + def test_dagrun_timeout_skips_running_task_with_end_date_and_duration(self, dag_maker): + """A task still running when its Dag run times out must be skipped with an end_date and duration.""" + session = settings.Session() + with dag_maker( + dag_id="test_scheduler_dagrun_timeout_running_task", + dagrun_timeout=datetime.timedelta(seconds=60), + session=session, + ): + EmptyOperator(task_id="dummy") + + now = timezone.utcnow().replace(microsecond=0) + dr = dag_maker.create_dagrun(start_date=now - datetime.timedelta(days=1)) + ti = dr.get_task_instance("dummy", session=session) + ti.state = TaskInstanceState.RUNNING + ti.start_date = now - datetime.timedelta(minutes=10) + session.flush() + + scheduler_job = Job() + self.job_runner = SchedulerJobRunner(job=scheduler_job) + + with time_machine.travel(now, tick=False): + self.job_runner._schedule_dag_run(dr, session) + session.flush() + + session.refresh(ti) + assert ti.state == TaskInstanceState.SKIPPED + assert ti.end_date == now + assert ti.duration == 600.0 + + session.rollback() + session.close() + + def test_dagrun_timeout_skips_unstarted_task_with_start_and_end_date(self, dag_maker): + """A task that never started gets start_date == end_date and a zero duration, like other skips.""" + session = settings.Session() + with dag_maker( + dag_id="test_scheduler_dagrun_timeout_unstarted_task", + dagrun_timeout=datetime.timedelta(seconds=60), + session=session, + ): + EmptyOperator(task_id="dummy") + + now = timezone.utcnow().replace(microsecond=0) + dr = dag_maker.create_dagrun(start_date=now - datetime.timedelta(days=1)) + ti = dr.get_task_instance("dummy", session=session) + ti.state = None + session.flush() + + scheduler_job = Job() + self.job_runner = SchedulerJobRunner(job=scheduler_job) + + with time_machine.travel(now, tick=False): + self.job_runner._schedule_dag_run(dr, session) + session.flush() + + session.refresh(ti) + assert ti.state == TaskInstanceState.SKIPPED + assert ti.start_date == now + assert ti.end_date == now + assert ti.duration == 0.0 + + session.rollback() + session.close() + @mock.patch("airflow._shared.observability.metrics.stats._get_backend") def test_dagrun_timeout_duration_metric_has_run_type(self, mock_get_backend, dag_maker): """