Disable separate Keycloak JWTs for Airflow v3.2-

providers/keycloak/docs/changelog.rst

@@ -29,12 +29,12 @@ Changelog
 .....
 
 .. note::
-    Keycloak's access and refresh tokens are now stored in dedicated ``_access_token`` and
-    ``_refresh_token`` cookies instead of being carried inside the Airflow JWT claims. Sessions
-    established before this release carry the tokens in the old form, so the first request after
-    the upgrade cannot be refreshed and the session is cleared -- **every logged-in user is signed
-    out once when you upgrade**. No action is required beyond logging back in; this is a one-time
-    effect of the move and does not recur.
+    For Airflow v3.3.0+ Keycloak's access and refresh tokens are now stored in dedicated
+    ``_access_token`` and ``_refresh_token`` cookies instead of being carried inside the
+    Airflow JWT claims. Sessions established before this release carry the tokens in the old form,
+    so the first request after the upgrade cannot be refreshed and the session is cleared --
+    **every logged-in user is signed out once when you upgrade**. No action is required
+    beyond logging back in; this is a one-time effect of the move and does not recur.
 
 Features
 ~~~~~~~~

providers/keycloak/src/airflow/providers/keycloak/auth_manager/keycloak_auth_manager.py

@@ -64,6 +64,7 @@
 )
 from airflow.providers.keycloak.auth_manager.resources import KeycloakResource
 from airflow.providers.keycloak.auth_manager.user import KeycloakAuthManagerUser
+from airflow.providers.keycloak.version_compat import AIRFLOW_V_3_3_PLUS
 from airflow.utils.helpers import prune_dict
 
 if TYPE_CHECKING:
@@ -139,13 +140,24 @@ def http_session(self) -> requests.Session:
 
     def deserialize_user(self, token: dict[str, Any]) -> KeycloakAuthManagerUser:
         return KeycloakAuthManagerUser(
-            user_id=token["user_id"], name=token["name"], access_token="", refresh_token=None
+            user_id=token["user_id"],
+            name=token["name"],
+            access_token=token.get("access_token", ""),
+            refresh_token=token.get("refresh_token"),
         )
 
     def serialize_user(self, user: KeycloakAuthManagerUser) -> dict[str, Any]:
