Use func.count to count rows (apache#23657)

NiklasBeierl · May 12, 2022 · 2111d73 · 2111d73
1 parent 1a8f1ee
commit 2111d73
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/airflow/models/trigger.py b/airflow/models/trigger.py
@@ -168,7 +168,7 @@ def assign_unassigned(cls, triggerer_id, capacity, session=None):
         """
         from airflow.jobs.base_job import BaseJob  # To avoid circular import
 
-        count = session.query(cls.id).filter(cls.triggerer_id == triggerer_id).count()
+        count = session.query(func.count(cls.id)).filter(cls.triggerer_id == triggerer_id).scalar()
         capacity -= count
 
         if capacity <= 0:

diff --git a/airflow/utils/db.py b/airflow/utils/db.py
@@ -984,7 +984,7 @@ def check_run_id_null(session: Session) -> Iterable[str]:
         dagrun_table.c.run_id.is_(None),
         dagrun_table.c.execution_date.is_(None),
     )
-    invalid_dagrun_count = session.query(dagrun_table.c.id).filter(invalid_dagrun_filter).count()
+    invalid_dagrun_count = session.query(func.count(dagrun_table.c.id)).filter(invalid_dagrun_filter).scalar()
     if invalid_dagrun_count > 0:
         dagrun_dangling_table_name = _format_airflow_moved_table_name(dagrun_table.name, '2.2', 'dangling')
         if dagrun_dangling_table_name in inspect(session.get_bind()).get_table_names():