From 4669fcff9dd291cd6a449113a0c6e47654ab2c3a Mon Sep 17 00:00:00 2001 From: Stephen Bracken Date: Mon, 3 Aug 2026 11:52:59 +0100 Subject: [PATCH] Disable separate Keycloak JWTs for Airflow v3.2- --- providers/keycloak/docs/changelog.rst | 12 +-- .../auth_manager/keycloak_auth_manager.py | 16 ++- .../keycloak/auth_manager/routes/login.py | 44 +++++---- .../providers/keycloak/version_compat.py | 1 + .../auth_manager/routes/test_login.py | 6 +- .../test_keycloak_auth_manager.py | 98 +++++++++++++++---- 6 files changed, 130 insertions(+), 47 deletions(-) diff --git a/providers/keycloak/docs/changelog.rst b/providers/keycloak/docs/changelog.rst index 35f18bbee0b93..c208d749903c6 100644 --- a/providers/keycloak/docs/changelog.rst +++ b/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 ~~~~~~~~ diff --git a/providers/keycloak/src/airflow/providers/keycloak/auth_manager/keycloak_auth_manager.py b/providers/keycloak/src/airflow/providers/keycloak/auth_manager/keycloak_auth_manager.py index 1bb0c9bda1b17..ce0fc9338076a 100644 --- a/providers/keycloak/src/airflow/providers/keycloak/auth_manager/keycloak_auth_manager.py +++ b/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 diff --git a/providers/keycloak/src/airflow/providers/keycloak/auth_manager/routes/login.py b/providers/keycloak/src/airflow/providers/keycloak/auth_manager/routes/login.py index 804733757c4f2..0728518ab3c0f 100644 --- a/providers/keycloak/src/airflow/providers/keycloak/auth_manager/routes/login.py +++ b/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 diff --git a/providers/keycloak/src/airflow/providers/keycloak/version_compat.py b/providers/keycloak/src/airflow/providers/keycloak/version_compat.py index e97df79b887d4..d87042a9edbb0 100644 --- a/providers/keycloak/src/airflow/providers/keycloak/version_compat.py +++ b/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) diff --git a/providers/keycloak/tests/unit/keycloak/auth_manager/routes/test_login.py b/providers/keycloak/tests/unit/keycloak/auth_manager/routes/test_login.py index 058bda6590d89..dcfbb3202aebb 100644 --- a/providers/keycloak/tests/unit/keycloak/auth_manager/routes/test_login.py +++ b/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): diff --git a/providers/keycloak/tests/unit/keycloak/auth_manager/test_keycloak_auth_manager.py b/providers/keycloak/tests/unit/keycloak/auth_manager/test_keycloak_auth_manager.py index 44ab60304bdae..d20c72e1aea5f 100644 --- a/providers/keycloak/tests/unit/keycloak/auth_manager/test_keycloak_auth_manager.py +++ b/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,26 +134,47 @@ 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( @@ -157,10 +182,19 @@ def test_serialize_user(self, auth_manager): 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(