safaehar opened a new pull request, #74856: URL: https://github.com/apache/airflow/pull/74856
Every request to the API server's core app goes through `JWTRefreshMiddleware`, including static UI assets, `/api/v2/monitor/health` and `/api/v2/version`. When the browser sends the `_token` cookie, which it always does once a user is logged in, the middleware decodes the JWT and checks the `revoked_token` table for each of these requests. That check is a synchronous SQLAlchemy query that runs straight on the event loop, inside the async `get_user_from_token`. The middleware is also a `BaseHTTPMiddleware`, which adds a task and a body stream to every request, including large streamed responses like the UI entry bundle. The effect is that requests on one worker wait for each other. The event loop stalls for the duration of every revocation query, and for longer if the sync connection pool is exhausted and the checkout blocks. Observed on a production deployment (4 api-server pods, CPU well under capacity) during a cold-cache load of the home page: - `GET /static/assets/<bundle>.js` took up to ~790 ms server-side, almost all of it inside the Python process. Each static request ran a pool pre-ping, a `SELECT EXISTS (... revoked_token ...)` and a commit, about 10 ms of DB time in total; the rest was queueing. - `GET /static` across the page: p50 ~95 ms, p95 ~380 ms. - The ~15 parallel home page API calls took 250–450 ms each under that burst, vs 70–145 ms when issued with less concurrency. - CPU profiles show `BaseHTTPMiddleware.__call__` → `JWTRefreshMiddleware.dispatch` → SQLAlchemy `Session._execute_internal` on the event loop thread. ### Changes - `JWTRefreshMiddleware` is now a pure ASGI middleware. It adds the `Set-Cookie` headers to `http.response.start` and never touches the response body. - It passes requests under the static mount, plus `/api/v2/monitor/health` and `/api/v2/version`, straight through. None of these routes read the user. Paths are matched relative to the root path, so a `[api] base_url` subpath works. - `refresh_user` runs in the threadpool, because auth managers may do blocking I/O there (Keycloak calls the identity provider when the access token has expired). - `BaseAuthManager.get_user_from_token` runs `RevokedToken.is_revoked` in the threadpool. That covers both the cookie path (middleware) and the bearer token path (`get_user` dependency). Cookie attributes, the trust sentinel on `request.state`, invalid/expired token handling and the stale root-path cookie cleanup are unchanged. ### Behaviour change An expired or invalid `_token` cookie is no longer cleared on responses to static assets, health or version requests. The next request to any other route clears it, as before. ### Risks - Middleware order is unchanged. Inner middlewares that raise `HTTPException` still get a JSON error response, as before. - Revocation lookups now take threadpool slots (anyio default: 40 per worker). With the static and health paths skipped, that is one lookup per authenticated API request, the same number of DB queries as before. ### Related PRs related: #70783, #74722 Both touch `JWTRefreshMiddleware` and will conflict textually with this rewrite, though not in intent. #74722's fail-open handling and #70783's auth-manager hooks both carry over to the ASGI form. I'm happy to rebase onto whichever lands first. ### Testing - `uv run --project airflow-core pytest` on `airflow-core/tests/unit/api_fastapi/auth`, `core_api/test_security.py`, `core_api/routes/public/test_auth.py`, `core_api/routes/public/test_monitor.py`, `core_api/routes/public/test_version.py` and `test_app.py`: 577 passed. - The middleware tests now run against a real ASGI stack (`TestClient`). They cover the skipped paths with and without a root path, static file ETag/304/range responses through the middleware, streaming passthrough, cookie set/clear paths, and both threadpool offloads. The two threadpool tests fail when the change is reverted. - `prek run mypy-airflow-core` and ruff pass. To verify on a running instance: load the UI with an authenticated session and confirm that no `revoked_token` query is issued for `/static/*` or `/api/v2/monitor/health` (e.g. with SQL echo or APM spans). --- ##### Was generative AI tooling used to co-author this PR? - [X] Yes (please specify the tool below) Generated-by: Claude Code (Opus 5.5) following [the guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions) -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
