Coverage for src/paperap/resources/tasks.py: 27%
96 statements
« prev ^ index » next coverage.py v7.6.12, created at 2025-03-22 16:02 -0400
1"""
2----------------------------------------------------------------------------
4 METADATA:
6 File: tasks.py
7 Project: paperap
8 Created: 2025-03-04
9 Version: 0.0.9
10 Author: Jess Mann
11 Email: jess@jmann.me
12 Copyright (c) 2025 Jess Mann
14----------------------------------------------------------------------------
16 LAST MODIFIED:
18 2025-03-04 By Jess Mann
20"""
22from __future__ import annotations
24import enum
25import logging
26import time
27from typing import Any, Callable, Generic, TypeVar, cast
29from paperap.exceptions import APIError, BadResponseError, ResourceNotFoundError
30from paperap.models.task import Task, TaskQuerySet
31from paperap.resources.base import BaseResource, StandardResource
33logger = logging.getLogger(__name__)
36class TaskStatus(enum.Enum):
37 """Status of a task."""
39 PENDING = "PENDING"
40 STARTED = "STARTED"
41 RETRY = "RETRY"
42 SUCCESS = "SUCCESS"
43 FAILURE = "FAILURE"
44 REVOKED = "REVOKED"
47T = TypeVar("T")
50class TaskResource(StandardResource[Task, TaskQuerySet]):
51 """Resource for managing tasks."""
53 model_class = Task
54 queryset_class = TaskQuerySet
56 def acknowledge(self, task_id: int) -> None:
57 """
58 Acknowledge a task.
60 Args:
61 task_id: ID of the task to acknowledge.
63 """
64 self.client.request("PUT", f"tasks/{task_id}/acknowledge/")
66 def bulk_acknowledge(self, task_ids: list[int]) -> None:
67 """
68 Acknowledge multiple tasks.
70 Args:
71 task_ids: list of task IDs to acknowledge.
73 """
74 self.client.request("POST", "tasks/bulk_acknowledge/", data={"tasks": task_ids})
76 def wait_for_task(
77 self,
78 task_id: str,
79 max_wait: int = 300,
80 poll_interval: float = 1.0,
81 success_callback: Callable[[Task], None] | None = None,
82 failure_callback: Callable[[Task], None] | None = None,
83 ) -> Task:
84 """
85 Wait for a task to complete.
87 Args:
88 task_id: The task ID to wait for.
89 max_wait: Maximum time (in seconds) to wait for completion.
90 poll_interval: Seconds between polling attempts.
91 success_callback: Optional callback to execute when task succeeds.
92 failure_callback: Optional callback to execute when task fails.
94 Returns:
95 The completed Task instance.
97 Raises:
98 APIError: If the task fails or times out.
99 ResourceNotFoundError: If the task cannot be found.
101 """
102 logger.debug("Waiting for task %s to complete", task_id)
103 end_time = time.monotonic() + max_wait
105 while time.monotonic() < end_time:
106 try:
107 task = self(task_id=task_id).first()
108 if task is None:
109 logger.debug("Task %s not found, retrying...", task_id)
110 time.sleep(poll_interval)
111 continue
113 # Check if task is complete
114 if task.status == TaskStatus.SUCCESS.value:
115 logger.debug("Task %s completed successfully", task_id)
116 if success_callback:
117 success_callback(task)
118 return task
120 if task.status == TaskStatus.FAILURE.value:
121 logger.error("Task %s failed: %s", task_id, task.result)
122 if failure_callback:
123 failure_callback(task)
124 raise APIError(f"Task {task_id} failed: {task.result}")
126 if task.status == TaskStatus.REVOKED.value:
127 logger.warning("Task %s was revoked", task_id)
128 raise APIError(f"Task {task_id} was revoked")
130 logger.debug("Task %s status: %s, waiting...", task_id, task.status)
132 except ResourceNotFoundError:
133 logger.debug("Task %s not found yet, retrying...", task_id)
135 time.sleep(poll_interval)
137 raise APIError(f"Timed out waiting for task {task_id} to complete")
139 def wait_for_tasks(self, task_ids: list[str], max_wait: int = 300, poll_interval: float = 1.0) -> dict[str, Task]:
140 """
141 Wait for multiple tasks to complete.
143 Args:
144 task_ids: List of task IDs to wait for.
145 max_wait: Maximum time (in seconds) to wait for all tasks.
146 poll_interval: Seconds between polling attempts.
148 Returns:
149 Dictionary mapping task IDs to completed Task instances.
151 Raises:
152 APIError: If any task fails or times out.
154 """
155 logger.debug("Waiting for %d tasks to complete", len(task_ids))
156 end_time = time.monotonic() + max_wait
157 completed_tasks: dict[str, Task] = {}
158 pending_tasks = list(task_ids)
160 while pending_tasks and time.monotonic() < end_time:
161 for task_id in list(pending_tasks): # Create a copy to safely modify during iteration
162 try:
163 task = self(task_id=task_id).first()
164 if task is None:
165 continue
167 if task.status == TaskStatus.SUCCESS.value:
168 logger.debug("Task %s completed successfully", task_id)
169 completed_tasks[task_id] = task
170 pending_tasks.remove(task_id)
172 elif task.status == TaskStatus.FAILURE.value:
173 logger.error("Task %s failed: %s", task_id, task.result)
174 raise APIError(f"Task {task_id} failed: {task.result}")
176 elif task.status == TaskStatus.REVOKED.value:
177 logger.warning("Task %s was revoked", task_id)
178 raise APIError(f"Task {task_id} was revoked")
180 except ResourceNotFoundError:
181 pass # Task not found yet, continue waiting
183 if pending_tasks:
184 time.sleep(poll_interval)
186 if pending_tasks:
187 raise APIError(f"Timed out waiting for tasks: {', '.join(pending_tasks)}")
189 return completed_tasks
191 def get_task_result(self, task_id: str, wait: bool = True, max_wait: int = 300) -> str | None:
192 """
193 Get the result of a task.
195 Args:
196 task_id: The task ID.
197 wait: Whether to wait for the task to complete if it's not already.
198 max_wait: Maximum time (in seconds) to wait if wait=True.
200 Returns:
201 The result of the task.
203 Raises:
204 APIError: If the task fails or times out.
205 ResourceNotFoundError: If the task cannot be found.
207 """
208 task = None
209 if wait:
210 task = self.wait_for_task(task_id, max_wait=max_wait)
211 else:
212 task = self(task_id=task_id).first()
214 if task is None:
215 raise ResourceNotFoundError(f"Task {task_id} not found")
217 if task.status != TaskStatus.SUCCESS.value:
218 raise APIError(f"Task {task_id} is not successful (status: {task.status})")
220 return task.result
222 def execute_task(self, method: str, endpoint: str, data: dict[str, Any] | None = None, max_wait: int = 300) -> Task:
223 """
224 Execute a task synchronously.
226 This is a helper method that executes a task and waits for its completion.
228 Args:
229 method: HTTP method (GET, POST, etc.)
230 endpoint: API endpoint to call
231 data: Optional data to send with the request
232 max_wait: Maximum time to wait for task completion
234 Returns:
235 The task object, once completed.
237 Raises:
238 APIError: If the task fails or times out
240 """
241 response = self.client.request(method, endpoint, data=data)
242 if not response or not isinstance(response, str):
243 raise BadResponseError("Expected task ID in response")
245 task_id = str(response)
246 return self.wait_for_task(task_id, max_wait=max_wait)