Add task instance management commands to airflowctl CLI

Author: Suraj-kumar00

airflow-ctl-tests/tests/airflowctl_tests/test_airflowctl_commands.py

@@ -96,6 +96,11 @@ def date_param():
     "dags update example_bash_operator --no-is-paused",
     # Dag Run commands
     "dagrun list --dag-id example_bash_operator --state success --limit=1",
+    # Task instance commands
+    'tasks list example_bash_operator "manual__{date_param}"',
+    'tasks get example_bash_operator "manual__{date_param}" runme_0',
+    "tasks clear example_bash_operator --dry-run",
+    'tasks update example_bash_operator "manual__{date_param}" runme_0 --new-state=failed',
     # 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}',

airflow-ctl/docs/images/command_hashes.txt

@@ -1,4 +1,4 @@
-main:27a22c00dcf32e7a1a4f06672dc8e3c8
+main:164bc97843d5be583c0b48f7a34dc8c8
 assets:6419e20452692f577c4c6f570b74be0c
 auth:d79e9c7d00c432bdbcbc2a86e2e32053
 backfill:74c8737b0a62a86ed3605fa9e6165874
@@ -12,4 +12,5 @@ providers:34502fe09dc0b8b0a13e7e46efdffda6
 variables:f8fc76d3d398b2780f4e97f7cd816646
 version:31f4efdf8de0dbaaa4fac71ff7efecc3
 plugins:4864fd8f356704bd2b3cd1aec3567e35
+tasks:7ab24cac521242b6b6012e2bcd317831
 auth login:9fe2bb1dd5c602beea2eefb33a2b20a8

airflow-ctl/docs/images/output_main.svg

@@ -57,6 +57,7 @@
     PoolsOperations,
     ProvidersOperations,
     ServerResponseError,
+    TasksOperations,
     VariablesOperations,
     VersionOperations,
     XComOperations,
@@ -474,6 +475,12 @@ def plugins(self):
         """Operations related to plugins."""
         return PluginsOperations(self)
 
+    @lru_cache()  # type: ignore[prop-decorator]
+    @property
+    def tasks(self):
+        """Operations related to tasks."""
+        return TasksOperations(self)
+
 
 # API Client Decorator for CLI Actions
 @contextlib.contextmanager

airflow-ctl/docs/images/output_tasks.svg

@@ -39,6 +39,7 @@
     BulkBodyPoolBody,
     BulkBodyVariableBody,
     BulkResponse,
+    ClearTaskInstancesBody,
     Config,
     ConnectionBody,
     ConnectionCollectionResponse,
@@ -59,6 +60,7 @@
     ImportErrorCollectionResponse,
     ImportErrorResponse,
     JobCollectionResponse,
+    PatchTaskInstanceBody,
     PluginCollectionResponse,
     PluginImportErrorCollectionResponse,
     PoolBody,
@@ -68,6 +70,8 @@
     ProviderCollectionResponse,
     QueuedEventCollectionResponse,
     QueuedEventResponse,
+    TaskInstanceCollectionResponse,
+    TaskInstanceResponse,
     TriggerDAGRunPostBody,
     VariableBody,
     VariableCollectionResponse,
@@ -950,3 +954,45 @@ def list_import_errors(self) -> PluginImportErrorCollectionResponse | ServerResp
             return PluginImportErrorCollectionResponse.model_validate_json(self.response.content)
         except ServerResponseError as e:
             raise e
+
+
+class TasksOperations(BaseOperations):
+    """Tasks operations."""
+
+    def get(self, dag_id: str, dag_run_id: str, task_id: str) -> TaskInstanceResponse | ServerResponseError:
+        """Get a task instance."""
+        self.response = self.client.get(f"dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances/{task_id}")
+        return TaskInstanceResponse.model_validate_json(self.response.content)
+
+    def list(self, dag_id: str, dag_run_id: str) -> TaskInstanceCollectionResponse | ServerResponseError:
+        """List task instances."""
+        return super().execute_list(
+            path=f"dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances",
+            data_model=TaskInstanceCollectionResponse,
+        )
+
+    def clear(
+        self, dag_id: str, body: ClearTaskInstancesBody
+    ) -> TaskInstanceCollectionResponse | ServerResponseError:
+        """Clear task instances."""
+        self.response = self.client.post(
+            f"dags/{dag_id}/clearTaskInstances",
+            json=body.model_dump(mode="json", exclude_unset=True),
+        )
+        return TaskInstanceCollectionResponse.model_validate_json(self.response.content)
+
+    def update(
+        self,
+        dag_id: str,
+        dag_run_id: str,
+        task_id: str,
+        body: PatchTaskInstanceBody,
+        map_index: int | None = None,
+    ) -> TaskInstanceCollectionResponse | ServerResponseError:
+        """Update task instance state. When map_index is given, only that mapped instance is affected."""
+        if map_index is not None:
+            path = f"dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances/{task_id}/{map_index}"
+        else:
+            path = f"dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances/{task_id}"
+        self.response = self.client.patch(path, json=body.model_dump(mode="json", exclude_unset=True))
+        return TaskInstanceCollectionResponse.model_validate_json(self.response.content)

