henry3260 commented on code in PR #72746:
URL: https://github.com/apache/airflow/pull/72746#discussion_r4377856854
##########
scripts/ci/prek/check_airflowctl_command_coverage.py:
##########
@@ -101,7 +112,12 @@ def main():
missing = []
for group, subcommands in sorted(available.items()):
for subcommand in sorted(subcommands):
- cmd = f"{group} {subcommand}"
+ subgroup = next((sg for sg in SUBGROUPS.get(group, ()) if
subcommand.endswith(sg)), None)
Review Comment:
This duplicates the registry in
airflow-ctl/src/airflowctl/ctl/cli_config.py:542 and, more importantly, matches
differently: cli_config.py:716 uses `endswith(f"_{suffix}") `with the
separator, but this line drops it while the `removesuffix(f'-{subgroup}')`
below still expects it. Can we use `endswith(f"-{sg}") `here?
##########
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:
This block now appears five times in this file (:349, :797, :818, :966,
:991), two of them added here. Could we extract something like
`_parse_json_or_str(value)` alongside _build_query_params?
--
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]
