From 675ed7b0bb5c55fb21cd12bf20bbafae5b178f9c Mon Sep 17 00:00:00 2001 From: Dev-iL <6509619+Dev-iL@users.noreply.github.com> Date: Tue, 21 Oct 2025 10:21:43 +0300 Subject: [PATCH] SQLA2/FAB: fix some type hints --- .../providers/fab/auth_manager/models/__init__.py | 6 +++--- .../tests/unit/fab/auth_manager/test_security.py | 14 ++++++++------ 2 files changed, 11 insertions(+), 9 deletions(-) diff --git a/providers/fab/src/airflow/providers/fab/auth_manager/models/__init__.py b/providers/fab/src/airflow/providers/fab/auth_manager/models/__init__.py index 62789c02f7d39..375a8a6d9526b 100644 --- a/providers/fab/src/airflow/providers/fab/auth_manager/models/__init__.py +++ b/providers/fab/src/airflow/providers/fab/auth_manager/models/__init__.py @@ -216,9 +216,9 @@ class Group(Model): Sequence("ab_group_id_seq", start=1, increment=1, minvalue=1, cycle=False), primary_key=True, ) - name: Mapped[str] = Column(String(100), unique=True, nullable=False) - label: Mapped[str] = Column(String(150)) - description: Mapped[str] = Column(String(512)) + name: Mapped[str] = mapped_column(String(100), unique=True, nullable=False) + label: Mapped[str | None] = mapped_column(String(150)) + description: Mapped[str | None] = mapped_column(String(512)) users: Mapped[list[User]] = relationship( "User", secondary=assoc_user_group, backref="groups", passive_deletes=True ) diff --git a/providers/fab/tests/unit/fab/auth_manager/test_security.py b/providers/fab/tests/unit/fab/auth_manager/test_security.py index d91e2a2de0629..09e8008ab032e 100644 --- a/providers/fab/tests/unit/fab/auth_manager/test_security.py +++ b/providers/fab/tests/unit/fab/auth_manager/test_security.py @@ -29,13 +29,15 @@ from flask_appbuilder import Model, expose, has_access from flask_appbuilder.models.sqla.interface import SQLAInterface from flask_appbuilder.views import BaseView, ModelView -from sqlalchemy import Column, Date, Float, Integer, String, delete +from sqlalchemy import Date, Float, Integer, String, delete +from sqlalchemy.orm import Mapped from airflow.api_fastapi.app import get_auth_manager from airflow.exceptions import AirflowException from airflow.models import DagModel from airflow.models.dag import DAG from airflow.models.dagbundle import DagBundleModel +from airflow.providers.common.compat.sqlalchemy.orm import mapped_column from airflow.providers.fab.auth_manager.fab_auth_manager import FabAuthManager from airflow.providers.fab.auth_manager.models.anonymous_user import AnonymousUser from airflow.providers.fab.auth_manager.security_manager.override import FabAirflowSecurityManagerOverride @@ -91,11 +93,11 @@ class MockSecurityManager(FabAirflowSecurityManagerOverride): class SomeModel(Model): __tablename__ = "some_model" - id = Column(Integer, primary_key=True) - field_string = Column(String(50), unique=True, nullable=False) - field_integer = Column(Integer()) - field_float = Column(Float()) - field_date = Column(Date()) + id: Mapped[int] = mapped_column(Integer, primary_key=True) + field_string: Mapped[str] = mapped_column(String(50), unique=True, nullable=False) + field_integer: Mapped[int | None] = mapped_column(Integer()) + field_float: Mapped[float | None] = mapped_column(Float()) + field_date: Mapped[datetime.date | None] = mapped_column(Date()) def __repr__(self): return str(self.field_string)