diff --git a/airflow/datasets/manager.py b/airflow/datasets/manager.py index 306781daba210..dd6a761b32614 100644 --- a/airflow/datasets/manager.py +++ b/airflow/datasets/manager.py @@ -37,6 +37,7 @@ DatasetModel, ) from airflow.stats import Stats +from airflow.utils import timezone from airflow.utils.log.logging_mixin import LoggingMixin from airflow.utils.session import NEW_SESSION, provide_session @@ -184,7 +185,7 @@ def _slow_path_queue_dagruns( cls, dataset_id: int, dags_to_queue: set[DagModel], session: Session ) -> None: def _queue_dagrun_if_needed(dag: DagModel) -> str | None: - item = DatasetDagRunQueue(target_dag_id=dag.dag_id, dataset_id=dataset_id) + item = DatasetDagRunQueue(target_dag_id=dag.dag_id, dataset_id=dataset_id, created_at=timezone.utcnow()) # Don't error whole transaction when a single RunQueue item conflicts. # https://docs.sqlalchemy.org/en/14/orm/session_transaction.html#using-savepoint try: @@ -202,8 +203,13 @@ def _queue_dagrun_if_needed(dag: DagModel) -> str | None: def _postgres_queue_dagruns(cls, dataset_id: int, dags_to_queue: set[DagModel], session: Session) -> None: from sqlalchemy.dialects.postgresql import insert - values = [{"target_dag_id": dag.dag_id} for dag in dags_to_queue] - stmt = insert(DatasetDagRunQueue).values(dataset_id=dataset_id).on_conflict_do_nothing() + values = [{"target_dag_id": dag.dag_id, "created_at": timezone.utcnow()} for dag in dags_to_queue] + stmt = insert(DatasetDagRunQueue).values(dataset_id=dataset_id) + stmt = stmt.on_conflict_do_update( + index_elements=["dataset_id", "target_dag_id"], + set_={"created_at": stmt.excluded.created_at}, + where=(DatasetDagRunQueue.created_at < stmt.excluded.created_at), + ) session.execute(stmt, values) @classmethod diff --git a/airflow/jobs/scheduler_job_runner.py b/airflow/jobs/scheduler_job_runner.py index 2725dd71d9f3b..0c76c1e42fe02 100644 --- a/airflow/jobs/scheduler_job_runner.py +++ b/airflow/jobs/scheduler_job_runner.py @@ -1489,21 +1489,25 @@ def _create_dag_runs_dataset_triggered( events=dataset_events, ) - dag_run = dag.create_dagrun( - run_id=run_id, - run_type=DagRunType.DATASET_TRIGGERED, - execution_date=exec_date, - data_interval=data_interval, - state=DagRunState.QUEUED, - external_trigger=False, - session=session, - dag_hash=dag_hash, - creating_job_id=self.job.id, - ) - Stats.incr("dataset.triggered_dagruns") - dag_run.consumed_dataset_events.extend(dataset_events) + if dataset_events: + dag_run = dag.create_dagrun( + run_id=run_id, + run_type=DagRunType.DATASET_TRIGGERED, + execution_date=exec_date, + data_interval=data_interval, + state=DagRunState.QUEUED, + external_trigger=False, + session=session, + dag_hash=dag_hash, + creating_job_id=self.job.id, + ) + Stats.incr("dataset.triggered_dagruns") + dag_run.consumed_dataset_events.extend(dataset_events) session.execute( - delete(DatasetDagRunQueue).where(DatasetDagRunQueue.target_dag_id == dag_run.dag_id) + delete(DatasetDagRunQueue).where( + DatasetDagRunQueue.target_dag_id == dag.dag_id, + DatasetDagRunQueue.created_at <= exec_date, + ) ) def _should_update_dag_next_dagruns( diff --git a/tests/datasets/test_manager.py b/tests/datasets/test_manager.py index cf36c8a8e3a9d..1cc4696588a4e 100644 --- a/tests/datasets/test_manager.py +++ b/tests/datasets/test_manager.py @@ -228,3 +228,97 @@ def test_create_datasets_notifies_dataset_listener(self, session): # Ensure the listener was notified assert len(dataset_listener.created) == 1 assert dataset_listener.created[0].uri == dsm.uri + + @pytest.mark.skip_if_database_isolation_mode + def test_slow_path_queue_dagruns_updates_timestamp(self, session): + """ + On non-Postgres backends, _slow_path_queue_dagruns should update created_at + when a DatasetDagRunQueue record already exists for the same (dataset_id, target_dag_id). + """ + from datetime import timedelta + + from airflow.utils import timezone + + dsem = DatasetManager() + ds = DatasetModel(uri="test_slow_path_update_ts") + dag = DagModel(dag_id="test_slow_path_dag", is_active=True) + session.add_all([ds, dag]) + session.flush() + + old_ts = timezone.utcnow() - timedelta(days=1) + session.add(DatasetDagRunQueue(dataset_id=ds.id, target_dag_id=dag.dag_id, created_at=old_ts)) + session.commit() + + before_update = timezone.utcnow() + dsem._slow_path_queue_dagruns(dataset_id=ds.id, dags_to_queue={dag}, session=session) + session.commit() + + updated = ( + session.query(DatasetDagRunQueue) + .filter_by(dataset_id=ds.id, target_dag_id=dag.dag_id) + .one() + ) + assert updated.created_at >= before_update + + @pytest.mark.skip_if_database_isolation_mode + @pytest.mark.backend("postgres") + def test_postgres_queue_dagruns_updates_timestamp(self, session): + """ + On PostgreSQL, _postgres_queue_dagruns should update created_at via ON CONFLICT DO UPDATE + when a DatasetDagRunQueue record already exists for the same (dataset_id, target_dag_id). + """ + from datetime import timedelta + + from airflow.utils import timezone + + dsem = DatasetManager() + ds = DatasetModel(uri="test_postgres_update_ts") + dag = DagModel(dag_id="test_postgres_dag", is_active=True) + session.add_all([ds, dag]) + session.flush() + + old_ts = timezone.utcnow() - timedelta(days=1) + session.add(DatasetDagRunQueue(dataset_id=ds.id, target_dag_id=dag.dag_id, created_at=old_ts)) + session.commit() + + before_update = timezone.utcnow() + dsem._postgres_queue_dagruns(dataset_id=ds.id, dags_to_queue={dag}, session=session) + session.commit() + + updated = ( + session.query(DatasetDagRunQueue) + .filter_by(dataset_id=ds.id, target_dag_id=dag.dag_id) + .one() + ) + assert updated.created_at >= before_update + + @pytest.mark.skip_if_database_isolation_mode + @pytest.mark.backend("postgres") + def test_postgres_queue_dagruns_where_guard_prevents_backwards_drift(self, session): + """ + The WHERE guard on the Postgres upsert should prevent a slower transaction + from overwriting a newer created_at with an older one. + """ + from datetime import timedelta + + from airflow.utils import timezone + + dsem = DatasetManager() + ds = DatasetModel(uri="test_postgres_where_guard") + dag = DagModel(dag_id="test_pg_guard_dag", is_active=True) + session.add_all([ds, dag]) + session.flush() + + newer_ts = timezone.utcnow() + session.add(DatasetDagRunQueue(dataset_id=ds.id, target_dag_id=dag.dag_id, created_at=newer_ts)) + session.commit() + + dsem._postgres_queue_dagruns(dataset_id=ds.id, dags_to_queue={dag}, session=session) + session.commit() + + updated = ( + session.query(DatasetDagRunQueue) + .filter_by(dataset_id=ds.id, target_dag_id=dag.dag_id) + .one() + ) + assert updated.created_at >= newer_ts