44 lines
1.5 KiB
Python
44 lines
1.5 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
from typing import Any, Dict, Optional
|
||
|
|
|
||
|
|
from app.schemas.tasks import TaskStatusResponse
|
||
|
|
from app.services.supabase_rest import supabase_rest
|
||
|
|
|
||
|
|
|
||
|
|
class SupabaseTaskTracker:
|
||
|
|
"""利用 Supabase background_tasks 表追踪进度。"""
|
||
|
|
|
||
|
|
def create_task(self, *, user_id: str, document_id: str, task_type: str = "ocr") -> TaskStatusResponse:
|
||
|
|
payload = {
|
||
|
|
"user_id": user_id,
|
||
|
|
"document_id": document_id,
|
||
|
|
"task_type": task_type,
|
||
|
|
"status": "pending",
|
||
|
|
"progress": 0,
|
||
|
|
}
|
||
|
|
record = supabase_rest.insert("background_tasks", payload)
|
||
|
|
return self._to_response(record)
|
||
|
|
|
||
|
|
def get_task(self, *, task_id: str, user_id: str) -> Optional[TaskStatusResponse]:
|
||
|
|
record = supabase_rest.select_one("background_tasks", {"id": task_id, "user_id": user_id})
|
||
|
|
if not record:
|
||
|
|
return None
|
||
|
|
return self._to_response(record)
|
||
|
|
|
||
|
|
def update_task(self, *, task_id: str, **kwargs: Any) -> None:
|
||
|
|
supabase_rest.update("background_tasks", {"id": task_id}, kwargs)
|
||
|
|
|
||
|
|
def _to_response(self, record: Optional[Dict[str, Any]]) -> TaskStatusResponse:
|
||
|
|
if not record:
|
||
|
|
raise ValueError("Task record missing")
|
||
|
|
return TaskStatusResponse(
|
||
|
|
task_id=str(record["id"]),
|
||
|
|
status=record.get("status", "pending"),
|
||
|
|
progress=record.get("progress", 0),
|
||
|
|
message=record.get("message"),
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
task_tracker = SupabaseTaskTracker()
|