Refactor models tests to use SQLA2 (#59837)

Migrate deprecated SQLAlchemy Query object usage to SQLAlchemy 2.0 style
in the following test files:
- test_dag.py
- test_dagcode.py
- test_mappedoperator.py
- test_taskinstance.py
- test_variable.py

Closes #59402

Author: rickychen-infinirc

.pre-commit-config.yaml

@@ -433,6 +433,11 @@ repos:
             ^airflow-core/tests/unit/models/test_cleartasks.py$|
             ^airflow-core/tests/unit/models/test_xcom.py$|
             ^airflow-core/tests/unit/models/test_dagrun.py$|
+            ^airflow-core/tests/unit/models/test_dag\.py$|
+            ^airflow-core/tests/unit/models/test_dagcode\.py$|
+            ^airflow-core/tests/unit/models/test_mappedoperator\.py$|
+            ^airflow-core/tests/unit/models/test_taskinstance\.py$|
+            ^airflow-core/tests/unit/models/test_variable\.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/models/test_dag.py

@@ -33,7 +33,7 @@
 import pendulum
 import pytest
 import time_machine
-from sqlalchemy import inspect, select
+from sqlalchemy import delete, inspect, select, update
 
 from airflow import settings
 from airflow._shared.module_loading import qualname
@@ -386,7 +386,8 @@ def test_dagtag_repr(self, testing_dag_bundle):
         sync_dag_to_db(dag)
         with create_session() as session:
             assert {"tag-1", "tag-2"} == {
-                repr(t) for t in session.query(DagTag).filter(DagTag.dag_id == "dag-test-dagtag").all()
+                repr(t)
+                for t in session.scalars(select(DagTag).where(DagTag.dag_id == "dag-test-dagtag")).all()
             }
 
     def test_bulk_write_to_db(self, testing_dag_bundle):
@@ -402,16 +403,16 @@ def test_bulk_write_to_db(self, testing_dag_bundle):
             SerializedDAG.bulk_write_to_db("testing", None, dags)
         with create_session() as session:
             assert {"dag-bulk-sync-0", "dag-bulk-sync-1", "dag-bulk-sync-2", "dag-bulk-sync-3"} == {
-                row[0] for row in session.query(DagModel.dag_id).all()
+                row[0] for row in session.execute(select(DagModel.dag_id)).all()
             }
             assert {
                 ("dag-bulk-sync-0", "test-dag"),
                 ("dag-bulk-sync-1", "test-dag"),
                 ("dag-bulk-sync-2", "test-dag"),
                 ("dag-bulk-sync-3", "test-dag"),
-            } == set(session.query(DagTag.dag_id, DagTag.name).all())
+            } == set(session.execute(select(DagTag.dag_id, DagTag.name)).all())
 
-            for row in session.query(DagModel.last_parsed_time).all():
+            for row in session.execute(select(DagModel.last_parsed_time)).all():
                 assert row[0] is not None
 
         # Re-sync should do fewer queries
@@ -426,7 +427,7 @@ def test_bulk_write_to_db(self, testing_dag_bundle):
             SerializedDAG.bulk_write_to_db("testing", None, dags)
         with create_session() as session:
             assert {"dag-bulk-sync-0", "dag-bulk-sync-1", "dag-bulk-sync-2", "dag-bulk-sync-3"} == {
-                row[0] for row in session.query(DagModel.dag_id).all()
+                row[0] for row in session.execute(select(DagModel.dag_id)).all()
             }
             assert {
                 ("dag-bulk-sync-0", "test-dag"),
@@ -437,24 +438,24 @@ def test_bulk_write_to_db(self, testing_dag_bundle):
                 ("dag-bulk-sync-2", "test-dag2"),
                 ("dag-bulk-sync-3", "test-dag"),
                 ("dag-bulk-sync-3", "test-dag2"),
-            } == set(session.query(DagTag.dag_id, DagTag.name).all())
+            } == set(session.execute(select(DagTag.dag_id, DagTag.name)).all())
         # Removing tags
         for dag in dags:
             dag.tags.remove("test-dag")
         with assert_queries_count(10):
             SerializedDAG.bulk_write_to_db("testing", None, dags)
         with create_session() as session:
             assert {"dag-bulk-sync-0", "dag-bulk-sync-1", "dag-bulk-sync-2", "dag-bulk-sync-3"} == {
-                row[0] for row in session.query(DagModel.dag_id).all()
+                row[0] for row in session.execute(select(DagModel.dag_id)).all()
             }
             assert {
                 ("dag-bulk-sync-0", "test-dag2"),
                 ("dag-bulk-sync-1", "test-dag2"),
                 ("dag-bulk-sync-2", "test-dag2"),
                 ("dag-bulk-sync-3", "test-dag2"),
-            } == set(session.query(DagTag.dag_id, DagTag.name).all())
+            } == set(session.execute(select(DagTag.dag_id, DagTag.name)).all())
 
