diff --git a/changelog/+task-retry-cancel.added.md b/changelog/+task-retry-cancel.added.md new file mode 100644 index 000000000..7b2ca3308 --- /dev/null +++ b/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. diff --git a/infrahub_sdk/task/__init__.py b/infrahub_sdk/task/__init__.py index 601803158..6cd642aa9 100644 --- a/infrahub_sdk/task/__init__.py +++ b/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", diff --git a/infrahub_sdk/task/manager.py b/infrahub_sdk/task/manager.py index 913c6b753..dbd2671d3 100644 --- a/infrahub_sdk/task/manager.py +++ b/infrahub_sdk/task/manager.py @@ -4,7 +4,7 @@ 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 @@ -12,6 +12,8 @@ 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( diff --git a/infrahub_sdk/task/models.py b/infrahub_sdk/task/models.py index 2525bda21..7ef3e4ed1 100644 --- a/infrahub_sdk/task/models.py +++ b/infrahub_sdk/task/models.py @@ -18,12 +18,23 @@ class TaskState(str, Enum): CANCELLING = "CANCELLING" +class TaskActionName(str, Enum): + RETRY = "RETRY" + CANCEL = "CANCEL" + + class TaskLog(BaseModel): message: str severity: str timestamp: datetime +class TaskAction(BaseModel): + action: TaskActionName + available: bool + unavailability_reason: str | None = None + + class TaskRelatedNode(BaseModel): id: str kind: str @@ -43,11 +54,23 @@ class Task(BaseModel): tags: list[str] | None = None related_nodes: list[TaskRelatedNode] = Field(default_factory=list) logs: list[TaskLog] = Field(default_factory=list) + available_actions: list[TaskAction] = Field(default_factory=list) + + @property + def can_retry(self) -> bool: + """Whether this task can currently be retried.""" + return any(action.action is TaskActionName.RETRY and action.available for action in self.available_actions) + + @property + def can_cancel(self) -> bool: + """Whether this task can currently be cancelled.""" + return any(action.action is TaskActionName.CANCEL and action.available for action in self.available_actions) @classmethod def from_graphql(cls, data: dict) -> Task: related_nodes: list[TaskRelatedNode] = [] logs: list[TaskLog] = [] + available_actions: list[TaskAction] = [] if "related_nodes" in data: if data.get("related_nodes"): @@ -59,7 +82,12 @@ def from_graphql(cls, data: dict) -> Task: logs = [TaskLog(**item["node"]) for item in data["logs"]["edges"]] del data["logs"] - return cls(**data, related_nodes=related_nodes, logs=logs) + if "available_actions" in data: + if data.get("available_actions"): + available_actions = [TaskAction(**item) for item in data["available_actions"]] + del data["available_actions"] + + return cls(**data, related_nodes=related_nodes, logs=logs, available_actions=available_actions) class TaskFilter(BaseModel): diff --git a/tests/unit/sdk/test_task.py b/tests/unit/sdk/test_task.py index dd029003f..f34cf66af 100644 --- a/tests/unit/sdk/test_task.py +++ b/tests/unit/sdk/test_task.py @@ -1,12 +1,14 @@ from __future__ import annotations +import json from datetime import datetime, timezone from typing import TYPE_CHECKING import pytest +from infrahub_sdk.graphql import Mutation from infrahub_sdk.task.exceptions import TaskNotFoundError, TooManyTasksError -from infrahub_sdk.task.manager import InfraHubTaskManagerBase +from infrahub_sdk.task.manager import MUTATION_TASK_QUERY, InfraHubTaskManagerBase from infrahub_sdk.task.models import Task, TaskFilter, TaskState if TYPE_CHECKING: @@ -39,6 +41,82 @@ async def test_method_all_full(clients: BothClients, mock_query_tasks_01: HTTPXM assert isinstance(tasks[0], Task) +@pytest.mark.parametrize("client_type", client_types) +async def test_method_retry(clients: BothClients, httpx_mock: HTTPXMock, client_type: str) -> None: + httpx_mock.add_response( + method="POST", + json={"data": {"InfrahubTaskRetry": {"ok": True, "task": {"id": "b71f5542-7b54-562f-9053-8fd6ec0b0481"}}}}, + match_headers={"X-Infrahub-Tracker": "mutation-task-retry"}, + ) + + if client_type == "standard": + new_id = await clients.standard.task.retry(id="a60f4431-6a43-451e-8f42-9ec5db9a9370") + else: + new_id = clients.sync.task.retry(id="a60f4431-6a43-451e-8f42-9ec5db9a9370") + + assert new_id == "b71f5542-7b54-562f-9053-8fd6ec0b0481" + sent_query = json.loads(httpx_mock.get_requests()[-1].content)["query"] + assert "InfrahubTaskRetry(" in sent_query + assert 'id: "a60f4431-6a43-451e-8f42-9ec5db9a9370"' in sent_query + + +@pytest.mark.parametrize("client_type", client_types) +async def test_method_cancel(clients: BothClients, httpx_mock: HTTPXMock, client_type: str) -> None: + httpx_mock.add_response( + method="POST", + json={"data": {"InfrahubTaskCancel": {"ok": True, "task": {"id": "a60f4431-6a43-451e-8f42-9ec5db9a9370"}}}}, + match_headers={"X-Infrahub-Tracker": "mutation-task-cancel"}, + ) + + if client_type == "standard": + cancelled = await clients.standard.task.cancel(id="a60f4431-6a43-451e-8f42-9ec5db9a9370") + else: + cancelled = clients.sync.task.cancel(id="a60f4431-6a43-451e-8f42-9ec5db9a9370") + + assert cancelled is True + sent_query = json.loads(httpx_mock.get_requests()[-1].content)["query"] + assert "InfrahubTaskCancel(" in sent_query + assert 'id: "a60f4431-6a43-451e-8f42-9ec5db9a9370"' in sent_query + + +@pytest.mark.parametrize("client_type", client_types) +async def test_filter_limit_forwards_include_actions( + clients: BothClients, mock_query_tasks_03: HTTPXMock, client_type: str +) -> None: + if client_type == "standard": + await clients.standard.task.filter(limit=5, include_actions=True) + else: + clients.sync.task.filter(limit=5, include_actions=True) + + sent_query = json.loads(mock_query_tasks_03.get_requests()[-1].content)["query"] + assert "available_actions" in sent_query + + +async def test_action_mutation_render() -> None: + query = Mutation( + mutation="InfrahubTaskRetry", + input_data={"data": {"id": "a60f4431-6a43-451e-8f42-9ec5db9a9370"}}, + query=MUTATION_TASK_QUERY, + ) + assert ( + query.render() + == """ +mutation { + InfrahubTaskRetry( + data: { + id: "a60f4431-6a43-451e-8f42-9ec5db9a9370" + } + ){ + ok + task { + id + } + } +} +""" + ) + + async def test_generate_count_query() -> None: query = InfraHubTaskManagerBase._generate_count_query() assert query @@ -121,6 +199,7 @@ async def test_method_get_full(clients: BothClients, mock_query_tasks_05: HTTPXM assert len(task.logs) == 4 assert len(task.related_nodes) == 2 assert task.model_dump() == { + "available_actions": [], "branch": "main", "created_at": datetime(2025, 1, 18, 22, 12, 20, 228112, tzinfo=timezone.utc), "id": "32116fcd-9071-43a7-9f14-777901020b5b", @@ -161,3 +240,37 @@ async def test_method_get_full(clients: BothClients, mock_query_tasks_05: HTTPXM "updated_at": datetime(2025, 1, 18, 22, 12, 22, 44921, tzinfo=timezone.utc), "workflow": "import-python-files", } + + +def _base_task_data() -> dict: + return { + "id": "a60f4431-6a43-451e-8f42-9ec5db9a9370", + "title": "Webhook delivery", + "state": "COMPLETED", + "created_at": "2025-01-18T22:12:20.228112+00:00", + "updated_at": "2025-01-18T22:12:22.044921+00:00", + } + + +async def test_available_actions_parsed() -> None: + task = Task.from_graphql( + { + **_base_task_data(), + "available_actions": [ + {"action": "RETRY", "available": True, "unavailability_reason": None}, + {"action": "CANCEL", "available": False, "unavailability_reason": "the task has already settled"}, + ], + } + ) + + assert task.can_retry is True + assert task.can_cancel is False + assert task.available_actions[1].unavailability_reason == "the task has already settled" + + +async def test_available_actions_absent_defaults_empty() -> None: + task = Task.from_graphql(_base_task_data()) + + assert task.available_actions == [] + assert task.can_retry is False + assert task.can_cancel is False