diff --git a/airflow/models/dag.py b/airflow/models/dag.py index 1ec78b64c3791..31bcfec5a0e88 100644 --- a/airflow/models/dag.py +++ b/airflow/models/dag.py @@ -2925,27 +2925,30 @@ def bulk_write_to_db( session.add(orm_dag) orm_dags.append(orm_dag) - # Get the latest dag run for each existing dag as a single query (avoid n+1 query) - most_recent_subq = ( - select(DagRun.dag_id, func.max(DagRun.execution_date).label("max_execution_date")) - .where( - DagRun.dag_id.in_(existing_dags), - or_(DagRun.run_type == DagRunType.BACKFILL_JOB, DagRun.run_type == DagRunType.SCHEDULED), + most_recent_runs: dict[str, DagRun] = {} + num_active_runs: dict[str, int] = {} + # Skip these queries entirely if no DAGs can be scheduled to save time. + if any(dag.timetable.can_be_scheduled for dag in dags): + # Get the latest dag run for each existing dag as a single query (avoid n+1 query) + most_recent_subq = ( + select(DagRun.dag_id, func.max(DagRun.execution_date).label("max_execution_date")) + .where( + DagRun.dag_id.in_(existing_dags), + or_(DagRun.run_type == DagRunType.BACKFILL_JOB, DagRun.run_type == DagRunType.SCHEDULED), + ) + .group_by(DagRun.dag_id) + .subquery() ) - .group_by(DagRun.dag_id) - .subquery() - ) - most_recent_runs_iter = session.scalars( - select(DagRun).where( - DagRun.dag_id == most_recent_subq.c.dag_id, - DagRun.execution_date == most_recent_subq.c.max_execution_date, + most_recent_runs_iter = session.scalars( + select(DagRun).where( + DagRun.dag_id == most_recent_subq.c.dag_id, + DagRun.execution_date == most_recent_subq.c.max_execution_date, + ) ) - ) - most_recent_runs = {run.dag_id: run for run in most_recent_runs_iter} - - # Get number of active dagruns for all dags we are processing as a single query. + most_recent_runs = {run.dag_id: run for run in most_recent_runs_iter} - num_active_runs = DagRun.active_runs_of_dags(dag_ids=existing_dags, session=session) + # Get number of active dagruns for all dags we are processing as a single query. + num_active_runs = DagRun.active_runs_of_dags(dag_ids=existing_dags, session=session) filelocs = [] diff --git a/tests/models/test_dag.py b/tests/models/test_dag.py index bd460e6cc7cbd..2d0dde711762b 100644 --- a/tests/models/test_dag.py +++ b/tests/models/test_dag.py @@ -942,6 +942,20 @@ def test_bulk_write_to_db(self): for row in session.query(DagModel.last_parsed_time).all(): assert row[0] is not None + @pytest.mark.parametrize("interval", [None, "@daily"]) + def test_bulk_write_to_db_interval_save_runtime(self, interval): + mock_active_runs_of_dags = mock.MagicMock(side_effect=DagRun.active_runs_of_dags) + with mock.patch.object(DagRun, "active_runs_of_dags", mock_active_runs_of_dags): + dags_null_timetable = [ + DAG("dag-interval-None", schedule_interval=None), + DAG("dag-interval-test", schedule_interval=interval), + ] + DAG.bulk_write_to_db(dags_null_timetable, session=settings.Session()) + if interval: + mock_active_runs_of_dags.assert_called_once() + else: + mock_active_runs_of_dags.assert_not_called() + @pytest.mark.parametrize("state", [DagRunState.RUNNING, DagRunState.QUEUED]) def test_bulk_write_to_db_max_active_runs(self, state): """