-            for row in session.query(DagModel.last_parsed_time).all():
+            for row in session.execute(select(DagModel.last_parsed_time)).all():
                 assert row[0] is not None
 
         # Removing all tags
@@ -464,11 +465,11 @@ def test_bulk_write_to_db(self, testing_dag_bundle):
             SerializedDAG.bulk_write_to_db("testing", None, dags)
         with create_session() as session:
             assert {"dag-bulk-sync-0", "dag-bulk-sync-1", "dag-bulk-sync-2", "dag-bulk-sync-3"} == {
-                row[0] for row in session.query(DagModel.dag_id).all()
+                row[0] for row in session.execute(select(DagModel.dag_id)).all()
             }
-            assert not set(session.query(DagTag.dag_id, DagTag.name).all())
+            assert not set(session.execute(select(DagTag.dag_id, DagTag.name)).all())
 
-            for row in session.query(DagModel.last_parsed_time).all():
+            for row in session.execute(select(DagModel.last_parsed_time)).all():
                 assert row[0] is not None
 
     def test_bulk_write_to_db_single_dag(self, testing_dag_bundle):
@@ -486,12 +487,12 @@ def test_bulk_write_to_db_single_dag(self, testing_dag_bundle):
         with assert_queries_count(6):
             SerializedDAG.bulk_write_to_db("testing", None, dags)
         with create_session() as session:
-            assert {"dag-bulk-sync-0"} == {row[0] for row in session.query(DagModel.dag_id).all()}
+            assert {"dag-bulk-sync-0"} == {row[0] for row in session.execute(select(DagModel.dag_id)).all()}
             assert {
                 ("dag-bulk-sync-0", "test-dag"),
-            } == set(session.query(DagTag.dag_id, DagTag.name).all())
+            } == set(session.execute(select(DagTag.dag_id, DagTag.name)).all())
 
-            for row in session.query(DagModel.last_parsed_time).all():
+            for row in session.execute(select(DagModel.last_parsed_time)).all():
                 assert row[0] is not None
 
         # Re-sync should do fewer queries
@@ -516,16 +517,16 @@ def test_bulk_write_to_db_multiple_dags(self, testing_dag_bundle):
             SerializedDAG.bulk_write_to_db("testing", None, dags)
         with create_session() as session:
             assert {"dag-bulk-sync-0", "dag-bulk-sync-1", "dag-bulk-sync-2", "dag-bulk-sync-3"} == {
-                row[0] for row in session.query(DagModel.dag_id).all()
+                row[0] for row in session.execute(select(DagModel.dag_id)).all()
             }
             assert {
                 ("dag-bulk-sync-0", "test-dag"),
                 ("dag-bulk-sync-1", "test-dag"),
                 ("dag-bulk-sync-2", "test-dag"),
                 ("dag-bulk-sync-3", "test-dag"),
-            } == set(session.query(DagTag.dag_id, DagTag.name).all())
+            } == set(session.execute(select(DagTag.dag_id, DagTag.name)).all())
 
-            for row in session.query(DagModel.last_parsed_time).all():
+            for row in session.execute(select(DagModel.last_parsed_time)).all():
                 assert row[0] is not None
 
         # Re-sync should do fewer queries
