amoghrajesh commented on code in PR #66160:
URL: https://github.com/apache/airflow/pull/66160#discussion_r3180667498
##########
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:
Handled in 7815e8c430
##########
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:
Handled in
https://github.com/apache/airflow/commit/7815e8c4307b7cc10450407e6f7c72926525285c
##########
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:
Handled in
https://github.com/apache/airflow/commit/7815e8c4307b7cc10450407e6f7c72926525285c
##########
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:
Handled in
https://github.com/apache/airflow/commit/7815e8c4307b7cc10450407e6f7c72926525285c
--
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]
