134 lines
4.3 KiB
Python
134 lines
4.3 KiB
Python
import sqlite3
|
|
import threading
|
|
from dataclasses import dataclass
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
|
|
SCHEMA = """
|
|
CREATE TABLE IF NOT EXISTS jobs (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
title TEXT,
|
|
image_file TEXT NOT NULL,
|
|
audio_file TEXT NOT NULL,
|
|
output_file TEXT,
|
|
status TEXT DEFAULT 'waiting' NOT NULL,
|
|
error_message TEXT,
|
|
created_at TEXT DEFAULT CURRENT_TIMESTAMP NOT NULL,
|
|
updated_at TEXT DEFAULT CURRENT_TIMESTAMP NOT NULL,
|
|
position INTEGER DEFAULT 0 NOT NULL
|
|
);
|
|
"""
|
|
|
|
COLUMNS = {
|
|
'id': 'INTEGER',
|
|
'title': 'TEXT',
|
|
'image_file': 'TEXT',
|
|
'audio_file': 'TEXT',
|
|
'output_file': 'TEXT',
|
|
'status': 'TEXT',
|
|
'error_message': 'TEXT',
|
|
'created_at': 'TEXT',
|
|
'updated_at': 'TEXT',
|
|
'position': 'INTEGER',
|
|
}
|
|
|
|
_lock = threading.Lock()
|
|
|
|
|
|
def _connect(db_path: Path) -> sqlite3.Connection:
|
|
conn = sqlite3.connect(str(db_path), check_same_thread=False)
|
|
conn.row_factory = sqlite3.Row
|
|
return conn
|
|
|
|
|
|
class JobStore:
|
|
def __init__(self, db_path: Path):
|
|
self.db_path = Path(db_path)
|
|
self.local = threading.local()
|
|
self.init_db()
|
|
|
|
def _conn(self) -> sqlite3.Connection:
|
|
if not hasattr(self.local, 'conn') or self.local.conn is None:
|
|
self.local.conn = _connect(self.db_path)
|
|
return self.local.conn
|
|
|
|
def init_db(self):
|
|
with _lock:
|
|
self.db_path.parent.mkdir(parents=True, exist_ok=True)
|
|
conn = self._conn()
|
|
conn.executescript(SCHEMA)
|
|
conn.commit()
|
|
self._migrate_columns()
|
|
|
|
def _migrate_columns(self):
|
|
conn = self._conn()
|
|
existing = {
|
|
row['name']: row['type']
|
|
for row in conn.execute("PRAGMA table_info(jobs)")
|
|
}
|
|
for col, col_type in COLUMNS.items():
|
|
if col not in existing:
|
|
conn.execute(f"ALTER TABLE jobs ADD COLUMN {col} {col_type}")
|
|
conn.commit()
|
|
|
|
def create_job(self, title: str, image_file: str, audio_file: str) -> int:
|
|
conn = self._conn()
|
|
max_pos = conn.execute("SELECT COALESCE(MAX(position), 0) FROM jobs").fetchone()[0]
|
|
cur = conn.execute(
|
|
"""INSERT INTO jobs (title, image_file, audio_file, status, position)
|
|
VALUES (?, ?, ?, 'waiting', ?)""",
|
|
(title, image_file, audio_file, int(max_pos) + 1)
|
|
)
|
|
conn.commit()
|
|
return cur.lastrowid
|
|
|
|
def list_jobs(self) -> list[sqlite3.Row]:
|
|
conn = self._conn()
|
|
return conn.execute(
|
|
"SELECT * FROM jobs ORDER BY position ASC, id ASC"
|
|
).fetchall()
|
|
|
|
def get_job(self, job_id: int) -> sqlite3.Row | None:
|
|
conn = self._conn()
|
|
row = conn.execute("SELECT * FROM jobs WHERE id = ?", (job_id,)).fetchone()
|
|
return row
|
|
|
|
def update_status(self, job_id: int, status: str, error_message: str | None = None, output_file: str | None = None):
|
|
conn = self._conn()
|
|
conn.execute(
|
|
"""UPDATE jobs
|
|
SET status = ?, error_message = ?, output_file = ?, updated_at = ?
|
|
WHERE id = ?""",
|
|
(status, error_message, output_file, datetime.utcnow().isoformat(), job_id)
|
|
)
|
|
conn.commit()
|
|
|
|
def delete_job(self, job_id: int) -> bool:
|
|
conn = self._conn()
|
|
cur = conn.execute("DELETE FROM jobs WHERE id = ?", (job_id,))
|
|
conn.commit()
|
|
return cur.rowcount > 0
|
|
|
|
def move_to_end(self, job_id: int) -> bool:
|
|
conn = self._conn()
|
|
max_pos = conn.execute("SELECT COALESCE(MAX(position), 0) FROM jobs").fetchone()[0]
|
|
cur = conn.execute(
|
|
"UPDATE jobs SET status = 'waiting', error_message = NULL, position = ? WHERE id = ?",
|
|
(int(max_pos) + 1, job_id)
|
|
)
|
|
conn.commit()
|
|
return cur.rowcount > 0
|
|
|
|
def pop_next_waiting(self) -> sqlite3.Row | None:
|
|
conn = self._conn()
|
|
with _lock:
|
|
row = conn.execute(
|
|
"SELECT * FROM jobs WHERE status = 'waiting' ORDER BY position ASC, id ASC LIMIT 1"
|
|
).fetchone()
|
|
if row:
|
|
conn.execute(
|
|
"UPDATE jobs SET status = 'processing', updated_at = ? WHERE id = ?",
|
|
(datetime.utcnow().isoformat(), row['id'])
|
|
)
|
|
conn.commit()
|
|
return row
|