potiuk commented on code in PR #66141:
URL: https://github.com/apache/airflow/pull/66141#discussion_r3215312365
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -1261,17 +1261,30 @@ def process_executor_events(
)
continue
+ # Callback state transitions are now driven by the supervisor
through
+ # the Execution API (POST /callbacks/{id}/run, PATCH
/callbacks/{id}/state).
+ # The in-process events from the executor are kept as a fallback
safety
+ # net for cases where the supervisor crashed before reporting a
terminal state
+
+ need_to_modify = False
+
if state == CallbackState.RUNNING:
Review Comment:
**Major (open thread, not yet addressed).** @amoghrajesh flagged this in the
prior review pass: the `if state == CallbackState.RUNNING:` branch only emits a
log and modifies nothing, because the API is now the source of truth for
RUNNING transitions. With the API-driven design, executor RUNNING events are
informational at best.
Two cleanups worth picking from:
1. Drop the branch entirely (and the corresponding entry from the `state in
(...)` filter at the top of `process_executor_events`). The fallback only needs
to advance terminal states.
2. Keep the branch, but downgrade the log to `cls.logger().debug(...)` and
add a one-line comment that the API is authoritative for RUNNING transitions
and this log is informational.
Either is fine; the current shape sits in the middle and reads as dead code
at first glance.
##########
airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_04_30.py:
##########
@@ -0,0 +1,31 @@
+# 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 cadwyn import VersionChange, endpoint
+
+
+class AddCallbackEndpoints(VersionChange):
Review Comment:
**Minor / question.** The version label `2026-04-30` is the date this PR's
first commit was authored. The convention in the rest of the bundle (and what
an archeologist reading the file later would expect) is the date the change
LANDED. With the PR opened in late April and now mid-May, a `2026-05-XX` label
is closer to reality.
Not a blocker — if `apache/airflow` has settled on a different convention
(dev-start date, milestone date, etc.), happy to be corrected. The mechanical
change is just renaming the file and the version-bundle entry. cadwyn's
`Version("YYYY-MM-DD", AddCallbackEndpoints)` doesn't care about the actual
date semantically; it's a sortable label.
##########
task-sdk/src/airflow/sdk/api/datamodels/_generated.py:
##########
@@ -85,6 +85,23 @@ class AssetStateResponse(BaseModel):
value: Annotated[str, Field(title="Value")]
+class State(str, Enum):
Review Comment:
**Major.** Generic class name `State` in a shared `_generated.py` is
collision-prone (it's literally the most-used identifier in the codebase). The
cause is the source: `CallbackTerminalState = Literal["success", "failed"]` in
`execution_api/datamodels/callback.py` — Cadwyn's OpenAPI generation can't see
through type aliases and falls back to naming the enum after the field (`state`
→ `State`).
Fix is in the source, not here. Replace the `Literal` alias with a real Enum:
```python
# execution_api/datamodels/callback.py
class CallbackTerminalState(str, Enum):
SUCCESS = "success"
FAILED = "failed"
class CallbackTerminalStatePayload(StrictBaseModel):
state: CallbackTerminalState
output: str | None = None
```
Then the generated module gets `class CallbackTerminalState(str, Enum)`,
which is collision-safe and self-documenting. The prek hook regenerates
`_generated.py` automatically.
##########
task-sdk/src/airflow/sdk/api/client.py:
##########
@@ -1008,6 +1009,36 @@ def get_detail_response(self, ti_id: uuid.UUID) ->
HITLDetailResponse:
return HITLDetailResponse.model_validate_json(resp.read())
+class CallbackOperations:
+ __slots__ = ("client",)
+
+ def __init__(self, client: Client):
+ self.client = client
+
+ def start(self, id: uuid.UUID | str) -> None:
+ """
+ Mark a callback as RUNNING and exchange a workload token for an
execution token.
+
+ Mirrors ``TaskInstanceOperations.start``: this is the single API call
that
+ accepts a workload-scoped token; the server returns the new execution
token
+ via the ``Refreshed-API-Token`` response header and the Client's
response
+ hook automatically swaps it onto subsequent requests.
+ """
+ self.client.post(f"callbacks/{id}/run")
+
+ def finish(
+ self,
+ id: uuid.UUID | str,
+ state: str,
Review Comment:
**Minor.** `state: str` is too loose — the only valid values are `"success"`
and `"failed"`. A typo like `state="successs"` or `state="FAILED"` would only
surface as a 422 from the API. Tighten with `Literal` (or import
`CallbackTerminalState` once #2 lands):
```python
from typing import Literal
def finish(
self,
id: uuid.UUID | str,
state: Literal["success", "failed"],
output: str | None = None,
) -> None:
```
Matches the API's actual contract and gives the supervisor's caller a
compile-time check.
##########
airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_callbacks.py:
##########
@@ -0,0 +1,264 @@
+# 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 Literal
+from unittest import mock
+from uuid import UUID, uuid4
+
+import pytest
+from fastapi import Request
+
+from airflow.api_fastapi.auth.tokens import JWTGenerator, JWTValidator
+from airflow.api_fastapi.execution_api.app import lifespan
+from airflow.api_fastapi.execution_api.datamodels.token import TIClaims,
TIToken
+from airflow.api_fastapi.execution_api.security import require_auth
+from airflow.executors.workloads.callback import CallbackFetchMethod
+from airflow.models.callback import ExecutorCallback
+from airflow.utils.state import CallbackState
+
+from tests_common.test_utils.db import clear_db_callbacks
+
+pytestmark = pytest.mark.db_test
+
+
+class _FakeCallbackDef:
+ """Minimal CallbackDefinitionProtocol stand-in for tests."""
+
+ path: str = "tests.fake.callback"
+ kwargs: dict = {}
+ executor: str | None = None
+
+ def serialize(self) -> dict:
+ return {"path": self.path, "kwargs": self.kwargs, "executor":
self.executor}
+
+
+def _make_callback(state: CallbackState, session) -> ExecutorCallback:
+ cb = ExecutorCallback(callback_def=_FakeCallbackDef(),
fetch_method=CallbackFetchMethod.IMPORT_PATH)
+ cb.state = state
+ session.add(cb)
+ session.commit()
+ return cb
+
+
+def _override_require_auth(exec_app, scope: Literal["execution", "workload"] =
"execution") -> None:
+ """Override require_auth to return a token whose sub matches the path
callback_id."""
+
+ async def _token(request: Request) -> TIToken:
+ path_id = request.path_params.get("callback_id")
+ cb_id = UUID(path_id) if path_id else uuid4()
+ return TIToken(id=cb_id, claims=TIClaims(scope=scope))
+
+ exec_app.dependency_overrides[require_auth] = _token
+
+
[email protected]
+def _use_real_jwt_bearer(exec_app):
+ """Remove the mock require_auth override so the real JWT validation runs
end-to-end."""
+ exec_app.dependency_overrides.pop(require_auth, None)
+
+
+class TestCallbackRun:
+ def setup_method(self):
+ clear_db_callbacks()
+
+ def teardown_method(self):
+ clear_db_callbacks()
+
+ @pytest.mark.parametrize("scope", ["workload", "execution"])
+ def test_run_marks_callback_running_and_swaps_workload_token(self, client,
exec_app, session, scope):
Review Comment:
**Major — missing test for the highest-value behaviour change in this PR.**
The race-condition fix (RUNNING-idempotent retry on `POST /callbacks/{id}/run`,
with the workload-token swap still firing) is the exact path @amoghrajesh's
earlier comment was about. The fix landed in the route, but no test covers it.
Worth adding alongside the existing run test:
```python
def test_run_idempotent_when_already_running_with_workload_token(
self, client, exec_app, session
):
cb = _make_callback(CallbackState.RUNNING, session)
mock_gen = mock.MagicMock(spec=JWTGenerator)
mock_gen.generate.return_value = "mock-execution-token"
lifespan.registry.register_value(JWTGenerator, mock_gen)
_override_require_auth(exec_app, scope="workload")
response = client.post(f"/execution/callbacks/{cb.id}/run")
exec_app.dependency_overrides.pop(require_auth, None)
assert response.status_code == 204
assert response.headers["Refreshed-API-Token"] == "mock-execution-token"
session.expire_all()
assert session.get(ExecutorCallback, cb.id).state ==
CallbackState.RUNNING
```
Locks in (a) idempotent 204, (b) fresh execution token issued for the retry,
(c) state stays RUNNING. Also worth a sibling case for execution-scoped retry
where no swap should happen.
--
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]
