anishgirianish commented on code in PR #60108:
URL: https://github.com/apache/airflow/pull/60108#discussion_r3125232646
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py:
##########
@@ -289,13 +295,24 @@ def ti_run(
context.next_method = ti.next_method
context.next_kwargs = ti.next_kwargs
context.start_date = ti.start_date
- 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"
)
+ try:
+ generator: JWTGenerator = services.get(JWTGenerator)
+ execution_token = generator.generate(extras={"sub":
str(task_instance_id), "scope": "execution"})
+ except jwt.PyJWTError:
+ log.exception("Failed to generate execution token for task instance
%s", task_instance_id)
+ raise HTTPException(
+ status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="Token
generation failed"
+ )
+ response.headers["Refreshed-API-Token"] = execution_token
Review Comment:
Done, `ti_run` now accepts `token: TIToken = CurrentTIToken` and the
`generator.generate(...)` call is gated on `if token.claims.scope ==
"workload":`.
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py:
##########
@@ -289,13 +295,24 @@ def ti_run(
context.next_method = ti.next_method
context.next_kwargs = ti.next_kwargs
context.start_date = ti.start_date
- 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"
)
+ try:
+ generator: JWTGenerator = services.get(JWTGenerator)
+ execution_token = generator.generate(extras={"sub":
str(task_instance_id), "scope": "execution"})
+ except jwt.PyJWTError:
+ log.exception("Failed to generate execution token for task instance
%s", task_instance_id)
+ raise HTTPException(
+ status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="Token
generation failed"
+ )
+ response.headers["Refreshed-API-Token"] = execution_token
Review Comment:
Added a comment on the `response.headers["Refreshed-API-Token"]` assignment.
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py:
##########
@@ -289,13 +295,24 @@ def ti_run(
context.next_method = ti.next_method
context.next_kwargs = ti.next_kwargs
context.start_date = ti.start_date
- 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"
)
+ try:
+ generator: JWTGenerator = services.get(JWTGenerator)
+ execution_token = generator.generate(extras={"sub":
str(task_instance_id), "scope": "execution"})
+ except jwt.PyJWTError:
Review Comment:
Broadened to `except (jwt.PyJWTError, ValueError, KeyError)` as suggested.
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py:
##########
@@ -289,13 +295,24 @@ def ti_run(
context.next_method = ti.next_method
context.next_kwargs = ti.next_kwargs
context.start_date = ti.start_date
- 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"
)
+ try:
+ generator: JWTGenerator = services.get(JWTGenerator)
+ execution_token = generator.generate(extras={"sub":
str(task_instance_id), "scope": "execution"})
+ except jwt.PyJWTError:
+ log.exception("Failed to generate execution token for task instance
%s", task_instance_id)
+ raise HTTPException(
+ status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="Token
generation failed"
+ )
+ response.headers["Refreshed-API-Token"] = execution_token
+
+ return context
Review Comment:
Moved token generation before `session.execute(query)` so a `generate()`
failure leaves the TI in QUEUED rather than stuck in RUNNING with no recovery
path.
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py:
##########
@@ -289,13 +295,24 @@ def ti_run(
context.next_method = ti.next_method
context.next_kwargs = ti.next_kwargs
context.start_date = ti.start_date
- 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"
)
+ try:
+ generator: JWTGenerator = services.get(JWTGenerator)
+ execution_token = generator.generate(extras={"sub":
str(task_instance_id), "scope": "execution"})
+ except jwt.PyJWTError:
+ log.exception("Failed to generate execution token for task instance
%s", task_instance_id)
+ raise HTTPException(
+ status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="Token
generation failed"
+ )
+ response.headers["Refreshed-API-Token"] = execution_token
+
+ return context
Review Comment:
Moved token generation before `session.execute(query)` ,
##########
airflow-core/tests/unit/api_fastapi/execution_api/conftest.py:
##########
@@ -55,6 +59,10 @@ async def mock_jwt_bearer(request: Request):
exec_app.dependency_overrides[_jwt_bearer] = mock_jwt_bearer
with TestClient(app, headers={"Authorization": "Bearer fake"}) as client:
+ mock_generator = MagicMock(spec=JWTGenerator)
+ mock_generator.generate.return_value = "mock-execution-token"
+ lifespan.registry.register_value(JWTGenerator, mock_generator)
Review Comment:
Removed the `register_value` call from the shared fixture entirely; the mock
generator is now registered inside `test_ti_run_returns_execution_token` only,
where it's needed.
##########
airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py:
##########
@@ -281,6 +281,36 @@ def test_ti_run_state_to_running(
)
assert response.status_code == 409
+ def test_ti_run_returns_execution_token(self, client, session,
create_task_instance, time_machine):
+ """PATCH /run should return a Refreshed-API-Token header on success."""
+ 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()
+
+ response = client.patch(
+ f"/execution/task-instances/{ti.id}/run",
+ json={
+ "state": "running",
+ "hostname": "test-host",
+ "unixname": "test-user",
+ "pid": 100,
+ "start_date": "2024-10-31T12:00:00Z",
+ },
+ )
+
+ assert response.status_code == 200
+ assert "Refreshed-API-Token" in response.headers
+ assert response.headers["Refreshed-API-Token"] ==
"mock-execution-token"
Review Comment:
Added `mock_gen.generate.assert_called_once()` plus assertions on
`extras["scope"] == "execution"` and `extras["sub"] == str(ti.id)`.
--
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]