@@ -682,21 +683,21 @@ def test_bulk_write_to_db_assets(self, testing_dag_bundle):
         create_scheduler_dag(dag1).clear()
         SerializedDAG.bulk_write_to_db("testing", None, [dag1, dag2], session=session)
         session.commit()
-        stored_assets = {x.uri: x for x in session.query(AssetModel).all()}
+        stored_assets = {x.uri: x for x in session.scalars(select(AssetModel)).all()}
         asset1_orm = stored_assets[a1.uri]
         asset2_orm = stored_assets[a2.uri]
         asset3_orm = stored_assets[a3.uri]
         assert stored_assets[uri1].extra == {"should": "be used"}
         assert [x.dag_id for x in asset1_orm.scheduled_dags] == [dag_id1]
         assert [(x.task_id, x.dag_id) for x in asset1_orm.producing_tasks] == [(task_id, dag_id2)]
         assert set(
-            session.query(
-                TaskOutletAssetReference.task_id,
-                TaskOutletAssetReference.dag_id,
-                TaskOutletAssetReference.asset_id,
-            )
-            .filter(TaskOutletAssetReference.dag_id.in_((dag_id1, dag_id2)))
-            .all()
+            session.execute(
+                select(
+                    TaskOutletAssetReference.task_id,
+                    TaskOutletAssetReference.dag_id,
+                    TaskOutletAssetReference.asset_id,
+                ).where(TaskOutletAssetReference.dag_id.in_((dag_id1, dag_id2)))
+            ).all()
         ) == {
             (task_id, dag_id1, asset2_orm.id),
             (task_id, dag_id1, asset3_orm.id),
@@ -715,18 +716,18 @@ def test_bulk_write_to_db_assets(self, testing_dag_bundle):
         SerializedDAG.bulk_write_to_db("testing", None, [dag1, dag2], session=session)
         session.commit()
         session.expunge_all()
-        stored_assets = {x.uri: x for x in session.query(AssetModel).all()}
+        stored_assets = {x.uri: x for x in session.scalars(select(AssetModel)).all()}
         asset1_orm = stored_assets[a1.uri]
         asset2_orm = stored_assets[a2.uri]
         assert [x.dag_id for x in asset1_orm.scheduled_dags] == []
         assert set(
-            session.query(
-                TaskOutletAssetReference.task_id,
-                TaskOutletAssetReference.dag_id,
-                TaskOutletAssetReference.asset_id,
-            )
-            .filter(TaskOutletAssetReference.dag_id.in_((dag_id1, dag_id2)))
-            .all()
+            session.execute(
+                select(
+                    TaskOutletAssetReference.task_id,
+                    TaskOutletAssetReference.dag_id,
+                    TaskOutletAssetReference.asset_id,
+                ).where(TaskOutletAssetReference.dag_id.in_((dag_id1, dag_id2)))
+            ).all()
         ) == {(task_id, dag_id1, asset2_orm.id)}
 
     def test_bulk_write_to_db_asset_aliases(self, testing_dag_bundle):
@@ -749,7 +750,7 @@ def test_bulk_write_to_db_asset_aliases(self, testing_dag_bundle):
         SerializedDAG.bulk_write_to_db("testing", None, [dag1, dag2], session=session)
         session.commit()
 
-        stored_asset_alias_models = {x.name: x for x in session.query(AssetAliasModel).all()}
+        stored_asset_alias_models = {x.name: x for x in session.scalars(select(AssetAliasModel)).all()}
         asset_alias_1_orm = stored_asset_alias_models[asset_alias_1.name]
         asset_alias_2_orm = stored_asset_alias_models[asset_alias_2.name]
         asset_alias_3_orm = stored_asset_alias_models[asset_alias_3.name]
@@ -818,7 +819,7 @@ def test_dag_is_deactivated_upon_dagfile_deletion(self, dag_maker):
         session = settings.Session()
         sync_dag_to_db(dag, session=session)
 
-        orm_dag = session.query(DagModel).filter(DagModel.dag_id == dag_id).one()
+        orm_dag = session.scalar(select(DagModel).where(DagModel.dag_id == dag_id))
 
         assert not orm_dag.is_stale
 
@@ -827,10 +828,10 @@ def test_dag_is_deactivated_upon_dagfile_deletion(self, dag_maker):
             rel_filelocs=list_py_file_paths(settings.DAGS_FOLDER),
         )
 
