Make JWTRefreshMiddleware extensible via BaseAuthManager.get_jwt_refresh_middleware()

[stephen-bracken]

closes: #70720

Make JWTRefreshMiddleware extensible via BaseAuthManager.get_jwt_refresh_middleware()

Refactor the initialisation flow for the JWTRefreshMiddleware to give greater control to auth managers over how this middleware gets initialised.

The current behaviour is preserved by setting the BaseAuthManager.use_jwt_middleware attribute to True (default True).

By moving the JWTRefreshMiddleware initialisation to BaseAuthManager.get_jwt_refresh_middleware(), this allows auth managers to override the JWTRefreshMiddleware behaviour by inheriting from it and using the inherited class in get_fastapi_middlewares instead.

By factoring out the _set_new_token() method from dispatch(), this gives an interface for the inherited token refresh middleware to alter the behaviour of the middleware when setting the tokens. The interface has access to the response object to set any cookies as needed.

Changed the _refresh_user() method to accept a request object to allow accessing any cookies or state information from the request.

Also adds the airflow.api_fastapi.app.request_cookie_is_secure() helper to standardise setting HTTP secure cookies

Was generative AI tooling used to co-author this PR?

• Yes (please specify the tool below)
• No

---

[vincbeck]

CI is failing

[pierrejeambrun]

How does this work for old provider version (fab) and new core (removing the hardcoded JWTRefreshMiddleware to move it to BaseAuthManager ?
New fab calls super().get_fastapi_middlewares(), old fab do not so the refresh token middleware disappears?

[stephen-bracken]

How does this work for old provider version (fab) and new core (removing the hardcoded JWTRefreshMiddleware to move it to BaseAuthManager ?
New fab calls super().get_fastapi_middlewares(), old fab do not so the refresh token middleware disappears?

In 3.3.1 and prior versions of airflow the JWTRefreshMiddleware will be injected in the hardcoded initialisation in airflow.api_fastapi.core_api.app.init_middlewares:

airflow/airflow-core/src/airflow/api_fastapi/core_api/app.py

Line 178 in d4581ae

app.add_middleware(JWTRefreshMiddleware)

In Airflow v3.3.0 and v3.3.1 BaseAuthManager.get_fastapi_middlewares() returns a [], but after this change it will return [(JWTRefreshMiddleware,{})].

The super().get_fastapi_middlewares() is there to include any middlewares we add to BaseAuthManager in FabAuthManager. This will need to be included in any auth managers that alter get_fastapi_middlewares() if they also use a JWTRefreshMiddleware. the if AIRFLOW_V_3_3_PLUS guard is there because get_fastapi_middlewares() was added in 3.3.0

[pierrejeambrun]

In your latest reply, you answer "old core <-> new provider", but my question was more about "old provider <-> new core".

I believe this is breaking the because core no longer hardcode JWTRefreshMiddleware, but old provider doesn't call super().get_fastapi_middlewares().

Can you try this setup loacally and report back.

[stephen-bracken]

@pierrejeambrun

I believe this is breaking the because core no longer hardcode JWTRefreshMiddleware, but old provider doesn't call super().get_fastapi_middlewares().

I have added BaseAuthManager.use_jwt_middleware (default True) that will trigger the old refresh behaviour where necessary. This allows auth managers to control the order in which the middleware from get_jwt_refresh_middleware() gets initialised without breaking old auth flows.