kaxil commented on code in PR #66160:
URL: https://github.com/apache/airflow/pull/66160#discussion_r3171192441
##########
task-sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -248,6 +250,14 @@ def get_template_context(self) -> Context:
"value": VariableAccessor(deserialize_json=False),
},
"conn": ConnectionAccessor(),
+ "task_state": TaskStateAccessor(ti_id=self.id),
+ # TODO: right now, asset_state is only provided for
single-inlet tasks.
+ # Multi-inlet support is deferred — see AssetStateAccessor
docstring.
+ **(
+ {"asset_state":
AssetStateAccessor(name=self.task.inlets[0].name)}
Review Comment:
`self.task.inlets[0].name` crashes when the single inlet is
`Asset.ref(uri="...")`. `AssetUriRef` only has `.uri` (no `.name`), so any task
with an inlet like `Asset.ref(uri="s3://x")` raises `AttributeError:
'AssetUriRef' object has no attribute 'name'` at template-context construction,
before user code runs. For `AssetAlias(name=...)` it doesn't crash, but the
alias name then gets sent to `/state/asset/{name}` and the route's
`_resolve_asset_id` looks for a real asset by that name and 404s, which is
silently confusing.
Worth gating on `isinstance(self.task.inlets[0], Asset)` (and probably
`AssetNameRef`) so non-Asset / non-NameRef inlets fall through to "no
`asset_state` in context", matching the multi-inlet behavior. Existing tests
only cover the `Asset(name=..., uri=...)` shape, so this isn't caught.
##########
task-sdk/src/airflow/sdk/execution_time/context.py:
##########
@@ -406,6 +405,137 @@ def get(self, key, default: Any = NOTSET) -> Any:
raise
+class TaskStateAccessor:
+ """Accessor for task state scoped to the current task instance. Available
as ``context['task_state']`` at task execution time."""
+
+ def __init__(self, ti_id: UUID) -> None:
+ self._ti_id = ti_id
+
+ def __eq__(self, other: object) -> bool:
+ if not isinstance(other, TaskStateAccessor):
+ return False
+ return self._ti_id == other._ti_id
+
+ def __hash__(self) -> int:
+ return hash(self._ti_id)
+
+ def __repr__(self) -> str:
+ return f"<TaskStateAccessor ti_id={self._ti_id}>"
+
+ # TODO: ``__getattr__`` for jinja template access like ``{{
task_state.job_id }}``
+ # is not implemented yet cos it's unclear whether task state values will be
+ # used in templates.
+
+ def get(self, key: str) -> str | None:
Review Comment:
`get(key)` returns `None` on missing but takes no `default` parameter.
`VariableAccessor.get(key, default=NOTSET)` lets users distinguish "not set"
from "set to None". Same shape would help here, especially since the stored
value is `str` and `None` is a meaningful sentinel users may want to override.
Same applies to `AssetStateAccessor.get` at line 500.
##########
task-sdk/src/airflow/sdk/execution_time/context.py:
##########
@@ -406,6 +405,137 @@ def get(self, key, default: Any = NOTSET) -> Any:
raise
+class TaskStateAccessor:
+ """Accessor for task state scoped to the current task instance. Available
as ``context['task_state']`` at task execution time."""
+
+ def __init__(self, ti_id: UUID) -> None:
+ self._ti_id = ti_id
+
+ def __eq__(self, other: object) -> bool:
+ if not isinstance(other, TaskStateAccessor):
+ return False
+ return self._ti_id == other._ti_id
+
+ def __hash__(self) -> int:
+ return hash(self._ti_id)
+
+ def __repr__(self) -> str:
+ return f"<TaskStateAccessor ti_id={self._ti_id}>"
+
+ # TODO: ``__getattr__`` for jinja template access like ``{{
task_state.job_id }}``
+ # is not implemented yet cos it's unclear whether task state values will be
+ # used in templates.
+
+ def get(self, key: str) -> str | None:
+ """Return the stored value, or ``None`` if the key does not exist."""
+ from airflow.sdk.exceptions import AirflowRuntimeError, ErrorType
+ from airflow.sdk.execution_time.comms import GetTaskState,
TaskStateResult
+ from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
+
+ try:
+ resp = SUPERVISOR_COMMS.send(GetTaskState(ti_id=self._ti_id,
key=key))
+ except AirflowRuntimeError as e:
+ if e.error.error == ErrorType.API_SERVER_ERROR and
isinstance(e.error.detail, dict):
+ if e.error.detail.get("status_code") == 404:
Review Comment:
Detecting "not found" via `e.error.detail.get("status_code") == 404` couples
the accessor to the supervisor's error-wrapping shape (the magic `status_code`
key inside `detail`). Every other accessor in this module
(`ConnectionAccessor.get`, `VariableAccessor.get`) translates 404 into a typed
`ErrorType.X_NOT_FOUND` at the client layer (see `client.py:425`,
`client.py:448`) and matches on that.
Worth adding `TASK_STATE_NOT_FOUND` / `ASSET_STATE_NOT_FOUND` to
`ErrorType`, doing the translation in `client.task_state.get` /
`client.asset_state.get`, and matching on the typed error here. Same change
applies to the duplicated branch at line 510 in `AssetStateAccessor.get`.
##########
task-sdk/src/airflow/sdk/execution_time/context.py:
##########
@@ -406,6 +405,137 @@ def get(self, key, default: Any = NOTSET) -> Any:
raise
+class TaskStateAccessor:
+ """Accessor for task state scoped to the current task instance. Available
as ``context['task_state']`` at task execution time."""
+
+ def __init__(self, ti_id: UUID) -> None:
+ self._ti_id = ti_id
+
+ def __eq__(self, other: object) -> bool:
+ if not isinstance(other, TaskStateAccessor):
+ return False
+ return self._ti_id == other._ti_id
+
+ def __hash__(self) -> int:
+ return hash(self._ti_id)
+
+ def __repr__(self) -> str:
+ return f"<TaskStateAccessor ti_id={self._ti_id}>"
+
+ # TODO: ``__getattr__`` for jinja template access like ``{{
task_state.job_id }}``
+ # is not implemented yet cos it's unclear whether task state values will be
+ # used in templates.
+
+ def get(self, key: str) -> str | None:
+ """Return the stored value, or ``None`` if the key does not exist."""
+ from airflow.sdk.exceptions import AirflowRuntimeError, ErrorType
Review Comment:
`AirflowRuntimeError` and `ErrorType` are already imported at module top
(line 45), so this inline re-import is redundant. Same in
`AssetStateAccessor.get` at line 502. The `comms` and `task_runner` imports do
need to stay inline (circular), but the exceptions ones don't.
##########
task-sdk/src/airflow/sdk/definitions/context.py:
##########
@@ -72,6 +72,8 @@ class Context(TypedDict, total=False):
task_reschedule_count: int
task_instance: RuntimeTaskInstanceProtocol
task_instance_key_str: str
+ task_state: TaskStateAccessor
Review Comment:
`task_state: TaskStateAccessor` is unannotated, but `asset_state:
NotRequired[AssetStateAccessor]` on the next line is. With `Context(TypedDict,
total=False)` the two are functionally equivalent, but the asymmetry reads as
if `task_state` is required and `asset_state` isn't. The PR description says
`task_state` is "always available", which matches your intent. Either drop
`NotRequired` from `asset_state` for symmetry, or wrap `task_state` in
`Required[...]` if you really want to enforce it (and switch the dict to
`total=True`).
##########
task-sdk/src/airflow/sdk/api/client.py:
##########
@@ -627,6 +631,63 @@ def get_sequence_slice(
return XComSequenceSliceResponse.model_validate_json(resp.read())
+class TaskStateOperations:
+ __slots__ = ("client",)
+
+ def __init__(self, client: Client):
+ self.client = client
+
+ def get(self, ti_id: uuid.UUID, key: str) -> TaskStateResponse:
Review Comment:
Other client `get` methods translate 404 into a typed
`ErrorResponse(error=ErrorType.X_NOT_FOUND, ...)`: `connections.get` at line
417/425, `variables.get` at 440/448, `xcoms.get` at 585/598, `assets.get` at
710/718. `task_state.get` and `asset_state.get` (line 669) skip that step and
let `ServerResponseError` propagate, which is what forces the brittle
`e.error.detail.get("status_code") == 404` check in `TaskStateAccessor.get` /
`AssetStateAccessor.get`.
If you adopt the typed-error approach here (wrap in `try/except
ServerResponseError`, return
`ErrorResponse(error=ErrorType.TASK_STATE_NOT_FOUND, detail={...})` on 404),
the accessor's "is this a 404?" logic becomes a clean enum match against
`resp.error`.
--
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]