airflow-ctl/src/airflowctl/api/client.py

@@ -26,6 +26,8 @@
 import inspect
 import os
 import sys
+import types as builtin_types
+import typing
 from argparse import Namespace
 from collections.abc import Callable, Iterable
 from enum import Enum
@@ -52,6 +54,22 @@
 BUILD_DOCS = "BUILDING_AIRFLOW_DOCS" in os.environ
 
 
+def _is_list_annotation(annotation: Any) -> bool:
+    """Check whether a Pydantic field annotation is a list type (including Optional[list[...]])."""
+    origin = typing.get_origin(annotation)
+    if origin is list:
+        return True
+    # Handle both typing.Union (Optional[list[...]]) and PEP-604 X | Y (types.UnionType)
+    if origin is typing.Union or isinstance(annotation, builtin_types.UnionType):
+        return any(_is_list_annotation(arg) for arg in typing.get_args(annotation) if arg is not type(None))
+    return False
+
+
+def _get_bool_flag_default(field_info: Any) -> bool:
+    """Return the CLI default for a generated bool flag: the datamodel field's default (the API default), or False."""
+    return field_info.default if isinstance(field_info.default, bool) else False
+
+
 def lazy_load_command(import_path: str) -> Callable:
     """Create a lazy loader for command."""
     _, _, name = import_path.rpartition(".")
@@ -399,7 +417,17 @@ def __init__(self, file_path: str | Path | None = None):
         # Exclude parameters that are not needed for CLI from datamodels
         self.excluded_parameters = ["schema_"]
         # This list is used to determine if the command/operation needs to output data
-        self.output_command_list = ["list", "get", "create", "delete", "update", "trigger", "add", "edit"]
+        self.output_command_list = [
+            "list",
+            "get",
+            "create",
+            "delete",
+            "update",
+            "trigger",
+            "add",
+            "edit",
+            "clear",
+        ]
         self.exclude_operation_names = ["LoginOperations", "VersionOperations", "BaseOperations"]
         self.exclude_method_names = [
             "error",
@@ -411,6 +439,8 @@ def __init__(self, file_path: str | Path | None = None):
         ]
         self.excluded_output_keys = [
             "total_entries",
+            "next_cursor",
+            "previous_cursor",
         ]
 
     def _inspect_operations(self) -> None:
@@ -587,7 +617,9 @@ def _create_arg_for_non_primitive_type(
                         arg_type=self._python_type_from_string(field_type.annotation),
                         arg_action=argparse.BooleanOptionalAction if field_type.annotation is bool else None,  # type: ignore
                         arg_help=f"{field} for {parameter_key} operation",
-                        arg_default=False if field_type.annotation is bool else None,
+                        arg_default=_get_bool_flag_default(field_type)
+                        if field_type.annotation is bool
+                        else None,
                     )
                 )
             else:
@@ -602,7 +634,7 @@ def _create_arg_for_non_primitive_type(
                         arg_type=self._python_type_from_string(annotation),
                         arg_action=argparse.BooleanOptionalAction if annotation is bool else None,  # type: ignore
                         arg_help=f"{field} for {parameter_key} operation",
-                        arg_default=False if annotation is bool else None,
+                        arg_default=_get_bool_flag_default(field_type) if annotation is bool else None,
                     )
                 )
         return commands
@@ -717,10 +749,19 @@ def _get_func(args: Namespace, api_operation: dict, api_client: Client = NEW_API
                                 datamodel_param_name = parameter_key
                             if expanded_parameter in self.excluded_parameters:
                                 continue
-                            if expanded_parameter in args_dict.keys():
+                            if (
+                                expanded_parameter in args_dict.keys()
+                                and args_dict[expanded_parameter] is not None
+                            ):
+                                val = args_dict[expanded_parameter]
+                                if isinstance(val, str) and expanded_parameter in datamodel.model_fields:
+                                    if _is_list_annotation(
+                                        datamodel.model_fields[expanded_parameter].annotation
+                                    ):
+                                        val = [v.strip() for v in val.split(",") if v.strip()]
                                 method_params[parameter_key][
                                     self._sanitize_method_param_key(expanded_parameter)
-                                ] = args_dict[expanded_parameter]
+                                ] = val
 
             if datamodel:
                 if datamodel_param_name:

airflow-ctl/src/airflowctl/api/operations.py

@@ -102,3 +102,9 @@ xcom:
 plugins:
   list: "List all installed Airflow plugins"
   list-import-errors: "List all plugin import errors"
+
+tasks:
+  get: "Retrieve a task instance by Dag ID, run ID, and task ID"
+  list: "List all task instances for a given Dag run"
+  clear: "Clear task instance state, optionally filtering by task IDs or run scope"
+  update: "Update task instance state or note; use --map-index to target a single mapped instance"