amoghrajesh commented on code in PR #60108:
URL: https://github.com/apache/airflow/pull/60108#discussion_r2980465748
##########
airflow-core/src/airflow/config_templates/config.yml:
##########
@@ -2069,6 +2069,15 @@ execution_api:
type: integer
example: ~
default: "600"
+ jwt_workload_token_expiration_time:
+ description: |
+ Seconds until workload JWT tokens expire. These long-lived tokens are
sent
+ with task workloads to executors and can only call the /run endpoint.
+ Set long enough to cover maximum expected queue wait time.
+ version_added: 3.2.0
Review Comment:
+1
##########
task-sdk/src/airflow/sdk/api/client.py:
##########
@@ -952,7 +952,10 @@ def __init__(self, *, base_url: str | None, dry_run: bool
= False, token: str, *
)
def _update_auth(self, response: httpx.Response):
- if new_token := response.headers.get("Refreshed-API-Token"):
+ if new_token := response.headers.get("X-Execution-Token"):
+ log.debug("Received execution token, swapping auth")
+ self.auth = BearerAuth(new_token)
+ elif new_token := response.headers.get("Refreshed-API-Token"):
Review Comment:
I agree too.
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py:
##########
@@ -281,14 +286,18 @@ def ti_run(
if ti.next_method:
context.next_method = ti.next_method
context.next_kwargs = ti.next_kwargs
-
- return context
except SQLAlchemyError:
log.exception("Error marking Task Instance state as running")
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail="Database error occurred"
)
+ generator: JWTGenerator = services.get(JWTGenerator)
+ execution_token = generator.generate(extras={"sub": str(task_instance_id)})
Review Comment:
Do we need to set the scope as `execution` here to be explicit?
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py:
##########
@@ -281,14 +286,18 @@ def ti_run(
if ti.next_method:
context.next_method = ti.next_method
context.next_kwargs = ti.next_kwargs
-
- return context
except SQLAlchemyError:
log.exception("Error marking Task Instance state as running")
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail="Database error occurred"
)
+ generator: JWTGenerator = services.get(JWTGenerator)
+ execution_token = generator.generate(extras={"sub": str(task_instance_id)})
+ response.headers["X-Execution-Token"] = execution_token
Review Comment:
If `services.get(JWTGenerator)` raises or fails for whatever reason / if
`generate` raises, the API would return 500 here, but we committed the TI as
running earlier. To avoid this, I suggest to move this up in the try/except
itself.
--
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]
