Dev-iL commented on code in PR #75396:
URL: https://github.com/apache/airflow/pull/75396#discussion_r4557099653
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py:
##########
@@ -129,8 +131,11 @@ def has_xcom_access(
_team_name_for_ti_stmt,
)
- requester_team = session.scalar(_team_name_for_ti_stmt(token.id))
- target_team = session.scalar(_team_name_for_dag_stmt(dag_id))
+ # Cadwyn solves this router dependency in two passes, and the first pass's
request session would stay
+ # checked out until the response; a session of its own returns the
connection before the producer lookup.
+ async with create_session_async() as session:
Review Comment:
Done: with `[core] multi_team` on, the check now runs once per request, so a
request takes one access-check connection and runs the two team queries once.
I didn't copy `_resolve_trigger_token`'s guard, though. There, skipping the
first pass falls back to the caller's own identity, which is safe. Here, a
skipped pass grants access, so a request solved without Cadwyn's pass would go
unchecked. Instead, the first pass records the grant in the request scope and
the second reuses it. A denial still raises the 403 in the first pass.
`test_access_is_checked_once_per_request` (`GET`, `POST`, `DELETE`) counts
the Dag-team query per request. Without the change, all three cases fail with 2
queries instead of 1.
---
Drafted-by: Claude Code (Opus 5.5) (no human review before posting)
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py:
##########
@@ -178,88 +183,144 @@ def get_order(self, entity: AliasedClass, *, descending:
bool = False) -> tuple[
return tuple(column.desc() if descending else column.asc() for column
in columns)
+def _render_json(content: JsonValue) -> JSONResponse:
+ """
+ Render an XCom read's response on a worker thread.
+
+ The values come from a JSON column, so returning a ``JSONResponse`` skips
FastAPI's re-validation of
+ them as ``JsonValue``, which for large values costs several times more
than encoding them. Cadwyn still
+ decodes and re-encodes a returned ``JSONResponse`` body on the event loop
+ (https://github.com/zmievsa/cadwyn/issues/51), so only the first encoding
runs on the worker thread.
+ """
+ return JSONResponse(content)
+
+
+def _reject_archived_caller(
+ token: ExecutionToken, *, dag_id: str, run_id: str, task_id: str,
map_index: int | None, session: Session
+) -> None:
+ """Refuse a task's read of its own XCom once its attempt has been
archived."""
+ if token.claims.scope != "execution":
+ return
+ caller = session.execute(
+ select(TaskInstance.id, TaskInstance.working_set)
+ .where(
+ TaskInstance.id == token.id,
+ TaskInstance.dag_id == dag_id,
+ TaskInstance.run_id == run_id,
+ TaskInstance.task_id == task_id,
+ *build_coordinate_filters(TaskInstance, map_index=map_index),
+ )
+ .limit(1)
+ .execution_options(include_all_attempts=True)
+ ).first()
+ if caller is None or caller.working_set is not None:
+ return
+ identify_archived = IdentifyArchivedTaskStateUpdates.is_applied
+ raise HTTPException(
+ status_code=status.HTTP_410_GONE if identify_archived else
status.HTTP_404_NOT_FOUND,
+ detail={
+ "reason": "not_found",
+ "message": (
+ "Task Instance not found in the working set; its attempt has
been archived"
+ if identify_archived
+ else "Task Instance not found"
+ ),
+ },
+ )
+
+
def _build_xcom_read(
*,
dag_id: str,
run_id: str,
task_id: str,
key: str,
- session: Session,
dag_bag: DBDagBag,
token: ExecutionToken,
map_index: int | None = None,
include_prior_dates: bool = False,
previous_iteration: bool = False,
sequence: bool = False,
+ reject_archived_caller: bool = False,
) -> _XComRead:
"""
Select the XCom rows of the producers visible to the calling task instance.
A ``sequence`` read returns several values in order. When it targets a
loop task from outside that
loop, it covers every iteration of the task rather than the one a caller
inside the loop would see.
+
+ Routes run this on a worker thread: resolving producers can deserialize a
Dag on a DagBag cache miss,
+ CPU work that would stall the event loop. A session cannot cross threads,
so this opens its own.
+
+ ``reject_archived_caller`` makes a caller reading its own XCom get the
archived-attempt error first. The
+ check runs here rather than on the route's session, so it leaves no route
connection idle during the lookup.
"""
- resolver = TaskCoordinateResolver(dag_bag, session)
- read = partial(
- XComModel.get_many,
- dag_ids=dag_id,
- run_id=run_id,
- task_ids=task_id,
- key=key,
- include_prior_dates=include_prior_dates,
- )
- try:
- if previous_iteration and include_prior_dates:
- raise ValueError("previous_iteration cannot be combined with prior
dates")
- if not previous_iteration and not resolver.has_regions(
- dag_id, None if include_prior_dates else run_id, task_id
- ):
- return _XComRead(read(region_id=SENTINEL_REGION_ID,
map_indexes=map_index))
- caller = session.get(TaskInstance, token.id) if token.claims.scope ==
"execution" else None
- if include_prior_dates:
- task = None
- if session.scalar(select(DagRun.id).where(DagRun.dag_id == dag_id,
DagRun.run_id == run_id)):
- with suppress(TaskNotFound):
- task = resolver.get_task(
- dag_id,
- run_id,
- task_id,
- dag_version_id=caller.dag_version_id
- if caller and (caller.dag_id, caller.run_id) ==
(dag_id, run_id)
- else None,
- )
- if task is not None and enclosing_loop(task) is not None:
- raise ValueError("Prior-date lookups are not supported for
tasks inside a loop")
- return _XComRead(
- read(region_id=SENTINEL_REGION_ID, include_node_regions=True,
map_indexes=map_index)
+ with create_session(scoped=False) as session:
Review Comment:
You're right. I checked the path: for a task outside any loop or mapped
region, without `previous_iteration`, `_build_xcom_read` runs only the
archived-caller check and `has_regions` before returning a statement.
`_find_writer_id` returns after `session.get`, and `get_public_map_index` runs
no query for such a task.
I'll keep this PR a plain conversion and do the fast path separately. It
reopens the trade-off from the other thread: once the pre-check runs on the
route's session, a task with regions either keeps that connection checked out
across the lookup or pays a second checkout. That version needs its own
benchmarks against this one. The description now says that the common path pays
for the bridge, and lists the fast path under Debt left.
---
Drafted-by: Claude Code (Opus 5.5) (no human review before posting)
--
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]
