Restructure Execution API security to better use FastAPI's Security scopes#62582
Restructure Execution API security to better use FastAPI's Security scopes#62582
Conversation
…copes
Before this change, `JWTBearer` in deps.py does everything: crypto validation,
sub-claim matching, and it runs twice per request on ti:self routes because
FastAPI includes scopes in dependency cache keys for `HTTPBearer` subclasses,
defeating dedup.
In a PR that is already created (but not yet merged) we want per-endpoint
token type policies (e.g. the /run endpoint will need to accept workload
tokens while other routes stay execution-only). This changes is the
"foundation" that enables that to work in a nice clear fashion
`SecurityScopes` can't express this directly because FastAPI resolves outer
router deps before inner ones -- a `token:workload` scope on an endpoint needs
to *relax* the default restriction, but `SecurityScopes` only accumulate
additively.
The fix is a new security.py with a three-layer split:
- `JWTBearer` (`_jwt_bearer`) now does only crypto validation and caches the
result on the ASGI request scope. It never looks at scopes or token types.
- `require_auth` is a plain function (not an `HTTPBearer` subclass) used via
`Security(require_auth)` on routers. Because plain functions have
`_uses_scopes=False` in FastAPI's dependency system, `_jwt_bearer` (its
sub-dep) deduplicates correctly across multiple Security resolutions. It
enforces `ti:self` via `SecurityScopes` and reads allowed token types from
the matched route object.
- `ExecutionAPIRoute` is a custom `APIRoute` subclass that precomputes
`allowed_token_types` from `token:*` Security scopes at route registration
time — after `include_router` has merged all parent and child
dependencies. This sidesteps the resolution ordering problem entirely.
To opt a route into workload tokens, it's now a one-liner:
```python
@ti_id_router.patch(
"/{task_instance_id}/run",
dependencies=[Security(require_auth, scopes=["token:execution", "token:workload"])],
)
```
Nothing uses the workload-scoped tokens just yet -- this PR lays the
foundation; a follow-up PR will add token:workload to /run.
Also cleaned up the module boundaries: security.py owns all auth-related deps
(CurrentTIToken, get_team_name_dep, require_auth); deps.py is just the svcs
DepContainer. Renamed JWTBearerDep to CurrentTIToken to match the FastAPI
current_user convention.
I tried _lots_ of different approaches to get this merge/override behaviour,
and the cleanest was a custom route class
There was a problem hiding this comment.
Pull request overview
This PR restructures the Execution API security architecture to properly leverage FastAPI's Security scopes mechanism, laying the foundation for PR #60108 which will add two-token support (workload + execution tokens) to prevent token expiration while tasks wait in executor queues.
Purpose: The previous implementation had JWTBearer handling both crypto validation AND authorization logic, running twice per request on ti:self routes due to FastAPI's dependency cache keys. The new three-layer architecture separates concerns: JWTBearer (crypto only, cached on ASGI request scope), require_auth (route-level authorization via Security scopes), and ExecutionAPIRoute (precomputes allowed token types at route registration time).
Changes:
- Introduced three-layer security architecture with
JWTBearer,require_auth, andExecutionAPIRoutein new security.py module - Cleaned up deps.py to only contain
DepContainer, moving all auth logic to security.py - Updated all route files to use
CurrentTITokeninstead of deprecatedJWTBearerDep/JWTBearerTIPathDep - Updated test fixtures to mock
_jwt_bearerdirectly with proper scope claims - Added comprehensive test coverage for token type validation (workload vs execution scopes)
Reviewed changes
Copilot reviewed 16 out of 16 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| airflow-core/src/airflow/api_fastapi/execution_api/security.py | New module implementing three-layer security: JWTBearer for crypto validation, require_auth for authorization, ExecutionAPIRoute for precomputing allowed token types |
| airflow-core/src/airflow/api_fastapi/execution_api/deps.py | Cleaned up to only contain DepContainer for dependency injection; all auth code moved to security.py |
| airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py | Updated ti_id_router to use ExecutionAPIRoute and Security(require_auth, scopes=["ti:self"]) |
| airflow-core/src/airflow/api_fastapi/execution_api/routes/init.py | Updated authenticated_router to use Security(require_auth) wrapper |
| airflow-core/src/airflow/api_fastapi/execution_api/routes/{connections,variables,xcoms,health}.py | Updated imports from JWTBearerDep to CurrentTIToken |
| airflow-core/src/airflow/api_fastapi/execution_api/app.py | Updated dependency overrides to use _jwt_bearer; added OpenAPI x-airflow-* field cleanup |
| airflow-core/tests/unit/api_fastapi/execution_api/conftest.py | Refactored client fixture to mock _jwt_bearer directly with TIToken including scope claims |
| airflow-core/tests/unit/api_fastapi/execution_api/test_security.py | Added comprehensive unit tests for ExecutionAPIRoute and token type enforcement |
| airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py | Added TestTokenTypeValidation class with tests for workload/execution scope enforcement |
| airflow-core/tests/unit/api_fastapi/execution_api/versions/head/{test_xcoms,test_variables}.py | Updated test fixtures to use CurrentTIToken instead of JWTBearerDep |
| airflow-core/tests/unit/api_fastapi/execution_api/test_app.py | Added test_ti_self_routes_have_task_instance_id_param to validate ti:self scope convention |
| airflow-core/src/airflow/api_fastapi/execution_api/AGENTS.md | Added documentation pointer to security.py for token scope infrastructure |
airflow-core/src/airflow/api_fastapi/execution_api/routes/health.py
Outdated
Show resolved
Hide resolved
airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py
Outdated
Show resolved
Hide resolved
|
Thanks for this, Ash, learned a lot going through the three-layer split. I'll rebase #60108 on top once this lands. |
|
cc @o-nikolas (Just for the moved import for |
Ack, thanks for the ping. Also looping in @vincbeck |
amoghrajesh
left a comment
There was a problem hiding this comment.
Looking at the code change and the tests, it looks fine to me.
pierrejeambrun
left a comment
There was a problem hiding this comment.
This broke main. It might have been behind compared to #62886
…copes (apache#62582) Before this change, `JWTBearer` in deps.py does everything: crypto validation, sub-claim matching, and it runs twice per request on ti:self routes because FastAPI includes scopes in dependency cache keys for `HTTPBearer` subclasses, defeating dedup. In a PR that is already created (but not yet merged) we want per-endpoint token type policies (e.g. the /run endpoint will need to accept workload tokens while other routes stay execution-only). This changes is the "foundation" that enables that to work in a nice clear fashion `SecurityScopes` can't express this directly because FastAPI resolves outer router deps before inner ones -- a `token:workload` scope on an endpoint needs to *relax* the default restriction, but `SecurityScopes` only accumulate additively. The fix is a new security.py with a three-layer split: - `JWTBearer` (`_jwt_bearer`) now does only crypto validation and caches the result on the ASGI request scope. It never looks at scopes or token types. - `require_auth` is a plain function (not an `HTTPBearer` subclass) used via `Security(require_auth)` on routers. Because plain functions have `_uses_scopes=False` in FastAPI's dependency system, `_jwt_bearer` (its sub-dep) deduplicates correctly across multiple Security resolutions. It enforces `ti:self` via `SecurityScopes` and reads allowed token types from the matched route object. - `ExecutionAPIRoute` is a custom `APIRoute` subclass that precomputes `allowed_token_types` from `token:*` Security scopes at route registration time — after `include_router` has merged all parent and child dependencies. This sidesteps the resolution ordering problem entirely. To opt a route into workload tokens, it's now a one-liner: ```python @ti_id_router.patch( "/{task_instance_id}/run", dependencies=[Security(require_auth, scopes=["token:execution", "token:workload"])], ) ``` Nothing uses the workload-scoped tokens just yet -- this PR lays the foundation; a follow-up PR will add token:workload to /run. Also cleaned up the module boundaries: security.py owns all auth-related deps (CurrentTIToken, get_team_name_dep, require_auth); deps.py is just the svcs DepContainer. Renamed JWTBearerDep to CurrentTIToken to match the FastAPI current_user convention. I tried _lots_ of different approaches to get this merge/override behaviour, and the cleanest was a custom route class
…copes (apache#62582) Before this change, `JWTBearer` in deps.py does everything: crypto validation, sub-claim matching, and it runs twice per request on ti:self routes because FastAPI includes scopes in dependency cache keys for `HTTPBearer` subclasses, defeating dedup. In a PR that is already created (but not yet merged) we want per-endpoint token type policies (e.g. the /run endpoint will need to accept workload tokens while other routes stay execution-only). This changes is the "foundation" that enables that to work in a nice clear fashion `SecurityScopes` can't express this directly because FastAPI resolves outer router deps before inner ones -- a `token:workload` scope on an endpoint needs to *relax* the default restriction, but `SecurityScopes` only accumulate additively. The fix is a new security.py with a three-layer split: - `JWTBearer` (`_jwt_bearer`) now does only crypto validation and caches the result on the ASGI request scope. It never looks at scopes or token types. - `require_auth` is a plain function (not an `HTTPBearer` subclass) used via `Security(require_auth)` on routers. Because plain functions have `_uses_scopes=False` in FastAPI's dependency system, `_jwt_bearer` (its sub-dep) deduplicates correctly across multiple Security resolutions. It enforces `ti:self` via `SecurityScopes` and reads allowed token types from the matched route object. - `ExecutionAPIRoute` is a custom `APIRoute` subclass that precomputes `allowed_token_types` from `token:*` Security scopes at route registration time — after `include_router` has merged all parent and child dependencies. This sidesteps the resolution ordering problem entirely. To opt a route into workload tokens, it's now a one-liner: ```python @ti_id_router.patch( "/{task_instance_id}/run", dependencies=[Security(require_auth, scopes=["token:execution", "token:workload"])], ) ``` Nothing uses the workload-scoped tokens just yet -- this PR lays the foundation; a follow-up PR will add token:workload to /run. Also cleaned up the module boundaries: security.py owns all auth-related deps (CurrentTIToken, get_team_name_dep, require_auth); deps.py is just the svcs DepContainer. Renamed JWTBearerDep to CurrentTIToken to match the FastAPI current_user convention. I tried _lots_ of different approaches to get this merge/override behaviour, and the cleanest was a custom route class
Before this change,
JWTBearerin deps.py does everything: crypto validation, sub-claim matching, and it runs twice per request on ti:self routes because FastAPI includes scopes in dependency cache keys forHTTPBearersubclasses, defeating dedup.In #60108 we want per-endpoint token type policies (e.g. the /run endpoint will need to accept workload tokens while other routes stay execution-only). This changes is the "foundation" that enables that to work in a nice clear fashion
SecurityScopescan't express this directly because FastAPI resolves outer router deps before inner ones -- atoken:workloadscope on an endpoint needs to relax the default restriction, butSecurityScopesonly accumulate additively.The fix is a new security.py with a three-layer split:
JWTBearer(_jwt_bearer) now does only crypto validation and caches the result on the ASGI request scope. It never looks at scopes or token types.require_authis a plain function (not anHTTPBearersubclass) used viaSecurity(require_auth)on routers. Because plain functions have_uses_scopes=Falsein FastAPI's dependency system,_jwt_bearer(its sub-dep) deduplicates correctly across multiple Security resolutions. It enforcesti:selfviaSecurityScopesand reads allowed token types from the matched route object.ExecutionAPIRouteis a customAPIRoutesubclass that precomputesallowed_token_typesfromtoken:*Security scopes at route registration time — afterinclude_routerhas merged all parent and child dependencies. This sidesteps the resolution ordering problem entirely.To opt a route into workload tokens, it's now a one-liner:
Nothing uses the workload-scoped tokens just yet -- this PR lays the foundation; a follow-up PR will add token:workload to /run.
Also cleaned up the module boundaries: security.py owns all auth-related deps (CurrentTIToken, get_team_name_dep, require_auth); deps.py is just the svcs DepContainer. Renamed JWTBearerDep to CurrentTIToken to match the FastAPI current_user convention.
I tried lots of different approaches to get this merge/override behaviour, and the cleanest was a custom route class
Was generative AI tooling used to co-author this PR?
{pr_number}.significant.rstor{issue_number}.significant.rst, in airflow-core/newsfragments.