Repository: apache/airflow
File: airflow-core/tests/unit/models/test_dag.py (Python)

Code before:

    def test_dags_needing_dagruns_asset_aliases(self, dag_maker, session):
        # link asset_alias hello_alias to asset hello
        asset_model = AssetModel(uri="hello")
        asset_alias_model = AssetAliasModel(name="hello_alias")
        asset_alias_model.assets.append(asset_model)
        session.add_all([asset_model, asset_alias_model])
        session.commit()

        with dag_maker(
            session=session,
            dag_id="my_dag",
            max_active_runs=1,
            schedule=[AssetAlias(name="hello_alias")],
            start_date=pendulum.now().add(days=-2),
        ):
            EmptyOperator(task_id="dummy")

        # there's no queue record yet, so no runs needed at this time.
        query, _ = DagModel.dags_needing_dagruns(session)
        dag_models = query.all()
        assert dag_models == []

        # add queue records so we'll need a run
        dag_model = dag_maker.dag_model
        session.add(AssetDagRunQueue(asset_id=asset_model.id, target_dag_id=dag_model.dag_id))
        session.flush()
        query, _ = DagModel.dags_needing_dagruns(session)
        dag_models = query.all()
        assert dag_models == [dag_model]

        # create run so we don't need a run anymore (due to max active runs)
        dag_maker.create_dagrun(
            run_type=DagRunType.ASSET_TRIGGERED,
            state=DagRunState.QUEUED,
            logical_date=pendulum.now("UTC"),
        )
        query, _ = DagModel.dags_needing_dagruns(session)
        dag_models = query.all()
        assert dag_models == []

        # increase max active runs and we should now need another run
        dag_maker.dag_model.max_active_runs = 2
        session.flush()
        query, _ = DagModel.dags_needing_dagruns(session)
        dag_models = query.all()
        assert dag_models == [dag_model]

    @pytest.mark.parametrize("ref", [Asset.ref(name="1"), Asset.ref(uri="s3://bucket/assets/1")])
    @pytest.mark.want_activate_assets
    @pytest.mark.need_serialized_dag

Review comment (suggestion):
The expected query count of 6 doesn't match the PR description which states 'After: 4 queries'. Either the assertion should be updated to 4, or the PR description should be updated to explain why 6 queries are expected.
```suggestion
        with assert_queries_count(4):
```

Code after:

    def test_dags_needing_dagruns_query_count(self, dag_maker, session):
        """Test that dags_needing_dagruns avoids N+1 on adrq.asset access."""
        num_assets = 10
        assets = [Asset(uri=f"test://asset{i}", group="test-group") for i in range(num_assets)]

        with dag_maker(
            session=session,
            dag_id="my_dag",
            max_active_runs=10,
            schedule=assets,
            start_date=pendulum.now().add(days=-2),
        ):
            EmptyOperator(task_id="dummy")

        dag_model = dag_maker.dag_model
        asset_models = dag_model.schedule_assets
        assert len(asset_models) == num_assets
        for asset_model in asset_models:
            session.add(AssetDagRunQueue(asset_id=asset_model.id, target_dag_id=dag_model.dag_id))
        session.flush()

        # Clear identity map so N+1 on adrq.asset is exposed
        session.expire_all()

        with assert_queries_count(6):
            query, _ = DagModel.dags_needing_dagruns(session)
            query.all()

    def test_dags_needing_dagruns_asset_aliases(self, dag_maker, session):
        # link asset_alias hello_alias to asset hello
        asset_model = AssetModel(uri="hello")
        asset_alias_model = AssetAliasModel(name="hello_alias")
        asset_alias_model.assets.append(asset_model)
        session.add_all([asset_model, asset_alias_model])
        session.commit()

        with dag_maker(
            session=session,
            dag_id="my_dag",
            max_active_runs=1,
            schedule=[AssetAlias(name="hello_alias")],
            start_date=pendulum.now().add(days=-2),
        ):
            EmptyOperator(task_id="dummy")

        # there's no queue record yet, so no runs needed at this time.
        query, _ = DagModel.dags_needing_dagruns(session)
        dag_models = query.all()
        assert dag_models == []