Refactor devel-common/tests_common to use SQLA2 (#59849)

Author: rickychen-infinirc

.pre-commit-config.yaml

@@ -472,7 +472,10 @@ repos:
             ^providers/edge3/.*\.py$|
             ^providers/mysql/.*\.py$|
             ^providers/openlineage/.*\.py$|
-            ^task_sdk.*\.py$
+            ^task_sdk.*\.py$|
+            ^devel-common/src/tests_common/pytest_plugin\.py$|
+            ^devel-common/src/tests_common/test_utils/.*\.py$
+
         pass_filenames: true
       - id: update-supported-versions
         name: Updates supported versions in documentation

devel-common/src/tests_common/pytest_plugin.py

@@ -1252,10 +1252,16 @@ def __call__(
             self.bundle_name = bundle_name or "dag_maker"
             self.bundle_version = bundle_version
             if AIRFLOW_V_3_0_PLUS:
+                from sqlalchemy import func, select
+
                 from airflow.models.dagbundle import DagBundleModel
 
                 if (
-                    self.session.query(DagBundleModel).filter(DagBundleModel.name == self.bundle_name).count()
+                    self.session.scalar(
+                        select(func.count())
+                        .select_from(DagBundleModel)
+                        .where(DagBundleModel.name == self.bundle_name)
+                    )
                     == 0
                 ):
                     self.session.add(DagBundleModel(name=self.bundle_name))
@@ -1285,39 +1291,25 @@ def cleanup(self):
                     self.session.rollback()
 
                     if AIRFLOW_V_3_0_PLUS:
+                        from sqlalchemy import delete
+
                         from airflow.models.dag_version import DagVersion
 
-                        self.session.query(DagRun).filter(DagRun.dag_id.in_(dag_ids)).delete(
-                            synchronize_session=False,
-                        )
-                        self.session.query(TaskInstance).filter(TaskInstance.dag_id.in_(dag_ids)).delete(
-                            synchronize_session=False,
-                        )
-                        self.session.query(DagVersion).filter(DagVersion.dag_id.in_(dag_ids)).delete(
-                            synchronize_session=False
-                        )
+                        self.session.execute(delete(DagRun).where(DagRun.dag_id.in_(dag_ids)))
+                        self.session.execute(delete(TaskInstance).where(TaskInstance.dag_id.in_(dag_ids)))
+                        self.session.execute(delete(DagVersion).where(DagVersion.dag_id.in_(dag_ids)))
                     else:
-                        self.session.query(SerializedDagModel).filter(
-                            SerializedDagModel.dag_id.in_(dag_ids)
-                        ).delete(synchronize_session=False)
-                        self.session.query(DagRun).filter(DagRun.dag_id.in_(dag_ids)).delete(
-                            synchronize_session=False,
-                        )
-                        self.session.query(TaskInstance).filter(TaskInstance.dag_id.in_(dag_ids)).delete(
-                            synchronize_session=False,
+                        from sqlalchemy import delete
+
+                        self.session.execute(
+                            delete(SerializedDagModel).where(SerializedDagModel.dag_id.in_(dag_ids))
                         )
-                    self.session.query(XCom).filter(XCom.dag_id.in_(dag_ids)).delete(
-                        synchronize_session=False,
-                    )
-                    self.session.query(DagModel).filter(DagModel.dag_id.in_(dag_ids)).delete(
-                        synchronize_session=False,
-                    )
-                    self.session.query(TaskMap).filter(TaskMap.dag_id.in_(dag_ids)).delete(
-                        synchronize_session=False,
-                    )
-                    self.session.query(AssetEvent).filter(AssetEvent.source_dag_id.in_(dag_ids)).delete(
-                        synchronize_session=False,
-                    )
+                        self.session.execute(delete(DagRun).where(DagRun.dag_id.in_(dag_ids)))
+                        self.session.execute(delete(TaskInstance).where(TaskInstance.dag_id.in_(dag_ids)))
+                    self.session.execute(delete(XCom).where(XCom.dag_id.in_(dag_ids)))
+                    self.session.execute(delete(DagModel).where(DagModel.dag_id.in_(dag_ids)))
+                    self.session.execute(delete(TaskMap).where(TaskMap.dag_id.in_(dag_ids)))
+                    self.session.execute(delete(AssetEvent).where(AssetEvent.source_dag_id.in_(dag_ids)))
                     self.session.commit()
                     if self._own_session:
                         self.session.expunge_all()
@@ -1714,10 +1706,12 @@ def _create_log_template(filename_template, elasticsearch_id=""):
         session.commit()
 
         def _delete_log_template():
+            from sqlalchemy import delete
+
             from airflow.models import DagRun, TaskInstance
 
-            session.query(TaskInstance).delete()
-            session.query(DagRun).delete()
+            session.execute(delete(TaskInstance))
+            session.execute(delete(DagRun))
             session.delete(log_template)
             session.commit()
 
@@ -2753,11 +2747,18 @@ def testing_dag_bundle():
     from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
 
     if AIRFLOW_V_3_0_PLUS:
+        from sqlalchemy import func, select
+
         from airflow.models.dagbundle import DagBundleModel
         from airflow.utils.session import create_session
 
         with create_session() as session:
-            if session.query(DagBundleModel).filter(DagBundleModel.name == "testing").count() == 0:
+            if (
+                session.scalar(
+                    select(func.count()).select_from(DagBundleModel).where(DagBundleModel.name == "testing")
+                )
+                == 0
+            ):
                 testing = DagBundleModel(name="testing")
                 session.add(testing)
 
@@ -2767,11 +2768,13 @@ def testing_team():
     from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
 
     if AIRFLOW_V_3_0_PLUS:
+        from sqlalchemy import select
+
         from airflow.models.team import Team
         from airflow.utils.session import create_session
 
         with create_session() as session:
-            team = session.query(Team).filter_by(name="testing").one_or_none()
+            team = session.scalar(select(Team).where(Team.name == "testing"))
             if not team:
                 team = Team(name="testing")
                 session.add(team)

devel-common/src/tests_common/test_utils/api_client_helpers.py

@@ -114,14 +114,15 @@ def create_airflow_connection(connection_id: str, connection_conf: dict[str, Any
             print(f"Connection '{connection_id}' does not exist. A new one will be created")
         create_connection_request(connection_id=connection_id, connection=connection_conf)
     else:
+        from sqlalchemy import delete
+
         from airflow.models import Connection
         from airflow.settings import Session
 
         if Session is None:
             raise RuntimeError("Session not configured. Call configure_orm() first.")
         session = Session()
-        query = session.query(Connection).filter(Connection.conn_id == connection_id)
-        query.delete()
+        session.execute(delete(Connection).where(Connection.conn_id == connection_id))
         connection = Connection(conn_id=connection_id, **connection_conf)
         session.add(connection)
         session.commit()
@@ -135,12 +136,13 @@ def delete_airflow_connection(connection_id: str) -> None:
     if AIRFLOW_V_3_0_PLUS:
         delete_connection_request(connection_id=connection_id)
     else:
+        from sqlalchemy import delete
+
         from airflow.models import Connection
         from airflow.settings import Session
 
         if Session is None:
             raise RuntimeError("Session not configured. Call configure_orm() first.")
         session = Session()
-        query = session.query(Connection).filter(Connection.conn_id == connection_id)
-        query.delete()
+        session.execute(delete(Connection).where(Connection.conn_id == connection_id))
         session.commit()

devel-common/src/tests_common/test_utils/api_fastapi.py

@@ -49,24 +49,25 @@ def _masked_value_check(data, sensitive_fields):
 
 
 def _check_last_log(session, dag_id, event, logical_date, expected_extra=None, check_masked=False):
-    logs = (
-        session.query(
+    from sqlalchemy import select
+
+    logs = session.execute(
+        select(
             Log.dag_id,
             Log.task_id,
             Log.event,
             Log.logical_date,
             Log.owner,
             Log.extra,
         )
-        .filter(
+        .where(
             Log.dag_id == dag_id,
             Log.event == event,
             Log.logical_date == logical_date,
         )
         .order_by(Log.dttm.desc())
         .limit(1)
-        .all()
-    )
+    ).all()
     assert len(logs) == 1
     assert logs[0].extra
     if expected_extra:
@@ -77,11 +78,10 @@ def _check_last_log(session, dag_id, event, logical_date, expected_extra=None, c
 
 
 def _check_dag_run_note(session, dr_id, note_data):
-    dr_note = (
-        session.query(DagRunNote)
-        .join(DagRun, DagRunNote.dag_run_id == DagRun.id)
-        .filter(DagRun.run_id == dr_id)
-        .one_or_none()
+    from sqlalchemy import select
+
+    dr_note = session.scalar(
+        select(DagRunNote).join(DagRun, DagRunNote.dag_run_id == DagRun.id).where(DagRun.run_id == dr_id)
     )
     if note_data is None:
         assert dr_note is None
@@ -91,7 +91,9 @@ def _check_dag_run_note(session, dr_id, note_data):
 
 
 def _check_task_instance_note(session, ti_id, note_data):
-    ti_note = session.query(TaskInstanceNote).filter_by(ti_id=ti_id).one_or_none()
+    from sqlalchemy import select
+
+    ti_note = session.scalar(select(TaskInstanceNote).where(TaskInstanceNote.ti_id == ti_id))
     if note_data is None:
         assert ti_note is None
     else: