279 lines
10 KiB
Python
279 lines
10 KiB
Python
"""
|
|
TaskService — creates tasks, dispatches them as asyncio background tasks,
|
|
and persists every state change to the DB so restarts don't lose history.
|
|
|
|
Flow:
|
|
create() → INSERT tasks (status=pending) → asyncio.create_task(run) → return task_id
|
|
run() → UPDATE running → per-device execute → UPDATE success/failed
|
|
"""
|
|
|
|
import asyncio
|
|
from datetime import datetime, timezone
|
|
from uuid import uuid4
|
|
|
|
from sqlalchemy import select
|
|
from sqlalchemy.orm import selectinload
|
|
|
|
from app.models.device import Device
|
|
from app.models.object import Object
|
|
from app.models.rack import Rack
|
|
from app.models.task import Task, TaskEvent, TaskResult
|
|
from app.models.user import User
|
|
from app.plugins.registry import ActionRegistry, action_registry
|
|
from app.services.database import AsyncSessionLocal
|
|
from app.services.sse import SSEEvent, sse_manager
|
|
|
|
# Max parallel devices within a single task
|
|
_TASK_DEVICE_CONCURRENCY = 20
|
|
|
|
|
|
async def _emit(task_id: str, message: str, level: str = "info") -> None:
|
|
"""Write a progress line to DB and push to SSE subscribers."""
|
|
now = datetime.now(timezone.utc)
|
|
async with AsyncSessionLocal() as db:
|
|
async with db.begin():
|
|
db.add(TaskEvent(task_id=task_id, message=message, level=level, created_at=now))
|
|
|
|
await sse_manager.publish(SSEEvent(
|
|
type="task_event",
|
|
task_id=task_id,
|
|
data={"message": message, "level": level, "ts": now.isoformat()},
|
|
))
|
|
|
|
|
|
async def _resolve_scope(scope: dict) -> list[tuple[Device, Object]]:
|
|
"""Expand target_scope to a list of (device, parent_object) pairs."""
|
|
async with AsyncSessionLocal() as db:
|
|
query = select(Device).options(
|
|
selectinload(Device.object).selectinload(Object.ssh_credentials),
|
|
selectinload(Device.object).selectinload(Object.category_credentials),
|
|
).where(Device.is_active == True)
|
|
|
|
scope_type = scope.get("type")
|
|
if scope_type == "device":
|
|
query = query.where(Device.id == scope["id"])
|
|
elif scope_type == "device_list":
|
|
query = query.where(Device.id.in_(scope["ids"]))
|
|
elif scope_type == "object":
|
|
query = query.where(Device.object_id == scope["id"])
|
|
elif scope_type == "rack":
|
|
query = query.where(Device.rack_id == scope["id"])
|
|
elif scope_type == "zone":
|
|
# Expand zone → rack IDs → devices
|
|
rack_rows = await db.execute(
|
|
select(Rack.id).where(Rack.zone_id == scope["id"])
|
|
)
|
|
rack_ids = [r[0] for r in rack_rows.all()]
|
|
if not rack_ids:
|
|
return []
|
|
query = query.where(Device.rack_id.in_(rack_ids))
|
|
elif scope_type == "category":
|
|
# support both "value" (legacy) and "category" key
|
|
cat = scope.get("category") or scope.get("value")
|
|
if cat:
|
|
query = query.where(Device.category == cat)
|
|
if "object_ids" in scope:
|
|
query = query.where(Device.object_id.in_(scope["object_ids"]))
|
|
elif "object_id" in scope:
|
|
query = query.where(Device.object_id == scope["object_id"])
|
|
elif scope_type == "role":
|
|
role = scope.get("role") or scope.get("value")
|
|
if role:
|
|
query = query.where(Device.role == role)
|
|
if "object_ids" in scope:
|
|
query = query.where(Device.object_id.in_(scope["object_ids"]))
|
|
elif "object_id" in scope:
|
|
query = query.where(Device.object_id == scope["object_id"])
|
|
elif scope_type == "filter":
|
|
# Cross-object filter scope (from WorkDevices global page)
|
|
from app.models.object import Object as Obj
|
|
if scope.get("category"):
|
|
cats = scope["category"]
|
|
if isinstance(cats, list):
|
|
query = query.where(Device.category.in_(cats))
|
|
else:
|
|
query = query.where(Device.category == cats)
|
|
if scope.get("object_id"):
|
|
query = query.where(Device.object_id == scope["object_id"])
|
|
if scope.get("city"):
|
|
query = query.where(Obj.city == scope["city"])
|
|
if scope.get("client"):
|
|
query = query.where(Obj.client == scope["client"])
|
|
else:
|
|
return []
|
|
|
|
result = await db.execute(query)
|
|
devices = result.scalars().all()
|
|
return [(d, d.object) for d in devices]
|
|
|
|
|
|
async def _run_on_device(
|
|
task_id: str,
|
|
device: Device,
|
|
obj: Object,
|
|
action,
|
|
params: dict,
|
|
sem: asyncio.Semaphore,
|
|
) -> bool:
|
|
async with sem:
|
|
emit = lambda msg, level="info": _emit(task_id, msg, level)
|
|
|
|
# Create task_result row
|
|
result_id: int | None = None
|
|
started_at = datetime.now(timezone.utc)
|
|
async with AsyncSessionLocal() as db:
|
|
async with db.begin():
|
|
tr = TaskResult(
|
|
task_id=task_id,
|
|
device_id=device.id,
|
|
status="running",
|
|
executed_at=started_at,
|
|
)
|
|
db.add(tr)
|
|
await db.flush()
|
|
result_id = tr.id
|
|
|
|
try:
|
|
action_result = await action.execute(device, obj, params, emit)
|
|
except Exception as exc:
|
|
action_result_success = False
|
|
action_stdout = ""
|
|
action_stderr = ""
|
|
action_exit_code = -1
|
|
action_data = {}
|
|
action_error = str(exc)
|
|
else:
|
|
action_result_success = action_result.success
|
|
action_stdout = action_result.stdout
|
|
action_stderr = action_result.stderr
|
|
action_exit_code = action_result.exit_code
|
|
action_data = action_result.data
|
|
action_error = action_result.error
|
|
|
|
finished_at = datetime.now(timezone.utc)
|
|
duration_ms = int((finished_at - started_at).total_seconds() * 1000)
|
|
|
|
async with AsyncSessionLocal() as db:
|
|
async with db.begin():
|
|
tr = await db.get(TaskResult, result_id)
|
|
if tr:
|
|
tr.status = "success" if action_result_success else "failed"
|
|
tr.stdout = action_stdout
|
|
tr.stderr = action_stderr
|
|
tr.exit_code = action_exit_code
|
|
tr.data = action_data or None
|
|
tr.duration_ms = duration_ms
|
|
|
|
await sse_manager.publish(SSEEvent(
|
|
type="task_result",
|
|
task_id=task_id,
|
|
data={
|
|
"device_id": device.id,
|
|
"device_ip": device.ip,
|
|
"status": "success" if action_result_success else "failed",
|
|
"duration_ms": duration_ms,
|
|
"error": action_error or None,
|
|
},
|
|
))
|
|
|
|
return action_result_success
|
|
|
|
|
|
async def _run_task(task_id: str) -> None:
|
|
await _emit(task_id, "Task started", "info")
|
|
|
|
async with AsyncSessionLocal() as db:
|
|
async with db.begin():
|
|
task = await db.get(Task, task_id)
|
|
if task is None:
|
|
return
|
|
task.status = "running"
|
|
task.started_at = datetime.now(timezone.utc)
|
|
scope = dict(task.target_scope)
|
|
action_type = task.type
|
|
params = dict(task.params)
|
|
|
|
action = action_registry.get(action_type)
|
|
if action is None:
|
|
async with AsyncSessionLocal() as db:
|
|
async with db.begin():
|
|
task = await db.get(Task, task_id)
|
|
if task:
|
|
task.status = "failed"
|
|
task.error_message = f"Unknown action: {action_type}"
|
|
task.finished_at = datetime.now(timezone.utc)
|
|
await _emit(task_id, f"Unknown action: {action_type}", "error")
|
|
return
|
|
|
|
targets = await _resolve_scope(scope)
|
|
if not targets:
|
|
async with AsyncSessionLocal() as db:
|
|
async with db.begin():
|
|
task = await db.get(Task, task_id)
|
|
if task:
|
|
task.status = "success"
|
|
task.result = {"total": 0, "ok": 0, "failed": 0}
|
|
task.finished_at = datetime.now(timezone.utc)
|
|
await _emit(task_id, "No devices matched the scope", "warn")
|
|
return
|
|
|
|
await _emit(task_id, f"Running on {len(targets)} device(s)...", "info")
|
|
sem = asyncio.Semaphore(_TASK_DEVICE_CONCURRENCY)
|
|
|
|
results = await asyncio.gather(
|
|
*[_run_on_device(task_id, device, obj, action, params, sem) for device, obj in targets],
|
|
return_exceptions=True,
|
|
)
|
|
|
|
ok = sum(1 for r in results if r is True)
|
|
failed = len(results) - ok
|
|
final_status = "success" if failed == 0 else ("failed" if ok == 0 else "partial")
|
|
|
|
async with AsyncSessionLocal() as db:
|
|
async with db.begin():
|
|
task = await db.get(Task, task_id)
|
|
if task:
|
|
task.status = final_status
|
|
task.result = {"total": len(results), "ok": ok, "failed": failed}
|
|
task.finished_at = datetime.now(timezone.utc)
|
|
|
|
await _emit(task_id, f"Done: {ok}/{len(results)} succeeded", "info")
|
|
await sse_manager.publish(SSEEvent(
|
|
type="task_status",
|
|
task_id=task_id,
|
|
data={"status": final_status, "total": len(results), "ok": ok, "failed": failed},
|
|
))
|
|
|
|
|
|
async def create_task(
|
|
action_type: str,
|
|
target_scope: dict,
|
|
params: dict,
|
|
user_id: int | None = None,
|
|
) -> Task:
|
|
if action_registry.get(action_type) is None:
|
|
raise ValueError(f"Unknown action: {action_type}")
|
|
|
|
now = datetime.now(timezone.utc)
|
|
task_id = str(uuid4())
|
|
|
|
async with AsyncSessionLocal() as db:
|
|
async with db.begin():
|
|
task = Task(
|
|
id=task_id,
|
|
type=action_type,
|
|
status="pending",
|
|
target_scope=target_scope,
|
|
params=params,
|
|
created_by=user_id,
|
|
created_at=now,
|
|
)
|
|
db.add(task)
|
|
|
|
# Dispatch — fire and forget; state is already persisted
|
|
asyncio.create_task(_run_task(task_id))
|
|
|
|
# Return a fresh copy to the caller
|
|
async with AsyncSessionLocal() as db:
|
|
result = await db.get(Task, task_id)
|
|
return result
|