118 lines
3.6 KiB
Python
118 lines
3.6 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
from datetime import datetime, timezone
|
|
from typing import Optional
|
|
|
|
import aiosqlite
|
|
from pydantic import BaseModel
|
|
|
|
|
|
VALID_STATUSES = {"pending", "processing", "completed", "failed"}
|
|
|
|
|
|
class JobRecord(BaseModel):
|
|
id: str
|
|
status: str = "pending"
|
|
prompt: Optional[str] = None
|
|
params: dict | None = None
|
|
error: Optional[str] = None
|
|
created_at: str
|
|
completed_at: Optional[str] = None
|
|
|
|
|
|
class JobDB:
|
|
def __init__(self, db_path: str) -> None:
|
|
self.db_path = db_path
|
|
|
|
async def init(self) -> None:
|
|
os.makedirs(os.path.dirname(self.db_path) or ".", exist_ok=True)
|
|
async with aiosqlite.connect(self.db_path) as db:
|
|
await db.execute(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS jobs (
|
|
id TEXT PRIMARY KEY,
|
|
status TEXT NOT NULL DEFAULT 'pending',
|
|
prompt TEXT,
|
|
params TEXT,
|
|
error TEXT,
|
|
created_at TEXT NOT NULL,
|
|
completed_at TEXT
|
|
)
|
|
"""
|
|
)
|
|
await db.commit()
|
|
|
|
async def create_job(self, job: JobRecord) -> None:
|
|
async with aiosqlite.connect(self.db_path) as db:
|
|
await db.execute(
|
|
"INSERT INTO jobs (id, status, prompt, params, created_at) VALUES (?, ?, ?, ?, ?)",
|
|
(
|
|
job.id,
|
|
job.status,
|
|
job.prompt,
|
|
_json_encode(job.params),
|
|
job.created_at,
|
|
),
|
|
)
|
|
await db.commit()
|
|
|
|
async def get_job(self, job_id: str) -> JobRecord | None:
|
|
async with aiosqlite.connect(self.db_path) as db:
|
|
db.row_factory = aiosqlite.Row
|
|
async with db.execute("SELECT * FROM jobs WHERE id = ?", (job_id,)) as cursor:
|
|
row = await cursor.fetchone()
|
|
if row is None:
|
|
return None
|
|
return JobRecord(
|
|
id=row["id"],
|
|
status=row["status"],
|
|
prompt=row["prompt"],
|
|
params=_json_decode(row["params"]),
|
|
error=row["error"],
|
|
created_at=row["created_at"],
|
|
completed_at=row["completed_at"],
|
|
)
|
|
|
|
async def update_status(
|
|
self,
|
|
job_id: str,
|
|
status: str,
|
|
error: Optional[str] = None,
|
|
) -> None:
|
|
now = datetime.now(timezone.utc).isoformat()
|
|
async with aiosqlite.connect(self.db_path) as db:
|
|
if status in {"completed", "failed"}:
|
|
await db.execute(
|
|
"UPDATE jobs SET status = ?, error = ?, completed_at = ? WHERE id = ?",
|
|
(status, error, now, job_id),
|
|
)
|
|
else:
|
|
await db.execute(
|
|
"UPDATE jobs SET status = ?, error = ? WHERE id = ?",
|
|
(status, error, job_id),
|
|
)
|
|
await db.commit()
|
|
|
|
async def delete_job(self, job_id: str) -> bool:
|
|
async with aiosqlite.connect(self.db_path) as db:
|
|
cur = await db.execute("DELETE FROM jobs WHERE id = ?", (job_id,))
|
|
await db.commit()
|
|
return cur.rowcount > 0
|
|
|
|
|
|
def _json_encode(obj: dict | None) -> str | None:
|
|
if obj is None:
|
|
return None
|
|
import json
|
|
|
|
return json.dumps(obj)
|
|
|
|
|
|
def _json_decode(s: str | None) -> dict | None:
|
|
if s is None:
|
|
return None
|
|
import json
|
|
|
|
return json.loads(s)
|