mirror of
https://github.com/DrizzleTime/Foxel.git
synced 2026-09-05 15:47:00 +08:00
Initial commit
This commit is contained in:
shiyu
@@ -0,0 +1,26 @@
|
||||
from .service import TaskService
|
||||
from .scheduler import task_scheduler
|
||||
from .task_queue import Task, TaskProgress, TaskStatus, task_queue_service
|
||||
from .types import (
|
||||
AutomationTaskBase,
|
||||
AutomationTaskCreate,
|
||||
AutomationTaskRead,
|
||||
AutomationTaskUpdate,
|
||||
TaskQueueSettings,
|
||||
TaskQueueSettingsResponse,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"TaskService",
|
||||
"Task",
|
||||
"TaskProgress",
|
||||
"TaskStatus",
|
||||
"task_queue_service",
|
||||
"task_scheduler",
|
||||
"AutomationTaskBase",
|
||||
"AutomationTaskCreate",
|
||||
"AutomationTaskRead",
|
||||
"AutomationTaskUpdate",
|
||||
"TaskQueueSettings",
|
||||
"TaskQueueSettingsResponse",
|
||||
]
|
||||
@@ -0,0 +1,110 @@
|
||||
from fastapi import APIRouter, Depends, Request
|
||||
|
||||
from api.response import success
|
||||
from domain.audit import AuditAction, audit
|
||||
from domain.auth import get_current_active_user
|
||||
from .service import TaskService
|
||||
from .types import (
|
||||
AutomationTaskCreate,
|
||||
AutomationTaskUpdate,
|
||||
TaskQueueSettings,
|
||||
)
|
||||
|
||||
CurrentUser = TaskService.current_user_dep
|
||||
|
||||
router = APIRouter(
|
||||
prefix="/api/tasks",
|
||||
tags=["Tasks"],
|
||||
dependencies=[Depends(get_current_active_user)],
|
||||
responses={404: {"description": "Not found"}},
|
||||
)
|
||||
|
||||
|
||||
@router.get("/queue")
|
||||
@audit(action=AuditAction.READ, description="获取任务队列状态")
|
||||
async def get_task_queue_status(request: Request, current_user: CurrentUser):
|
||||
payload = TaskService.get_queue_tasks()
|
||||
return success(payload)
|
||||
|
||||
|
||||
@router.get("/queue/settings")
|
||||
@audit(action=AuditAction.READ, description="获取任务队列设置")
|
||||
async def get_task_queue_settings(request: Request, current_user: CurrentUser):
|
||||
payload = TaskService.get_queue_settings()
|
||||
return success(payload.model_dump())
|
||||
|
||||
|
||||
@router.post("/queue/settings")
|
||||
@audit(
|
||||
action=AuditAction.UPDATE,
|
||||
description="更新任务队列设置",
|
||||
body_fields=["concurrency"],
|
||||
)
|
||||
async def update_task_queue_settings(request: Request, settings: TaskQueueSettings, current_user: CurrentUser):
|
||||
payload = await TaskService.update_queue_settings(settings, getattr(current_user, "id", None))
|
||||
return success(payload.model_dump())
|
||||
|
||||
|
||||
@router.get("/queue/{task_id}")
|
||||
@audit(action=AuditAction.READ, description="获取队列任务状态")
|
||||
async def get_task_status(task_id: str, request: Request, current_user: CurrentUser):
|
||||
payload = TaskService.get_queue_task(task_id)
|
||||
return success(payload)
|
||||
|
||||
|
||||
@router.post("/")
|
||||
@audit(
|
||||
action=AuditAction.CREATE,
|
||||
description="创建自动化任务",
|
||||
body_fields=[
|
||||
"name",
|
||||
"event",
|
||||
"trigger_config",
|
||||
"processor_type",
|
||||
"processor_config",
|
||||
"enabled",
|
||||
],
|
||||
user_kw="user",
|
||||
)
|
||||
async def create_task(request: Request, task_in: AutomationTaskCreate, user: CurrentUser):
|
||||
task = await TaskService.create_task(task_in, user)
|
||||
return success(task)
|
||||
|
||||
|
||||
@router.get("/{task_id}")
|
||||
@audit(action=AuditAction.READ, description="获取自动化任务详情")
|
||||
async def get_task(task_id: int, request: Request, current_user: CurrentUser):
|
||||
task = await TaskService.get_task(task_id)
|
||||
return success(task)
|
||||
|
||||
|
||||
@router.get("/")
|
||||
@audit(action=AuditAction.READ, description="获取自动化任务列表")
|
||||
async def list_tasks(request: Request, current_user: CurrentUser):
|
||||
tasks = await TaskService.list_tasks()
|
||||
return success(tasks)
|
||||
|
||||
|
||||
@router.put("/{task_id}")
|
||||
@audit(
|
||||
action=AuditAction.UPDATE,
|
||||
description="更新自动化任务",
|
||||
body_fields=[
|
||||
"name",
|
||||
"event",
|
||||
"trigger_config",
|
||||
"processor_type",
|
||||
"processor_config",
|
||||
"enabled",
|
||||
],
|
||||
)
|
||||
async def update_task(request: Request, current_user: CurrentUser, task_id: int, task_in: AutomationTaskUpdate):
|
||||
task = await TaskService.update_task(task_id, task_in, current_user)
|
||||
return success(task)
|
||||
|
||||
|
||||
@router.delete("/{task_id}")
|
||||
@audit(action=AuditAction.DELETE, description="删除自动化任务", user_kw="user")
|
||||
async def delete_task(task_id: int, request: Request, user: CurrentUser):
|
||||
await TaskService.delete_task(task_id, user)
|
||||
return success(msg="Task deleted")
|
||||
@@ -0,0 +1,102 @@
|
||||
import asyncio
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
|
||||
from croniter import croniter
|
||||
|
||||
from models.database import AutomationTask
|
||||
from .task_queue import task_queue_service
|
||||
|
||||
|
||||
@dataclass
|
||||
class CronTaskItem:
|
||||
task_id: int
|
||||
processor_type: str
|
||||
path: str
|
||||
cron: croniter
|
||||
next_run: datetime
|
||||
|
||||
|
||||
class AutomationTaskScheduler:
|
||||
def __init__(self):
|
||||
self._items: list[CronTaskItem] = []
|
||||
self._worker: asyncio.Task | None = None
|
||||
self._reload_event = asyncio.Event()
|
||||
self._stop_event = asyncio.Event()
|
||||
|
||||
async def start(self) -> None:
|
||||
if self._worker and not self._worker.done():
|
||||
return
|
||||
self._stop_event.clear()
|
||||
await self._load_tasks()
|
||||
self._worker = asyncio.create_task(self._run_loop())
|
||||
|
||||
async def stop(self) -> None:
|
||||
if not self._worker:
|
||||
return
|
||||
self._stop_event.set()
|
||||
self._reload_event.set()
|
||||
await self._worker
|
||||
self._worker = None
|
||||
|
||||
def refresh(self) -> None:
|
||||
if self._worker and not self._worker.done():
|
||||
self._reload_event.set()
|
||||
|
||||
async def _load_tasks(self) -> None:
|
||||
tasks = await AutomationTask.filter(event="cron", enabled=True)
|
||||
items: list[CronTaskItem] = []
|
||||
now = datetime.now()
|
||||
for task in tasks:
|
||||
trigger = task.trigger_config or {}
|
||||
if not isinstance(trigger, dict):
|
||||
continue
|
||||
cron_expr = trigger.get("cron_expr")
|
||||
path = trigger.get("path")
|
||||
if not cron_expr or not path:
|
||||
continue
|
||||
cron = self._build_cron(cron_expr, now)
|
||||
if not cron:
|
||||
continue
|
||||
next_run = cron.get_next(datetime)
|
||||
items.append(
|
||||
CronTaskItem(
|
||||
task_id=task.id,
|
||||
processor_type=task.processor_type,
|
||||
path=path,
|
||||
cron=cron,
|
||||
next_run=next_run,
|
||||
)
|
||||
)
|
||||
self._items = items
|
||||
|
||||
def _build_cron(self, expr: str, base_time: datetime) -> croniter | None:
|
||||
expr = str(expr or "").strip()
|
||||
if not expr:
|
||||
return None
|
||||
parts = [p for p in expr.split() if p]
|
||||
if len(parts) not in (5, 6):
|
||||
return None
|
||||
second_at_beginning = len(parts) == 6
|
||||
try:
|
||||
return croniter(expr, base_time, second_at_beginning=second_at_beginning)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
async def _run_loop(self) -> None:
|
||||
while not self._stop_event.is_set():
|
||||
if self._reload_event.is_set():
|
||||
self._reload_event.clear()
|
||||
await self._load_tasks()
|
||||
now = datetime.now()
|
||||
for item in list(self._items):
|
||||
if item.next_run <= now:
|
||||
await task_queue_service.add_task(
|
||||
item.processor_type,
|
||||
{"task_id": item.task_id, "path": item.path},
|
||||
)
|
||||
item.next_run = item.cron.get_next(datetime)
|
||||
await asyncio.sleep(1)
|
||||
|
||||
|
||||
task_scheduler = AutomationTaskScheduler()
|
||||
@@ -0,0 +1,117 @@
|
||||
import re
|
||||
from typing import Annotated, Any, Dict, Optional
|
||||
|
||||
from fastapi import Depends, HTTPException
|
||||
|
||||
from domain.auth import User, get_current_active_user
|
||||
from domain.config import ConfigService
|
||||
from .scheduler import task_scheduler
|
||||
from .task_queue import task_queue_service
|
||||
from .types import (
|
||||
AutomationTaskCreate,
|
||||
AutomationTaskUpdate,
|
||||
TaskQueueSettings,
|
||||
TaskQueueSettingsResponse,
|
||||
)
|
||||
from models.database import AutomationTask
|
||||
|
||||
|
||||
class TaskService:
|
||||
current_user_dep = Annotated[User, Depends(get_current_active_user)]
|
||||
|
||||
@classmethod
|
||||
def get_queue_tasks(cls) -> list[dict[str, Any]]:
|
||||
tasks = task_queue_service.get_all_tasks()
|
||||
return [task.dict() for task in tasks]
|
||||
|
||||
@classmethod
|
||||
def get_queue_settings(cls) -> TaskQueueSettingsResponse:
|
||||
return TaskQueueSettingsResponse(
|
||||
concurrency=task_queue_service.get_concurrency(),
|
||||
active_workers=task_queue_service.get_active_worker_count(),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def update_queue_settings(cls, settings: TaskQueueSettings, user_id: Optional[int]) -> TaskQueueSettingsResponse:
|
||||
await task_queue_service.set_concurrency(settings.concurrency)
|
||||
await ConfigService.set("TASK_QUEUE_CONCURRENCY", str(task_queue_service.get_concurrency()))
|
||||
return cls.get_queue_settings()
|
||||
|
||||
@classmethod
|
||||
def get_queue_task(cls, task_id: str) -> dict[str, Any]:
|
||||
task = task_queue_service.get_task(task_id)
|
||||
if not task:
|
||||
raise HTTPException(status_code=404, detail="Task not found")
|
||||
return task.dict()
|
||||
|
||||
@classmethod
|
||||
async def create_task(cls, payload: AutomationTaskCreate, user: Optional[User]) -> AutomationTask:
|
||||
task = await AutomationTask.create(**payload.model_dump())
|
||||
task_scheduler.refresh()
|
||||
return task
|
||||
|
||||
@classmethod
|
||||
async def get_task(cls, task_id: int) -> AutomationTask:
|
||||
task = await AutomationTask.get_or_none(id=task_id)
|
||||
if not task:
|
||||
raise HTTPException(status_code=404, detail=f"Task {task_id} not found")
|
||||
return task
|
||||
|
||||
@classmethod
|
||||
async def list_tasks(cls) -> list[AutomationTask]:
|
||||
tasks = await AutomationTask.all()
|
||||
return tasks
|
||||
|
||||
@classmethod
|
||||
async def update_task(cls, task_id: int, payload: AutomationTaskUpdate, current_user: User) -> AutomationTask:
|
||||
task = await AutomationTask.get_or_none(id=task_id)
|
||||
if not task:
|
||||
raise HTTPException(status_code=404, detail=f"Task {task_id} not found")
|
||||
update_data = payload.model_dump(exclude_unset=True)
|
||||
for key, value in update_data.items():
|
||||
setattr(task, key, value)
|
||||
await task.save()
|
||||
task_scheduler.refresh()
|
||||
return task
|
||||
|
||||
@classmethod
|
||||
async def delete_task(cls, task_id: int, user: Optional[User]) -> None:
|
||||
deleted_count = await AutomationTask.filter(id=task_id).delete()
|
||||
if not deleted_count:
|
||||
raise HTTPException(status_code=404, detail=f"Task {task_id} not found")
|
||||
task_scheduler.refresh()
|
||||
|
||||
@classmethod
|
||||
async def trigger_tasks(cls, event: str, path: str):
|
||||
tasks = await AutomationTask.filter(event=event, enabled=True)
|
||||
for task in tasks:
|
||||
if cls.match(task, path):
|
||||
await cls.execute(task, path)
|
||||
|
||||
@classmethod
|
||||
def match(cls, task: AutomationTask, path: str) -> bool:
|
||||
trigger_config = task.trigger_config or {}
|
||||
if not isinstance(trigger_config, dict):
|
||||
trigger_config = {}
|
||||
path_prefix = trigger_config.get("path_prefix")
|
||||
filename_regex = trigger_config.get("filename_regex")
|
||||
if path_prefix and not path.startswith(path_prefix):
|
||||
return False
|
||||
if filename_regex:
|
||||
filename = path.split("/")[-1]
|
||||
if not re.match(filename_regex, filename):
|
||||
return False
|
||||
return True
|
||||
|
||||
@classmethod
|
||||
async def execute(cls, task: AutomationTask, path: str):
|
||||
await task_queue_service.add_task(
|
||||
task.processor_type,
|
||||
{
|
||||
"task_id": task.id,
|
||||
"path": path,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
task_service = TaskService
|
||||
@@ -0,0 +1,214 @@
|
||||
import asyncio
|
||||
from typing import Dict, Any
|
||||
from pydantic import BaseModel, Field
|
||||
import uuid
|
||||
from enum import Enum
|
||||
|
||||
|
||||
class TaskStatus(str, Enum):
|
||||
PENDING = "pending"
|
||||
RUNNING = "running"
|
||||
SUCCESS = "success"
|
||||
FAILED = "failed"
|
||||
|
||||
|
||||
class TaskProgress(BaseModel):
|
||||
stage: str | None = None
|
||||
percent: float | None = None
|
||||
bytes_total: int | None = None
|
||||
bytes_done: int | None = None
|
||||
detail: str | None = None
|
||||
|
||||
|
||||
class Task(BaseModel):
|
||||
id: str = Field(default_factory=lambda: uuid.uuid4().hex)
|
||||
name: str
|
||||
status: TaskStatus = TaskStatus.PENDING
|
||||
result: Any = None
|
||||
error: str | None = None
|
||||
task_info: Dict[str, Any] = {}
|
||||
progress: TaskProgress | None = None
|
||||
meta: Dict[str, Any] | None = None
|
||||
|
||||
|
||||
_SENTINEL = object()
|
||||
|
||||
|
||||
class TaskQueueService:
|
||||
def __init__(self):
|
||||
self._queue: asyncio.Queue[Task | object] = asyncio.Queue()
|
||||
self._tasks: Dict[str, Task] = {}
|
||||
self._worker_tasks: list[asyncio.Task] = []
|
||||
self._concurrency: int = 1
|
||||
self._worker_seq: int = 0
|
||||
|
||||
async def add_task(self, name: str, task_info: Dict[str, Any]) -> Task:
|
||||
task = Task(name=name, task_info=task_info)
|
||||
self._tasks[task.id] = task
|
||||
await self._queue.put(task)
|
||||
return task
|
||||
|
||||
def get_task(self, task_id: str) -> Task | None:
|
||||
return self._tasks.get(task_id)
|
||||
|
||||
def get_all_tasks(self) -> list[Task]:
|
||||
return list(self._tasks.values())
|
||||
|
||||
async def update_progress(self, task_id: str, progress: TaskProgress | Dict[str, Any]):
|
||||
task = self._tasks.get(task_id)
|
||||
if not task:
|
||||
return
|
||||
if isinstance(progress, TaskProgress):
|
||||
task.progress = progress
|
||||
else:
|
||||
task.progress = TaskProgress(**progress)
|
||||
|
||||
async def update_meta(self, task_id: str, meta: Dict[str, Any]):
|
||||
task = self._tasks.get(task_id)
|
||||
if not task:
|
||||
return
|
||||
task.meta = (task.meta or {}) | meta
|
||||
|
||||
async def _execute_task(self, task: Task):
|
||||
task.status = TaskStatus.RUNNING
|
||||
|
||||
try:
|
||||
# Local import to avoid circular dependency during module load.
|
||||
from domain.virtual_fs import VirtualFSService
|
||||
|
||||
if task.name == "process_file":
|
||||
params = task.task_info
|
||||
result = await VirtualFSService.process_file(
|
||||
path=params["path"],
|
||||
processor_type=params["processor_type"],
|
||||
config=params["config"],
|
||||
save_to=params.get("save_to"),
|
||||
overwrite=params.get("overwrite", False),
|
||||
)
|
||||
task.result = result
|
||||
elif task.name == "process_directory_scan":
|
||||
from domain.processors import ProcessDirectoryRequest, ProcessorService
|
||||
|
||||
params = task.task_info or {}
|
||||
req = ProcessDirectoryRequest(**params)
|
||||
task.result = await ProcessorService.scan_directory(req)
|
||||
elif task.name == "automation_task" or self._is_processor_task(task.name):
|
||||
from models.database import AutomationTask
|
||||
|
||||
params = task.task_info
|
||||
auto_task = await AutomationTask.get(id=params["task_id"])
|
||||
path = params["path"]
|
||||
|
||||
processor_type = auto_task.processor_type
|
||||
config = auto_task.processor_config or {}
|
||||
save_to = config.get("save_to") if isinstance(config, dict) else None
|
||||
overwrite = bool(config.get("overwrite")) if isinstance(config, dict) else False
|
||||
try:
|
||||
if await VirtualFSService.path_is_directory(path):
|
||||
overwrite = True
|
||||
except Exception:
|
||||
pass
|
||||
await VirtualFSService.process_file(
|
||||
path=path,
|
||||
processor_type=processor_type,
|
||||
config=config if isinstance(config, dict) else {},
|
||||
save_to=save_to,
|
||||
overwrite=overwrite,
|
||||
)
|
||||
task.result = "Automation task completed"
|
||||
elif task.name == "offline_http_download":
|
||||
from domain.offline_downloads import OfflineDownloadService
|
||||
|
||||
result_path = await OfflineDownloadService.run_http_download(task)
|
||||
task.result = {"path": result_path}
|
||||
elif task.name == "cross_mount_transfer":
|
||||
result = await VirtualFSService.run_cross_mount_transfer_task(task)
|
||||
task.result = result
|
||||
elif task.name == "send_email":
|
||||
from domain.email import EmailService
|
||||
await EmailService.send_from_task(task.id, task.task_info)
|
||||
task.result = "Email sent"
|
||||
else:
|
||||
raise ValueError(f"Unknown task name: {task.name}")
|
||||
task.status = TaskStatus.SUCCESS
|
||||
|
||||
except Exception as e:
|
||||
task.status = TaskStatus.FAILED
|
||||
task.error = str(e)
|
||||
|
||||
def _cleanup_workers(self):
|
||||
self._worker_tasks = [task for task in self._worker_tasks if not task.done()]
|
||||
|
||||
def _is_processor_task(self, task_name: str) -> bool:
|
||||
try:
|
||||
from domain.processors import get_processor
|
||||
|
||||
return get_processor(task_name) is not None
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
async def _ensure_worker_count(self):
|
||||
self._cleanup_workers()
|
||||
current = len(self._worker_tasks)
|
||||
if current < self._concurrency:
|
||||
for _ in range(self._concurrency - current):
|
||||
self._worker_seq += 1
|
||||
worker_id = self._worker_seq
|
||||
worker_task = asyncio.create_task(self._worker_loop(worker_id))
|
||||
self._worker_tasks.append(worker_task)
|
||||
elif current > self._concurrency:
|
||||
for _ in range(current - self._concurrency):
|
||||
await self._queue.put(_SENTINEL)
|
||||
|
||||
async def _worker_loop(self, worker_id: int):
|
||||
current_task = asyncio.current_task()
|
||||
try:
|
||||
while True:
|
||||
job = await self._queue.get()
|
||||
if job is _SENTINEL:
|
||||
self._queue.task_done()
|
||||
break
|
||||
try:
|
||||
await self._execute_task(job)
|
||||
except Exception as e:
|
||||
pass
|
||||
finally:
|
||||
self._queue.task_done()
|
||||
finally:
|
||||
if current_task in self._worker_tasks:
|
||||
self._worker_tasks.remove(current_task) # type: ignore[arg-type]
|
||||
|
||||
async def start_worker(self, concurrency: int | None = None):
|
||||
if concurrency is None:
|
||||
from domain.config import ConfigService
|
||||
|
||||
stored_value = await ConfigService.get("TASK_QUEUE_CONCURRENCY", self._concurrency)
|
||||
try:
|
||||
concurrency = int(stored_value)
|
||||
except (TypeError, ValueError):
|
||||
concurrency = self._concurrency
|
||||
await self.set_concurrency(concurrency)
|
||||
|
||||
async def set_concurrency(self, value: int):
|
||||
value = max(1, int(value))
|
||||
if value != self._concurrency:
|
||||
self._concurrency = value
|
||||
await self._ensure_worker_count()
|
||||
|
||||
async def stop_worker(self):
|
||||
self._cleanup_workers()
|
||||
for _ in range(len(self._worker_tasks)):
|
||||
await self._queue.put(_SENTINEL)
|
||||
if self._worker_tasks:
|
||||
await asyncio.gather(*self._worker_tasks, return_exceptions=True)
|
||||
self._worker_tasks.clear()
|
||||
|
||||
def get_concurrency(self) -> int:
|
||||
return self._concurrency
|
||||
|
||||
def get_active_worker_count(self) -> int:
|
||||
self._cleanup_workers()
|
||||
return len(self._worker_tasks)
|
||||
|
||||
|
||||
task_queue_service = TaskQueueService()
|
||||
@@ -0,0 +1,40 @@
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class AutomationTaskBase(BaseModel):
|
||||
name: str
|
||||
event: str
|
||||
trigger_config: Dict[str, Any] = {}
|
||||
processor_type: str
|
||||
processor_config: Dict[str, Any] = {}
|
||||
enabled: bool = True
|
||||
|
||||
|
||||
class AutomationTaskCreate(AutomationTaskBase):
|
||||
pass
|
||||
|
||||
|
||||
class AutomationTaskUpdate(AutomationTaskBase):
|
||||
name: Optional[str] = None
|
||||
event: Optional[str] = None
|
||||
processor_type: Optional[str] = None
|
||||
processor_config: Optional[Dict[str, Any]] = None
|
||||
trigger_config: Optional[Dict[str, Any]] = None
|
||||
enabled: Optional[bool] = None
|
||||
|
||||
|
||||
class AutomationTaskRead(AutomationTaskBase):
|
||||
id: int
|
||||
|
||||
class Config:
|
||||
from_attributes = True
|
||||
|
||||
|
||||
class TaskQueueSettings(BaseModel):
|
||||
concurrency: int = Field(..., ge=1, description="Desired number of concurrent task workers")
|
||||
|
||||
|
||||
class TaskQueueSettingsResponse(TaskQueueSettings):
|
||||
active_workers: int = Field(..., ge=0, description="Currently running worker count")
|
||||
Reference in New Issue
Block a user