diff --git a/airflow/jobs/scheduler_job_runner.py b/airflow/jobs/scheduler_job_runner.py index 2999ed391bc94..f6821f57aa83c 100644 --- a/airflow/jobs/scheduler_job_runner.py +++ b/airflow/jobs/scheduler_job_runner.py @@ -30,7 +30,7 @@ from pathlib import Path from typing import TYPE_CHECKING, Any, Callable, Collection, Iterable, Iterator -from sqlalchemy import and_, delete, func, not_, or_, select, text, update +from sqlalchemy import and_, delete, exists, func, not_, or_, select, text, update from sqlalchemy.exc import OperationalError from sqlalchemy.orm import lazyload, load_only, make_transient, selectinload from sqlalchemy.sql import expression @@ -51,6 +51,7 @@ DagScheduleAssetReference, TaskOutletAssetReference, ) +from airflow.models.backfill import Backfill from airflow.models.dag import DAG, DagModel from airflow.models.dagbag import DagBag from airflow.models.dagrun import DagRun @@ -1063,6 +1064,11 @@ def _run_scheduler_loop(self) -> None: self.check_trigger_timeouts, ) + timers.call_regular_interval( + 30, + self._mark_backfills_complete, + ) + timers.call_regular_interval( conf.getfloat("scheduler", "pool_metrics_interval", fallback=5.0), self._emit_pool_metrics, @@ -1288,6 +1294,28 @@ def _create_dagruns_for_dags(self, guard: CommitProhibitorGuard, session: Sessio guard.commit() # END: create dagruns + @provide_session + def _mark_backfills_complete(self, session: Session = NEW_SESSION) -> None: + """Mark completed backfills as completed.""" + 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 + query = select(Backfill).where( + Backfill.completed_at.is_(None), + ~exists( + select(DagRun.id).where( + and_(DagRun.backfill_id == Backfill.id, DagRun.state.in_(unfinished_states)) + ) + ), + ) + backfills = session.scalars(query).all() + if not backfills: + return + self.log.info("marking %s backfills as complete", len(backfills)) + for b in backfills: + b.completed_at = now + @add_span 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.""" diff --git a/tests/jobs/test_scheduler_job.py b/tests/jobs/test_scheduler_job.py index a5748ebaeef9c..92fba901b2c76 100644 --- a/tests/jobs/test_scheduler_job.py +++ b/tests/jobs/test_scheduler_job.py @@ -52,7 +52,7 @@ from airflow.jobs.local_task_job_runner import LocalTaskJobRunner from airflow.jobs.scheduler_job_runner import SchedulerJobRunner from airflow.models.asset import AssetDagRunQueue, AssetEvent, AssetModel -from airflow.models.backfill import _create_backfill +from airflow.models.backfill import Backfill, _create_backfill from airflow.models.dag import DAG, DagModel from airflow.models.dagbag import DagBag from airflow.models.dagrun import DagRun @@ -6449,3 +6449,32 @@ def test_process_dags_queries_count( prefix = "Collected database query count mismatches:" joined = "\n\n".join(failures) raise AssertionError(f"{prefix}\n\n{joined}") + + +def test_mark_backfills_completed(dag_maker, session): + clear_db_backfills() + with dag_maker(serialized=True, dag_id="test_mark_backfills_completed", schedule="@daily") as dag: + BashOperator(task_id="hi", bash_command="echo hi") + b = _create_backfill( + dag_id=dag.dag_id, + from_date=pendulum.parse("2021-01-01"), + to_date=pendulum.parse("2021-01-03"), + max_active_runs=10, + reverse=False, + dag_run_conf={}, + ) + session.expunge_all() + runner = SchedulerJobRunner( + job=Job(job_type=SchedulerJobRunner.job_type, executor=MockExecutor(do_update=False)) + ) + runner._mark_backfills_complete() + b = session.get(Backfill, b.id) + assert b.completed_at is None + session.expunge_all() + drs = session.scalars(select(DagRun).where(DagRun.dag_id == dag.dag_id)) + for dr in drs: + dr.state = DagRunState.SUCCESS + session.commit() + runner._mark_backfills_complete() + b = session.get(Backfill, b.id) + assert b.completed_at.timestamp() > 0