ferruzzi commented on code in PR #66141:
URL: https://github.com/apache/airflow/pull/66141#discussion_r3222821358
##########
task-sdk/src/airflow/sdk/execution_time/callback_supervisor.py:
##########
@@ -352,6 +352,15 @@ def supervise_callback(
logger = structlog.get_logger(logger_name="callback").bind()
with _ensure_client(server, token, client=client) as client:
+ # Mark the callback as RUNNING via the API. This is the single endpoint
+ # that accepts a workload-scoped token; it returns a fresh execution
+ # token via the Refreshed-API-Token header which the Client's response
+ # hook adopts automatically for the rest of the run.
+ client.callbacks.start(id)
+
+ terminal_state = "failed"
Review Comment:
Here and below, shouldn't we be using `CallbackTerminalState.FAILED`,
`CallbackTerminalState.SUCCESS`, etc instead of the strings? This technically
works, but only until someone changes one of those enum values. I suggest that
we tie them together just to future-proof this.
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/callbacks.py:
##########
@@ -0,0 +1,153 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Annotated
+from uuid import UUID
+
+import structlog
+from cadwyn import VersionedAPIRouter
+from fastapi import Body, HTTPException, Response, Security, status
+from structlog.contextvars import bind_contextvars
+
+from airflow.api_fastapi.auth.tokens import JWTGenerator
+from airflow.api_fastapi.common.db.common import SessionDep
+from airflow.api_fastapi.execution_api.datamodels.callback import
CallbackTerminalStatePayload
+from airflow.api_fastapi.execution_api.datamodels.token import TIToken
+from airflow.api_fastapi.execution_api.deps import DepContainer
+from airflow.api_fastapi.execution_api.security import CurrentTIToken,
ExecutionAPIRoute, require_auth
+from airflow.models.callback import Callback
+from airflow.utils.state import CallbackState
+
+log = structlog.get_logger(__name__)
+
+router = VersionedAPIRouter(route_class=ExecutionAPIRoute)
+
+
+def _require_self(token: TIToken, callback_id: UUID) -> None:
+ """Mirror the ``ti:self`` enforcement from security.py for callback
routes."""
+ if str(token.id) != str(callback_id):
+ raise HTTPException(
+ status_code=status.HTTP_403_FORBIDDEN,
+ detail="Token subject does not match callback id",
+ )
+
+
[email protected](
+ "/{callback_id}/run",
+ status_code=status.HTTP_204_NO_CONTENT,
+ dependencies=[Security(require_auth, scopes=["token:execution",
"token:workload"])],
+ responses={
+ status.HTTP_403_FORBIDDEN: {"description": "Token subject does not
match callback id"},
+ status.HTTP_404_NOT_FOUND: {"description": "Callback not found"},
+ status.HTTP_409_CONFLICT: {"description": "Callback is not in a state
that can be marked running"},
+ },
+)
+def callback_run(
+ callback_id: UUID,
+ response: Response,
+ session: SessionDep,
+ services=DepContainer,
+ token: TIToken = CurrentTIToken,
+) -> Response:
+ """
+ Mark a callback as RUNNING.
+
+ Mirrors ``PATCH /task-instances/{id}/run``: this is the single endpoint
that
+ accepts a workload-scoped token and atomically (a) transitions the callback
+ from QUEUED to RUNNING and (b) issues a fresh execution-scoped token via
the
+ ``Refreshed-API-Token`` response header. All subsequent supervisor calls
hit
+ execution-only routes.
+ """
+ bind_contextvars(callback_id=str(callback_id))
+ _require_self(token, callback_id)
+
+ callback = session.get(Callback, callback_id)
+ if callback is None:
+ raise HTTPException(
+ status_code=status.HTTP_404_NOT_FOUND,
+ detail={"reason": "not_found", "message": "Callback not found"},
+ )
+
+ # Allow QUEUED → RUNNING transition; treat RUNNING as idempotent so a
retried
+ # supervisor start does not 409. Anything else (PENDING / SCHEDULED /
terminal) rejects.
+ if callback.state == CallbackState.RUNNING:
+ log.info("Duplicate start request received from %s", callback.id)
Review Comment:
I think the rest of the file is using the "new" structlog format, if you
don't mind a nitpick for consistency:
```suggestion
log.info("Duplicate start request received",
callback_id=str(callback.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]
