This is an automated email from the ASF dual-hosted git repository.
shahar1 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 675b02a2f9b Make airflow-core typing hold for SQLAlchemy 2.0 and 2.1
(#74629)
675b02a2f9b is described below
commit 675b02a2f9b21b77777b1577b43e37dbf2513e37
Author: Dev-iL <[email protected]>
AuthorDate: Sun Oct 18 12:23:13 2026 +0300
Make airflow-core typing hold for SQLAlchemy 2.0 and 2.1 (#74629)
SQLAlchemy 2.1 turns Select/Row/Result into TypeVarTuple generics, so
`Select[tuple[X]]` now means a one-column select of a tuple, and rows
unpack to the real (possibly nullable) column types instead of Any.
No parameterised `Select` spelling is valid on both versions, so helpers
that only refine a select they receive are generic over it, and helpers
that build their own select return an unparameterised `Select`. Nullable
columns are narrowed with `is None` checks rather than cast; the pool and
asset-event queries already exclude NULL there, while get_authorized_pools
now skips pools whose name is NULL. The 12-column task instance row, beyond
SQLAlchemy's typed select() overloads, is unpacked into locals declared
with their column types.
Type-checks with mypy on both SQLAlchemy 2.0.54 and 2.1.4.
---
.../api_fastapi/auth/managers/base_auth_manager.py | 3 ++-
.../core_api/services/ui/dependencies.py | 6 +++++-
.../execution_api/routes/task_instances.py | 20 ++++++++++++++++---
.../src/airflow/jobs/scheduler_job_runner.py | 2 +-
airflow-core/src/airflow/models/dynamic_region.py | 23 +++++++++++++---------
airflow-core/src/airflow/models/pool.py | 2 +-
.../src/airflow/models/task_coordinates.py | 9 +++++----
airflow-core/src/airflow/models/xcom.py | 7 +++++--
airflow-core/src/airflow/utils/sqlalchemy.py | 3 +++
.../auth/managers/test_base_auth_manager.py | 17 ++++++++++++++++
.../tests/unit/models/test_taskinstance.py | 4 +++-
11 files changed, 73 insertions(+), 23 deletions(-)
diff --git
a/airflow-core/src/airflow/api_fastapi/auth/managers/base_auth_manager.py
b/airflow-core/src/airflow/api_fastapi/auth/managers/base_auth_manager.py
index c1d901a8bc8..77c0ef9a141 100644
--- a/airflow-core/src/airflow/api_fastapi/auth/managers/base_auth_manager.py
+++ b/airflow-core/src/airflow/api_fastapi/auth/managers/base_auth_manager.py
@@ -786,7 +786,8 @@ class BaseAuthManager(Generic[T], LoggingMixin,
metaclass=ABCMeta):
rows = session.execute(stmt).all()
pools_by_team: dict[str | None, set[str]] = defaultdict(set)
for pool_name, team_name in rows:
- pools_by_team[team_name].add(pool_name)
+ if pool_name is not None:
+ pools_by_team[team_name].add(pool_name)
pool_names: set[str] = set()
for team_name, team_pool_names in pools_by_team.items():
diff --git
a/airflow-core/src/airflow/api_fastapi/core_api/services/ui/dependencies.py
b/airflow-core/src/airflow/api_fastapi/core_api/services/ui/dependencies.py
index b48a2480502..eb2e665e889 100644
--- a/airflow-core/src/airflow/api_fastapi/core_api/services/ui/dependencies.py
+++ b/airflow-core/src/airflow/api_fastapi/core_api/services/ui/dependencies.py
@@ -417,7 +417,9 @@ def get_data_dependencies(
select(TaskInletAssetReference.dag_id).where(TaskInletAssetReference.asset_id
== asset_id),
select(DagScheduleAssetReference.dag_id).where(DagScheduleAssetReference.asset_id
== asset_id),
)
- connected_dag_ids =
set(session.scalars(select(connected_dag_ids_query.subquery().c.dag_id)))
+ connected_dag_ids: set[str] = set(
+
session.scalars(select(connected_dag_ids_query.subquery().c.dag_id))
+ )
if not connected_dag_ids & readable_dag_ids:
return {"nodes": [], "edges": []}
@@ -631,6 +633,8 @@ def get_data_dependencies(
.distinct()
)
for event_asset_id, source_dag_id, source_task_id in alias_events:
+ if source_dag_id is None or source_task_id is None:
+ continue
if readable_dag_ids is not None and source_dag_id not in
readable_dag_ids:
continue
asset_node_id = f"asset:{event_asset_id}"
diff --git
a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py
b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py
index 5b49827ca29..c27e2f74f6a 100644
---
a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py
+++
b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py
@@ -24,6 +24,7 @@ import itertools
import json
from collections import defaultdict
from collections.abc import Callable, Collection, Iterator, Sequence
+from datetime import datetime
from typing import Annotated, Any, NoReturn, cast
from uuid import UUID
@@ -347,7 +348,7 @@ def ti_run(
# backend.
# However, do not clear it for deferral
- xcom_keys = []
+ xcom_keys: list[str] = []
if not ti.next_method:
read =
build_xcom_read_query(producer_ids=select(TI.id).where(TI.id ==
task_instance_id))
entity = xcom_entity(read)
@@ -569,6 +570,19 @@ def ti_update_state(
.with_for_update(of=TI)
.execution_options(include_all_attempts=True)
)
+ # A 12-column select exceeds SQLAlchemy's typed ``select()`` overloads, so
the row is untyped.
+ previous_state: str | None
+ try_number: int
+ max_tries: int
+ dag_id: str
+ task_id: str
+ run_id: str
+ region_index: int
+ region_id: UUID
+ hostname: str | None
+ logical_date: datetime | None
+ owners: str | None
+ working_set: bool | None
try:
(
previous_state,
@@ -583,7 +597,7 @@ def ti_update_state(
logical_date,
owners,
working_set,
- ) = session.execute(old).one()
+ ) = tuple(session.execute(old).one())
log.debug(
"Retrieved current task instance state",
previous_state=previous_state,
@@ -1122,7 +1136,7 @@ def ti_skip_downstream(
raise HTTPException(status.HTTP_400_BAD_REQUEST, str(error)) from error
if not selects:
return
- selected_ids = set(session.scalars(union(*selects)))
+ selected_ids: set[UUID] = set(session.scalars(union(*selects)))
# Don't overwrite tasks that are already executing or finished.
# See: https://github.com/apache/airflow/issues/59378
diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py
b/airflow-core/src/airflow/jobs/scheduler_job_runner.py
index c5e7732f6d8..19662ae5a07 100644
--- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py
+++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py
@@ -3588,7 +3588,7 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
count_result: int | None = session.scalar(statement)
return count_result if count_result is not None else 0
- previous_ti_metrics: dict[TaskInstanceState, dict[tuple[str, str, str],
int]] = {}
+ previous_ti_metrics: dict[TaskInstanceState, dict[tuple[str, str, str |
None], int]] = {}
@provide_session
def _emit_ti_metrics(self, *, session: Session = NEW_SESSION) -> None:
diff --git a/airflow-core/src/airflow/models/dynamic_region.py
b/airflow-core/src/airflow/models/dynamic_region.py
index 494cd09a926..b2100d95734 100644
--- a/airflow-core/src/airflow/models/dynamic_region.py
+++ b/airflow-core/src/airflow/models/dynamic_region.py
@@ -21,7 +21,7 @@ import struct
from collections.abc import Collection, Iterable, Iterator
from datetime import datetime
from functools import partial
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, TypeVar
from uuid import UUID
import attrs
@@ -72,6 +72,11 @@ if TYPE_CHECKING:
from sqlalchemy.sql.elements import ColumnElement
from airflow.models.taskinstance import TaskInstance
+ from airflow.utils.sqlalchemy import UUIDSelect
+
+# Bound unparameterised: no ``Select[...]`` spelling accepts every arity on
both SQLAlchemy 2.0 and 2.1.
+# TODO: parameterise the bound with the variadic Select once the minimum
SQLAlchemy is 2.1.
+_SelectT = TypeVar("_SelectT", bound="Select")
def is_regional(region_id: UUID, region_index: int) -> bool:
@@ -566,8 +571,8 @@ def loop_position(
def build_loop_family_legs(
- query: Select[Any], family_member: DynamicRegion, *, iteration: int | None
= None
-) -> tuple[Select[Any], Select[Any]]:
+ query: _SelectT, family_member: DynamicRegion, *, iteration: int | None =
None
+) -> tuple[_SelectT, _SelectT]:
"""
Split ``query``, a selection from ``task_instance``, into the rows that
belong to one loop family.
@@ -648,7 +653,7 @@ def _match_any(column, value: str | Collection[str]) ->
ColumnElement[bool]:
def _filter_producers(
- query: Select,
+ query: _SelectT,
*,
dag_id: str,
run_id: str,
@@ -658,7 +663,7 @@ def _filter_producers(
region_id: UUID | None,
region_index: int | None,
top_level_only: bool,
-) -> Select:
+) -> _SelectT:
"""Push every coordinate predicate that does not need loop ancestry into
SQL."""
from airflow.models.taskinstance import TaskInstance
@@ -707,7 +712,7 @@ def _validate_producer_request(
def _build_candidate_query(
- query: Select[Any],
+ query: _SelectT,
*,
dag_id: str,
run_id: str,
@@ -718,7 +723,7 @@ def _build_candidate_query(
region_id: UUID | None,
region_index: int | None,
session: Session,
-) -> Select[Any] | CompoundSelect | None:
+) -> _SelectT | CompoundSelect | None:
"""
Narrow ``query`` to the candidate producers of the caller's scope, or
``None`` if there are none.
@@ -823,7 +828,7 @@ def select_current_producer_ids(
region_id: UUID | None = None,
region_index: int | None = None,
session: Session,
-) -> Select[tuple[UUID]]:
+) -> UUIDSelect:
"""
Select the ids of the live producers :func:`resolve_current_producers`
would return.
@@ -928,7 +933,7 @@ def _build_pass_ownership_filter(
def select_loop_producer_ids(
*, dag_id: str, run_id: str, loop_node_id: str, task_id: str, is_mapped:
bool, session: Session
-) -> Select[tuple[UUID]]:
+) -> UUIDSelect:
"""
Select every live execution of ``task_id`` across all passes of the loop
``loop_node_id``.
diff --git a/airflow-core/src/airflow/models/pool.py
b/airflow-core/src/airflow/models/pool.py
index 3a2bc8f0922..e791704c95c 100644
--- a/airflow-core/src/airflow/models/pool.py
+++ b/airflow-core/src/airflow/models/pool.py
@@ -398,4 +398,4 @@ class Pool(Base):
pool_names: list[str], *, session: Session = NEW_SESSION
) -> dict[str, str | None]:
stmt = select(Pool.pool,
Pool.team_name).where(Pool.pool.in_(pool_names))
- return {pool: team_name for pool, team_name in session.execute(stmt)}
+ return {pool: team_name for pool, team_name in session.execute(stmt)
if pool is not None}
diff --git a/airflow-core/src/airflow/models/task_coordinates.py
b/airflow-core/src/airflow/models/task_coordinates.py
index 7ce280a31bd..3d97ef3eade 100644
--- a/airflow-core/src/airflow/models/task_coordinates.py
+++ b/airflow-core/src/airflow/models/task_coordinates.py
@@ -49,6 +49,7 @@ if TYPE_CHECKING:
from sqlalchemy.sql.elements import ColumnElement
from airflow.serialization.definitions.dag import SerializedDAG,
SerializedOperator
+ from airflow.utils.sqlalchemy import UUIDSelect
LOOP_GATE_OPERATOR = "LoopGateOperator"
@@ -361,7 +362,7 @@ class TaskCoordinateResolver:
run_id: str,
task_ids: Collection[str],
slots: Collection[tuple[str, int]],
- ) -> Select[tuple[UUID]]:
+ ) -> UUIDSelect:
"""Select, in one statement, the region-less task instances that the
named tasks or slots address."""
targets = []
if task_ids:
@@ -730,7 +731,7 @@ class TaskCoordinateResolver:
def select_skip_target_ids(
self, *, caller: TaskInstance, task_id: str, map_indexes: int | None =
None
- ) -> Select[tuple[UUID]]:
+ ) -> UUIDSelect:
"""
Select the task instance ids a skip request from ``caller`` reaches
for ``task_id``.
@@ -765,7 +766,7 @@ class TaskCoordinateResolver:
map_indexes=map_indexes,
)
- def _select_loop_producer_ids(self, request: _ProducerRequest) ->
Select[tuple[UUID]]:
+ def _select_loop_producer_ids(self, request: _ProducerRequest) ->
UUIDSelect:
if TYPE_CHECKING:
assert request.loop_node_id is not None
return select_loop_producer_ids(
@@ -789,7 +790,7 @@ class TaskCoordinateResolver:
map_indexes: int | Collection[int] | None = None,
previous_iteration: bool = False,
all_iterations: bool = False,
- ) -> Select[tuple[UUID]]:
+ ) -> UUIDSelect:
"""Select the task instance ids :meth:`resolve` returns, without
loading the producers."""
if region_index is not None and region_id is None:
raise ValueError("region_index requires an explicit producer
region_id")
diff --git a/airflow-core/src/airflow/models/xcom.py
b/airflow-core/src/airflow/models/xcom.py
index 544545fa04d..2054ba1368a 100644
--- a/airflow-core/src/airflow/models/xcom.py
+++ b/airflow-core/src/airflow/models/xcom.py
@@ -68,6 +68,9 @@ if TYPE_CHECKING:
from airflow.models.dagrun import DagRun
from airflow.models.taskinstance import TaskInstance
+ # TODO: use Select[XComModel] once the minimum SQLAlchemy is 2.1; 2.0 uses
Select[tuple[XComModel]].
+ XComModelSelect = Select
+
XCOM_RETURN_KEY = "return_value"
@@ -328,7 +331,7 @@ class _XComOperations:
include_prior_dates: bool = False,
limit: int | None = None,
try_number: int | None = None,
- ) -> Select[tuple[XComModel]]:
+ ) -> XComModelSelect:
"""
Composes a query to get one or more XCom entries.
@@ -596,7 +599,7 @@ def _reject_projection_write(mapper, connection, target):
raise TypeError("XCom read projection is read-only")
-def build_xcom_read_query(*, producer_ids: Select, key: str | None = None) ->
Select[tuple[XComModel]]:
+def build_xcom_read_query(*, producer_ids: Select, key: str | None = None) ->
XComModelSelect:
model = get_xcom_model()
entity = aliased(model, _filter_rows(model.__table__,
producer_ids=producer_ids, key=key))
return select(entity).execution_options(populate_existing=True)
diff --git a/airflow-core/src/airflow/utils/sqlalchemy.py
b/airflow-core/src/airflow/utils/sqlalchemy.py
index 966217a68d0..18c5da08214 100644
--- a/airflow-core/src/airflow/utils/sqlalchemy.py
+++ b/airflow-core/src/airflow/utils/sqlalchemy.py
@@ -52,6 +52,9 @@ if TYPE_CHECKING:
from airflow.typing_compat import Self
+ # TODO: use Select[UUID] once the minimum SQLAlchemy is 2.1; 2.0 uses
Select[tuple[UUID]].
+ UUIDSelect = Select
+
log = logging.getLogger(__name__)
diff --git
a/airflow-core/tests/unit/api_fastapi/auth/managers/test_base_auth_manager.py
b/airflow-core/tests/unit/api_fastapi/auth/managers/test_base_auth_manager.py
index f25f4dbc9fd..49e3699f203 100644
---
a/airflow-core/tests/unit/api_fastapi/auth/managers/test_base_auth_manager.py
+++
b/airflow-core/tests/unit/api_fastapi/auth/managers/test_base_auth_manager.py
@@ -38,6 +38,7 @@ from
airflow.api_fastapi.auth.managers.models.resource_details import (
from airflow.api_fastapi.auth.tokens import JWTGenerator, JWTValidator
from airflow.api_fastapi.common.types import MenuItem
from airflow.exceptions import RemovedInAirflow4Warning
+from airflow.models.pool import Pool
from airflow.models.team import Team
from tests_common.test_utils.config import conf_vars
@@ -799,6 +800,22 @@ class TestBaseAuthManager:
result = auth_manager.get_authorized_pools(user=user, session=session)
assert result == expected
+ @pytest.mark.db_test
+ def test_get_authorized_pools_skips_null_names(self, auth_manager,
session):
+ session.add_all(
+ [
+ Pool(pool=None, slots=1, include_deferred=False),
+ Pool(pool="authorized_pool", slots=1, include_deferred=False),
+ ]
+ )
+ session.flush()
+ auth_manager.is_authorized_pool =
create_autospec(auth_manager.is_authorized_pool, return_value=True)
+
+ result =
auth_manager.get_authorized_pools(user=Mock(spec=BaseAuthManagerUserTest),
session=session)
+
+ assert None not in result
+ assert "authorized_pool" in result
+
@pytest.mark.parametrize(
("authorized_uri_prefix", "rows", "expected"),
[
diff --git a/airflow-core/tests/unit/models/test_taskinstance.py
b/airflow-core/tests/unit/models/test_taskinstance.py
index 4a7721b3a41..d2bec4dc393 100644
--- a/airflow-core/tests/unit/models/test_taskinstance.py
+++ b/airflow-core/tests/unit/models/test_taskinstance.py
@@ -137,6 +137,8 @@ from tests_common.test_utils.taskinstance import (
from unit.models import DEFAULT_DATE
if TYPE_CHECKING:
+ from collections.abc import Sequence
+
from tests_common.pytest_plugin import DagMaker
pytestmark = [pytest.mark.db_test, pytest.mark.need_serialized_dag,
pytest.mark.want_activate_assets]
@@ -650,7 +652,7 @@ class TestTaskInstance:
)
ti.run(session=session)
- executor_configs = session.scalars(
+ executor_configs: Sequence[dict | None] = session.scalars(
select(TaskInstance.executor_config).where(TaskInstance.dag_id ==
ti.dag_id)
).all()
assert executor_configs == [{"foo": "bar"}]
