Flush dataset events before queuing dagruns (#26276)

When we go to schedule dagruns from the dataset dagrun queue, we assume
the events will happen before the queue records, which isn't the case
unless we explicitly flush them first. This ensures that dagruns are
properly related to their upstream dataset events.

airflow/datasets/manager.py

@@ -61,6 +61,7 @@ def register_dataset_change(
                 extra=extra,
             )
         )
+        session.flush()
         if dataset_model.consuming_dags:
             self._queue_dagruns(dataset_model, session)
         session.flush()

airflow/jobs/scheduler_job.py

@@ -1135,6 +1135,7 @@ def _create_dag_runs_dataset_triggered(
                 ]
                 if previous_dag_run:
                     dataset_event_filters.append(DatasetEvent.timestamp > previous_dag_run.execution_date)
+
                 dataset_events = (
                     session.query(DatasetEvent)
                     .join(

tests/models/test_taskinstance.py

@@ -1732,6 +1732,12 @@ def test_outlet_datasets(self, create_task_instance):
             DatasetEvent.source_task_instance == ti
         ).one() == ('s3://dag1/output_1.txt',)
 
+        # check that the dataset event has an earlier timestamp than the DDRQ's
+        ddrq_timestamps = (
+            session.query(DatasetDagRunQueue.created_at).filter_by(dataset_id=event.dataset.id).all()
+        )
+        assert all([event.timestamp < ddrq_timestamp for (ddrq_timestamp,) in ddrq_timestamps])
+
     def test_outlet_datasets_failed(self, create_task_instance):
         """
         Verify that when we have an outlet dataset on a task, and the task