sqla2 upgrade (test_dag_sources, test_hitl, test_assets, test_structure, test_deadline)

.pre-commit-config.yaml

@@ -425,10 +425,15 @@ repos:
             ^airflow-ctl.*\.py$|
             ^airflow-core/src/airflow/models/.*\.py$|
             ^airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_assets.py$|
+            ^airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_sources.py$|
+            ^airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_hitl.py$|
+            ^airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_assets.py$|
+            ^airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_structure.py$|
             ^airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_variables.py$|
             ^airflow-core/tests/unit/cli/commands/test_task_command.py$|
             ^airflow-core/tests/unit/dag_processing/bundles/test_dag_bundle_manager.py$|
             ^airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py$|
+            ^airflow-core/tests/unit/models/test_deadline.py$|
             ^airflow-core/tests/unit/models/test_renderedtifields.py$|
             ^airflow-core/tests/unit/models/test_timestamp.py$|
             ^airflow-core/tests/unit/utils/test_cli_util.py$|

airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_sources.py

@@ -137,12 +137,9 @@ def test_should_respond_200_version(self, test_client, accept, session, test_dag
         # force reserialization
         test_dag.doc_md = "new doc"
         force_reserialization(test_dag.dag_id, "dag-folder")
-        dagcode = (
-            session.query(DagCode)
-            .filter(DagCode.fileloc == test_dag.fileloc)
-            .order_by(DagCode.id.desc())
-            .first()
-        )
+        dagcode = session.scalars(
+            select(DagCode).where(DagCode.fileloc == test_dag.fileloc).order_by(DagCode.id.desc())
+        ).first()
         assert dagcode.dag_version.version_number == 2
         # populate the latest dagcode with a value
         dag_content2 = "new source code"

airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_hitl.py

@@ -25,7 +25,7 @@
 
 import pytest
 import time_machine
-from sqlalchemy import select
+from sqlalchemy import delete, select
 from sqlalchemy.orm import Session
 
 from airflow._shared.timezones.timezone import utc, utcnow
@@ -275,7 +275,7 @@ def expected_sample_hitl_detail_dict(sample_ti: TaskInstance) -> dict[str, Any]:
 
 @pytest.fixture(autouse=True)
 def cleanup_audit_log(session: Session) -> None:
-    session.query(Log).delete()
+    session.execute(delete(Log))
     session.commit()
 
 

airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_assets.py

@@ -20,6 +20,7 @@
 
 import pendulum
 import pytest
+from sqlalchemy import select
 
 from airflow.models.asset import AssetDagRunQueue, AssetEvent, AssetModel
 from airflow.providers.standard.operators.empty import EmptyOperator
@@ -95,7 +96,9 @@ def test_should_set_last_update_only_for_queued_and_hide_flag(self, test_client,
 
         assets = {
             a.uri: a
-            for a in session.query(AssetModel).filter(AssetModel.uri.in_(["s3://bucket/A", "s3://bucket/B"]))
+            for a in session.scalars(
+                select(AssetModel).where(AssetModel.uri.in_(["s3://bucket/A", "s3://bucket/B"]))
+            )
         }
         # Queue and add an event only for A
         session.add(AssetDagRunQueue(asset_id=assets["s3://bucket/A"].id, target_dag_id="two_assets_equal"))
@@ -145,7 +148,7 @@ def test_last_update_respects_latest_run_filter(self, test_client, dag_maker, se
         dr = dag_maker.create_dagrun()
         dag_maker.sync_dagbag_to_db()
 
-        asset = session.query(AssetModel).filter(AssetModel.uri == "s3://bucket/F").one()
+        asset = session.scalars(select(AssetModel).where(AssetModel.uri == "s3://bucket/F")).one()
         session.add(AssetDagRunQueue(asset_id=asset.id, target_dag_id="filter_run"))
         # event before latest_run should be ignored
         ts_base = dr.logical_date or pendulum.now()

airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_structure.py

@@ -544,7 +544,7 @@ def test_should_return_200_with_resolved_asset_alias_attached_to_the_corrrect_pr
         self, test_client, session
     ):
         resolved_asset = session.scalar(
-            session.query(AssetModel).filter_by(name="resolved_example_asset_alias")
+            select(AssetModel).where(AssetModel.name == "resolved_example_asset_alias")
         )
         params = {
             "dag_id": DAG_ID_RESOLVED_ASSET_ALIAS,

airflow-core/tests/unit/models/test_deadline.py

@@ -23,6 +23,7 @@
 
 import pytest
 import time_machine
+from sqlalchemy import select
 from sqlalchemy.exc import SQLAlchemyError
 
 from airflow.api_fastapi.core_api.datamodels.dag_run import DAGRunResponse
@@ -97,9 +98,9 @@ def dagrun(session, dag_maker):
         dag_maker.create_dagrun(state=DagRunState.QUEUED, logical_date=DEFAULT_DATE)
 
         session.commit()
-        assert session.query(DagRun).count() == 1
-
-        return session.query(DagRun).one()
+        dag_runs = session.scalars(select(DagRun)).all()
+        assert len(dag_runs) == 1
+        return dag_runs[0]
 
 
 @pytest.fixture