feat: initial Tauri desktop shell for BatchTube v1.0
This commit is contained in:
Agent Zero
371
backend/database.py
Normal file
371
backend/database.py
Normal file
@@ -0,0 +1,371 @@
|
||||
"""
|
||||
database.py - SQLite state management for YouTube SIPOC Tool
|
||||
Handles jobs, videos, artifacts, and logs with resume capability.
|
||||
"""
|
||||
import sqlite3
|
||||
import json
|
||||
import uuid
|
||||
from datetime import datetime
|
||||
from contextlib import contextmanager
|
||||
from typing import Optional, List, Dict, Any
|
||||
|
||||
import config
|
||||
|
||||
|
||||
@contextmanager
|
||||
def get_conn():
|
||||
"""Context manager providing a thread-safe SQLite connection."""
|
||||
conn = sqlite3.connect(config.DB_PATH, check_same_thread=False)
|
||||
conn.row_factory = sqlite3.Row
|
||||
conn.execute("PRAGMA journal_mode=WAL") # WAL mode for concurrent reads
|
||||
conn.execute("PRAGMA foreign_keys=ON")
|
||||
try:
|
||||
yield conn
|
||||
conn.commit()
|
||||
except Exception:
|
||||
conn.rollback()
|
||||
raise
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def init_db():
|
||||
"""Create all tables if they don't exist."""
|
||||
with get_conn() as conn:
|
||||
conn.executescript("""
|
||||
CREATE TABLE IF NOT EXISTS jobs (
|
||||
job_id TEXT PRIMARY KEY,
|
||||
status TEXT NOT NULL DEFAULT 'pending',
|
||||
created_at TEXT NOT NULL,
|
||||
updated_at TEXT NOT NULL,
|
||||
output_folder TEXT NOT NULL,
|
||||
selections TEXT NOT NULL,
|
||||
resolution TEXT NOT NULL DEFAULT 'best',
|
||||
delay_min REAL NOT NULL DEFAULT 2.0,
|
||||
delay_max REAL NOT NULL DEFAULT 7.0,
|
||||
input_type TEXT NOT NULL DEFAULT 'url',
|
||||
input_value TEXT NOT NULL DEFAULT '',
|
||||
summary_prompt TEXT NOT NULL DEFAULT '',
|
||||
openrouter_model TEXT NOT NULL DEFAULT '',
|
||||
total_videos INTEGER NOT NULL DEFAULT 0,
|
||||
done_videos INTEGER NOT NULL DEFAULT 0
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS videos (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
job_id TEXT NOT NULL REFERENCES jobs(job_id),
|
||||
video_id TEXT NOT NULL,
|
||||
title TEXT NOT NULL DEFAULT '',
|
||||
url TEXT NOT NULL,
|
||||
position INTEGER NOT NULL DEFAULT 0,
|
||||
status TEXT NOT NULL DEFAULT 'pending',
|
||||
UNIQUE(job_id, video_id)
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS artifacts (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
job_id TEXT NOT NULL REFERENCES jobs(job_id),
|
||||
video_id TEXT NOT NULL,
|
||||
artifact_type TEXT NOT NULL,
|
||||
status TEXT NOT NULL DEFAULT 'pending',
|
||||
file_path TEXT,
|
||||
error_msg TEXT,
|
||||
attempts INTEGER NOT NULL DEFAULT 0,
|
||||
started_at TEXT,
|
||||
finished_at TEXT,
|
||||
UNIQUE(job_id, video_id, artifact_type)
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS logs (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
job_id TEXT NOT NULL REFERENCES jobs(job_id),
|
||||
video_id TEXT,
|
||||
artifact_type TEXT,
|
||||
level TEXT NOT NULL DEFAULT 'INFO',
|
||||
message TEXT NOT NULL,
|
||||
timestamp TEXT NOT NULL
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_artifacts_job ON artifacts(job_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_videos_job ON videos(job_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_logs_job ON logs(job_id);
|
||||
""")
|
||||
|
||||
|
||||
# ── Job Operations ────────────────────────────────────────────────────────────
|
||||
|
||||
def create_job(
|
||||
output_folder: str,
|
||||
selections: List[str],
|
||||
resolution: str,
|
||||
delay_min: float,
|
||||
delay_max: float,
|
||||
input_type: str,
|
||||
input_value: str,
|
||||
summary_prompt: str,
|
||||
openrouter_model: str,
|
||||
) -> str:
|
||||
"""Create a new job and return its job_id."""
|
||||
job_id = str(uuid.uuid4())
|
||||
now = datetime.utcnow().isoformat()
|
||||
with get_conn() as conn:
|
||||
conn.execute(
|
||||
"""INSERT INTO jobs
|
||||
(job_id, status, created_at, updated_at, output_folder, selections,
|
||||
resolution, delay_min, delay_max, input_type, input_value,
|
||||
summary_prompt, openrouter_model)
|
||||
VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?)""",
|
||||
(
|
||||
job_id, "pending", now, now, output_folder,
|
||||
json.dumps(selections), resolution,
|
||||
delay_min, delay_max, input_type, input_value,
|
||||
summary_prompt, openrouter_model,
|
||||
),
|
||||
)
|
||||
return job_id
|
||||
|
||||
|
||||
def get_job(job_id: str) -> Optional[Dict]:
|
||||
with get_conn() as conn:
|
||||
row = conn.execute("SELECT * FROM jobs WHERE job_id=?", (job_id,)).fetchone()
|
||||
if row:
|
||||
d = dict(row)
|
||||
d["selections"] = json.loads(d["selections"])
|
||||
return d
|
||||
return None
|
||||
|
||||
|
||||
def list_jobs() -> List[Dict]:
|
||||
with get_conn() as conn:
|
||||
rows = conn.execute(
|
||||
"SELECT * FROM jobs ORDER BY created_at DESC"
|
||||
).fetchall()
|
||||
result = []
|
||||
for row in rows:
|
||||
d = dict(row)
|
||||
d["selections"] = json.loads(d["selections"])
|
||||
result.append(d)
|
||||
return result
|
||||
|
||||
|
||||
def update_job_status(job_id: str, status: str):
|
||||
now = datetime.utcnow().isoformat()
|
||||
with get_conn() as conn:
|
||||
conn.execute(
|
||||
"UPDATE jobs SET status=?, updated_at=? WHERE job_id=?",
|
||||
(status, now, job_id),
|
||||
)
|
||||
|
||||
|
||||
def update_job_progress(job_id: str, total: int, done: int):
|
||||
now = datetime.utcnow().isoformat()
|
||||
with get_conn() as conn:
|
||||
conn.execute(
|
||||
"UPDATE jobs SET total_videos=?, done_videos=?, updated_at=? WHERE job_id=?",
|
||||
(total, done, now, job_id),
|
||||
)
|
||||
|
||||
|
||||
def set_job_field(job_id: str, field: str, value: Any):
|
||||
"""Update a single field on the jobs table (safe whitelist)."""
|
||||
allowed = {"status", "total_videos", "done_videos", "updated_at"}
|
||||
if field not in allowed:
|
||||
raise ValueError(f"Field '{field}' not allowed")
|
||||
now = datetime.utcnow().isoformat()
|
||||
with get_conn() as conn:
|
||||
conn.execute(f"UPDATE jobs SET {field}=?, updated_at=? WHERE job_id=?", (value, now, job_id))
|
||||
|
||||
|
||||
# ── Video Operations ──────────────────────────────────────────────────────────
|
||||
|
||||
def upsert_video(job_id: str, video_id: str, url: str, title: str = "", position: int = 0):
|
||||
with get_conn() as conn:
|
||||
conn.execute(
|
||||
"""INSERT INTO videos (job_id, video_id, title, url, position, status)
|
||||
VALUES (?,?,?,?,?,'pending')
|
||||
ON CONFLICT(job_id, video_id) DO UPDATE SET
|
||||
title=excluded.title,
|
||||
position=excluded.position""",
|
||||
(job_id, video_id, title, url, position),
|
||||
)
|
||||
|
||||
|
||||
def update_video_title(job_id: str, video_id: str, title: str):
|
||||
with get_conn() as conn:
|
||||
conn.execute(
|
||||
"UPDATE videos SET title=? WHERE job_id=? AND video_id=?",
|
||||
(title, job_id, video_id),
|
||||
)
|
||||
|
||||
|
||||
def update_video_status(job_id: str, video_id: str, status: str):
|
||||
with get_conn() as conn:
|
||||
conn.execute(
|
||||
"UPDATE videos SET status=? WHERE job_id=? AND video_id=?",
|
||||
(status, job_id, video_id),
|
||||
)
|
||||
|
||||
|
||||
def get_videos(job_id: str) -> List[Dict]:
|
||||
with get_conn() as conn:
|
||||
rows = conn.execute(
|
||||
"SELECT * FROM videos WHERE job_id=? ORDER BY position ASC",
|
||||
(job_id,),
|
||||
).fetchall()
|
||||
return [dict(r) for r in rows]
|
||||
|
||||
|
||||
# ── Artifact Operations ───────────────────────────────────────────────────────
|
||||
|
||||
def upsert_artifact(job_id: str, video_id: str, artifact_type: str):
|
||||
"""Create artifact record if not exists (skip if already done/skipped)."""
|
||||
with get_conn() as conn:
|
||||
conn.execute(
|
||||
"""INSERT INTO artifacts (job_id, video_id, artifact_type, status)
|
||||
VALUES (?,?,?,'pending')
|
||||
ON CONFLICT(job_id, video_id, artifact_type) DO NOTHING""",
|
||||
(job_id, video_id, artifact_type),
|
||||
)
|
||||
|
||||
|
||||
def get_artifact(job_id: str, video_id: str, artifact_type: str) -> Optional[Dict]:
|
||||
with get_conn() as conn:
|
||||
row = conn.execute(
|
||||
"SELECT * FROM artifacts WHERE job_id=? AND video_id=? AND artifact_type=?",
|
||||
(job_id, video_id, artifact_type),
|
||||
).fetchone()
|
||||
return dict(row) if row else None
|
||||
|
||||
|
||||
def artifact_is_done(job_id: str, video_id: str, artifact_type: str) -> bool:
|
||||
"""Return True if artifact already succeeded or was skipped."""
|
||||
art = get_artifact(job_id, video_id, artifact_type)
|
||||
return art is not None and art["status"] in ("done", "skipped")
|
||||
|
||||
|
||||
def update_artifact_status(
|
||||
job_id: str,
|
||||
video_id: str,
|
||||
artifact_type: str,
|
||||
status: str,
|
||||
file_path: Optional[str] = None,
|
||||
error_msg: Optional[str] = None,
|
||||
):
|
||||
now = datetime.utcnow().isoformat()
|
||||
with get_conn() as conn:
|
||||
conn.execute(
|
||||
"""UPDATE artifacts
|
||||
SET status=?, file_path=?, error_msg=?, finished_at=?,
|
||||
attempts = attempts + CASE WHEN ? IN ('done','failed','skipped') THEN 1 ELSE 0 END
|
||||
WHERE job_id=? AND video_id=? AND artifact_type=?""",
|
||||
(status, file_path, error_msg, now, status, job_id, video_id, artifact_type),
|
||||
)
|
||||
|
||||
|
||||
def mark_artifact_started(job_id: str, video_id: str, artifact_type: str):
|
||||
now = datetime.utcnow().isoformat()
|
||||
with get_conn() as conn:
|
||||
conn.execute(
|
||||
"UPDATE artifacts SET status='running', started_at=? WHERE job_id=? AND video_id=? AND artifact_type=?",
|
||||
(now, job_id, video_id, artifact_type),
|
||||
)
|
||||
|
||||
|
||||
def get_artifacts_for_video(job_id: str, video_id: str) -> List[Dict]:
|
||||
with get_conn() as conn:
|
||||
rows = conn.execute(
|
||||
"SELECT * FROM artifacts WHERE job_id=? AND video_id=? ORDER BY id ASC",
|
||||
(job_id, video_id),
|
||||
).fetchall()
|
||||
return [dict(r) for r in rows]
|
||||
|
||||
|
||||
def get_all_artifacts(job_id: str) -> List[Dict]:
|
||||
with get_conn() as conn:
|
||||
rows = conn.execute(
|
||||
"""SELECT a.*, v.title as video_title
|
||||
FROM artifacts a
|
||||
JOIN videos v ON a.job_id=v.job_id AND a.video_id=v.video_id
|
||||
WHERE a.job_id=?
|
||||
ORDER BY v.position ASC, a.artifact_type ASC""",
|
||||
(job_id,),
|
||||
).fetchall()
|
||||
return [dict(r) for r in rows]
|
||||
|
||||
|
||||
# ── Logging ───────────────────────────────────────────────────────────────────
|
||||
|
||||
def add_log(
|
||||
job_id: str,
|
||||
message: str,
|
||||
level: str = "INFO",
|
||||
video_id: Optional[str] = None,
|
||||
artifact_type: Optional[str] = None,
|
||||
):
|
||||
now = datetime.utcnow().isoformat()
|
||||
with get_conn() as conn:
|
||||
conn.execute(
|
||||
"""INSERT INTO logs (job_id, video_id, artifact_type, level, message, timestamp)
|
||||
VALUES (?,?,?,?,?,?)""",
|
||||
(job_id, video_id, artifact_type, level, message, now),
|
||||
)
|
||||
|
||||
|
||||
def get_logs(job_id: str, limit: int = 500) -> List[Dict]:
|
||||
with get_conn() as conn:
|
||||
rows = conn.execute(
|
||||
"SELECT * FROM logs WHERE job_id=? ORDER BY id DESC LIMIT ?",
|
||||
(job_id, limit),
|
||||
).fetchall()
|
||||
return [dict(r) for r in reversed(rows)]
|
||||
|
||||
|
||||
def get_log_csv(job_id: str) -> str:
|
||||
"""Export artifacts as CSV string for download."""
|
||||
import io, csv
|
||||
artifacts = get_all_artifacts(job_id)
|
||||
output = io.StringIO()
|
||||
writer = csv.writer(output)
|
||||
writer.writerow(["video_id", "video_title", "artifact_type", "status", "file_path", "error_msg", "attempts", "started_at", "finished_at"])
|
||||
for a in artifacts:
|
||||
writer.writerow([
|
||||
a.get("video_id", ""),
|
||||
a.get("video_title", ""),
|
||||
a.get("artifact_type", ""),
|
||||
a.get("status", ""),
|
||||
a.get("file_path", ""),
|
||||
a.get("error_msg", ""),
|
||||
a.get("attempts", 0),
|
||||
a.get("started_at", ""),
|
||||
a.get("finished_at", ""),
|
||||
])
|
||||
return output.getvalue()
|
||||
|
||||
|
||||
def get_job_stats(job_id: str) -> Dict:
|
||||
"""Return aggregate counts for a job."""
|
||||
with get_conn() as conn:
|
||||
total_videos = conn.execute(
|
||||
"SELECT COUNT(*) FROM videos WHERE job_id=?", (job_id,)
|
||||
).fetchone()[0]
|
||||
done_videos = conn.execute(
|
||||
"SELECT COUNT(*) FROM videos WHERE job_id=? AND status='done'", (job_id,)
|
||||
).fetchone()[0]
|
||||
artifact_counts = conn.execute(
|
||||
"""SELECT status, COUNT(*) as cnt FROM artifacts WHERE job_id=? GROUP BY status""",
|
||||
(job_id,),
|
||||
).fetchall()
|
||||
counts = {row["status"]: row["cnt"] for row in artifact_counts}
|
||||
return {
|
||||
"total_videos": total_videos,
|
||||
"done_videos": done_videos,
|
||||
"artifact_done": counts.get("done", 0),
|
||||
"artifact_failed": counts.get("failed", 0),
|
||||
"artifact_pending": counts.get("pending", 0),
|
||||
"artifact_running": counts.get("running", 0),
|
||||
"artifact_skipped": counts.get("skipped", 0),
|
||||
}
|
||||
|
||||
|
||||
# Initialize on import
|
||||
init_db()
|
||||
Reference in New Issue
Block a user