justinpakzad commented on code in PR #72746:
URL: https://github.com/apache/airflow/pull/72746#discussion_r4383151986
##########
airflow-ctl/src/airflowctl/api/operations.py:
##########
@@ -759,6 +763,93 @@ def list(self, dag_id: str, dag_run_id: str) ->
TaskInstanceCollectionResponse |
data_model=TaskInstanceCollectionResponse,
)
+ def list_state_store(
+ self, dag_id: str, dag_run_id: str, task_id: str, map_index: int |
None = None
+ ) -> TaskStateStoreCollectionResponse | ServerResponseError:
+ """List all state store entries for a task instance."""
+ return super().execute_list(
+
path=f"dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances/{task_id}/state-store",
+ data_model=TaskStateStoreCollectionResponse,
+ params=_build_query_params(map_index=map_index),
+ )
+
+ def get_state_store(
+ self, dag_id: str, dag_run_id: str, task_id: str, key: str, map_index:
int | None = None
+ ) -> TaskStateStoreResponse | ServerResponseError:
+ """Get a single task state store entry."""
+ self.response = self.client.get(
+
f"dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances/{task_id}/state-store/{key}",
+ params=_build_query_params(map_index=map_index),
+ )
+ return
TaskStateStoreResponse.model_validate_json(self.response.content)
+
+ def set_state_store(
+ self,
+ dag_id: str,
+ dag_run_id: str,
+ task_id: str,
+ key: str,
+ value: str,
+ map_index: int | None = None,
+ ) -> str | ServerResponseError:
+ """Set a task state store value. Creates or overwrites the key."""
+ try:
+ parsed_value = json.loads(value)
+ except (ValueError, TypeError):
+ parsed_value = value
Review Comment:
Yep, good shout. I've added a helper.
--
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]
