feat(sdk): add task retry/cancel methods and expose available actions

changelog/+task-retry-cancel.added.md

@@ -0,0 +1 @@
+Added `retry()` and `cancel()` methods to the task manager. The `Task` model now exposes `available_actions` along with `can_retry` / `can_cancel` helpers.

infrahub_sdk/task/__init__.py

@@ -1,9 +1,11 @@
 from __future__ import annotations
 
-from .models import Task, TaskFilter, TaskLog, TaskRelatedNode, TaskState
+from .models import Task, TaskAction, TaskActionName, TaskFilter, TaskLog, TaskRelatedNode, TaskState
 
 __all__ = [
     "Task",
+    "TaskAction",
+    "TaskActionName",
     "TaskFilter",
     "TaskLog",
     "TaskRelatedNode",

infrahub_sdk/task/manager.py

@@ -4,14 +4,16 @@
 import time
 from typing import TYPE_CHECKING, Any
 
-from ..graphql import Query
+from ..graphql import Mutation, Query
 from .constants import FINAL_STATES
 from .exceptions import TaskNotCompletedError, TaskNotFoundError, TooManyTasksError
 from .models import Task, TaskFilter
 
 if TYPE_CHECKING:
     from ..client import InfrahubClient, InfrahubClientSync
 
+MUTATION_TASK_QUERY = {"ok": None, "task": {"id": None}}
+
 
 class InfraHubTaskManagerBase:
     @classmethod
@@ -20,6 +22,7 @@ def _generate_query(
         filters: TaskFilter | None = None,
         include_logs: bool = False,
         include_related_nodes: bool = False,
+        include_actions: bool = False,
         offset: int | None = None,
         limit: int | None = None,
         count: bool = False,
@@ -68,6 +71,13 @@ def _generate_query(
         if include_related_nodes:
             query["InfrahubTask"]["edges"]["node"]["related_nodes"] = {"id": None, "kind": None}
 
+        if include_actions:
+            query["InfrahubTask"]["edges"]["node"]["available_actions"] = {
+                "action": None,
+                "available": None,
+                "unavailability_reason": None,
+            }
+
         return Query(query=query)
 
     @classmethod
@@ -113,6 +123,7 @@ async def all(
         parallel: bool = False,
         include_logs: bool = False,
         include_related_nodes: bool = False,
+        include_actions: bool = False,
     ) -> list[Task]:
         """Get all tasks.
 
@@ -123,6 +134,7 @@ async def all(
             parallel: Whether to query the tasks in parallel. Defaults to False.
             include_logs: Whether to include the logs in the tasks. Defaults to False.
             include_related_nodes: Whether to include the related nodes in the tasks. Defaults to False.
+            include_actions: Whether to include the available actions in the tasks. Defaults to False.
 
         Returns:
             A list of tasks.
@@ -135,6 +147,7 @@ async def all(
             parallel=parallel,
             include_logs=include_logs,
             include_related_nodes=include_related_nodes,
+            include_actions=include_actions,
         )
 
     async def filter(
@@ -146,6 +159,7 @@ async def filter(
         parallel: bool = False,
         include_logs: bool = False,
         include_related_nodes: bool = False,
+        include_actions: bool = False,
     ) -> list[Task]:
         """Filter tasks.
 
@@ -157,6 +171,7 @@ async def filter(
             parallel: Whether to query the tasks in parallel. Defaults to False.
             include_logs: Whether to include the logs in the tasks. Defaults to False.
             include_related_nodes: Whether to include the related nodes in the tasks. Defaults to False.
+            include_actions: Whether to include the available actions in the tasks. Defaults to False.
 
         Returns:
             A list of tasks.
@@ -167,7 +182,18 @@ async def filter(
 
         if limit:
             tasks, _ = await self.process_page(
-                self.client, self._generate_query(filters=filter, offset=offset, limit=limit, count=False), 1, timeout
+                self.client,
+                self._generate_query(
+                    filters=filter,
+                    offset=offset,
+                    limit=limit,
+                    include_logs=include_logs,
+                    include_related_nodes=include_related_nodes,
+                    include_actions=include_actions,
+                    count=False,
+                ),
+                1,
+                timeout,
             )
             return tasks
 
@@ -177,6 +203,7 @@ async def filter(
                 timeout=timeout,
                 include_logs=include_logs,
                 include_related_nodes=include_related_nodes,
+                include_actions=include_actions,
             )
 
         return await self.process_non_batch(
@@ -186,13 +213,17 @@ async def filter(
             timeout=timeout,
             include_logs=include_logs,
             include_related_nodes=include_related_nodes,
+            include_actions=include_actions,
         )
 
-    async def get(self, id: str, include_logs: bool = False, include_related_nodes: bool = False) -> Task:
+    async def get(
+        self, id: str, include_logs: bool = False, include_related_nodes: bool = False, include_actions: bool = False
+    ) -> Task:
         tasks = await self.filter(
             filter=TaskFilter(ids=[id]),
             include_logs=include_logs,
             include_related_nodes=include_related_nodes,
+            include_actions=include_actions,
             parallel=False,
         )
         if not tasks:
@@ -225,6 +256,34 @@ async def wait_for_completion(self, id: str, interval: int = 1, timeout: int = 6
             await asyncio.sleep(interval)
         raise TaskNotCompletedError(id=id, message=f"Task {id} did not complete in {timeout} seconds")
 
+    async def retry(self, id: str) -> str:
+        """Retry a settled task by replaying it as a new, independent task.
+
+        Args:
+            id: The id of the task to retry.
+
+        Returns:
+            The id of the new task created by the retry.
+
+        """
+        query = Mutation(mutation="InfrahubTaskRetry", input_data={"data": {"id": id}}, query=MUTATION_TASK_QUERY)
+        response = await self.client.execute_graphql(query=query.render(), tracker="mutation-task-retry")
+        return response["InfrahubTaskRetry"]["task"]["id"]
+
+    async def cancel(self, id: str) -> bool:
+        """Cancel an in-flight task, stopping any remaining retries.
+
+        Args:
+            id: The id of the task to cancel.
+
+        Returns:
+            Whether the task was successfully cancelled.
+
+        """
+        query = Mutation(mutation="InfrahubTaskCancel", input_data={"data": {"id": id}}, query=MUTATION_TASK_QUERY)
+        response = await self.client.execute_graphql(query=query.render(), tracker="mutation-task-cancel")
+        return response["InfrahubTaskCancel"]["ok"]
+
     @staticmethod
     async def process_page(
         client: InfrahubClient, query: Query, page_number: int, timeout: int | None = None
@@ -255,6 +314,7 @@ async def process_batch(
         timeout: int | None = None,
         include_logs: bool = False,
         include_related_nodes: bool = False,
+        include_actions: bool = False,
     ) -> list[Task]:
         """Process queries in parallel mode."""
         pagination_size = self.client.pagination_size
@@ -271,6 +331,7 @@ async def process_batch(
                 limit=pagination_size,
                 include_logs=include_logs,
                 include_related_nodes=include_related_nodes,
+                include_actions=include_actions,
                 count=False,
             )
             batch_process.add(
@@ -290,6 +351,7 @@ async def process_non_batch(
         timeout: int | None = None,
         include_logs: bool = False,
         include_related_nodes: bool = False,
+        include_actions: bool = False,
     ) -> list[Task]:
         """Process queries without parallel mode.
 
@@ -309,6 +371,7 @@ async def process_non_batch(
                 limit=self.client.pagination_size,
                 include_logs=include_logs,
                 include_related_nodes=include_related_nodes,
+                include_actions=include_actions,
                 count=True,
             )
             new_tasks, count = await self.process_page(
@@ -353,6 +416,7 @@ def all(
         parallel: bool = False,
         include_logs: bool = False,
         include_related_nodes: bool = False,
+        include_actions: bool = False,
     ) -> list[Task]:
         """Get all tasks.
 
@@ -363,6 +427,7 @@ def all(
             parallel: Whether to query the tasks in parallel. Defaults to False.
             include_logs: Whether to include the logs in the tasks. Defaults to False.
             include_related_nodes: Whether to include the related nodes in the tasks. Defaults to False.
+            include_actions: Whether to include the available actions in the tasks. Defaults to False.
 
         Returns:
             A list of tasks.
@@ -375,6 +440,7 @@ def all(
             parallel=parallel,
             include_logs=include_logs,
             include_related_nodes=include_related_nodes,
+            include_actions=include_actions,
         )
 
     def filter(
@@ -386,6 +452,7 @@ def filter(
         parallel: bool = False,
         include_logs: bool = False,
         include_related_nodes: bool = False,
+        include_actions: bool = False,
     ) -> list[Task]:
         """Filter tasks.
 
@@ -397,6 +464,7 @@ def filter(
             parallel: Whether to query the tasks in parallel. Defaults to False.
             include_logs: Whether to include the logs in the tasks. Defaults to False.
             include_related_nodes: Whether to include the related nodes in the tasks. Defaults to False.
+            include_actions: Whether to include the available actions in the tasks. Defaults to False.
 
         Returns:
             A list of tasks.
@@ -407,7 +475,18 @@ def filter(
 
         if limit:
             tasks, _ = self.process_page(
-                self.client, self._generate_query(filters=filter, offset=offset, limit=limit, count=False), 1, timeout
+                self.client,
+                self._generate_query(
+                    filters=filter,
+                    offset=offset,
+                    limit=limit,
+                    include_logs=include_logs,
+                    include_related_nodes=include_related_nodes,
+                    include_actions=include_actions,
+                    count=False,
+                ),
+                1,
+                timeout,
             )
             return tasks
 
@@ -417,6 +496,7 @@ def filter(
                 timeout=timeout,
                 include_logs=include_logs,
                 include_related_nodes=include_related_nodes,
+                include_actions=include_actions,
             )
 
         return self.process_non_batch(
@@ -426,13 +506,17 @@ def filter(
             timeout=timeout,
             include_logs=include_logs,
             include_related_nodes=include_related_nodes,
+            include_actions=include_actions,
         )
 
-    def get(self, id: str, include_logs: bool = False, include_related_nodes: bool = False) -> Task:
+    def get(
+        self, id: str, include_logs: bool = False, include_related_nodes: bool = False, include_actions: bool = False
+    ) -> Task:
         tasks = self.filter(
             filter=TaskFilter(ids=[id]),
             include_logs=include_logs,
             include_related_nodes=include_related_nodes,
+            include_actions=include_actions,
             parallel=False,
         )
         if not tasks:
@@ -465,6 +549,34 @@ def wait_for_completion(self, id: str, interval: int = 1, timeout: int = 60) ->
             time.sleep(interval)
         raise TaskNotCompletedError(id=id, message=f"Task {id} did not complete in {timeout} seconds")
 
+    def retry(self, id: str) -> str:
+        """Retry a settled task by replaying it as a new, independent task.
+
+        Args:
+            id: The id of the task to retry.
+
+        Returns:
+            The id of the new task created by the retry.
+
+        """
+        query = Mutation(mutation="InfrahubTaskRetry", input_data={"data": {"id": id}}, query=MUTATION_TASK_QUERY)
+        response = self.client.execute_graphql(query=query.render(), tracker="mutation-task-retry")
+        return response["InfrahubTaskRetry"]["task"]["id"]
+
+    def cancel(self, id: str) -> bool:
+        """Cancel an in-flight task, stopping any remaining retries.
+
+        Args:
+            id: The id of the task to cancel.
+
+        Returns:
+            Whether the task was successfully cancelled.
+
+        """
+        query = Mutation(mutation="InfrahubTaskCancel", input_data={"data": {"id": id}}, query=MUTATION_TASK_QUERY)
+        response = self.client.execute_graphql(query=query.render(), tracker="mutation-task-cancel")
+        return response["InfrahubTaskCancel"]["ok"]
+
     @staticmethod
     def process_page(
         client: InfrahubClientSync, query: Query, page_number: int, timeout: int | None = None
@@ -495,6 +607,7 @@ def process_batch(
         timeout: int | None = None,
         include_logs: bool = False,
         include_related_nodes: bool = False,
+        include_actions: bool = False,
     ) -> list[Task]:
         """Process queries in parallel mode."""
         pagination_size = self.client.pagination_size
@@ -511,6 +624,7 @@ def process_batch(
                 limit=pagination_size,
                 include_logs=include_logs,
                 include_related_nodes=include_related_nodes,
+                include_actions=include_actions,
                 count=False,
             )
             batch_process.add(
@@ -530,6 +644,7 @@ def process_non_batch(
         timeout: int | None = None,
         include_logs: bool = False,
         include_related_nodes: bool = False,
+        include_actions: bool = False,
     ) -> list[Task]:
         """Process queries without parallel mode.
 
@@ -549,6 +664,7 @@ def process_non_batch(
                 limit=self.client.pagination_size,
                 include_logs=include_logs,
                 include_related_nodes=include_related_nodes,
+                include_actions=include_actions,
                 count=True,
             )
             new_tasks, count = self.process_page(