DatasetDagRunQueue(dataset_id=ds1_id, target_dag_id=dag.dag_id, created_at=DEFAULT_DATE),
                DatasetDagRunQueue(
                    dataset_id=ds2_id, target_dag_id=dag.dag_id, created_at=DEFAULT_DATE + timedelta(hours=1)
                ),
            ]
        )
        session.flush()
        query, dataset_triggered_dag_info = DagModel.dags_needing_dagruns(session)
        assert 1 == len(dataset_triggered_dag_info)
        assert dag.dag_id in dataset_triggered_dag_info
        first_queued_time, last_queued_time = dataset_triggered_dag_info[dag.dag_id]
        assert first_queued_time == DEFAULT_DATE
        assert last_queued_time == DEFAULT_DATE + timedelta(hours=1)
class TestQueries:
    def setup_method(self) -> None:
        clear_db_runs()
    def teardown_method(self) -> None:
        clear_db_runs()
    @pytest.mark.parametrize("tasks_count", [3, 12])
    def test_count_number_queries(self, tasks_count):
        dag = DAG("test_dagrun_query_count", start_date=DEFAULT_DATE)
        for i in range(tasks_count):
            EmptyOperator(task_id=f"dummy_task_{i}", owner="test", dag=dag)
        with assert_queries_count(2):
            dag.create_dagrun(
                run_id="test_dagrun_query_count",
                state=State.RUNNING,
                execution_date=TEST_DATE,
            )
class TestDagDecorator: