123 lines
3.2 KiB
Python
123 lines
3.2 KiB
Python
"""
|
|
AngelaMos | 2026
|
|
tasking.py
|
|
"""
|
|
|
|
import asyncio
|
|
|
|
import aiosqlite
|
|
|
|
from app.core.models import TaskRecord, TaskResult
|
|
|
|
|
|
class TaskManager:
|
|
"""
|
|
Per-beacon async task queues backed by SQLite persistence
|
|
"""
|
|
def __init__(self) -> None:
|
|
"""
|
|
Initialize empty queue registry
|
|
"""
|
|
self._queues: dict[str, asyncio.Queue[TaskRecord]] = {}
|
|
|
|
def _ensure_queue(self, beacon_id: str) -> asyncio.Queue[TaskRecord]:
|
|
"""
|
|
Create a task queue for a beacon if one does not exist
|
|
"""
|
|
if beacon_id not in self._queues:
|
|
self._queues[beacon_id] = asyncio.Queue()
|
|
return self._queues[beacon_id]
|
|
|
|
async def submit(
|
|
self,
|
|
task: TaskRecord,
|
|
db: aiosqlite.Connection,
|
|
) -> None:
|
|
"""
|
|
Persist a task to SQLite and enqueue it for beacon delivery
|
|
"""
|
|
await db.execute(
|
|
"""
|
|
INSERT INTO tasks (id, beacon_id, command, args, status, created_at)
|
|
VALUES (?, ?, ?, ?, ?, ?)
|
|
""",
|
|
(
|
|
task.id,
|
|
task.beacon_id,
|
|
task.command,
|
|
task.args,
|
|
task.status,
|
|
task.created_at,
|
|
),
|
|
)
|
|
await db.commit()
|
|
|
|
queue = self._ensure_queue(task.beacon_id)
|
|
await queue.put(task)
|
|
|
|
async def get_next(self, beacon_id: str) -> TaskRecord:
|
|
"""
|
|
Await the next task for a beacon (blocks until available)
|
|
"""
|
|
queue = self._ensure_queue(beacon_id)
|
|
return await queue.get()
|
|
|
|
async def store_result(
|
|
self,
|
|
result: TaskResult,
|
|
db: aiosqlite.Connection,
|
|
) -> None:
|
|
"""
|
|
Persist a task result and mark the parent task as completed
|
|
"""
|
|
await db.execute(
|
|
"""
|
|
INSERT INTO task_results (id, task_id, output, error, created_at)
|
|
VALUES (?, ?, ?, ?, ?)
|
|
""",
|
|
(
|
|
result.id,
|
|
result.task_id,
|
|
result.output,
|
|
result.error,
|
|
result.created_at,
|
|
),
|
|
)
|
|
await db.execute(
|
|
"UPDATE tasks SET status = 'completed', completed_at = ? WHERE id = ?",
|
|
(result.created_at,
|
|
result.task_id),
|
|
)
|
|
await db.commit()
|
|
|
|
async def get_history(
|
|
self,
|
|
beacon_id: str,
|
|
db: aiosqlite.Connection,
|
|
) -> list[dict[str,
|
|
str | None]]:
|
|
"""
|
|
Retrieve task history with results for a specific beacon
|
|
"""
|
|
cursor = await db.execute(
|
|
"""
|
|
SELECT
|
|
t.id, t.command, t.args, t.status, t.created_at, t.completed_at,
|
|
tr.output, tr.error
|
|
FROM tasks t
|
|
LEFT JOIN task_results tr ON t.id = tr.task_id
|
|
WHERE t.beacon_id = ?
|
|
ORDER BY t.created_at DESC
|
|
""",
|
|
(beacon_id,
|
|
),
|
|
)
|
|
rows = await cursor.fetchall()
|
|
return [dict(row) for row in rows]
|
|
|
|
def remove_queue(self, beacon_id: str) -> None:
|
|
"""
|
|
Clean up the task queue when a beacon disconnects
|
|
"""
|
|
self._queues.pop(beacon_id, None)
|