Skip to content

Restructure Execution API security to better use FastAPI's Security scopes#62582

Merged
ashb merged 4 commits intomainfrom
execution-api-scope-infra
Mar 10, 2026
Merged

Restructure Execution API security to better use FastAPI's Security scopes#62582
ashb merged 4 commits intomainfrom
execution-api-scope-infra

Conversation

@ashb
Copy link
Member

@ashb ashb commented Feb 27, 2026

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 #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

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:

@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


Was generative AI tooling used to co-author this PR?
  • Yes (please specify the tool below)

  • Read the Pull Request Guidelines for more information. Note: commit author/co-author name and email in commits become permanently public when merged.
  • For fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
  • When adding dependency, check compliance with the ASF 3rd Party License Policy.
  • For significant user-facing changes create newsfragment: {pr_number}.significant.rst or {issue_number}.significant.rst, in airflow-core/newsfragments.

…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
Copy link
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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, and ExecutionAPIRoute in 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 CurrentTIToken instead of deprecated JWTBearerDep/JWTBearerTIPathDep
  • Updated test fixtures to mock _jwt_bearer directly 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

@anishgirianish
Copy link
Contributor

Thanks for this, Ash, learned a lot going through the three-layer split. I'll rebase #60108 on top once this lands.

@ashb
Copy link
Member Author

ashb commented Feb 27, 2026

cc @o-nikolas (Just for the moved import for get_team_name going forward)

Copy link
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 15 out of 15 changed files in this pull request and generated 3 comments.

@o-nikolas
Copy link
Contributor

cc @o-nikolas (Just for the moved import for get_team_name going forward)

Ack, thanks for the ping. Also looping in @vincbeck

Copy link
Contributor

@ephraimbuddy ephraimbuddy left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

Copy link
Contributor

@amoghrajesh amoghrajesh left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looking at the code change and the tests, it looks fine to me.

@ashb ashb merged commit df4cb30 into main Mar 10, 2026
129 checks passed
@ashb ashb deleted the execution-api-scope-infra branch March 10, 2026 13:33
Copy link
Member

@pierrejeambrun pierrejeambrun left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This broke main. It might have been behind compared to #62886

dominikhei pushed a commit to dominikhei/airflow that referenced this pull request Mar 11, 2026
…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
Pyasma pushed a commit to Pyasma/airflow that referenced this pull request Mar 13, 2026
…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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:API Airflow's REST/HTTP API area:task-sdk

Projects

None yet

Development

Successfully merging this pull request may close these issues.

8 participants