apache-airflow-providers-fab 3.1.0rc1__py3-none-any.whl → 3.2.0rc1__py3-none-any.whl

airflow/providers/fab/auth_manager/security_manager/override.py

@@ -44,6 +44,7 @@ from flask_appbuilder.const import (
 from flask_appbuilder.models.sqla.interface import SQLAInterface
 from flask_appbuilder.security.api import SecurityApi
 from flask_appbuilder.security.registerviews import (
+    BaseRegisterUser,
     RegisterUserDBView,
     RegisterUserOAuthView,
 )
@@ -52,8 +53,10 @@ from flask_appbuilder.security.views import (
     AuthLDAPView,
     AuthOAuthView,
     AuthRemoteUserView,
+    AuthView,
     RegisterUserModelView,
     UserGroupModelView,
+    UserModelView,
 )
 from flask_babel import lazy_gettext
 from flask_jwt_extended import JWTManager
@@ -67,7 +70,6 @@ from sqlalchemy.orm import joinedload
 from werkzeug.security import check_password_hash, generate_password_hash
 
 from airflow.configuration import conf
-from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.fab.auth_manager.models import (
     Action,
     Group,
@@ -101,15 +103,16 @@ from airflow.providers.fab.version_compat import AIRFLOW_V_3_1_PLUS
 from airflow.providers.fab.www.security import permissions
 from airflow.providers.fab.www.security_manager import AirflowSecurityManagerV2
 from airflow.providers.fab.www.session import AirflowDatabaseSessionInterface
-from airflow.security.permissions import RESOURCE_BACKFILL
 
 if TYPE_CHECKING:
+    from authlib.integrations.flask_client import OAuth
+
     from airflow.providers.fab.www.security.permissions import (
         RESOURCE_ASSET,
         RESOURCE_ASSET_ALIAS,
     )
     from airflow.sdk import DAG
-    from airflow.serialization.serialized_objects import SerializedDAG
+    from airflow.serialization.definitions.dag import SerializedDAG
 else:
     from airflow.providers.common.compat.security.permissions import (
         RESOURCE_ASSET,
@@ -149,6 +152,10 @@ log = logging.getLogger(__name__)
 MAX_NUM_DATABASE_USER_SESSIONS = 50000
 
 
+class FabException(Exception):
+    """Custom exception for FAB security manager."""
+
+
 class FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
     """
     This security manager overrides the default AirflowSecurityManager security manager.
@@ -160,11 +167,11 @@ class FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
     :param appbuilder: The appbuilder.
     """
 
-    auth_view = None
+    auth_view: AuthView | None = None
     """ The obj instance for authentication view """
-    registeruser_view = None
+    registeruser_view: BaseRegisterUser | None = None
     """ The obj instance for registering user view """
-    user_view = None
+    user_view: type[UserModelView] | None = None
     """ The obj instance for user view """
 
     """ Models """
@@ -208,9 +215,9 @@ class FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
     security_api = SecurityApi
     """ Override if you want your own Security API login endpoint """
 
-    jwt_manager = None
+    jwt_manager: JWTManager | None = None
     """ Flask-JWT-Extended """
-    oauth = None
+    oauth: OAuth | None = None
     oauth_remotes: dict[str, Any]
     """ Initialized (remote_app) providers dict {'provider_name', OBJ } """
 
@@ -236,7 +243,7 @@ class FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
         (permissions.ACTION_CAN_READ, permissions.RESOURCE_DAG_WARNING),
         (permissions.ACTION_CAN_READ, RESOURCE_ASSET),
         (permissions.ACTION_CAN_READ, RESOURCE_ASSET_ALIAS),
-        (permissions.ACTION_CAN_READ, RESOURCE_BACKFILL),
+        (permissions.ACTION_CAN_READ, permissions.RESOURCE_BACKFILL),
         (permissions.ACTION_CAN_READ, permissions.RESOURCE_CLUSTER_ACTIVITY),
         (permissions.ACTION_CAN_READ, permissions.RESOURCE_POOL),
         (permissions.ACTION_CAN_READ, permissions.RESOURCE_IMPORT_ERROR),
@@ -305,12 +312,14 @@ class FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
         (permissions.ACTION_CAN_READ, permissions.RESOURCE_VARIABLE),
         (permissions.ACTION_CAN_EDIT, permissions.RESOURCE_VARIABLE),
         (permissions.ACTION_CAN_DELETE, permissions.RESOURCE_VARIABLE),
+        (permissions.ACTION_CAN_CREATE, permissions.RESOURCE_XCOM),
+        (permissions.ACTION_CAN_EDIT, permissions.RESOURCE_XCOM),
         (permissions.ACTION_CAN_DELETE, permissions.RESOURCE_XCOM),
         (permissions.ACTION_CAN_CREATE, RESOURCE_ASSET),
         (permissions.ACTION_CAN_DELETE, RESOURCE_ASSET),
-        (permissions.ACTION_CAN_CREATE, RESOURCE_BACKFILL),
-        (permissions.ACTION_CAN_EDIT, RESOURCE_BACKFILL),
-        (permissions.ACTION_CAN_DELETE, RESOURCE_BACKFILL),
+        (permissions.ACTION_CAN_CREATE, permissions.RESOURCE_BACKFILL),
+        (permissions.ACTION_CAN_EDIT, permissions.RESOURCE_BACKFILL),
+        (permissions.ACTION_CAN_DELETE, permissions.RESOURCE_BACKFILL),
     ]
     # [END security_op_perms]
 
@@ -418,9 +427,9 @@ class FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
             log.warning("JWT token is not validated!")
             return me
 
-        raise AirflowException("OAuth signature verify failed")
+        raise FabException("OAuth signature verify failed")
 
-    def register_views(self):
+    def register_views(self) -> None:
         """Register FAB auth manager related views."""
         if not current_app.config.get("FAB_ADD_SECURITY_VIEWS", True):
             return
@@ -457,7 +466,7 @@ class FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
 
         # this needs to be done after the view is added, otherwise the blueprint
         # is not initialized
-        if self.is_auth_limited:
+        if self.is_auth_limited and self.auth_view:
             self.limiter.limit(self.auth_rate_limit, methods=["POST"])(self.auth_view.blueprint)
 
         self.user_view = self.appbuilder.add_view(
@@ -480,14 +489,13 @@ class FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
         )
         role_view.related_views = [self.user_view.__class__]
 
-        if self.userstatschartview:
-            self.appbuilder.add_view(
-                self.userstatschartview,
-                "User's Statistics",
-                icon="fa-bar-chart-o",
-                label=lazy_gettext("User's Statistics"),
-                category="Security",
-            )
+        self.appbuilder.add_view(
+            self.userstatschartview,
+            "User's Statistics",
+            icon="fa-bar-chart-o",
+            label=lazy_gettext("User's Statistics"),
+            category="Security",
+        )
         if self.auth_user_registration:
             self.appbuilder.add_view(
                 self.registerusermodelview,
@@ -534,7 +542,7 @@ class FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
         lm.user_loader(self.load_user)
         return lm
 
-    def create_jwt_manager(self):
+    def create_jwt_manager(self) -> None:
         """Create the JWT manager."""
         jwt_manager = JWTManager()
         jwt_manager.init_app(current_app)
@@ -550,6 +558,8 @@ class FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
         :param password: the clear text password to reset and save hashed on the db
         """
         user = self.get_user_by_id(userid)
+        if not user:
+            return False
         user.password = generate_password_hash(password)
         self.reset_user_sessions(user)
         return self.update_user(user)
@@ -591,16 +601,17 @@ class FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
                 "warning",
             )
 
-    def load_user_jwt(self, _jwt_header, jwt_data):
+    def load_user_jwt(self, _jwt_header, jwt_data) -> User | None:
         identity = jwt_data["sub"]
         user = self.load_user(identity)
         if user and user.is_active:
             # Set flask g.user to JWT user, we can't do it on before request
             g.user = user
             return user
+        return None
 
     @property
-    def auth_type(self):
+    def auth_type(self) -> int:
         """Get the auth type."""
         return current_app.config["AUTH_TYPE"]
 
@@ -737,11 +748,6 @@ class FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
     def auth_remote_user_env_var(self) -> str:
         return current_app.config["AUTH_REMOTE_USER_ENV_VAR"]
 
-    @property
-    def auth_username_ci(self):
-        """Get the auth username for CI."""
-        return current_app.config.get("AUTH_USERNAME_CI", True)
-
     @property
     def auth_user_registration(self):
         """Will user self registration be allowed."""
@@ -763,15 +769,15 @@ class FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
         return current_app.config["AUTH_ROLE_ADMIN"]
 
     @property
-    def oauth_whitelists(self):
+    def oauth_whitelists(self) -> dict[str, list]:
         return self.oauth_allow_list
 
     @staticmethod
-    def create_builtin_roles():
+    def create_builtin_roles() -> dict:
         return current_app.config.get("FAB_ROLES", {})
 
     @property
-    def builtin_roles(self):
+    def builtin_roles(self) -> dict:
         """Get the builtin roles."""
         return self._builtin_roles
 
@@ -780,11 +786,11 @@ class FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
         return current_app.config["AUTH_API_LOGIN_ALLOW_MULTIPLE_PROVIDERS"]
 
     @property
-    def auth_type_provider_name(self):
+    def auth_type_provider_name(self) -> str | None:
         provider_to_auth_type = {AUTH_DB: "db", AUTH_LDAP: "ldap"}
         return provider_to_auth_type.get(self.auth_type)
 
-    def _init_config(self):
+    def _init_config(self) -> None:
         """
         Initialize config.
 
@@ -854,7 +860,7 @@ class FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
         current_app.config.setdefault("AUTH_RATE_LIMITED", True)
         current_app.config.setdefault("AUTH_RATE_LIMIT", "5 per 40 second")
 
-    def _init_auth(self):
+    def _init_auth(self) -> None:
         """
         Initialize authentication configuration.
 
@@ -877,7 +883,7 @@ class FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
                     self.oauth_allow_list[provider_name] = provider["whitelist"]
                 self.oauth_remotes[provider_name] = obj_provider
 
-    def _init_data_model(self):
+    def _init_data_model(self) -> None:
         user_data_model = SQLAInterface(self.user_model)
         if self.auth_type == const.AUTH_DB:
             self.userdbmodelview.datamodel = user_data_model
@@ -888,8 +894,7 @@ class FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
         elif self.auth_type == const.AUTH_REMOTE_USER:
             self.userremoteusermodelview.datamodel = user_data_model
 
-        if self.userstatschartview:
-            self.userstatschartview.datamodel = user_data_model
+        self.userstatschartview.datamodel = user_data_model
         if self.auth_user_registration:
             self.registerusermodelview.datamodel = SQLAInterface(self.registeruser_model)
 
@@ -899,7 +904,7 @@ class FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
         self.resourcemodelview.datamodel = SQLAInterface(self.resource_model)
         self.permissionmodelview.datamodel = SQLAInterface(self.permission_model)
 
-    def create_db(self):
+    def create_db(self) -> None:
         """
         Create the database.
 
@@ -1052,7 +1057,7 @@ class FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
         for rolename, resource_actions_raw in access_control.items():
             role = self.find_role(rolename)
             if not role:
-                raise AirflowException(
+                raise FabException(
                     f"The access_control mapping for DAG '{dag_id}' includes a role named "
                     f"'{rolename}', but that role does not exist"
                 )
@@ -1066,7 +1071,7 @@ class FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
 
             for resource_name, actions in resource_actions.items():
                 if resource_name not in self.RESOURCE_DETAILS_MAP:
-                    raise AirflowException(
+                    raise FabException(
                         f"The access_control map for DAG '{dag_id}' includes the following invalid "
                         f"resource name: '{resource_name}'; "
                         f"The set of valid resource names is: {self.RESOURCE_DETAILS_MAP.keys()}"
@@ -1078,7 +1083,7 @@ class FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
                 invalid_actions = set(actions) - self.RESOURCE_DETAILS_MAP[resource_name]["actions"]
 
                 if invalid_actions:
-                    raise AirflowException(
+                    raise FabException(
                         f"The access_control map for DAG '{dag_resource_name}' includes "
                         f"the following invalid permissions: {invalid_actions}; "
                         f"The set of valid permissions is: {self.RESOURCE_DETAILS_MAP[resource_name]['actions']}"
@@ -1089,7 +1094,9 @@ class FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
                     if dag_perm:
                         self.add_permission_to_role(role, dag_perm)
 
-    def add_permissions_view(self, base_action_names, resource_name):  # Keep name for compatibility with FAB.
+    def add_permissions_view(
+        self, base_action_names, resource_name
+    ) -> None:  # Keep name for compatibility with FAB.
         """
         Add an action on a resource to the backend.
 
@@ -1114,6 +1121,8 @@ class FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
         else:
             # Permissions on this view exist but....
             admin_role = self.find_role(self.auth_role_admin)
+            if not admin_role:
+                admin_role = self.add_role(self.auth_role_admin)
             for action_name in base_action_names:
                 # Check if base view permissions exist
                 if not self.perms_include_action(perms, action_name):
@@ -1137,7 +1146,7 @@ class FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
                     # Role Admin must have all permissions
                     self.add_permission_to_role(admin_role, perm)
 
-    def add_permissions_menu(self, resource_name):
+    def add_permissions_menu(self, resource_name) -> None:
         """
         Add menu_access to resource on permission_resource.
 
@@ -1150,6 +1159,8 @@ class FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
             perm = self.create_permission("menu_access", resource_name)
         if self.auth_role_admin not in self.builtin_roles:
             role_admin = self.find_role(self.auth_role_admin)
+            if not role_admin:
+                role_admin = self.add_role(self.auth_role_admin)
             self.add_permission_to_role(role_admin, perm)
 
     def sync_roles(self) -> None:
@@ -1204,6 +1215,8 @@ class FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
         perms = [p for p in perms if p.action and p.resource]
 
         admin = self.find_role("Admin")
+        if not admin:
+            admin = self.add_role("Admin")
         admin.permissions = list(set(admin.permissions) | set(perms))
 
         self.session.commit()
@@ -1231,7 +1244,7 @@ class FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
         if deleted_count:
             self.log.info("Deleted %s faulty permissions", deleted_count)
 
-    def perms_include_action(self, perms, action_name):
+    def perms_include_action(self, perms, action_name) -> bool:
         return any(perm.action and perm.action.name == action_name for perm in perms)
 
     def bulk_sync_roles(self, roles: Iterable[dict[str, Any]]) -> None:
@@ -1288,9 +1301,10 @@ class FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
             except Exception as e:
                 log.error(const.LOGMSG_ERR_SEC_ADD_ROLE, e)
                 self.session.rollback()
+                raise FabException(const.LOGMSG_ERR_SEC_ADD_ROLE) from e
         return role
 
-    def find_role(self, name):
+    def find_role(self, name) -> Role | None:
         """
         Find a role in the database.
 
@@ -1298,7 +1312,7 @@ class FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
         """
         return self.session.scalars(select(self.role_model).filter_by(name=name)).unique().one_or_none()
 
-    def get_all_roles(self):
+    def get_all_roles(self) -> list[Role]:
         return self.session.scalars(select(self.role_model)).unique().all()
 
     def delete_role(self, role_name: str) -> None:
@@ -1313,7 +1327,7 @@ class FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
             self.session.execute(delete(Role).where(Role.name == role_name))
             self.session.commit()
         else:
-            raise AirflowException(f"Role named '{role_name}' does not exist")
+            raise FabException(f"Role named '{role_name}' does not exist")
 
     def get_roles_from_keys(self, role_keys: list[str]) -> set[Role]:
         """
@@ -1340,7 +1354,7 @@ class FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
                         )
         return _roles
 
-    def get_public_role(self):
+    def get_public_role(self) -> Role | None:
         return (
             self.session.scalars(select(self.role_model).filter_by(name=self.auth_role_public))
             .unique()
@@ -1363,7 +1377,7 @@ class FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
         password: str = "",
         hashed_password: str = "",
         groups: list[Group] | None = None,
-    ):
+    ) -> User:
         """Create a user."""
         roles: list[Role] = []
         if role:
@@ -1390,22 +1404,24 @@ class FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
         except Exception as e:
             log.error(const.LOGMSG_ERR_SEC_ADD_USER, e)
             self.session.rollback()
-            return False
+            raise FabException(const.LOGMSG_ERR_SEC_ADD_USER) from e
 
-    def load_user(self, pk: int) -> Any | None:
+    def load_user(self, pk: int) -> User | None:
         user = self.get_user_by_id(int(pk))
         if user and user.is_active:
             return user
         return None
 
-    def get_user_by_id(self, pk):
+    def get_user_by_id(self, pk) -> User | None:
         return self.session.get(self.user_model, pk)
 
-    def count_users(self):
+    def count_users(self) -> int:
         """Return the number of users in the database."""
         return self.session.scalar(select(func.count(self.user_model.id)))
 
-    def add_register_user(self, username, first_name, last_name, email, password="", hashed_password=""):
+    def add_register_user(
+        self, username, first_name, last_name, email, password="", hashed_password=""
+    ) -> RegisterUser | None:
         """
         Add a registration request for the user.
 
@@ -1430,16 +1446,10 @@ class FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
             self.session.rollback()
             return None
 
-    def find_user(self, username=None, email=None):
+    def find_user(self, username=None, email=None) -> User | None:
         """Find user by username or email."""
         if username:
             try:
-                if self.auth_username_ci:
-                    return self.session.scalars(
-                        select(self.user_model).where(
-                            func.lower(self.user_model.username) == func.lower(username)
-                        )
-                    ).one_or_none()
                 return self.session.scalars(
                     select(self.user_model).where(
                         func.lower(self.user_model.username) == func.lower(username)
@@ -1447,13 +1457,12 @@ class FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
                 ).one_or_none()
             except MultipleResultsFound:
                 log.error("Multiple results found for user %s", username)
-                return None
         elif email:
             try:
                 return self.session.scalars(select(self.user_model).filter_by(email=email)).one_or_none()
             except MultipleResultsFound:
                 log.error("Multiple results found for user with email %s", email)
-                return None
+        return None
 
     def update_user(self, user: User) -> bool:
         try:
@@ -1466,7 +1475,7 @@ class FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
             return False
         return True
 
-    def del_register_user(self, register_user):
+    def del_register_user(self, register_user) -> bool:
         """
         Delete registration object from database.
 
@@ -1481,10 +1490,10 @@ class FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
             self.session.rollback()
             return False
 
-    def get_all_users(self):
+    def get_all_users(self) -> list[User]:
         return self.session.scalars(select(self.user_model)).all()
 
-    def update_user_auth_stat(self, user, success=True):
+    def update_user_auth_stat(self, user, success=True) -> None:
         """
         Update user authentication stats.
 
@@ -1524,7 +1533,7 @@ class FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
         """
         return self.session.scalars(select(self.action_model).filter_by(name=name)).one_or_none()
 
-    def create_action(self, name):
+    def create_action(self, name) -> Action:
         """
         Add an action to the backend, model action.
 
@@ -1583,7 +1592,7 @@ class FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
         """
         return self.session.scalars(select(self.resource_model).filter_by(name=name)).one_or_none()
 
-    def create_resource(self, name) -> Resource | None:
+    def create_resource(self, name) -> Resource:
         """
         Create a resource with the given name.
 
@@ -1600,6 +1609,7 @@ class FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
             except Exception as e:
                 log.error(const.LOGMSG_ERR_SEC_ADD_VIEWMENU, e)
                 self.session.rollback()
+                raise FabException(const.LOGMSG_ERR_SEC_ADD_VIEWMENU) from e
         return resource
 
     """
@@ -1738,7 +1748,7 @@ class FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
                 self.session.rollback()
 
     @staticmethod
-    def get_user_roles(user=None):
+    def get_user_roles(user=None) -> list[Role]:
         """
         Get all the roles associated with the user.
 
@@ -1755,7 +1765,7 @@ class FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
     --------------------
     """
 
-    def auth_user_ldap(self, username, password, rotate_session_id=True):
+    def auth_user_ldap(self, username, password, rotate_session_id=True) -> User | None:
         """
         Authenticate user with LDAP.
 
@@ -1948,11 +1958,11 @@ class FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
         user = self.find_user(username=username)
         if user is None:
             user = self.find_user(email=username)
-        if user is None:
+        if user is None or user.password is None:
             return False
         return check_password_hash(user.password, password)
 
-    def auth_user_db(self, username, password, rotate_session_id=True):
+    def auth_user_db(self, username, password, rotate_session_id=True) -> User | None:
         """
         Authenticate user, auth db style.
 
@@ -1976,6 +1986,8 @@ class FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
             )
             log.info(LOGMSG_WAR_SEC_LOGIN_FAILED, username)
             return None
+        if user.password is None:
+            return None
         if check_password_hash(user.password, password):
             if rotate_session_id:
                 self._rotate_session_id()
@@ -1985,7 +1997,7 @@ class FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
         log.info(LOGMSG_WAR_SEC_LOGIN_FAILED, username)
         return None
 
-    def set_oauth_session(self, provider, oauth_response):
+    def set_oauth_session(self, provider, oauth_response) -> None:
         """Set the current session with OAuth user secrets."""
         # Get this provider key names for token_key and token_secret
         token_key = self.get_oauth_token_key_name(provider)
@@ -2019,7 +2031,7 @@ class FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
             if _provider["name"] == provider:
                 return _provider.get("token_secret", "oauth_token_secret")
 
-    def auth_user_oauth(self, userinfo):
+    def auth_user_oauth(self, userinfo, rotate_session_id=True) -> User | None:
         """
         Authenticate user with OAuth.
 
@@ -2073,7 +2085,8 @@ class FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
 
         # LOGIN SUCCESS (only if user is now registered)
         if user:
-            self._rotate_session_id()
+            if rotate_session_id:
+                self._rotate_session_id()
             self.update_user_auth_stat(user)
             return user
         return None
@@ -2217,7 +2230,7 @@ class FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
         # decode - if empty string, default to fallback, otherwise take first element
         return raw_value[0].decode("utf-8") or fallback
 
-    def auth_user_remote_user(self, username):
+    def auth_user_remote_user(self, username) -> User | None:
         """
         REMOTE_USER user Authentication.
 
@@ -2253,7 +2266,7 @@ class FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
     ---------------
     """
 
-    def _rotate_session_id(self):
+    def _rotate_session_id(self) -> None:
         """
         Rotate the session ID.
 
@@ -2261,7 +2274,7 @@ class FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
         database session backend.
         """
         if conf.get("fab", "SESSION_BACKEND") == "database":
-            session.sid = str(uuid.uuid4())
+            session.sid = str(uuid.uuid4())  # type: ignore
 
     def _get_microsoft_jwks(self) -> list[dict[str, Any]]:
         import requests
@@ -2273,7 +2286,7 @@ class FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
         if verify_signature:
             from authlib.jose import JsonWebKey, jwt as authlib_jwt
 
-            keyset = JsonWebKey.import_key_set(self._get_microsoft_jwks())
+            keyset = JsonWebKey.import_key_set(self._get_microsoft_jwks())  # type: ignore
             claims = authlib_jwt.decode(id_token, keyset)
             claims.validate()
             return claims
@@ -2415,7 +2428,7 @@ class FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
         except ldap.INVALID_CREDENTIALS:
             return False
 
-    def _ldap_calculate_user_roles(self, user_attributes: dict[str, list[bytes]]) -> list[str]:
+    def _ldap_calculate_user_roles(self, user_attributes: dict[str, list[bytes]]) -> list[Role]:
         user_role_objects = set()
 
         # apply AUTH_ROLES_MAPPING
@@ -2495,7 +2508,7 @@ class FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
         else:
             getattr(log, level)(text.replace("<br>", "\n").replace("<b>", "*").replace("</b>", "*"))
 
-    def _oauth_calculate_user_roles(self, userinfo) -> list[str]:
+    def _oauth_calculate_user_roles(self, userinfo) -> list[Role]:
         user_role_objects = set()
 
         # apply AUTH_ROLES_MAPPING

airflow/providers/fab/auth_manager/views/user.py

@@ -26,6 +26,7 @@ from flask_appbuilder.security.views import (
     UserOAuthModelView,
     UserRemoteUserModelView,
 )
+from wtforms.validators import DataRequired
 
 from airflow.providers.fab.www.security import permissions
 
@@ -186,6 +187,17 @@ class CustomUserDBModelView(MultiResourceUserMixin, UserDBModelView):
         "conf_password",
     ]
 
+    edit_columns = [
+        "first_name",
+        "last_name",
+        "username",
+        "active",
+        "email",
+        "roles",
+    ]
+
+    validators_columns = {"roles": [DataRequired()]}
+
     base_permissions = [
         permissions.ACTION_CAN_CREATE,
         permissions.ACTION_CAN_READ,

airflow/providers/fab/cli/__init__.py

@@ -0,0 +1,18 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations