Support --map-index in airflowctl tasks state

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

@@ -94,6 +94,10 @@ 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",
+    # Tasks state command (manual command, uses flags) - needs a Dag run with completed tasks
+    'tasks state --dag-id=example_bash_operator --dag-run-id="manual__{date_param}" --task-id=runme_0',
+    # Task instance get (auto-generated command, uses positional args) - needs a Dag run with completed tasks
+    'taskinstances get 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}',

airflow-ctl/docs/images/command_hashes.txt

@@ -1,4 +1,4 @@
-main:27a22c00dcf32e7a1a4f06672dc8e3c8
+main:0460d9c03248bee26207b20b05aa36b9
 assets:70619a2d92bda80930cde2aefcd8e1cd
 auth:d79e9c7d00c432bdbcbc2a86e2e32053
 backfill:74c8737b0a62a86ed3605fa9e6165874

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

@@ -57,6 +57,7 @@
     PoolsOperations,
     ProvidersOperations,
     ServerResponseError,
+    TaskInstancesOperations,
     VariablesOperations,
     VersionOperations,
     XComOperations,
@@ -467,6 +468,12 @@ def xcom(self):
         """Operations related to XComs."""
         return XComOperations(self)
 
+    @lru_cache()  # type: ignore[prop-decorator]
+    @property
+    def task_instances(self):
+        """Operations related to task instances."""
+        return TaskInstancesOperations(self)
+
     @lru_cache()  # type: ignore[prop-decorator]
     @property
     def plugins(self):

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

@@ -68,6 +68,7 @@
     ProviderCollectionResponse,
     QueuedEventCollectionResponse,
     QueuedEventResponse,
+    TaskInstanceResponse,
     TriggerDAGRunPostBody,
     VariableBody,
     VariableCollectionResponse,
@@ -911,6 +912,32 @@ def delete(
             raise e
 
 
+class TaskInstancesOperations(BaseOperations):
+    """Task instance operations."""
+
+    def get(
+        self,
+        dag_id: str,
+        dag_run_id: str,
+        task_id: str,
+        map_index: int = None,  # type: ignore
+    ) -> TaskInstanceResponse | ServerResponseError:
+        """
+        Get a task instance.
+
+        When ``map_index`` is non-negative, the mapped task instance endpoint is
+        called; otherwise the standard (unmapped) endpoint is used.
+        """
+        path = f"dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances/{task_id}"
+        if map_index is not None and map_index >= 0:
+            path = f"{path}/{map_index}"
+        try:
+            self.response = self.client.get(path)
+            return TaskInstanceResponse.model_validate_json(self.response.content)
+        except ServerResponseError as e:
+            raise e
+
+
 class PluginsOperations(BaseOperations):
     """Plugins operations."""
 

airflow-ctl/src/airflowctl/ctl/cli_config.py

@@ -268,6 +268,36 @@ def _load_help_texts_yaml() -> dict[str, dict[str, str]]:
     help="The Dag ID of the Dag to pause or unpause",
 )
 
+# Task Commands Args
+ARG_TASK_DAG_ID = Arg(
+    flags=("--dag-id",),
+    type=str,
+    dest="dag_id",
+    required=True,
+    help="The Dag ID",
+)
+ARG_DAG_RUN_ID = Arg(
+    flags=("--dag-run-id",),
+    type=str,
+    dest="dag_run_id",
+    required=True,
+    help="The Dag Run ID",
+)
+ARG_TASK_ID = Arg(
+    flags=("--task-id",),
+    type=str,
+    dest="task_id",
+    required=True,
+    help="The Task ID",
+)
+ARG_MAP_INDEX = Arg(
+    flags=("--map-index",),
+    type=int,
+    dest="map_index",
+    default=-1,
+    help="If set, query the mapped task instance with this map index (negative means non-mapped)",
+)
+
 ARG_ACTION_ON_EXISTING_KEY = Arg(
     flags=("-a", "--action-on-existing-key"),
     type=str,
@@ -1015,6 +1045,21 @@ def merge_commands(
     ),
 )
 
+TASK_COMMANDS = (
+    ActionCommand(
+        name="state",
+        help="Get the state of a task instance",
+        func=lazy_load_command("airflowctl.ctl.commands.task_command.task_state"),
+        args=(
+            ARG_TASK_DAG_ID,
+            ARG_DAG_RUN_ID,
+            ARG_TASK_ID,
+            ARG_MAP_INDEX,
+            ARG_OUTPUT,
+        ),
+    ),
+)
+
 core_commands: list[CLICommand] = [
     GroupCommand(
         name="auth",
@@ -1057,6 +1102,11 @@ def merge_commands(
         help="Manage Airflow variables",
         subcommands=VARIABLE_COMMANDS,
     ),
+    GroupCommand(
+        name="tasks",
+        help="Manage Airflow tasks",
+        subcommands=TASK_COMMANDS,
+    ),
 ]
 # Add generated group commands
 core_commands = merge_commands(

airflow-ctl/src/airflowctl/ctl/commands/task_command.py

@@ -0,0 +1,33 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+from airflowctl.api.client import NEW_API_CLIENT, ClientKind, provide_api_client
+from airflowctl.ctl.console_formatting import AirflowConsole
+
+
+@provide_api_client(kind=ClientKind.CLI)
+def task_state(args, api_client=NEW_API_CLIENT) -> None:
+    """Get the state of a task instance."""
+    ti = api_client.task_instances.get(
+        dag_id=args.dag_id,
+        dag_run_id=args.dag_run_id,
+        task_id=args.task_id,
+        map_index=args.map_index,
+    )
+    AirflowConsole().print_as(data=[{"state": ti.state}], output=args.output)

airflow-ctl/src/airflowctl/ctl/help_texts.yaml

@@ -100,3 +100,6 @@ xcom:
 plugins:
   list: "List all installed Airflow plugins"
   list-import-errors: "List all plugin import errors"
+
+taskinstances:
+  get: "Retrieve a task instance by Dag ID, run ID, and task ID"

airflow-ctl/tests/airflow_ctl/api/test_operations.py

@@ -91,6 +91,8 @@
     QueuedEventCollectionResponse,
     QueuedEventResponse,
     ReprocessBehavior,
+    TaskInstanceResponse,
+    TaskInstanceState,
     TriggerDAGRunPostBody,
     VariableBody,
     VariableCollectionResponse,
@@ -100,7 +102,7 @@
     XComResponse,
     XComResponseNative,
 )
-from airflowctl.api.operations import BaseOperations
+from airflowctl.api.operations import BaseOperations, ServerResponseError
 from airflowctl.exceptions import AirflowCtlConnectionException
 
 if TYPE_CHECKING:
@@ -1907,6 +1909,102 @@ def handle_request(request: httpx.Request) -> httpx.Response:
         assert response == self.key
 
 
+class TestTaskInstancesOperations:
+    """Test suite for task instance operations."""
+
+    dag_id: str = "test_dag"
+    dag_run_id: str = "manual__2025-01-24T00:00:00+00:00"
+    task_id: str = "test_task"
+
+    task_instance_response = TaskInstanceResponse(
+        id=uuid.uuid4(),
+        task_id=task_id,
+        dag_id=dag_id,
+        dag_run_id=dag_run_id,
+        map_index=-1,
+        run_after=datetime.datetime(2025, 1, 24, 0, 0, 0),
+        try_number=1,
+        max_tries=1,
+        task_display_name=task_id,
+        dag_display_name=dag_id,
+        pool="default_pool",
+        pool_slots=1,
+        executor_config="{}",
+        state=TaskInstanceState.SUCCESS,
+    )
+
+    def test_get(self):
+        """Test fetching an unmapped task instance hits the standard endpoint."""
+
+        def handle_request(request: httpx.Request) -> httpx.Response:
+            assert request.url.path == (
+                f"/api/v2/dags/{self.dag_id}/dagRuns/{self.dag_run_id}/taskInstances/{self.task_id}"
+            )
+            return httpx.Response(200, json=json.loads(self.task_instance_response.model_dump_json()))
+
+        client = make_api_client(transport=httpx.MockTransport(handle_request))
+        response = client.task_instances.get(
+            dag_id=self.dag_id,
+            dag_run_id=self.dag_run_id,
+            task_id=self.task_id,
+        )
+        assert response == self.task_instance_response
+
+    @pytest.mark.parametrize("map_index", [-1, None])
+    def test_get_without_map_index_uses_unmapped_endpoint(self, map_index):
+        """A negative or omitted ``map_index`` must not append a map index to the path."""
+
+        def handle_request(request: httpx.Request) -> httpx.Response:
+            assert request.url.path == (
+                f"/api/v2/dags/{self.dag_id}/dagRuns/{self.dag_run_id}/taskInstances/{self.task_id}"
+            )
+            return httpx.Response(200, json=json.loads(self.task_instance_response.model_dump_json()))
+
+        client = make_api_client(transport=httpx.MockTransport(handle_request))
+        response = client.task_instances.get(
+            dag_id=self.dag_id,
+            dag_run_id=self.dag_run_id,
+            task_id=self.task_id,
+            map_index=map_index,
+        )
+        assert response == self.task_instance_response
+
+    @pytest.mark.parametrize("map_index", [0, 1, 7])
+    def test_get_with_map_index_uses_mapped_endpoint(self, map_index):
+        """A non-negative ``map_index`` must hit the mapped task instance endpoint."""
+        mapped_response = self.task_instance_response.model_copy(update={"map_index": map_index})
+
+        def handle_request(request: httpx.Request) -> httpx.Response:
+            assert request.url.path == (
+                f"/api/v2/dags/{self.dag_id}/dagRuns/{self.dag_run_id}/"
+                f"taskInstances/{self.task_id}/{map_index}"
+            )
+            return httpx.Response(200, json=json.loads(mapped_response.model_dump_json()))
+
+        client = make_api_client(transport=httpx.MockTransport(handle_request))
+        response = client.task_instances.get(
+            dag_id=self.dag_id,
+            dag_run_id=self.dag_run_id,
+            task_id=self.task_id,
+            map_index=map_index,
+        )
+        assert response == mapped_response
+
+    def test_get_not_found_raises(self):
+        """A 404 from the server must surface as a ServerResponseError."""
+
+        def handle_request(request: httpx.Request) -> httpx.Response:
+            return httpx.Response(404, json={"detail": "Task instance not found"})
+
+        client = make_api_client(transport=httpx.MockTransport(handle_request))
+        with pytest.raises(ServerResponseError):
+            client.task_instances.get(
+                dag_id=self.dag_id,
+                dag_run_id=self.dag_run_id,
+                task_id=self.task_id,
+            )
+
+
 class TestPluginsOperations:
     plugin_response = PluginResponse(
         name="test-plugin",