This is an automated email from the ASF dual-hosted git repository.
henry3260 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 3f0cea0e8ba Add task instance state store commands to airflowctl
(#72746)
3f0cea0e8ba is described below
commit 3f0cea0e8ba9074b13a2ce155d83925600864589
Author: Justin Pakzad <[email protected]>
AuthorDate: Mon Oct 26 22:59:01 2026 -0400
Add task instance state store commands to airflowctl (#72746)
---
.../airflowctl_tests/test_airflowctl_commands.py | 7 +
airflow-ctl/docs/images/command_hashes.txt | 2 +-
airflow-ctl/docs/images/output_taskinstances.svg | 82 ++++++------
airflow-ctl/src/airflowctl/api/operations.py | 107 +++++++++++++--
airflow-ctl/src/airflowctl/ctl/cli_config.py | 48 ++++++-
airflow-ctl/src/airflowctl/ctl/cli_parser.py | 2 +
airflow-ctl/src/airflowctl/ctl/help_texts.yaml | 9 ++
.../tests/airflow_ctl/api/test_operations.py | 143 +++++++++++++++++++++
.../tests/airflow_ctl/ctl/test_cli_config.py | 64 ++++++++-
.../ci/prek/check_airflowctl_command_coverage.py | 24 +++-
10 files changed, 426 insertions(+), 62 deletions(-)
diff --git
a/airflow-ctl-tests/tests/airflowctl_tests/test_airflowctl_commands.py
b/airflow-ctl-tests/tests/airflowctl_tests/test_airflowctl_commands.py
index 0e7ce34ed8f..02f96d5c038 100644
--- a/airflow-ctl-tests/tests/airflowctl_tests/test_airflowctl_commands.py
+++ b/airflow-ctl-tests/tests/airflowctl_tests/test_airflowctl_commands.py
@@ -115,6 +115,13 @@ TEST_COMMANDS = [
'taskinstances get example_bash_operator "manual__{date_param}" runme_0',
'taskinstances get-dependencies example_bash_operator
"manual__{date_param}" runme_0',
'taskinstances list example_bash_operator "manual__{date_param}"',
+ 'taskinstances state-store set example_bash_operator
"manual__{date_param}" runme_0 test_key test_value',
+ 'taskinstances state-store get example_bash_operator
"manual__{date_param}" runme_0 test_key',
+ 'taskinstances state-store list example_bash_operator
"manual__{date_param}" runme_0',
+ 'taskinstances state-store update example_bash_operator
"manual__{date_param}" runme_0 test_key new_value',
+ 'taskinstances state-store delete example_bash_operator
"manual__{date_param}" runme_0 test_key',
+ 'taskinstances state-store set example_bash_operator
"manual__{date_param}" runme_0 clear_key value',
+ 'taskinstances state-store clear example_bash_operator
"manual__{date_param}" runme_0',
# XCom commands - need a Dag run with completed tasks
'xcom add example_bash_operator "manual__{date_param}" runme_0 {xcom_key}
\'{{"test": "value"}}\'',
'xcom get example_bash_operator "manual__{date_param}" runme_0 {xcom_key}',
diff --git a/airflow-ctl/docs/images/command_hashes.txt
b/airflow-ctl/docs/images/command_hashes.txt
index 0c1bf556267..d9605381065 100644
--- a/airflow-ctl/docs/images/command_hashes.txt
+++ b/airflow-ctl/docs/images/command_hashes.txt
@@ -9,7 +9,7 @@ dagrun:07035226eaaff0a557d3ad9aab4b41b5
jobs:d4af478f28dae48ee18d43b998edc345
pools:19efe105b9515ab1926ebcaf0e028d71
providers:34502fe09dc0b8b0a13e7e46efdffda6
-taskinstances:bea84117114c2438eb7e7026f6bb7042
+taskinstances:af95eda24a2bc4c6bd87c9f1324dbc3b
tasks:eb70701dfe1b9baeda39d5eea07fde5f
variables:f8fc76d3d398b2780f4e97f7cd816646
version:31f4efdf8de0dbaaa4fac71ff7efecc3
diff --git a/airflow-ctl/docs/images/output_taskinstances.svg
b/airflow-ctl/docs/images/output_taskinstances.svg
index ccde77e2bb8..7c02563d8f2 100644
--- a/airflow-ctl/docs/images/output_taskinstances.svg
+++ b/airflow-ctl/docs/images/output_taskinstances.svg
@@ -1,4 +1,4 @@
-<svg class="rich-terminal" viewBox="0 0 933 367.2"
xmlns="http://www.w3.org/2000/svg";>
+<svg class="rich-terminal" viewBox="0 0 933 416.0"
xmlns="http://www.w3.org/2000/svg";>
<!-- Generated with Rich https://www.textualize.io -->
<style>
@@ -19,90 +19,98 @@
font-weight: 700;
}
- .terminal-909544278-matrix {
+ .terminal-944627566-matrix {
font-family: Fira Code, monospace;
font-size: 20px;
line-height: 24.4px;
font-variant-east-asian: full-width;
}
- .terminal-909544278-title {
+ .terminal-944627566-title {
font-size: 18px;
font-weight: bold;
font-family: arial;
}
- .terminal-909544278-r1 { fill: #ff8700 }
-.terminal-909544278-r2 { fill: #c5c8c6 }
-.terminal-909544278-r3 { fill: #808080 }
-.terminal-909544278-r4 { fill: #68a0b3 }
+ .terminal-944627566-r1 { fill: #ff8700 }
+.terminal-944627566-r2 { fill: #c5c8c6 }
+.terminal-944627566-r3 { fill: #808080 }
+.terminal-944627566-r4 { fill: #68a0b3 }
</style>
<defs>
- <clipPath id="terminal-909544278-clip-terminal">
- <rect x="0" y="0" width="914.0" height="316.2" />
+ <clipPath id="terminal-944627566-clip-terminal">
+ <rect x="0" y="0" width="914.0" height="365.0" />
</clipPath>
- <clipPath id="terminal-909544278-line-0">
+ <clipPath id="terminal-944627566-line-0">
<rect x="0" y="1.5" width="915" height="24.65"/>
</clipPath>
-<clipPath id="terminal-909544278-line-1">
+<clipPath id="terminal-944627566-line-1">
<rect x="0" y="25.9" width="915" height="24.65"/>
</clipPath>
-<clipPath id="terminal-909544278-line-2">
+<clipPath id="terminal-944627566-line-2">
<rect x="0" y="50.3" width="915" height="24.65"/>
</clipPath>
-<clipPath id="terminal-909544278-line-3">
+<clipPath id="terminal-944627566-line-3">
<rect x="0" y="74.7" width="915" height="24.65"/>
</clipPath>
-<clipPath id="terminal-909544278-line-4">
+<clipPath id="terminal-944627566-line-4">
<rect x="0" y="99.1" width="915" height="24.65"/>
</clipPath>
-<clipPath id="terminal-909544278-line-5">
+<clipPath id="terminal-944627566-line-5">
<rect x="0" y="123.5" width="915" height="24.65"/>
</clipPath>
-<clipPath id="terminal-909544278-line-6">
+<clipPath id="terminal-944627566-line-6">
<rect x="0" y="147.9" width="915" height="24.65"/>
</clipPath>
-<clipPath id="terminal-909544278-line-7">
+<clipPath id="terminal-944627566-line-7">
<rect x="0" y="172.3" width="915" height="24.65"/>
</clipPath>
-<clipPath id="terminal-909544278-line-8">
+<clipPath id="terminal-944627566-line-8">
<rect x="0" y="196.7" width="915" height="24.65"/>
</clipPath>
-<clipPath id="terminal-909544278-line-9">
+<clipPath id="terminal-944627566-line-9">
<rect x="0" y="221.1" width="915" height="24.65"/>
</clipPath>
-<clipPath id="terminal-909544278-line-10">
+<clipPath id="terminal-944627566-line-10">
<rect x="0" y="245.5" width="915" height="24.65"/>
</clipPath>
-<clipPath id="terminal-909544278-line-11">
+<clipPath id="terminal-944627566-line-11">
<rect x="0" y="269.9" width="915" height="24.65"/>
</clipPath>
+<clipPath id="terminal-944627566-line-12">
+ <rect x="0" y="294.3" width="915" height="24.65"/>
+ </clipPath>
+<clipPath id="terminal-944627566-line-13">
+ <rect x="0" y="318.7" width="915" height="24.65"/>
+ </clipPath>
</defs>
- <rect fill="#292929" stroke="rgba(255,255,255,0.35)" stroke-width="1"
x="1" y="1" width="931" height="365.2" rx="8"/>
+ <rect fill="#292929" stroke="rgba(255,255,255,0.35)" stroke-width="1"
x="1" y="1" width="931" height="414" rx="8"/>
<g transform="translate(26,22)">
<circle cx="0" cy="0" r="7" fill="#ff5f57"/>
<circle cx="22" cy="0" r="7" fill="#febc2e"/>
<circle cx="44" cy="0" r="7" fill="#28c840"/>
</g>
- <g transform="translate(9, 41)"
clip-path="url(#terminal-909544278-clip-terminal)">
+ <g transform="translate(9, 41)"
clip-path="url(#terminal-944627566-clip-terminal)">
- <g class="terminal-909544278-matrix">
- <text class="terminal-909544278-r1" x="0" y="20" textLength="73.2"
clip-path="url(#terminal-909544278-line-0)">Usage:</text><text
class="terminal-909544278-r3" x="85.4" y="20" textLength="292.8"
clip-path="url(#terminal-909544278-line-0)">airflowctl taskinstances</text><text
class="terminal-909544278-r2" x="378.2" y="20" textLength="24.4"
clip-path="url(#terminal-909544278-line-0)"> [</text><text
class="terminal-909544278-r4" x="402.6" y="20" textLength="24.4" clip-path="ur
[...]
-</text><text class="terminal-909544278-r2" x="915" y="44.4" textLength="12.2"
clip-path="url(#terminal-909544278-line-1)">
-</text><text class="terminal-909544278-r2" x="0" y="68.8" textLength="390.4"
clip-path="url(#terminal-909544278-line-2)">Perform TaskInstances operations</text><text
class="terminal-909544278-r2" x="915" y="68.8" textLength="12.2"
clip-path="url(#terminal-909544278-line-2)">
-</text><text class="terminal-909544278-r2" x="915" y="93.2" textLength="12.2"
clip-path="url(#terminal-909544278-line-3)">
-</text><text class="terminal-909544278-r1" x="0" y="117.6" textLength="256.2"
clip-path="url(#terminal-909544278-line-4)">Positional Arguments:</text><text
class="terminal-909544278-r2" x="915" y="117.6" textLength="12.2"
clip-path="url(#terminal-909544278-line-4)">
-</text><text class="terminal-909544278-r4" x="24.4" y="142" textLength="85.4"
clip-path="url(#terminal-909544278-line-5)">COMMAND</text><text
class="terminal-909544278-r2" x="915" y="142" textLength="12.2"
clip-path="url(#terminal-909544278-line-5)">
-</text><text class="terminal-909544278-r4" x="48.8" y="166.4"
textLength="36.6" clip-path="url(#terminal-909544278-line-6)">get</text><text
class="terminal-909544278-r2" x="244" y="166.4" textLength="475.8"
clip-path="url(#terminal-909544278-line-6)">Get a task instance for a given Dag run</text><text
class="terminal-909544278-r2" x="915" y="166.4" textLength="12.2"
clip-path="url(#terminal-909544278-line-6)">
-</text><text class="terminal-909544278-r4" x="48.8" y="190.8"
textLength="195.2"
clip-path="url(#terminal-909544278-line-7)">get-dependencies</text><text
class="terminal-909544278-r2" x="915" y="190.8" textLength="12.2"
clip-path="url(#terminal-909544278-line-7)">
-</text><text class="terminal-909544278-r2" x="244" y="215.2"
textLength="634.4"
clip-path="url(#terminal-909544278-line-8)">Get unmet scheduler dependencies for a task instance</text><text
class="terminal-909544278-r2" x="915" y="215.2" textLength="12.2"
clip-path="url(#terminal-909544278-line-8)">
-</text><text class="terminal-909544278-r4" x="48.8" y="239.6"
textLength="48.8" clip-path="url(#terminal-909544278-line-9)">list</text><text
class="terminal-909544278-r2" x="244" y="239.6" textLength="524.6"
clip-path="url(#terminal-909544278-line-9)">List all task instances for a given Dag run</text><text
class="terminal-909544278-r2" x="915" y="239.6" textLength="12.2"
clip-path="url(#terminal-909544278-line-9)">
-</text><text class="terminal-909544278-r2" x="915" y="264" textLength="12.2"
clip-path="url(#terminal-909544278-line-10)">
-</text><text class="terminal-909544278-r1" x="0" y="288.4" textLength="97.6"
clip-path="url(#terminal-909544278-line-11)">Options:</text><text
class="terminal-909544278-r2" x="915" y="288.4" textLength="12.2"
clip-path="url(#terminal-909544278-line-11)">
-</text><text class="terminal-909544278-r4" x="24.4" y="312.8"
textLength="24.4" clip-path="url(#terminal-909544278-line-12)">-h</text><text
class="terminal-909544278-r2" x="48.8" y="312.8" textLength="24.4"
clip-path="url(#terminal-909544278-line-12)">, </text><text
class="terminal-909544278-r4" x="73.2" y="312.8" textLength="73.2"
clip-path="url(#terminal-909544278-line-12)">--help</text><text
class="terminal-909544278-r2" x="244" y="312.8" textLength="378.2"
clip-path="url(#termin [...]
+ <g class="terminal-944627566-matrix">
+ <text class="terminal-944627566-r1" x="0" y="20" textLength="73.2"
clip-path="url(#terminal-944627566-line-0)">Usage:</text><text
class="terminal-944627566-r3" x="85.4" y="20" textLength="292.8"
clip-path="url(#terminal-944627566-line-0)">airflowctl taskinstances</text><text
class="terminal-944627566-r2" x="378.2" y="20" textLength="24.4"
clip-path="url(#terminal-944627566-line-0)"> [</text><text
class="terminal-944627566-r4" x="402.6" y="20" textLength="24.4" clip-path="ur
[...]
+</text><text class="terminal-944627566-r2" x="915" y="44.4" textLength="12.2"
clip-path="url(#terminal-944627566-line-1)">
+</text><text class="terminal-944627566-r2" x="0" y="68.8" textLength="390.4"
clip-path="url(#terminal-944627566-line-2)">Perform TaskInstances operations</text><text
class="terminal-944627566-r2" x="915" y="68.8" textLength="12.2"
clip-path="url(#terminal-944627566-line-2)">
+</text><text class="terminal-944627566-r2" x="915" y="93.2" textLength="12.2"
clip-path="url(#terminal-944627566-line-3)">
+</text><text class="terminal-944627566-r1" x="0" y="117.6" textLength="256.2"
clip-path="url(#terminal-944627566-line-4)">Positional Arguments:</text><text
class="terminal-944627566-r2" x="915" y="117.6" textLength="12.2"
clip-path="url(#terminal-944627566-line-4)">
+</text><text class="terminal-944627566-r4" x="24.4" y="142" textLength="85.4"
clip-path="url(#terminal-944627566-line-5)">COMMAND</text><text
class="terminal-944627566-r2" x="915" y="142" textLength="12.2"
clip-path="url(#terminal-944627566-line-5)">
+</text><text class="terminal-944627566-r4" x="48.8" y="166.4"
textLength="36.6" clip-path="url(#terminal-944627566-line-6)">get</text><text
class="terminal-944627566-r2" x="244" y="166.4" textLength="475.8"
clip-path="url(#terminal-944627566-line-6)">Get a task instance for a given Dag run</text><text
class="terminal-944627566-r2" x="915" y="166.4" textLength="12.2"
clip-path="url(#terminal-944627566-line-6)">
+</text><text class="terminal-944627566-r4" x="48.8" y="190.8"
textLength="195.2"
clip-path="url(#terminal-944627566-line-7)">get-dependencies</text><text
class="terminal-944627566-r2" x="915" y="190.8" textLength="12.2"
clip-path="url(#terminal-944627566-line-7)">
+</text><text class="terminal-944627566-r2" x="244" y="215.2"
textLength="634.4"
clip-path="url(#terminal-944627566-line-8)">Get unmet scheduler dependencies for a task instance</text><text
class="terminal-944627566-r2" x="915" y="215.2" textLength="12.2"
clip-path="url(#terminal-944627566-line-8)">
+</text><text class="terminal-944627566-r4" x="48.8" y="239.6"
textLength="48.8" clip-path="url(#terminal-944627566-line-9)">list</text><text
class="terminal-944627566-r2" x="244" y="239.6" textLength="524.6"
clip-path="url(#terminal-944627566-line-9)">List all task instances for a given Dag run</text><text
class="terminal-944627566-r2" x="915" y="239.6" textLength="12.2"
clip-path="url(#terminal-944627566-line-9)">
+</text><text class="terminal-944627566-r4" x="48.8" y="264" textLength="134.2"
clip-path="url(#terminal-944627566-line-10)">state-store</text><text
class="terminal-944627566-r2" x="244" y="264" textLength="500.2"
clip-path="url(#terminal-944627566-line-10)">Manage the state store of a task instance</text><text
class="terminal-944627566-r2" x="915" y="264" textLength="12.2"
clip-path="url(#terminal-944627566-line-10)">
+</text><text class="terminal-944627566-r2" x="915" y="288.4" textLength="12.2"
clip-path="url(#terminal-944627566-line-11)">
+</text><text class="terminal-944627566-r1" x="0" y="312.8" textLength="97.6"
clip-path="url(#terminal-944627566-line-12)">Options:</text><text
class="terminal-944627566-r2" x="915" y="312.8" textLength="12.2"
clip-path="url(#terminal-944627566-line-12)">
+</text><text class="terminal-944627566-r4" x="24.4" y="337.2"
textLength="24.4" clip-path="url(#terminal-944627566-line-13)">-h</text><text
class="terminal-944627566-r2" x="48.8" y="337.2" textLength="24.4"
clip-path="url(#terminal-944627566-line-13)">, </text><text
class="terminal-944627566-r4" x="73.2" y="337.2" textLength="73.2"
clip-path="url(#terminal-944627566-line-13)">--help</text><text
class="terminal-944627566-r2" x="244" y="337.2" textLength="378.2"
clip-path="url(#termin [...]
+</text><text class="terminal-944627566-r2" x="915" y="361.6" textLength="12.2"
clip-path="url(#terminal-944627566-line-14)">
</text>
</g>
</g>
diff --git a/airflow-ctl/src/airflowctl/api/operations.py
b/airflow-ctl/src/airflowctl/api/operations.py
index a28048d5a05..a2758740b41 100644
--- a/airflow-ctl/src/airflowctl/api/operations.py
+++ b/airflow-ctl/src/airflowctl/api/operations.py
@@ -75,6 +75,10 @@ from airflowctl.api.datamodels.generated import (
TaskDependencyCollectionResponse,
TaskInstanceCollectionResponse,
TaskInstanceResponse,
+ TaskStateStoreBody,
+ TaskStateStoreCollectionResponse,
+ TaskStateStorePatchBody,
+ TaskStateStoreResponse,
TriggerDAGRunPostBody,
VariableBody,
VariableCollectionResponse,
@@ -114,6 +118,13 @@ def _has_json_body(response: httpx.Response) -> bool:
return media_type == "application/json" or media_type.endswith("+json")
+def _parse_json_or_str(value: str) -> Any:
+ try:
+ return json.loads(value)
+ except (ValueError, TypeError):
+ return value
+
+
# Generic Server Response Error
class ServerResponseError(httpx.HTTPStatusError):
"""Server response error (Generic)."""
@@ -355,10 +366,7 @@ class AssetsOperations(BaseOperations):
def set_state_store(self, asset_id: str, key: str, value: str) -> str |
ServerResponseError:
"""Set an asset state store value. Creates or overwrites the key."""
- try:
- parsed_value = json.loads(value)
- except (ValueError, TypeError):
- parsed_value = value
+ parsed_value = _parse_json_or_str(value)
self.client.put(
f"assets/{asset_id}/state-store/{key}",
json=AssetStateStoreBody(value=parsed_value).model_dump(mode="json"),
@@ -773,6 +781,87 @@ class TaskInstancesOperations(BaseOperations):
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."""
+ parsed_value = _parse_json_or_str(value)
+ self.client.put(
+
f"dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances/{task_id}/state-store/{key}",
+
json=TaskStateStoreBody(value=parsed_value).model_dump(mode="json"),
+ params=_build_query_params(map_index=map_index),
+ )
+ return key
+
+ def update_state_store(
+ self,
+ dag_id: str,
+ dag_run_id: str,
+ task_id: str,
+ key: str,
+ value: str,
+ map_index: int | None = None,
+ ) -> str | ServerResponseError:
+ """Update the value of an existing task state store key."""
+ parsed_value = _parse_json_or_str(value)
+ self.client.patch(
+
f"dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances/{task_id}/state-store/{key}",
+
json=TaskStateStorePatchBody(value=parsed_value).model_dump(mode="json"),
+ params=_build_query_params(map_index=map_index),
+ )
+ return key
+
+ def delete_state_store(
+ self, dag_id: str, dag_run_id: str, task_id: str, key: str, map_index:
int | None = None
+ ) -> str | ServerResponseError:
+ """Delete a single task state store key."""
+ self.client.delete(
+
f"dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances/{task_id}/state-store/{key}",
+ params=_build_query_params(map_index=map_index),
+ )
+ return key
+
+ def clear_state_store(
+ self,
+ dag_id: str,
+ dag_run_id: str,
+ task_id: str,
+ map_index: int | None = None,
+ all_map_indices: bool = False,
+ ) -> str | ServerResponseError:
+ """Delete all task state store keys for a task instance."""
+ self.client.delete(
+
f"dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances/{task_id}/state-store",
+ params=_build_query_params(map_index=map_index,
all_map_indices=all_map_indices or None),
+ )
+ return task_id
+
class TasksOperations(BaseOperations):
"""Tasks operations."""
@@ -885,10 +974,7 @@ class XComOperations(BaseOperations):
map_index: int = None, # type: ignore
) -> XComResponseNative | ServerResponseError:
"""Add an XCom entry."""
- try:
- parsed_value = json.loads(value)
- except (ValueError, TypeError):
- parsed_value = value
+ parsed_value = _parse_json_or_str(value)
body_dict: dict[str, Any] = {"key": key, "value": parsed_value}
if map_index is not None:
@@ -910,10 +996,7 @@ class XComOperations(BaseOperations):
map_index: int = None, # type: ignore
) -> XComResponseNative | ServerResponseError:
"""Edit an XCom entry."""
- try:
- parsed_value = json.loads(value)
- except (ValueError, TypeError):
- parsed_value = value
+ parsed_value = _parse_json_or_str(value)
body_dict: dict[str, Any] = {"value": parsed_value}
if map_index is not None:
diff --git a/airflow-ctl/src/airflowctl/ctl/cli_config.py
b/airflow-ctl/src/airflowctl/ctl/cli_config.py
index f6deacfeb25..0f44714654e 100755
--- a/airflow-ctl/src/airflowctl/ctl/cli_config.py
+++ b/airflow-ctl/src/airflowctl/ctl/cli_config.py
@@ -492,7 +492,7 @@ class CommandFactory:
operations: list[dict]
args_map: dict[tuple, list[Arg]]
func_map: dict[tuple, Callable]
- commands_map: dict[str, list[ActionCommand]]
+ commands_map: dict[str, list[CLICommand]]
group_commands_list: list[CLICommand]
auth_environment_command_list: list[str]
exclude_operation_names: list[str]
@@ -539,6 +539,12 @@ class CommandFactory:
self.excluded_output_keys = [
"total_entries",
]
+ # Lookup for building nested subcommands. Each entry is the shared
method suffix
+ # (hyphenated into the subcommand name).
+ # e.g. list_state_store -> "taskinstances state-store list".
+ self.subgroups: dict[str, set[str]] = {
+ "TaskInstancesOperations": {"state_store"},
+ }
def _inspect_operations(self) -> None:
"""Parse file and return matching Operation Method with details."""
@@ -704,6 +710,13 @@ class CommandFactory:
help=arg_help,
)
+ def _find_subgroup_suffix(self, operation_group_name: str, operation_name:
str) -> str | None:
+ """Find the subgroup suffix an operation belongs to, if one exists."""
+ for suffix in self.subgroups.get(operation_group_name, set()):
+ if operation_name.endswith(f"_{suffix}"):
+ return suffix
+ return None
+
def _get_bool_arg_default(self, parameter_type: str, field_default: Any)
-> bool | None:
"""Get default for a generated bool flag: the datamodel field default
for datamodels in ``field_bool_default_datamodels``, otherwise False."""
if parameter_type in self.field_bool_default_datamodels and (
@@ -954,6 +967,7 @@ class CommandFactory:
def _create_group_commands_from_operation(self):
"""Create GroupCommand from Operation Methods."""
+ subgroup_commands_map: dict[tuple[str, str], list[ActionCommand]] = {}
for operation in self.operations:
operation_name = operation["name"]
operation_group_name = operation["parent"].name
@@ -962,12 +976,34 @@ class CommandFactory:
)
if operation_group_name not in self.commands_map:
self.commands_map[operation_group_name] = []
+ subgroup_suffix = self._find_subgroup_suffix(
+ operation_group_name=operation_group_name,
operation_name=operation_name
+ )
+ command_name = (
+ operation_name.removesuffix(f"_{subgroup_suffix}") if
subgroup_suffix else operation_name
+ )
+ action_command = ActionCommand(
+ name=command_name.replace("_", "-"),
+ help=help_text,
+ func=self.func_map[(operation_name, operation_group_name)],
+ args=self.args_map[(operation_name, operation_group_name)],
+ )
+ if subgroup_suffix:
+ subgroup_commands_map.setdefault((operation_group_name,
subgroup_suffix), []).append(
+ action_command
+ )
+ else:
+ self.commands_map[operation_group_name].append(action_command)
+
+ for (operation_group_name, subgroup_suffix), action_commands in
subgroup_commands_map.items():
+ subgroup_name = subgroup_suffix.replace("_", "-")
self.commands_map[operation_group_name].append(
- ActionCommand(
- name=operation["name"].replace("_", "-"),
- help=help_text,
- func=self.func_map[(operation_name, operation_group_name)],
- args=self.args_map[(operation_name, operation_group_name)],
+ GroupCommand(
+ name=subgroup_name,
+
help=self.help_texts.get(operation_group_name.replace("Operations",
"").lower(), {}).get(
+ subgroup_name
+ ),
+ subcommands=action_commands,
)
)
diff --git a/airflow-ctl/src/airflowctl/ctl/cli_parser.py
b/airflow-ctl/src/airflowctl/ctl/cli_parser.py
index d48c5688bdc..4e3c1fa3684 100644
--- a/airflow-ctl/src/airflowctl/ctl/cli_parser.py
+++ b/airflow-ctl/src/airflowctl/ctl/cli_parser.py
@@ -139,6 +139,8 @@ def _sort_args(args: Iterable[Arg]) -> Iterable[Arg]:
def _add_command(subparsers: argparse._SubParsersAction, sub: CLICommand) ->
None:
+ if isinstance(sub, GroupCommand):
+ sub = GroupCommandParser.from_group_command(sub)
if isinstance(sub, ActionCommand) and sub.hide:
sub_proc = subparsers.add_parser(sub.name, epilog=sub.epilog)
else:
diff --git a/airflow-ctl/src/airflowctl/ctl/help_texts.yaml
b/airflow-ctl/src/airflowctl/ctl/help_texts.yaml
index b2e5a4129cc..e0141373ef2 100644
--- a/airflow-ctl/src/airflowctl/ctl/help_texts.yaml
+++ b/airflow-ctl/src/airflowctl/ctl/help_texts.yaml
@@ -95,6 +95,15 @@ taskinstances:
list: "List all task instances for a given Dag run"
get: "Get a task instance for a given Dag run"
get-dependencies: "Get unmet scheduler dependencies for a task instance"
+ # state-store is the nested subcommand group. The *-state-store keys stay
because
+ # help is still looked up by the operation method name, not the CLI path.
+ state-store: "Manage the state store of a task instance"
+ list-state-store: "List all state store entries for a task instance"
+ get-state-store: "Retrieve a single task state store entry by its key"
+ set-state-store: "Set a task state store value. Parsed as JSON if valid,
otherwise stored as a string"
+ update-state-store: "Update an existing key's value. Parsed as JSON if
valid, otherwise stored as a string"
+ delete-state-store: "Delete a single task state store key by its key"
+ clear-state-store: "Delete all state store keys for a task instance"
tasks:
clear: "Clear task instances of a Dag by its ID"
diff --git a/airflow-ctl/tests/airflow_ctl/api/test_operations.py
b/airflow-ctl/tests/airflow_ctl/api/test_operations.py
index 7c369038e5b..b01f5231544 100644
--- a/airflow-ctl/tests/airflow_ctl/api/test_operations.py
+++ b/airflow-ctl/tests/airflow_ctl/api/test_operations.py
@@ -103,6 +103,8 @@ from airflowctl.api.datamodels.generated import (
TaskInstanceResponse,
TaskInstanceState,
TaskOutletAssetReference,
+ TaskStateStoreCollectionResponse,
+ TaskStateStoreResponse,
TriggerDAGRunPostBody,
VariableBody,
VariableCollectionResponse,
@@ -1977,6 +1979,12 @@ class TestTaskInstancesOperations:
task_dependency_collection_response = TaskDependencyCollectionResponse(
dependencies=[TaskDependencyResponse(name="Trigger Rule",
reason="upstream tasks not done")],
)
+ task_state_store_response = TaskStateStoreResponse(
+ key="my_key",
+ value={"my_val": 0}, # type: ignore[arg-type]
+ updated_at=datetime.datetime(2025, 1, 1, 0, 0, 0),
+ expires_at=None,
+ )
def _make_client_asserting_path(self, expected_path: str, response_model)
-> Client:
def handle_request(request: httpx.Request) -> httpx.Response:
@@ -2032,6 +2040,141 @@ class TestTaskInstancesOperations:
response = client.task_instances.list(dag_id="dag_id",
dag_run_id="dag_run_id")
assert response == self.task_instance_collection_response
+ @pytest.mark.parametrize("map_index", [None, 2])
+ def test_list_state_store(self, map_index):
+ collection_response = TaskStateStoreCollectionResponse(
+ task_state_store=[self.task_state_store_response],
+ total_entries=1,
+ )
+
+ def handle_request(request: httpx.Request) -> httpx.Response:
+ assert request.url.path == (
+
"/api/v2/dags/dag_id/dagRuns/dag_run_id/taskInstances/task_id/state-store"
+ )
+ if map_index is None:
+ assert "map_index" not in request.url.params
+ else:
+ assert request.url.params["map_index"] == str(map_index)
+ return httpx.Response(200,
json=json.loads(collection_response.model_dump_json()))
+
+ client = make_api_client(transport=httpx.MockTransport(handle_request))
+ response = client.task_instances.list_state_store(
+ "dag_id", "dag_run_id", "task_id", map_index=map_index
+ )
+ assert response == collection_response
+
+ @pytest.mark.parametrize("map_index", [None, 2])
+ def test_get_state_store(self, map_index):
+ key = self.task_state_store_response.key
+
+ def handle_request(request: httpx.Request) -> httpx.Response:
+ assert request.url.path == (
+
f"/api/v2/dags/dag_id/dagRuns/dag_run_id/taskInstances/task_id/state-store/{key}"
+ )
+ if map_index is None:
+ assert "map_index" not in request.url.params
+ else:
+ assert request.url.params["map_index"] == str(map_index)
+ return httpx.Response(200,
json=json.loads(self.task_state_store_response.model_dump_json()))
+
+ client = make_api_client(transport=httpx.MockTransport(handle_request))
+ response = client.task_instances.get_state_store(
+ "dag_id", "dag_run_id", "task_id", key, map_index=map_index
+ )
+ assert response == self.task_state_store_response
+
+ @pytest.mark.parametrize(
+ ("value", "expected"),
+ [
+ ('{"index": 0}', {"index": 0}),
+ ("hello", "hello"),
+ ],
+ )
+ def test_set_state_store(self, value, expected):
+ key = self.task_state_store_response.key
+
+ def handle_request(request: httpx.Request) -> httpx.Response:
+ assert request.method == "PUT"
+ assert request.url.path == (
+
f"/api/v2/dags/dag_id/dagRuns/dag_run_id/taskInstances/task_id/state-store/{key}"
+ )
+ assert json.loads(request.content) == {"value": expected,
"expires_at": "default"}
+ return httpx.Response(204)
+
+ client = make_api_client(transport=httpx.MockTransport(handle_request))
+ response = client.task_instances.set_state_store("dag_id",
"dag_run_id", "task_id", key, value)
+ assert response == key
+
+ @pytest.mark.parametrize(
+ ("value", "expected"),
+ [
+ ('{"index": 0}', {"index": 0}),
+ ("hello", "hello"),
+ ],
+ )
+ def test_update_state_store(self, value, expected):
+ key = self.task_state_store_response.key
+
+ def handle_request(request: httpx.Request) -> httpx.Response:
+ assert request.method == "PATCH"
+ assert request.url.path == (
+
f"/api/v2/dags/dag_id/dagRuns/dag_run_id/taskInstances/task_id/state-store/{key}"
+ )
+ assert json.loads(request.content) == {"value": expected}
+ return httpx.Response(200)
+
+ client = make_api_client(transport=httpx.MockTransport(handle_request))
+ response = client.task_instances.update_state_store("dag_id",
"dag_run_id", "task_id", key, value)
+ assert response == key
+
+ @pytest.mark.parametrize("map_index", [None, 2])
+ def test_delete_state_store(self, map_index):
+ key = self.task_state_store_response.key
+
+ def handle_request(request: httpx.Request) -> httpx.Response:
+ assert request.method == "DELETE"
+ assert request.url.path == (
+
f"/api/v2/dags/dag_id/dagRuns/dag_run_id/taskInstances/task_id/state-store/{key}"
+ )
+ if map_index is None:
+ assert "map_index" not in request.url.params
+ else:
+ assert request.url.params["map_index"] == str(map_index)
+ return httpx.Response(204)
+
+ client = make_api_client(transport=httpx.MockTransport(handle_request))
+ response = client.task_instances.delete_state_store(
+ "dag_id", "dag_run_id", "task_id", key, map_index=map_index
+ )
+ assert response == key
+
+ @pytest.mark.parametrize(
+ ("map_index", "all_map_indices", "expected_params"),
+ [
+ (None, False, {}),
+ (2, False, {"map_index": "2"}),
+ (None, True, {"all_map_indices": "true"}),
+ ],
+ )
+ def test_clear_state_store(self, map_index, all_map_indices,
expected_params):
+ def handle_request(request: httpx.Request) -> httpx.Response:
+ assert request.method == "DELETE"
+ assert request.url.path == (
+
"/api/v2/dags/dag_id/dagRuns/dag_run_id/taskInstances/task_id/state-store"
+ )
+ assert dict(request.url.params) == expected_params
+ return httpx.Response(204)
+
+ client = make_api_client(transport=httpx.MockTransport(handle_request))
+ response = client.task_instances.clear_state_store(
+ "dag_id",
+ "dag_run_id",
+ "task_id",
+ map_index=map_index,
+ all_map_indices=all_map_indices,
+ )
+ assert response == "task_id"
+
class TestTasksOperations:
dag_id = "dag_id"
diff --git a/airflow-ctl/tests/airflow_ctl/ctl/test_cli_config.py
b/airflow-ctl/tests/airflow_ctl/ctl/test_cli_config.py
index e357177263f..8417860fcc4 100644
--- a/airflow-ctl/tests/airflow_ctl/ctl/test_cli_config.py
+++ b/airflow-ctl/tests/airflow_ctl/ctl/test_cli_config.py
@@ -306,6 +306,44 @@ class TestCommandFactory:
assert arg.kwargs.get("default") ==
test_arg[1].get("default")
assert arg.kwargs["type"] == test_arg[1]["type"]
+ def test_command_factory_builds_nested_subgroup(self, tmp_path):
+ """Methods matching a declared suffix collapse into a nested
GroupCommand."""
+ temp_file = self._save_temp_operations_py(
+ tmp_path=tmp_path,
+ file_content="""
+ class TaskInstancesOperations(BaseOperations):
+ def list(self, dag_id: str, dag_run_id: str) ->
TaskInstanceCollectionResponse:
+ ...
+ def list_state_store(self, dag_id: str, dag_run_id: str,
task_id: str) -> TaskStateStoreCollectionResponse:
+ ...
+ def get_state_store(self, dag_id: str, dag_run_id: str,
task_id: str, key: str) -> TaskStateStoreResponse:
+ ...
+ """,
+ )
+
+ command_factory = CommandFactory(file_path=str(temp_file))
+ taskinstances = next(g for g in command_factory.group_commands if
g.name == "taskinstances")
+ subcommand_by_name = {c.name: c for c in taskinstances.subcommands}
+
+ assert isinstance(subcommand_by_name["list"], ActionCommand)
+
+ state_store = subcommand_by_name["state-store"]
+ assert isinstance(state_store, GroupCommand)
+ assert state_store.help == "Manage the state store of a task instance"
+ assert sorted(c.name for c in state_store.subcommands) == ["get",
"list"]
+ assert all(isinstance(c, ActionCommand) for c in
state_store.subcommands)
+
+ def test_find_subgroup_suffix(self):
+ """Suffix match is scoped to the declaring operations class."""
+ command_factory = CommandFactory()
+
+ assert (
+ command_factory._find_subgroup_suffix("TaskInstancesOperations",
"list_state_store")
+ == "state_store"
+ )
+ assert
command_factory._find_subgroup_suffix("TaskInstancesOperations", "list") is None
+ assert command_factory._find_subgroup_suffix("DagsOperations",
"list_state_store") is None
+
def test_command_factory_optional_bool_uses_boolean_optional_action(self,
tmp_path):
"""Optional bool parameters should support --flag and --no-flag
forms."""
temp_file = self._save_temp_operations_py(
@@ -492,10 +530,21 @@ class TestCommandFactory:
"""``_get_func`` always prints through ``args.output``, so every
generated command must declare it."""
command_factory = CommandFactory()
+ def iter_action_commands(commands, prefix):
+ """Yield (qualified_name, ActionCommand) for leaves, descending
into nested subgroups."""
+ for command in commands:
+ name = f"{prefix}{command.name}"
+ if isinstance(command, GroupCommand):
+ yield from iter_action_commands(command.subcommands,
prefix=f"{name} ")
+ else:
+ yield name, command
+
missing = [
- f"{group_command.name} {sub_command.name}"
+ name
for group_command in command_factory.group_commands
- for sub_command in group_command.subcommands
+ for name, sub_command in iter_action_commands(
+ group_command.subcommands, prefix=f"{group_command.name} "
+ )
if ARG_OUTPUT not in sub_command.args
]
@@ -947,6 +996,7 @@ class TestCliConfigMethods:
("taskinstances", "list", "List all task instances for a given Dag
run"),
("taskinstances", "get", "Get a task instance for a given Dag
run"),
("taskinstances", "get-dependencies", "Get unmet scheduler
dependencies for a task instance"),
+ ("taskinstances", "state-store", "Manage the state store of a task
instance"),
("tasks", "clear", "Clear task instances of a Dag by its ID"),
],
)
@@ -1116,3 +1166,13 @@ class TestCliConfigMethods:
rows = json.loads(capsys.readouterr().out)
assert [row["dag_run_id"] for row in rows] == ["manual_run"]
assert rows[0]["conf"] == ["a", "b"]
+
+ def test_parser_resolves_nested_subcommand(self):
+ """The parser descends into the state-store subgroup, binds the func,
and keeps the auth token."""
+ args = cli_parser.get_parser().parse_args(
+ ["taskinstances", "state-store", "list", "my_dag", "my_run",
"my_task"]
+ )
+
+ assert callable(args.func)
+ assert args.dag_id == "my_dag"
+ assert hasattr(args, "api_token")
diff --git a/scripts/ci/prek/check_airflowctl_command_coverage.py
b/scripts/ci/prek/check_airflowctl_command_coverage.py
index b5a73d86ea7..3e3be605f65 100755
--- a/scripts/ci/prek/check_airflowctl_command_coverage.py
+++ b/scripts/ci/prek/check_airflowctl_command_coverage.py
@@ -72,6 +72,11 @@ EXCLUDED_COMMANDS = {
"dags get-tags",
}
+SUBGROUPS = {
+ "taskinstances": {"state-store"},
+}
+NESTED_PREFIXES = {f"{group} {subgroup}" for group, subgroups in
SUBGROUPS.items() for subgroup in subgroups}
+
def parse_tested_commands() -> set[str]:
tested: set[str] = set()
@@ -81,11 +86,17 @@ def parse_tested_commands() -> set[str]:
# Match command patterns like "assets list", "dags list-import-errors",
etc.
# Also handles f-strings like f"dagrun get..." or f'dagrun get...'
- pattern = r'f?["\']([a-z]+(?:-[a-z]+)*\s+[a-z]+(?:-[a-z]+)*)'
+ # The third token is only kept for nested subgroups (NESTED_PREFIXES).
+ pattern = r'f?["\']([a-z]+(?:-[a-z]+)*(?:\s+[a-z]+(?:-[a-z]+)*){1,2})'
for match in re.findall(pattern, content):
parts = match.split()
- if len(parts) >= 2:
- tested.add(f"{parts[0]} {parts[1]}")
+ if len(parts) < 2:
+ continue
+ two = f"{parts[0]} {parts[1]}"
+ if two in NESTED_PREFIXES and len(parts) >= 3:
+ tested.add(f"{two} {parts[2]}")
+ else:
+ tested.add(two)
return tested
@@ -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(f"-{sg}")), None)
+ cmd = (
+ f"{group} {subgroup} {subcommand.removesuffix(f'-{subgroup}')}"
+ if subgroup is not None
+ else f"{group} {subcommand}"
+ )
if cmd not in tested and cmd not in EXCLUDED_COMMANDS:
missing.append(cmd)
