This is an automated email from the ASF dual-hosted git repository.
vincbeck pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 5ce21c5ad74 Fix mypy errors from SQLAlchemy 2.1 typing in fab sorting
and standard sensor_helper (#75075)
5ce21c5ad74 is described below
commit 5ce21c5ad7404871a619ecb8a7d7a1c4a266eb38
Author: Wei Lee <[email protected]>
AuthorDate: Mon Oct 19 23:09:57 2026 +0900
Fix mypy errors from SQLAlchemy 2.1 typing in fab sorting and standard
sensor_helper (#75075)
Use Column.asc()/desc() in fab build_ordering and make the sensor_helper
count helpers type-check on both SQLAlchemy 2.0 and 2.1.
---
.../fab/auth_manager/api_fastapi/sorting.py | 3 +--
.../providers/standard/utils/sensor_helper.py | 31 ++++++++++++----------
2 files changed, 18 insertions(+), 16 deletions(-)
diff --git
a/providers/fab/src/airflow/providers/fab/auth_manager/api_fastapi/sorting.py
b/providers/fab/src/airflow/providers/fab/auth_manager/api_fastapi/sorting.py
index b19d43cb9c8..217823037dc 100644
---
a/providers/fab/src/airflow/providers/fab/auth_manager/api_fastapi/sorting.py
+++
b/providers/fab/src/airflow/providers/fab/auth_manager/api_fastapi/sorting.py
@@ -20,7 +20,6 @@ from collections.abc import Mapping
from typing import TYPE_CHECKING, Any
from fastapi import HTTPException, status
-from sqlalchemy import asc, desc
if TYPE_CHECKING:
from sqlalchemy.orm import InstrumentedAttribute
@@ -46,4 +45,4 @@ def build_ordering(
)
col = allowed[key]
- return desc(col) if is_desc else asc(col)
+ return col.desc() if is_desc else col.asc()
diff --git
a/providers/standard/src/airflow/providers/standard/utils/sensor_helper.py
b/providers/standard/src/airflow/providers/standard/utils/sensor_helper.py
index f16e86bf34a..10ae4b478c2 100644
--- a/providers/standard/src/airflow/providers/standard/utils/sensor_helper.py
+++ b/providers/standard/src/airflow/providers/standard/utils/sensor_helper.py
@@ -57,11 +57,9 @@ def _get_count(
return 0
if external_task_ids:
- count = (
- session.scalar(
- _count_stmt(TI, states, dttm_filter,
external_dag_id).where(TI.task_id.in_(external_task_ids))
- )
- or 0
+ count = _get_scalar_count(
+ _count_stmt(TI, states, dttm_filter,
external_dag_id).where(TI.task_id.in_(external_task_ids)),
+ session,
) / len(external_task_ids)
elif external_task_group_id:
external_task_group_task_ids = _get_external_task_group_task_ids(
@@ -71,22 +69,26 @@ def _get_count(
count = 0
else:
count = (
- (
- session.scalar(
- _count_stmt(TI, states, dttm_filter,
external_dag_id).where(
- tuple_(TI.task_id,
TI.map_index).in_(external_task_group_task_ids)
- )
- )
- or 0
+ _get_scalar_count(
+ _count_stmt(TI, states, dttm_filter,
external_dag_id).where(
+ tuple_(TI.task_id,
TI.map_index).in_(external_task_group_task_ids)
+ ),
+ session,
)
/ len(external_task_group_task_ids)
* len(dttm_filter)
)
else:
- count = session.scalar(_count_stmt(DR, states, dttm_filter,
external_dag_id)) or 0
+ count = _get_scalar_count(_count_stmt(DR, states, dttm_filter,
external_dag_id), session)
return cast("int", count)
+def _get_scalar_count(stmt: Select[tuple[int]], session: Session) -> int:
+ # SQLAlchemy 2.0 and 2.1 type `session.scalar()` of a single-column select
differently.
+ result: Any = session.scalar(stmt)
+ return result or 0
+
+
def _count_stmt(
model: type[DagRun] | type[TaskInstance], states: list[str], dttm_filter:
list[Any], external_dag_id: str
) -> Select[tuple[int]]:
@@ -100,9 +102,10 @@ def _count_stmt(
"""
date_field = model.logical_date if AIRFLOW_V_3_0_PLUS else
model.execution_date
- return select(func.count()).where(
+ stmt: Any = select(func.count()).where(
model.dag_id == external_dag_id, model.state.in_(states),
date_field.in_(dttm_filter)
)
+ return stmt
def _get_external_task_group_task_ids(
