372 lines
14 KiB
Python
372 lines
14 KiB
Python
"""
|
|
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()
|