-        orm_dag = session.query(DagModel).filter(DagModel.dag_id == dag_id).one()
+        orm_dag = session.scalar(select(DagModel).where(DagModel.dag_id == dag_id))
         assert orm_dag.is_stale
 
-        session.execute(DagModel.__table__.delete().where(DagModel.dag_id == dag_id))
+        session.execute(delete(DagModel).where(DagModel.dag_id == dag_id))
         session.close()
 
     def test_dag_naive_default_args_start_date_with_timezone(self):
@@ -1107,7 +1108,7 @@ def test_get_paused_dag_ids(self, testing_dag_bundle):
         assert paused_dag_ids == {dag_id}
 
         with create_session() as session:
-            session.query(DagModel).filter(DagModel.dag_id == dag_id).delete(synchronize_session=False)
+            session.execute(delete(DagModel).where(DagModel.dag_id == dag_id))
 
     @pytest.mark.parametrize(
         ("schedule_arg", "expected_timetable", "interval_description"),
@@ -1293,7 +1294,7 @@ def consumer(value):
         assert upstream_ti.state is None  # cleared
         assert ti.state is None  # cleared
         assert ti2.state == State.SUCCESS  # not cleared
-        dagruns = session.query(DagRun).filter(DagRun.dag_id == dag_id).all()
+        dagruns = session.scalars(select(DagRun).where(DagRun.dag_id == dag_id)).all()
 
         assert len(dagruns) == 1
         dagrun: DagRun = dagruns[0]
@@ -1465,7 +1466,7 @@ def test_clear_dag(
             session=session,
         )
 
-        task_instances = session.query(TI).filter(TI.dag_id == dag_id).all()
+        task_instances = session.scalars(select(TI).where(TI.dag_id == dag_id)).all()
 
         assert len(task_instances) == 1
         task_instance: TI = task_instances[0]
@@ -1813,7 +1814,7 @@ def test_dag_owner_links(self, testing_dag_bundle):
         dag = DAG("dag", schedule=None, start_date=DEFAULT_DATE)
         sync_dag_to_db(dag, session=session)
 
-        orm_dag_owners = session.query(DagOwnerAttributes).all()
+        orm_dag_owners = session.scalars(select(DagOwnerAttributes)).all()
         assert not orm_dag_owners
 
     @pytest.mark.need_serialized_dag
@@ -1963,7 +1964,7 @@ def test_dags_needing_dagruns_assets(self, dag_maker, session):
         assert dag_models == []
 
         # add queue records so we'll need a run
-        dag_model = session.query(DagModel).filter(DagModel.dag_id == dag.dag_id).one()
+        dag_model = session.scalar(select(DagModel).where(DagModel.dag_id == dag.dag_id))
         asset_model: AssetModel = dag_model.schedule_assets[0]
         session.add(AssetDagRunQueue(asset_id=asset_model.id, target_dag_id=dag_model.dag_id))
         session.flush()
@@ -2250,7 +2251,7 @@ def test_dags_needing_dagruns_triggered_date_by_dag_queued_times(self, session,
                 EmptyOperator(task_id="task", outlets=[asset])
             dr = dag_maker.create_dagrun()
 
-            asset_id = session.query(AssetModel.id).filter_by(uri=asset.uri).scalar()
+            asset_id = session.scalar(select(AssetModel.id).where(AssetModel.uri == asset.uri))
 
             session.add(
                 AssetEvent(
@@ -2262,8 +2263,8 @@ def test_dags_needing_dagruns_triggered_date_by_dag_queued_times(self, session,
                 )
             )
 
-        asset1_id = session.query(AssetModel.id).filter_by(uri=asset1.uri).scalar()
-        asset2_id = session.query(AssetModel.id).filter_by(uri=asset2.uri).scalar()
+        asset1_id = session.scalar(select(AssetModel.id).where(AssetModel.uri == asset1.uri))
+        asset2_id = session.scalar(select(AssetModel.id).where(AssetModel.uri == asset2.uri))
 
         with dag_maker(dag_id="assets-consumer-multiple", schedule=[asset1, asset2]) as dag:
             pass
@@ -2314,7 +2315,9 @@ def test_asset_expression(self, session: Session, testing_dag_bundle) -> None:
         )
         SerializedDAG.bulk_write_to_db("testing", None, [dag], session=session)
 
-        expression = session.scalars(select(DagModel.asset_expression).filter_by(dag_id=dag.dag_id)).one()
+        expression = session.scalars(
+            select(DagModel.asset_expression).where(DagModel.dag_id == dag.dag_id)
+        ).one()
         assert expression == {
             "any": [
                 {
@@ -2506,14 +2509,12 @@ def test_set_task_instance_state(run_id, session, dag_maker):
     )
 
     def get_ti_from_db(task):
-        return (
-            session.query(TI)
-            .filter(
+        return session.scalar(
+            select(TI).where(
                 TI.dag_id == dag.dag_id,
                 TI.task_id == task.task_id,
                 TI.run_id == dagrun.run_id,
             )
-            .one()
         )
 
     get_ti_from_db(task_1).state = State.FAILED
@@ -2588,16 +2589,16 @@ def consumer(value):
     )
     expand_mapped_task(mapped, dr2.run_id, "make_arg_lists", length=2, session=session)
 
-    session.query(TI).filter_by(dag_id=dag.dag_id).update({"state": TaskInstanceState.FAILED})
+    session.execute(update(TI).where(TI.dag_id == dag.dag_id).values(state=TaskInstanceState.FAILED))
 
     ti_query = (
-        session.query(TI.task_id, TI.map_index, TI.run_id, TI.state)
-        .filter(TI.dag_id == dag.dag_id, TI.task_id.in_([task_id, "downstream"]))
+        select(TI.task_id, TI.map_index, TI.run_id, TI.state)
+        .where(TI.dag_id == dag.dag_id, TI.task_id.in_([task_id, "downstream"]))
         .order_by(TI.run_id, TI.task_id, TI.map_index)
     )
 
     # Check pre-conditions
-    assert ti_query.all() == [
+    assert session.execute(ti_query).all() == [
         ("downstream", -1, dr1.run_id, TaskInstanceState.FAILED),
         (task_id, 0, dr1.run_id, TaskInstanceState.FAILED),
         (task_id, 1, dr1.run_id, TaskInstanceState.FAILED),
@@ -2616,7 +2617,7 @@ def consumer(value):
     )
     assert dr1 in session, "Check session is passed down all the way"
 
-    assert ti_query.all() == [
+    assert session.execute(ti_query).all() == [
         ("downstream", -1, dr1.run_id, None),
         (task_id, 0, dr1.run_id, TaskInstanceState.FAILED),
         (task_id, 1, dr1.run_id, TaskInstanceState.SUCCESS),
@@ -2867,7 +2868,7 @@ def test_get_asset_triggered_next_run_info(dag_maker, clear_assets):
     dag3 = dag_maker.dag
 
     session = dag_maker.session
-    asset1_id = session.query(AssetModel.id).filter_by(uri=asset1.uri).scalar()
+    asset1_id = session.scalar(select(AssetModel.id).where(AssetModel.uri == asset1.uri))
     session.bulk_save_objects(
         [
             AssetDagRunQueue(asset_id=asset1_id, target_dag_id=dag2.dag_id),
@@ -2876,7 +2877,7 @@ def test_get_asset_triggered_next_run_info(dag_maker, clear_assets):
     )
     session.flush()
 
-    assets = session.query(AssetModel.uri).order_by(AssetModel.id).all()
+    assets = session.execute(select(AssetModel.uri).order_by(AssetModel.id)).all()
 
     info = get_asset_triggered_next_run_info([dag1.dag_id], session=session)
     assert info[dag1.dag_id] == {