+        if AIRFLOW_V_3_3_PLUS:
+            # Omit Keycloak JWTs from claims, they are stored in separate cookies
+            return {
+                "user_id": user.get_id(),
+                "name": user.get_name(),
+            }
         return {
             "user_id": user.get_id(),
             "name": user.get_name(),
+            "access_token": user.access_token,
+            "refresh_token": user.refresh_token,
         }
 
     async def get_user_from_token(
@@ -159,6 +171,8 @@ async def get_user_from_token(
         :param refresh_token: Keycloak refresh JWT
         """
         user = cast("KeycloakAuthManagerUser", await super().get_user_from_token(token))
+        if not AIRFLOW_V_3_3_PLUS:
+            return user
         if access_token:
             user.access_token = access_token
             user.refresh_token = refresh_token

providers/keycloak/src/airflow/providers/keycloak/auth_manager/routes/login.py

@@ -28,7 +28,11 @@
 
 from airflow.api_fastapi.app import AUTH_MANAGER_FASTAPI_APP_PREFIX, get_auth_manager
 from airflow.api_fastapi.auth.managers.base_auth_manager import COOKIE_NAME_JWT_TOKEN
-from airflow.providers.keycloak.version_compat import AIRFLOW_V_3_1_1_PLUS, AIRFLOW_V_3_1_8_PLUS
+from airflow.providers.keycloak.version_compat import (
+    AIRFLOW_V_3_1_1_PLUS,
+    AIRFLOW_V_3_1_8_PLUS,
+    AIRFLOW_V_3_3_PLUS,
+)
 
 if AIRFLOW_V_3_1_8_PLUS:
     from airflow.api_fastapi.app import get_cookie_path
@@ -139,13 +143,14 @@ def login_callback(request: Request):
         COOKIE_NAME_ID_TOKEN, tokens["id_token"], path=cookie_path, secure=secure, httponly=True
     )
 
-    response.set_cookie(
-        COOKIE_NAME_ACCESS_TOKEN, tokens["access_token"], path=cookie_path, secure=secure, httponly=True
-    )
+    if AIRFLOW_V_3_3_PLUS:
+        response.set_cookie(
+            COOKIE_NAME_ACCESS_TOKEN, tokens["access_token"], path=cookie_path, secure=secure, httponly=True
+        )
 
-    response.set_cookie(
-        COOKIE_NAME_REFRESH_TOKEN, tokens["refresh_token"], path=cookie_path, secure=secure, httponly=True
-    )
+        response.set_cookie(
+            COOKIE_NAME_REFRESH_TOKEN, tokens["refresh_token"], path=cookie_path, secure=secure, httponly=True
+        )
 
     return response
 
@@ -197,16 +202,17 @@ def logout_callback(request: Request):
         secure=secure,
         httponly=True,
     )
-    response.delete_cookie(
-        key=COOKIE_NAME_ACCESS_TOKEN,
-        path=cookie_path,
-        secure=secure,
-        httponly=True,
-    )
-    response.delete_cookie(
-        key=COOKIE_NAME_REFRESH_TOKEN,
-        path=cookie_path,
-        secure=secure,
-        httponly=True,
-    )
+    if AIRFLOW_V_3_3_PLUS:
+        response.delete_cookie(
+            key=COOKIE_NAME_ACCESS_TOKEN,
+            path=cookie_path,
+            secure=secure,
+            httponly=True,
+        )
+        response.delete_cookie(
+            key=COOKIE_NAME_REFRESH_TOKEN,
+            path=cookie_path,
+            secure=secure,
+            httponly=True,
+        )
     return response

providers/keycloak/src/airflow/providers/keycloak/version_compat.py

@@ -35,3 +35,4 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]:
 AIRFLOW_V_3_1_1_PLUS = get_base_airflow_version_tuple() >= (3, 1, 1)
 AIRFLOW_V_3_1_7_PLUS = get_base_airflow_version_tuple() >= (3, 1, 7)
 AIRFLOW_V_3_1_8_PLUS = get_base_airflow_version_tuple() >= (3, 1, 8)
+AIRFLOW_V_3_3_PLUS = get_base_airflow_version_tuple() >= (3, 3, 0)

providers/keycloak/tests/unit/keycloak/auth_manager/routes/test_login.py

@@ -23,6 +23,7 @@
 from airflow.api_fastapi.app import AUTH_MANAGER_FASTAPI_APP_PREFIX
 
 from tests_common.test_utils.config import conf_vars
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS
 
 
 class TestLoginRouter:
@@ -79,8 +80,9 @@ def test_login_callback(self, mock_get_keycloak_client, mock_get_auth_manager, c
         assert "_token" in response.cookies
         assert response.cookies["_token"] == token
         assert response.cookies["_id_token"] == "id_token"
-        assert response.cookies["_access_token"] == "access_token"
-        assert response.cookies["_refresh_token"] == "refresh_token"
+        if AIRFLOW_V_3_3_PLUS:
+            assert response.cookies["_access_token"] == "access_token"
+            assert response.cookies["_refresh_token"] == "refresh_token"
 
     @patch("airflow.providers.keycloak.auth_manager.routes.login.KeycloakAuthManager.get_keycloak_client")
     def test_login_sets_secure_state_cookie_behind_tls_proxy(self, mock_get_keycloak_client, client):

providers/keycloak/tests/unit/keycloak/auth_manager/test_keycloak_auth_manager.py

@@ -40,7 +40,11 @@
 )
 
 from tests_common.test_utils.config import conf_vars
-from tests_common.test_utils.version_compat import AIRFLOW_V_3_1_7_PLUS, AIRFLOW_V_3_2_PLUS
+from tests_common.test_utils.version_compat import (
+    AIRFLOW_V_3_1_7_PLUS,
+    AIRFLOW_V_3_2_PLUS,
+    AIRFLOW_V_3_3_PLUS,
+)
 
 if AIRFLOW_V_3_1_7_PLUS:
     from airflow.api_fastapi.auth.managers.exceptions import AuthManagerRefreshTokenExpiredException
@@ -130,37 +134,67 @@ def _clear_filter_cache():
 
 class TestKeycloakAuthManager:
     @pytest.mark.parametrize(
-        "token_data",
+        ("token_data", "exp"),
         [
-            {
-                "user_id": "user_id",
-                "name": "name",
-            },
-            {
-                "user_id": "user_id",
-                "name": "name",
-                "access_token": "access_token",
-                "refresh_token": "refresh_token",
-            },
+            pytest.param(
+                {
+                    "user_id": "user_id",
+                    "name": "name",
+                },
+                {
+                    "user_id": "user_id",
+                    "name": "name",
+                    "access_token": "",
+                    "refresh_token": None,
+                },
+                marks=pytest.mark.skipif(
+                    not AIRFLOW_V_3_3_PLUS, reason="Uses KeycloakJWTMiddleware and separate cookies"
+                ),
+                id="no-Keycloak-tokens",
+            ),
+            pytest.param(
+                {
+                    "user_id": "user_id",
+                    "name": "name",
+                    "access_token": "access_token",
+                    "refresh_token": "refresh_token",
+                },
+                {
+                    "user_id": "user_id",
+                    "name": "name",
+                    "access_token": "access_token",
+                    "refresh_token": "refresh_token",
+                },
+                id="with-Keycloak-tokens",
+            ),
         ],
     )
-    def test_deserialize_user(self, auth_manager, token_data):
+    def test_deserialize_user(self, auth_manager, token_data, exp):
         result = auth_manager.deserialize_user(token_data)
-        assert result.user_id == "user_id"
-        assert result.name == "name"
-        assert result.access_token == ""
-        assert result.refresh_token is None
+        assert result.user_id == exp["user_id"]
+        assert result.name == exp["name"]
+        assert result.access_token == exp["access_token"]
+        assert result.refresh_token == exp["refresh_token"]
 
     def test_serialize_user(self, auth_manager):
         result = auth_manager.serialize_user(
             KeycloakAuthManagerUser(
                 user_id="user_id", name="name", access_token="access_token", refresh_token="refresh_token"
             )
         )
-        assert result == {"user_id": "user_id", "name": "name"}
+        if AIRFLOW_V_3_3_PLUS:
+            assert result == {"user_id": "user_id", "name": "name"}
+        else:
+            assert result == {
+                "user_id": "user_id",
+                "name": "name",
+                "access_token": "access_token",
+                "refresh_token": "refresh_token",
+            }
 
+    @pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="Uses KeycloakJWTMiddleware and separate cookies")
     @pytest.mark.asyncio
-    async def test_get_user_from_token(self, auth_manager):
+    async def test_get_user_from_token_with_keycloak_tokens(self, auth_manager):
         mock_get_user_from_token = AsyncMock(
             return_value=KeycloakAuthManagerUser(
                 user_id="user_id", name="name", access_token="", refresh_token=None
@@ -180,6 +214,31 @@ async def test_get_user_from_token(self, auth_manager):
         assert user.access_token == "access_token"
         assert user.refresh_token == "refresh_token"
 
+    @pytest.mark.skipif(AIRFLOW_V_3_3_PLUS, reason="Testing Old Keycloak JWT flow.")
+    @pytest.mark.asyncio
+    async def test_get_user_from_token(self, auth_manager):
+        mock_token_validator = Mock()
+        mock_get_token_validator = Mock(return_value=mock_token_validator)
+        mock_token_validator.avalidated_claims = AsyncMock(
+            return_value=dict(
+                user_id="user_id", name="name", access_token="access_token", refresh_token="refresh_token"
+            )
+        )
+        with (
+            patch.object(
+                KeycloakAuthManager,
+                "_get_token_validator",
+                mock_get_token_validator,
+            ),
+        ):
+            user = await auth_manager.get_user_from_token("token")
+        mock_token_validator.avalidated_claims.assert_called_with("token")
+        assert user.get_id() == "user_id"
+        assert user.get_name() == "name"
+        assert user.access_token == "access_token"
+        assert user.refresh_token == "refresh_token"
+
+    @pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="Uses KeycloakJWTMiddleware and separate cookies")
     @pytest.mark.asyncio
     async def test_get_user_from_token_keycloak_jwts_missing(self, auth_manager):
         mock_get_user_from_token = AsyncMock(
@@ -196,6 +255,7 @@ async def test_get_user_from_token_keycloak_jwts_missing(self, auth_manager):
         ):
             assert await auth_manager.get_user_from_token("token") is None
 
+    @pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="Uses KeycloakJWTMiddleware and separate cookies")
     @pytest.mark.asyncio
     async def test_get_user_from_token_keycloak_jwt(self, auth_manager):
         mock_get_user_from_token = AsyncMock(