SQLA2 public tests (test_backfills, test_connections, test_pools) (#59733)

* sqla2 public tests & pre-commit

* scalars to scalar for one

Author: steveahnahn

.pre-commit-config.yaml

@@ -426,6 +426,9 @@ repos:
             ^airflow-core/src/airflow/models/.*\.py$|
             ^airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.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_backfills.py$|
+            ^airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py$|
+            ^airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_pools.py$|
             ^airflow-core/tests/unit/models/test_serialized_dag.py$|
             ^airflow-core/tests/unit/models/test_pool.py$|
             ^airflow-core/tests/unit/models/test_trigger.py$|

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

@@ -205,7 +205,7 @@ class TestCreateBackfill(TestBackfillEndpoint):
     def test_create_backfill(self, repro_act, repro_exp, session, dag_maker, test_client):
         with dag_maker(session=session, dag_id="TEST_DAG_1", schedule="0 * * * *") as dag:
             EmptyOperator(task_id="mytask")
-        session.query(DagModel).all()
+        session.scalars(select(DagModel)).all()
         session.commit()
         from_date = pendulum.parse("2024-01-01")
         from_date_iso = to_iso(from_date)
@@ -244,7 +244,7 @@ def test_create_backfill(self, repro_act, repro_exp, session, dag_maker, test_cl
         check_last_log(session, dag_id="TEST_DAG_1", event="create_backfill", logical_date=None)
 
     def test_dag_not_exist(self, session, test_client):
-        session.query(DagModel).all()
+        session.scalars(select(DagModel)).all()
         session.commit()
         from_date = pendulum.parse("2024-01-01")
         from_date_iso = to_iso(from_date)
@@ -270,7 +270,7 @@ def test_dag_not_exist(self, session, test_client):
     def test_no_schedule_dag(self, session, dag_maker, test_client):
         with dag_maker(session=session, dag_id="TEST_DAG_1", schedule="None") as dag:
             EmptyOperator(task_id="mytask")
-        session.query(DagModel).all()
+        session.scalars(select(DagModel)).all()
         session.commit()
         from_date = pendulum.parse("2024-01-01")
         from_date_iso = to_iso(from_date)
@@ -306,7 +306,7 @@ def test_create_backfill_with_depends_on_past(
     ):
         with dag_maker(session=session, dag_id="TEST_DAG_1", schedule="0 * * * *") as dag:
             EmptyOperator(task_id="mytask", depends_on_past=True)
-        session.query(DagModel).all()
+        session.scalars(select(DagModel)).all()
         session.commit()
         from_date = pendulum.parse("2024-01-01")
         from_date_iso = to_iso(from_date)
@@ -350,7 +350,7 @@ def test_create_backfill_with_depends_on_past(
     def test_create_backfill_future_dates(self, session, dag_maker, test_client, run_backwards):
         with dag_maker(session=session, dag_id="TEST_DAG_1", schedule="0 * * * *") as dag:
             EmptyOperator(task_id="mytask")
-        session.query(DagModel).all()
+        session.scalars(select(DagModel)).all()
         session.commit()
         from_date = timezone.utcnow() + timedelta(days=1)
         to_date = timezone.utcnow() + timedelta(days=1)
@@ -381,7 +381,7 @@ def test_create_backfill_future_dates(self, session, dag_maker, test_client, run
     def test_create_backfill_past_future_dates(self, session, dag_maker, test_client, run_backwards):
         with dag_maker(session=session, dag_id="TEST_DAG_1", schedule="@daily") as dag:
             EmptyOperator(task_id="mytask")
-        session.query(DagModel).all()
+        session.scalars(select(DagModel)).all()
         session.commit()
         from_date = timezone.utcnow() - timedelta(days=2)
         to_date = timezone.utcnow() + timedelta(days=1)
@@ -543,7 +543,7 @@ def test_create_backfill_with_existing_runs(
     def test_should_respond_401(self, unauthenticated_test_client, dag_maker, session):
         with dag_maker(session=session, dag_id="TEST_DAG_1", schedule="0 * * * *") as dag:
             EmptyOperator(task_id="mytask")
-        session.query(DagModel).all()
+        session.scalars(select(DagModel)).all()
         session.commit()
         from_date = pendulum.parse("2024-01-01")
         from_date_iso = to_iso(from_date)
@@ -564,7 +564,7 @@ def test_should_respond_401(self, unauthenticated_test_client, dag_maker, sessio
     def test_should_respond_403(self, unauthorized_test_client, dag_maker, session):
         with dag_maker(session=session, dag_id="TEST_DAG_1", schedule="0 * * * *") as dag:
             EmptyOperator(task_id="mytask")
-        session.query(DagModel).all()
+        session.scalars(select(DagModel)).all()
         session.commit()
         from_date = pendulum.parse("2024-01-01")
         from_date_iso = to_iso(from_date)
@@ -682,7 +682,7 @@ def test_create_backfill_dry_run_with_depends_on_past(
     ):
         with dag_maker(session=session, dag_id="TEST_DAG_1", schedule="0 * * * *") as dag:
             EmptyOperator(task_id="mytask", depends_on_past=True)
-        session.query(DagModel).all()
+        session.scalars(select(DagModel)).all()
         session.commit()
         from_date = pendulum.parse("2024-01-01")
         from_date_iso = to_iso(from_date)

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

@@ -21,6 +21,7 @@
 from unittest import mock
 
 import pytest
+from sqlalchemy import select
 from sqlalchemy.orm import Session
 
 from airflow.models import Connection
@@ -104,11 +105,11 @@ def create_connections(self):
 class TestDeleteConnection(TestConnectionEndpoint):
     def test_delete_should_respond_204(self, test_client, session):
         self.create_connection()
-        conns = session.query(Connection).all()
+        conns = session.scalars(select(Connection)).all()
         assert len(conns) == 1
         response = test_client.delete(f"/connections/{TEST_CONN_ID}")
         assert response.status_code == 204
-        connection = session.query(Connection).all()
+        connection = session.scalars(select(Connection)).all()
         assert len(connection) == 0
         _check_last_log(session, dag_id=None, event="delete_connection", logical_date=None)
 
@@ -153,7 +154,7 @@ def test_get_should_respond_404(self, test_client):
 
     def test_get_should_respond_200_with_extra(self, test_client, session):
         self.create_connection()
-        connection = session.query(Connection).first()
+        connection = session.scalars(select(Connection)).first()
         connection.extra = '{"extra_key": "extra_value"}'
         session.commit()
         response = test_client.get(f"/connections/{TEST_CONN_ID}")
@@ -166,7 +167,7 @@ def test_get_should_respond_200_with_extra(self, test_client, session):
     @pytest.mark.enable_redact
     def test_get_should_respond_200_with_extra_redacted(self, test_client, session):
         self.create_connection()
-        connection = session.query(Connection).first()
+        connection = session.scalars(select(Connection)).first()
         connection.extra = '{"password": "test-password"}'
         session.commit()
         response = test_client.get(f"/connections/{TEST_CONN_ID}")
@@ -277,7 +278,7 @@ class TestPostConnection(TestConnectionEndpoint):
     def test_post_should_respond_201(self, test_client, session, body):
         response = test_client.post("/connections", json=body)
         assert response.status_code == 201
-        connection = session.query(Connection).all()
+        connection = session.scalars(select(Connection)).all()
         assert len(connection) == 1
         _check_last_log(session, dag_id=None, event="post_connection", logical_date=None)
 
@@ -780,7 +781,7 @@ def test_patch_should_respond_200_with_update_mask(
         self.create_connection()
         response = test_client.patch(f"/connections/{TEST_CONN_ID}", json=body, params=update_mask)
         assert response.status_code == 200
-        connection = session.query(Connection).filter_by(conn_id=TEST_CONN_ID).first()
+        connection = session.scalars(select(Connection).where(Connection.conn_id == TEST_CONN_ID)).first()
         assert connection.password is None
         assert response.json() == updated_connection
 
@@ -1399,7 +1400,7 @@ def test_post_should_accept_empty_string_as_extra(self, test_client, session):
         response = test_client.post("/connections", json=body)
         assert response.status_code == 201
 
-        connection = session.query(Connection).filter_by(conn_id=TEST_CONN_ID).first()
+        connection = session.scalars(select(Connection).where(Connection.conn_id == TEST_CONN_ID)).first()
         assert connection is not None
         assert connection.extra == "{}"  # Backward compatibility: treat "" as empty JSON object
 

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

@@ -19,6 +19,7 @@
 from unittest import mock
 
 import pytest
+from sqlalchemy import func, select
 
 from airflow.models.pool import Pool
 from airflow.utils.session import provide_session
@@ -73,11 +74,11 @@ def create_pools(self):
 class TestDeletePool(TestPoolsEndpoint):
     def test_delete_should_respond_204(self, test_client, session):
         self.create_pools()
-        pools = session.query(Pool).all()
+        pools = session.scalars(select(Pool)).all()
         assert len(pools) == 4
         response = test_client.delete(f"/pools/{POOL1_NAME}")
         assert response.status_code == 204
-        pools = session.query(Pool).all()
+        pools = session.scalars(select(Pool)).all()
         assert len(pools) == 3
         check_last_log(session, dag_id=None, event="delete_pool", logical_date=None)
 
@@ -104,11 +105,11 @@ def test_delete_should_respond_404(self, test_client):
     def test_delete_pool3_should_respond_204(self, test_client, session):
         """Test deleting POOL3 with forward slash in name"""
         self.create_pools()
-        pools = session.query(Pool).all()
+        pools = session.scalars(select(Pool)).all()
         assert len(pools) == 4
         response = test_client.delete(f"/pools/{POOL3_NAME}")
         assert response.status_code == 204
-        pools = session.query(Pool).all()
+        pools = session.scalars(select(Pool)).all()
         assert len(pools) == 3
         check_last_log(session, dag_id=None, event="delete_pool", logical_date=None)
 
@@ -430,12 +431,12 @@ class TestPostPool(TestPoolsEndpoint):
     )
     def test_should_respond_200(self, test_client, session, body, expected_status_code, expected_response):
         self.create_pools()
-        n_pools = session.query(Pool).count()
+        n_pools = session.scalar(select(func.count()).select_from(Pool))
         response = test_client.post("/pools", json=body)
         assert response.status_code == expected_status_code
 
         assert response.json() == expected_response
-        assert session.query(Pool).count() == n_pools + 1
+        assert session.scalar(select(func.count()).select_from(Pool)) == n_pools + 1
         check_last_log(session, dag_id=None, event="post_pool", logical_date=None)
 
     def test_should_respond_401(self, unauthenticated_test_client):
@@ -486,11 +487,11 @@ def test_should_response_409(
         second_expected_response,
     ):
         self.create_pools()
-        n_pools = session.query(Pool).count()
+        n_pools = session.scalar(select(func.count()).select_from(Pool))
         response = test_client.post("/pools", json=body)
         assert response.status_code == first_expected_status_code
         assert response.json() == first_expected_response
-        assert session.query(Pool).count() == n_pools + 1
+        assert session.scalar(select(func.count()).select_from(Pool)) == n_pools + 1
         response = test_client.post("/pools", json=body)
         assert response.status_code == second_expected_status_code
         if second_expected_status_code == 201:
@@ -500,7 +501,7 @@ def test_should_response_409(
             assert "detail" in response_json
             assert list(response_json["detail"].keys()) == ["reason", "statement", "orig_error", "message"]
 
-        assert session.query(Pool).count() == n_pools + 1
+        assert session.scalar(select(func.count()).select_from(Pool)) == n_pools + 1
 
 
 class TestBulkPools(TestPoolsEndpoint):
@@ -990,7 +991,7 @@ def test_update_mask_preserves_other_fields(self, test_client, session):
         assert response_data["update"]["success"] == ["pool1"]
 
         # Assert: fetch from DB and check only masked field changed
-        updated_pool = session.query(Pool).filter_by(pool="pool1").one()
+        updated_pool = session.execute(select(Pool).where(Pool.pool == "pool1")).scalar_one()
         assert updated_pool.slots == 50  # updated
         assert updated_pool.description is None  # unchanged
         assert updated_pool.include_deferred is True  # unchanged