From c80c261b7208e6846c7bd0492213c6ef0235972b Mon Sep 17 00:00:00 2001
From: vincbeck <vincbeck@amazon.com>
Date: Mon, 10 Mar 2025 11:14:25 -0400
Subject: [PATCH] Add `logout` method in auth manager interface

---
 airflow/api_fastapi/auth/managers/base_auth_manager.py   | 9 +++++++++
 newsfragments/aip-79.significant.rst                     | 2 ++
 .../providers/fab/auth_manager/fab_auth_manager.py       | 5 +++++
 .../tests/unit/fab/auth_manager/test_fab_auth_manager.py | 6 ++++++
 .../api_fastapi/auth/managers/test_base_auth_manager.py  | 3 +++
 5 files changed, 25 insertions(+)

diff --git a/airflow/api_fastapi/auth/managers/base_auth_manager.py b/airflow/api_fastapi/auth/managers/base_auth_manager.py
index 30978e7c518c5..4e98bc56d4117 100644
--- a/airflow/api_fastapi/auth/managers/base_auth_manager.py
+++ b/airflow/api_fastapi/auth/managers/base_auth_manager.py
@@ -108,6 +108,15 @@ def get_jwt_token(
     def get_url_login(self, **kwargs) -> str:
         """Return the login page url."""
 
+    def logout(self) -> None:
+        """
+        Logout the user.
+
+        This method is called when the user is logging out. By default, it does nothing. Override it to
+        invalidate resources when logging out, such as a session.
+        """
+        return None
+
     @abstractmethod
     def is_authorized_configuration(
         self,
diff --git a/newsfragments/aip-79.significant.rst b/newsfragments/aip-79.significant.rst
index 497c2054b1435..b22290a45a083 100644
--- a/newsfragments/aip-79.significant.rst
+++ b/newsfragments/aip-79.significant.rst
@@ -24,6 +24,8 @@ As part of this change the following breaking changes have occurred:
     - ``get_api_endpoints``
     - ``register_views``
 
+  - A new optional method ``logout`` has been added to the interface
+
   - All the following method signatures changed to make the parameter ``user`` required (it was optional)
 
     - ``is_authorized_configuration``
diff --git a/providers/fab/src/airflow/providers/fab/auth_manager/fab_auth_manager.py b/providers/fab/src/airflow/providers/fab/auth_manager/fab_auth_manager.py
index 0c721e713a199..0d04baa911670 100644
--- a/providers/fab/src/airflow/providers/fab/auth_manager/fab_auth_manager.py
+++ b/providers/fab/src/airflow/providers/fab/auth_manager/fab_auth_manager.py
@@ -26,6 +26,7 @@
 from connexion import FlaskApi
 from fastapi import FastAPI
 from flask import Blueprint, g, url_for
+from flask_login import logout_user
 from sqlalchemy import select
 from sqlalchemy.orm import Session, joinedload
 from starlette.middleware.wsgi import WSGIMiddleware
@@ -419,6 +420,10 @@ def get_url_logout(self):
             raise AirflowException("`auth_view` not defined in the security manager.")
         return url_for(f"{self.security_manager.auth_view.endpoint}.logout")
 
+    def logout(self) -> None:
+        """Logout the user."""
+        logout_user()
+
     def register_views(self) -> None:
         self.security_manager.register_views()
 
diff --git a/providers/fab/tests/unit/fab/auth_manager/test_fab_auth_manager.py b/providers/fab/tests/unit/fab/auth_manager/test_fab_auth_manager.py
index a50489b1610d9..deaecc3f2b3b4 100644
--- a/providers/fab/tests/unit/fab/auth_manager/test_fab_auth_manager.py
+++ b/providers/fab/tests/unit/fab/auth_manager/test_fab_auth_manager.py
@@ -575,3 +575,9 @@ def test_get_url_logout(self, mock_url_for, auth_manager_with_appbuilder):
         auth_manager_with_appbuilder.security_manager.auth_view.endpoint = "test_endpoint"
         auth_manager_with_appbuilder.get_url_logout()
         mock_url_for.assert_called_once_with("test_endpoint.logout")
+
+    @pytest.mark.db_test
+    @mock.patch("airflow.providers.fab.auth_manager.fab_auth_manager.logout_user")
+    def test_logout(self, mock_logout_user, auth_manager_with_appbuilder):
+        auth_manager_with_appbuilder.logout()
+        mock_logout_user.assert_called_once()
diff --git a/tests/api_fastapi/auth/managers/test_base_auth_manager.py b/tests/api_fastapi/auth/managers/test_base_auth_manager.py
index 82b1ffbd832aa..53b6a93842d27 100644
--- a/tests/api_fastapi/auth/managers/test_base_auth_manager.py
+++ b/tests/api_fastapi/auth/managers/test_base_auth_manager.py
@@ -149,6 +149,9 @@ def test_get_cli_commands_return_empty_list(self, auth_manager):
     def test_get_fastapi_app_return_none(self, auth_manager):
         assert auth_manager.get_fastapi_app() is None
 
+    def test_logout_return_none(self, auth_manager):
+        assert auth_manager.logout() is None
+
     @patch("airflow.api_fastapi.auth.managers.base_auth_manager.JWTSigner")
     @patch.object(EmptyAuthManager, "deserialize_user")
     def test_get_user_from_token(self, mock_deserialize_user, mock_jwt_signer, auth_manager):