Make JWTRefreshMiddleware extensible via BaseAuthManager._get_jwt_refresh_middleware()

[stephen-bracken]

closes: #70720

Make JWTRefreshMiddleware extensible via BaseAuthManager.get_jwt_refresh_middleware()

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]

I would keep the PR focused around the target issue.

That's also updating/refactoring a whole bunch, making the PR harder to review and surface to test bigger.

1. Moves generate_jwt from before call_next to inside _set_new_token after call_next — introduces the fail-hard-vs-fail-soft semantic
2. Reorganizes the outer try scope so the "if current_token" gate lives inside _refresh_user — makes the flow harder to read for anyone tracing "when is _refresh_user called".
3. Changes the outer condition from if new_token is not None to if new_user or new_token is not None. Under a minimal refactor, the original condition
   still works.
4. Introduces the if new_user: else new_token = "" dead branch inside _set_new_token — dead because dispatch never calls it with new_user=None, but present because the extraction was over-scoped.
5. Swaps delete_cookie for set_cookie(max_age=0) — drops the expires=0 attribute. Small, but again: not needed for extensibility.
6. @classmethod async def on _set_new_token despite using neither cls nor await.

[stephen-bracken]

@pierrejeambrun

1. Moves generate_jwt from before call_next to inside _set_new_token after call_next — introduces the fail-hard-vs-fail-soft semantic

When call_next is called, the request.state.user attributes should be populated in the same way, allowing auth to take place via get_user(). If the user model was refreshed, the JWT cookies will be set on the response object from call_next like they normally would.

2. Reorganizes the outer try scope so the "if current_token" gate lives inside _refresh_user — makes the flow harder to read for anyone tracing "when is _refresh_user called".

Restored the if current_token gate

3. Changes the outer condition from if new_token is not None to if new_user or new_token is not None. Under a minimal refactor, the original condition
   still works.

In the new flow, new_token should only ever be None or "". The first condition checks whether we need to set any cookies and calls the necessary get_cookie_path and request_cookie_is_secure helpers, then clears the JWT cookie if new_token == "", or passes the user model and response to _set_new_token() to set the cookies if it is populated.

4. Introduces the if new_user: else new_token = "" dead branch inside _set_new_token — dead because dispatch never calls it with new_user=None, but present because the extraction was over-scoped.

Removed dead branch

5. Swaps delete_cookie for set_cookie(max_age=0) — drops the expires=0 attribute. Small, but again: not needed for extensibility.

Restored delete_cookie

6. @classmethod async def on _set_new_token despite using neither cls nor await.

Swapped to @staticmethod async def to match _refresh_user()

[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