kaxil commented on code in PR #60108:
URL: https://github.com/apache/airflow/pull/60108#discussion_r3131873700
##########
airflow-core/src/airflow/api_fastapi/auth/tokens.py:
##########
@@ -418,6 +418,10 @@ class JWTGenerator:
kid: str = attrs.field(default=attrs.Factory(_generate_kid,
takes_self=True))
valid_for: float
+ workload_valid_for: float = attrs.field(
Review Comment:
`JWTGenerator` is used by auth managers (FAB, Simple) via
`BaseAuthManager._get_token_signer`, not just the execution API. Adding a
class-level field that unconditionally reads `[scheduler] task_queued_timeout`
at construction couples a general auth primitive to scheduler config. In an
api-server-only or auth-manager-only deployment, this pulls a key that's
irrelevant to the caller and potentially absent.
Can you move `workload_valid_for` out of `JWTGenerator` and compute it at
the workload-generation site (e.g., in `BaseWorkloadSchema.generate_token`)?
`generate()` already accepts a per-call `valid_for` override, so the generator
doesn't need to know about this concept.
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py:
##########
@@ -230,6 +238,17 @@ def ti_run(
last_heartbeat_at=timezone.utcnow(),
)
+ execution_token: str | None = None
Review Comment:
Token minting runs before the state-transition executes. If the state update
fails (SQLAlchemyError at the `session.execute(query)` below, or the client
hits the 409 path further down), you've minted a token that's thrown away.
More importantly, in the duplicate-request branch earlier in this handler
(where we return early for a TI already in RUNNING), a workload-token call
still mints a new execution token even though the worker already has one from
its first request.
Can this block move down to just before `return context` so it only fires on
the success path?
##########
airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py:
##########
@@ -281,6 +284,54 @@ def test_ti_run_state_to_running(
)
assert response.status_code == 409
+ def test_ti_run_returns_execution_token(
+ self, client, exec_app, session, create_task_instance, time_machine
+ ):
+ """PATCH /run with a workload token should swap to an execution-scoped
token."""
+ instant = timezone.parse("2024-10-31T12:00:00Z")
+ time_machine.move_to(instant, tick=False)
+
+ ti = create_task_instance(
+ task_id="test_exec_token",
+ state=State.QUEUED,
+ dagrun_state=DagRunState.RUNNING,
+ session=session,
+ start_date=instant,
+ dag_id=str(uuid4()),
+ )
+ session.commit()
+
+ mock_gen = mock.MagicMock(spec=JWTGenerator)
+ mock_gen.generate.return_value = "mock-execution-token"
+ lifespan.registry.register_value(JWTGenerator, mock_gen)
Review Comment:
`lifespan.registry.register_value(JWTGenerator, mock_gen)` registers into
module-level state that persists past this test. I don't see a teardown that
removes it, so later tests in the session will see this mock generator instead
of the real one -- that's a hidden test-pollution source.
Same pattern exists in `TestTokenTypeValidation` with `JWTValidator`.
Can you add an autouse fixture (or `finally` block) that restores/clears the
registry entry after each test? The PR description mentions "registry teardown
to prevent test pollution" but I can't find the actual teardown in this file.
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py:
##########
@@ -230,6 +238,17 @@ def ti_run(
last_heartbeat_at=timezone.utcnow(),
)
+ execution_token: str | None = None
+ if token.claims.scope == "workload":
+ try:
+ generator: JWTGenerator = services.get(JWTGenerator)
+ execution_token = generator.generate(extras={"sub":
str(task_instance_id), "scope": "execution"})
+ except (jwt.PyJWTError, ValueError, KeyError):
Review Comment:
This catch is too broad. `generate()` is deterministic;
`ValueError`/`KeyError` here would indicate a programming bug (bad extras dict,
missing signing key) that should surface with a real stack trace, not a generic
500.
Also worth noting: `services.get(JWTGenerator)` raises
`svcs.exceptions.ServiceNotFoundError`, which isn't in this tuple -- so if the
registry is misconfigured you'll get an unhelpful 500 from the global handler
anyway.
Suggest narrowing to `jwt.PyJWTError` only, or dropping the try/except
entirely and letting the global handler log the real error.
##########
airflow-core/src/airflow/api_fastapi/execution_api/app.py:
##########
@@ -333,10 +342,30 @@ async def always_allow(request: Request):
claims = TIClaims(scope="execution")
return TIToken(id=ti_id, claims=claims)
+ # Override _container (the svcs service locator behind
DepContainer).
+ # The default _container reads request.app.state.svcs_registry, but
+ # Cadwyn's versioned sub-apps don't inherit the main app's state,
+ # so lookups raise ServiceNotFoundError. This registry provides
+ # services needed by routes called during dag.test().
+ # Note: tokens generated by this stub are never validated since
+ # _jwt_bearer is overridden with always_allow in dag.test() mode.
+ stub_generator = JWTGenerator(
Review Comment:
The stub generator uses a per-process random secret
(`secrets.token_urlsafe(32)`), so any token it signs is unverifiable against
the real validator. The comment above says tokens aren't validated because
`_jwt_bearer` is overridden, but `/run` now actively mints an execution token
via `services.get(JWTGenerator).generate(...)` and returns it in
`Refreshed-API-Token`.
Under `dag.test()` the current workload branch is skipped because
`always_allow` returns `scope="execution"`, but this is a fragile invariant.
Any future path that triggers the workload-token swap (or any test simulating a
workload token through the in-process path) would hand the client a token no
validator can verify.
Consider either reusing the real signing key from `get_signing_args()`, or
not registering a `JWTGenerator` under in-process mode at all so callers fail
loudly if they ever exercise that path.
##########
airflow-core/src/airflow/executors/workloads/base.py:
##########
@@ -76,7 +76,12 @@ class BaseWorkloadSchema(BaseModel):
@staticmethod
def generate_token(sub_id: str, generator: JWTGenerator | None = None) ->
str:
- return generator.generate({"sub": sub_id}) if generator else ""
+ if not generator:
+ return ""
+ return generator.generate(
Review Comment:
Since `JWTGenerator` ideally shouldn't know about `task_queued_timeout`,
this is the right place to read it:
```python
valid_for = conf.getfloat("scheduler", "task_queued_timeout")
return generator.generate(
extras={"sub": sub_id, "scope": "workload"},
valid_for=valid_for,
)
```
That also lets you drop `workload_valid_for` from `JWTGenerator` entirely.
##########
airflow-core/src/airflow/api_fastapi/execution_api/app.py:
##########
@@ -142,6 +145,12 @@ async def dispatch(self, request: Request, call_next):
validator: JWTValidator = await services.aget(JWTValidator)
claims = await validator.avalidated_claims(token, {})
+ # Workload tokens are long-lived and meant to survive queue
+ # wait times so avoid refreshing them. If avalidated_claims
+ # raises for a workload token, the outer except handles it.
+ if claims.get("scope") == "workload":
+ return response
+
now = int(time.time())
validity = conf.getint("execution_api",
"jwt_expiration_time")
refresh_when_less_than = max(int(validity * 0.20), 30)
Review Comment:
The refresh-threshold denominator is `jwt_expiration_time`, which only makes
sense for execution tokens. The new early-return at line 151 (`if
claims.get("scope") == "workload": return response`) guards workload tokens
today, but if someone removes that guard or adds a new scope with a different
lifetime, this threshold math silently becomes wrong (e.g. comparing a 5-minute
workload-token's remaining life against 20% of the 10-minute execution-token
lifetime).
Worth computing the threshold from the actual token lifetime instead:
```python
token_lifetime = int(claims.get("exp", 0)) - int(claims.get("iat", 0))
refresh_when_less_than = max(int(token_lifetime * 0.20), 30)
```
--
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]
