Lee-W commented on code in PR #66160:
URL: https://github.com/apache/airflow/pull/66160#discussion_r3181600658
##########
task-sdk/src/airflow/sdk/execution_time/context.py:
##########
@@ -406,6 +405,168 @@ 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.execution_time.comms import ErrorResponse,
GetTaskState, TaskStateResult
+ from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
+
+ resp = SUPERVISOR_COMMS.send(GetTaskState(ti_id=self._ti_id, key=key))
+ if isinstance(resp, ErrorResponse) and resp.error !=
ErrorType.TASK_STATE_NOT_FOUND:
+ raise AirflowRuntimeError(resp)
+ if isinstance(resp, TaskStateResult):
+ return resp.value
+ return None
+
+ def set(self, key: str, value: str) -> None:
+ """Write or overwrite the value for the given key."""
+ from airflow.sdk.execution_time.comms import SetTaskState
+ from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
+
+ SUPERVISOR_COMMS.send(SetTaskState(ti_id=self._ti_id, key=key,
value=value))
+
+ def delete(self, key: str) -> None:
+ """Delete a single key. No-op if the key does not exist."""
+ from airflow.sdk.execution_time.comms import DeleteTaskState
+ from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
+
+ SUPERVISOR_COMMS.send(DeleteTaskState(ti_id=self._ti_id, key=key))
+
+ def clear(self, all_map_indices: bool = False) -> None:
+ """
+ Delete all keys for this task instance.
+
+ Pass ``all_map_indices=True`` to wipe state across every mapped
+ instance of the task (fleet-wide reset). Defaults to clearing only
+ this task instance's own state.
+ """
+ from airflow.sdk.execution_time.comms import ClearTaskState
+ from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
+
+ SUPERVISOR_COMMS.send(ClearTaskState(ti_id=self._ti_id,
all_map_indices=all_map_indices))
+
+
+class AssetStateAccessor:
+ """
+ Accessor for asset state scoped to an asset.
+
+ Available as ``context['asset_state']`` when the task has exactly one inlet
+ asset - asset watcher pattern. The accessor is bound to that Asset at
context construction time.
+ """
+
+ # TODO: When the task has zero or more than one inlet — the key
+ # will not be present in ``context`` in those cases. Multi-inlet support
+ # ie: ``context['asset_state'][MY_ASSET].get(...)`` can be developed as
needed.
Review Comment:
After a second thought, I think supporting
`context['asset_state'][MY_ASSET].get(...)` first actually makes more sense to
the [current
design](https://airflow.apache.org/docs/apache-airflow/stable/authoring-and-scheduling/assets.html#dependency-between-asset-task-and-classic-operators).
We can make `context['asset_state'].get(...)` some kind of syntax sugar
afterward instead.
Would love to hear why we'd like to design a single fetch case first? Thanks!
##########
task-sdk/src/airflow/sdk/execution_time/context.py:
##########
@@ -406,6 +405,168 @@ 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.execution_time.comms import ErrorResponse,
GetTaskState, TaskStateResult
+ from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
+
+ resp = SUPERVISOR_COMMS.send(GetTaskState(ti_id=self._ti_id, key=key))
+ if isinstance(resp, ErrorResponse) and resp.error !=
ErrorType.TASK_STATE_NOT_FOUND:
+ raise AirflowRuntimeError(resp)
+ if isinstance(resp, TaskStateResult):
+ return resp.value
+ return None
+
+ def set(self, key: str, value: str) -> None:
+ """Write or overwrite the value for the given key."""
+ from airflow.sdk.execution_time.comms import SetTaskState
+ from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
+
+ SUPERVISOR_COMMS.send(SetTaskState(ti_id=self._ti_id, key=key,
value=value))
+
+ def delete(self, key: str) -> None:
+ """Delete a single key. No-op if the key does not exist."""
+ from airflow.sdk.execution_time.comms import DeleteTaskState
+ from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
+
+ SUPERVISOR_COMMS.send(DeleteTaskState(ti_id=self._ti_id, key=key))
+
+ def clear(self, all_map_indices: bool = False) -> None:
+ """
+ Delete all keys for this task instance.
+
+ Pass ``all_map_indices=True`` to wipe state across every mapped
+ instance of the task (fleet-wide reset). Defaults to clearing only
+ this task instance's own state.
+ """
+ from airflow.sdk.execution_time.comms import ClearTaskState
+ from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
+
+ SUPERVISOR_COMMS.send(ClearTaskState(ti_id=self._ti_id,
all_map_indices=all_map_indices))
+
+
+class AssetStateAccessor:
+ """
+ Accessor for asset state scoped to an asset.
+
+ Available as ``context['asset_state']`` when the task has exactly one inlet
+ asset - asset watcher pattern. The accessor is bound to that Asset at
context construction time.
+ """
+
+ # TODO: When the task has zero or more than one inlet — the key
+ # will not be present in ``context`` in those cases. Multi-inlet support
+ # ie: ``context['asset_state'][MY_ASSET].get(...)`` can be developed as
needed.
Review Comment:
```python
@task(inlets=[write_to_s3])
def post_process_s3_file(*, inlet_events):
events = inlet_events[example_s3_asset]
last_row_count = events[-1].extra["row_count"]
post_process_s3_file()
```
inlet event fetching works this way
```python
context["inlet_events"][example_s3_asset]
```
It's confusing if we only/first support
```python
context["asset_state"].get(...)
```
instead of
```python
context["asset_state"][example_s3_asset].get(...)
``
##########
task-sdk/src/airflow/sdk/execution_time/context.py:
##########
@@ -406,6 +404,172 @@ 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.execution_time.comms import ErrorResponse,
GetTaskState, TaskStateResult
+ from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
+
+ resp = SUPERVISOR_COMMS.send(GetTaskState(ti_id=self._ti_id, key=key))
+ if isinstance(resp, ErrorResponse):
+ if resp.error == ErrorType.TASK_STATE_NOT_FOUND:
+ return None
+ raise AirflowRuntimeError(resp)
+ if isinstance(resp, TaskStateResult):
+ return resp.value
+ return None
+
+ def set(self, key: str, value: str) -> None:
+ """Write or overwrite the value for the given key."""
+ from airflow.sdk.execution_time.comms import SetTaskState
+ from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
+
+ SUPERVISOR_COMMS.send(SetTaskState(ti_id=self._ti_id, key=key,
value=value))
+
+ def delete(self, key: str) -> None:
+ """Delete a single key. No-op if the key does not exist."""
+ from airflow.sdk.execution_time.comms import DeleteTaskState
+ from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
+
+ SUPERVISOR_COMMS.send(DeleteTaskState(ti_id=self._ti_id, key=key))
+
+ def clear(self, all_map_indices: bool = False) -> None:
+ """
+ Delete all keys for this task instance.
+
+ Pass ``all_map_indices=True`` to wipe state across every mapped
+ instance of the task (fleet-wide reset). Defaults to clearing only
+ this task instance's own state.
+ """
+ from airflow.sdk.execution_time.comms import ClearTaskState
+ from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
+
+ SUPERVISOR_COMMS.send(ClearTaskState(ti_id=self._ti_id,
all_map_indices=all_map_indices))
+
+
+class AssetStateAccessor:
+ """
+ Accessor for asset state scoped to an asset.
+
+ Available as ``context['asset_state']`` when the task has exactly one inlet
+ asset - asset watcher pattern. The accessor is bound to that Asset at
context construction time.
+ """
+
+ # TODO: When the task has zero or more than one inlet — the key
+ # will not be present in ``context`` in those cases. Multi-inlet support
+ # ie: ``context['asset_state'][MY_ASSET].get(...)`` can be developed as
needed.
+
+ def __init__(self, *, name: str | None = None, uri: str | None = None) ->
None:
Review Comment:
sounds good
--
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]
