amoghrajesh commented on code in PR #66160:
URL: https://github.com/apache/airflow/pull/66160#discussion_r3181066517
##########
task-sdk/src/airflow/sdk/api/client.py:
##########
@@ -639,6 +643,100 @@ 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 |
ErrorResponse:
+ """Get a task state value from the API server."""
+ try:
+ resp = self.client.get(f"state/ti/{ti_id}/{key}")
+ except ServerResponseError as e:
+ if e.response.status_code == HTTPStatus.NOT_FOUND:
+ log.debug("Task state key not found", ti_id=ti_id, key=key)
+ return ErrorResponse(error=ErrorType.TASK_STATE_NOT_FOUND,
detail={"key": key})
+ raise
+ return TaskStateResponse.model_validate_json(resp.read())
+
+ def set(self, ti_id: uuid.UUID, key: str, value: str) -> OKResponse:
+ """Set a task state value via the API server."""
+ body = TaskStatePutBody(value=value)
+ self.client.put(f"state/ti/{ti_id}/{key}",
content=body.model_dump_json())
+ return OKResponse(ok=True)
+
+ def delete(self, ti_id: uuid.UUID, key: str) -> OKResponse:
+ """Delete a single task state key via the API server."""
+ self.client.delete(f"state/ti/{ti_id}/{key}")
+ return OKResponse(ok=True)
+
+ def clear(self, ti_id: uuid.UUID, all_map_indices: bool = False) ->
OKResponse:
+ """Clear all task state keys for a task instance via the API server."""
+ params = {"all_map_indices": "true"} if all_map_indices else {}
+ self.client.delete(f"state/ti/{ti_id}", params=params)
+ return OKResponse(ok=True)
+
+
+class AssetStateOperations:
+ __slots__ = ("client",)
+
+ def __init__(self, client: Client):
+ self.client = client
+
+ def get(
+ self, key: str, *, name: str | None = None, uri: str | None = None
+ ) -> AssetStateResponse | ErrorResponse:
+ """Get an asset state value from the API server."""
+ if name:
+ endpoint, params = "state/asset/by-name/value", {"name": name,
"key": key}
+ elif uri:
+ endpoint, params = "state/asset/by-uri/value", {"uri": uri, "key":
key}
+ else:
+ raise ValueError("Either `name` or `uri` must be provided")
Review Comment:
Extracted as `_resolve_endpoint(op, *, key, name, uri)`
--
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]
