first commit
This commit is contained in:
gosangam
+752
@@ -0,0 +1,752 @@
|
||||
"""Persistent state: video records, state transitions, folder cache, cursors.
|
||||
|
||||
Two backends share one SQL dialect subset:
|
||||
|
||||
* PostgreSQL (``DATABASE_URL``) — durable production storage, e.g. Heroku Postgres.
|
||||
* SQLite (``DATABASE_PATH``) — local development. On Heroku the filesystem is
|
||||
ephemeral, so SQLite there would lose all history on every dyno restart; the
|
||||
configuration layer refuses that unless explicitly overridden.
|
||||
|
||||
SQL is written with ``?`` placeholders and translated to ``%s`` for psycopg.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import sqlite3
|
||||
import threading
|
||||
import time
|
||||
from contextlib import contextmanager
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from pathlib import Path
|
||||
from typing import Any, Iterator
|
||||
|
||||
from models import (
|
||||
IN_PROGRESS_STATES,
|
||||
MediaInfo,
|
||||
ParsedCaption,
|
||||
SourceMessage,
|
||||
State,
|
||||
TaskStatus,
|
||||
UploadTask,
|
||||
VideoRecord,
|
||||
can_transition,
|
||||
)
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
# Each entry is one schema version; {id_pk} differs between backends.
|
||||
MIGRATIONS: list[list[str]] = [
|
||||
[
|
||||
"""CREATE TABLE IF NOT EXISTS videos (
|
||||
channel_id BIGINT NOT NULL,
|
||||
message_id BIGINT NOT NULL,
|
||||
state TEXT NOT NULL,
|
||||
caption TEXT,
|
||||
indexes TEXT NOT NULL DEFAULT '[]',
|
||||
title TEXT,
|
||||
topic TEXT,
|
||||
batch TEXT,
|
||||
extracted_by TEXT,
|
||||
warnings TEXT NOT NULL DEFAULT '[]',
|
||||
media_kind TEXT,
|
||||
media_id BIGINT,
|
||||
mime_type TEXT,
|
||||
original_filename TEXT,
|
||||
expected_size BIGINT,
|
||||
message_date TEXT,
|
||||
sequence_no BIGINT,
|
||||
local_path TEXT,
|
||||
local_size BIGINT,
|
||||
sha256 TEXT,
|
||||
md5 TEXT,
|
||||
final_name TEXT,
|
||||
drive_folder_id TEXT,
|
||||
drive_file_id TEXT,
|
||||
upload_session_uri TEXT,
|
||||
upload_session_created_at TEXT,
|
||||
download_status TEXT,
|
||||
upload_status TEXT,
|
||||
attempts INTEGER NOT NULL DEFAULT 0,
|
||||
last_error TEXT,
|
||||
error_kind TEXT,
|
||||
next_retry_at TEXT,
|
||||
skip_reason TEXT,
|
||||
discovered_at TEXT NOT NULL,
|
||||
updated_at TEXT NOT NULL,
|
||||
download_started_at TEXT,
|
||||
downloaded_at TEXT,
|
||||
upload_started_at TEXT,
|
||||
completed_at TEXT,
|
||||
PRIMARY KEY (channel_id, message_id)
|
||||
)""",
|
||||
"CREATE INDEX IF NOT EXISTS idx_videos_state ON videos(state)",
|
||||
"CREATE INDEX IF NOT EXISTS idx_videos_media ON videos(media_id)",
|
||||
"""CREATE TABLE IF NOT EXISTS folders (
|
||||
parent_id TEXT NOT NULL,
|
||||
name_key TEXT NOT NULL,
|
||||
name TEXT NOT NULL,
|
||||
folder_id TEXT NOT NULL,
|
||||
created_at TEXT NOT NULL,
|
||||
PRIMARY KEY (parent_id, name_key)
|
||||
)""",
|
||||
"""CREATE TABLE IF NOT EXISTS transitions (
|
||||
id {id_pk},
|
||||
channel_id BIGINT NOT NULL,
|
||||
message_id BIGINT NOT NULL,
|
||||
from_state TEXT,
|
||||
to_state TEXT NOT NULL,
|
||||
at TEXT NOT NULL,
|
||||
note TEXT
|
||||
)""",
|
||||
"CREATE INDEX IF NOT EXISTS idx_transitions_msg ON transitions(channel_id, message_id)",
|
||||
],
|
||||
# v2: interactive /upload tasks
|
||||
[
|
||||
"ALTER TABLE videos ADD COLUMN task_id BIGINT",
|
||||
"CREATE INDEX IF NOT EXISTS idx_videos_task ON videos(task_id)",
|
||||
"""CREATE TABLE IF NOT EXISTS upload_tasks (
|
||||
id {id_pk},
|
||||
admin_user_id BIGINT NOT NULL,
|
||||
status TEXT NOT NULL,
|
||||
channel_ref TEXT,
|
||||
channel_id BIGINT,
|
||||
channel_title TEXT,
|
||||
start_message_id BIGINT,
|
||||
end_message_id BIGINT,
|
||||
created_at TEXT NOT NULL,
|
||||
updated_at TEXT NOT NULL,
|
||||
confirmed_at TEXT,
|
||||
started_at TEXT,
|
||||
scanned_at TEXT,
|
||||
completed_at TEXT,
|
||||
current_message_id BIGINT,
|
||||
examined_count INTEGER NOT NULL DEFAULT 0,
|
||||
eligible_count INTEGER NOT NULL DEFAULT 0,
|
||||
uploaded_count INTEGER NOT NULL DEFAULT 0,
|
||||
duplicate_count INTEGER NOT NULL DEFAULT 0,
|
||||
non_video_count INTEGER NOT NULL DEFAULT 0,
|
||||
skipped_count INTEGER NOT NULL DEFAULT 0,
|
||||
missing_count INTEGER NOT NULL DEFAULT 0,
|
||||
failed_count INTEGER NOT NULL DEFAULT 0,
|
||||
deleted_files_count INTEGER NOT NULL DEFAULT 0,
|
||||
last_error TEXT
|
||||
)""",
|
||||
"CREATE INDEX IF NOT EXISTS idx_tasks_status ON upload_tasks(status)",
|
||||
"CREATE INDEX IF NOT EXISTS idx_tasks_admin ON upload_tasks(admin_user_id)",
|
||||
"""CREATE TABLE IF NOT EXISTS task_items (
|
||||
task_id BIGINT NOT NULL,
|
||||
channel_id BIGINT NOT NULL,
|
||||
message_id BIGINT NOT NULL,
|
||||
result TEXT NOT NULL,
|
||||
folder TEXT,
|
||||
updated_at TEXT NOT NULL,
|
||||
PRIMARY KEY (task_id, channel_id, message_id)
|
||||
)""",
|
||||
],
|
||||
]
|
||||
|
||||
_COLUMNS = frozenset(VideoRecord.column_names())
|
||||
_MUTABLE = _COLUMNS - {"channel_id", "message_id", "state", "discovered_at"}
|
||||
|
||||
# download_status / upload_status mirrored from the state for the metadata record.
|
||||
_STAGE_STATUS = {
|
||||
State.PENDING: ("pending", "waiting"),
|
||||
State.DOWNLOADING: ("in_progress", "waiting"),
|
||||
State.DOWNLOADED: ("done", "waiting"),
|
||||
State.UPLOADING: ("done", "in_progress"),
|
||||
State.COMPLETED: ("done", "done"),
|
||||
State.SKIPPED: ("skipped", "skipped"),
|
||||
}
|
||||
|
||||
_TASK_MUTABLE = frozenset(UploadTask.column_names()) - {"id", "admin_user_id", "created_at"}
|
||||
_TASK_COUNTERS = frozenset({"deleted_files_count"})
|
||||
TERMINAL_VIDEO_STATES = ("COMPLETED", "SKIPPED", "FAILED")
|
||||
|
||||
_UPSERT_KV = "INSERT INTO kv (key, value) VALUES (?, ?) ON CONFLICT (key) DO UPDATE SET value = excluded.value"
|
||||
|
||||
|
||||
class InvalidTransition(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
class DatabaseError(RuntimeError):
|
||||
"""Backend-neutral wrapper so callers need not import sqlite3/psycopg."""
|
||||
|
||||
|
||||
class InstanceLockError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
def utcnow() -> datetime:
|
||||
return datetime.now(timezone.utc)
|
||||
|
||||
|
||||
def iso(dt: datetime) -> str:
|
||||
return dt.astimezone(timezone.utc).isoformat(timespec="seconds")
|
||||
|
||||
|
||||
def retry_time(seconds: float) -> str:
|
||||
return iso(utcnow() + timedelta(seconds=seconds))
|
||||
|
||||
|
||||
# -- backends --------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _SQLiteBackend:
|
||||
name = "sqlite"
|
||||
id_pk = "INTEGER PRIMARY KEY AUTOINCREMENT"
|
||||
for_update = ""
|
||||
errors: tuple[type[BaseException], ...] = (sqlite3.Error,)
|
||||
|
||||
def __init__(self, path: Path):
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
self.path = path
|
||||
self.conn = sqlite3.connect(str(path), isolation_level=None, check_same_thread=False, timeout=30)
|
||||
self.conn.row_factory = sqlite3.Row
|
||||
self.conn.execute("PRAGMA journal_mode=WAL")
|
||||
self.conn.execute("PRAGMA synchronous=FULL")
|
||||
|
||||
def execute(self, sql: str, params: tuple = ()):
|
||||
return self.conn.execute(sql, params)
|
||||
|
||||
@contextmanager
|
||||
def transaction(self):
|
||||
self.conn.execute("BEGIN IMMEDIATE")
|
||||
try:
|
||||
yield
|
||||
except BaseException:
|
||||
self.conn.execute("ROLLBACK")
|
||||
raise
|
||||
else:
|
||||
self.conn.execute("COMMIT")
|
||||
|
||||
def close(self) -> None:
|
||||
self.conn.close()
|
||||
|
||||
def describe(self) -> str:
|
||||
return f"SQLite ({self.path.name})"
|
||||
|
||||
def insert_returning_id(self, sql: str, params: tuple) -> int:
|
||||
return self.conn.execute(sql, params).lastrowid
|
||||
|
||||
|
||||
class _PostgresBackend:
|
||||
name = "postgres"
|
||||
id_pk = "BIGSERIAL PRIMARY KEY"
|
||||
for_update = " FOR UPDATE"
|
||||
|
||||
def __init__(self, url: str):
|
||||
import psycopg
|
||||
from psycopg.rows import dict_row
|
||||
|
||||
self._psycopg = psycopg
|
||||
self._dict_row = dict_row
|
||||
self.url = url
|
||||
self.errors = (psycopg.Error,)
|
||||
self.conn = self._connect()
|
||||
|
||||
def _connect(self):
|
||||
return self._psycopg.connect(self.url, autocommit=True, row_factory=self._dict_row, connect_timeout=15)
|
||||
|
||||
def _ensure_connected(self) -> None:
|
||||
if self.conn.closed or self.conn.broken:
|
||||
log.warning("PostgreSQL connection lost; reconnecting")
|
||||
self.conn = self._connect()
|
||||
|
||||
def execute(self, sql: str, params: tuple = ()):
|
||||
self._ensure_connected()
|
||||
return self.conn.execute(sql.replace("?", "%s"), params)
|
||||
|
||||
@contextmanager
|
||||
def transaction(self):
|
||||
self._ensure_connected()
|
||||
with self.conn.transaction():
|
||||
yield
|
||||
|
||||
def close(self) -> None:
|
||||
self.conn.close()
|
||||
|
||||
def describe(self) -> str:
|
||||
return "PostgreSQL"
|
||||
|
||||
def insert_returning_id(self, sql: str, params: tuple) -> int:
|
||||
return self.execute(sql + " RETURNING id", params).fetchone()["id"]
|
||||
|
||||
|
||||
# -- database --------------------------------------------------------------------------
|
||||
|
||||
|
||||
class Database:
|
||||
def __init__(self, *, path: Path | str | None = None, url: str | None = None):
|
||||
if url:
|
||||
self._backend: Any = _PostgresBackend(url)
|
||||
elif path:
|
||||
self._backend = _SQLiteBackend(Path(path))
|
||||
else:
|
||||
raise ValueError("either a SQLite path or a PostgreSQL URL is required")
|
||||
self.backend_name = self._backend.name
|
||||
self.errors = self._backend.errors
|
||||
self._lock = threading.RLock()
|
||||
self._migrate()
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config) -> "Database":
|
||||
return cls(url=config.database_url) if config.database_url else cls(path=config.database_path)
|
||||
|
||||
def describe(self) -> str:
|
||||
return self._backend.describe()
|
||||
|
||||
def close(self) -> None:
|
||||
with self._lock:
|
||||
self._backend.close()
|
||||
|
||||
def _migrate(self) -> None:
|
||||
with self._lock, self._backend.transaction():
|
||||
self._backend.execute("CREATE TABLE IF NOT EXISTS kv (key TEXT PRIMARY KEY, value TEXT)")
|
||||
row = self._backend.execute("SELECT value FROM kv WHERE key='schema_version'").fetchone()
|
||||
version = int(row["value"]) if row else 0
|
||||
for number, statements in enumerate(MIGRATIONS, start=1):
|
||||
if number <= version:
|
||||
continue
|
||||
for statement in statements:
|
||||
self._backend.execute(statement.format(id_pk=self._backend.id_pk))
|
||||
self._backend.execute(_UPSERT_KV, ("schema_version", str(number)))
|
||||
log.info("Database schema migrated to version %d", number)
|
||||
|
||||
@contextmanager
|
||||
def _tx(self) -> Iterator[Any]:
|
||||
with self._lock, self._backend.transaction():
|
||||
yield self._backend
|
||||
|
||||
def _query(self, sql: str, params: tuple = ()) -> list[Any]:
|
||||
with self._lock:
|
||||
return self._backend.execute(sql, params).fetchall()
|
||||
|
||||
# -- videos ---------------------------------------------------------------------
|
||||
|
||||
def insert_video(
|
||||
self,
|
||||
msg: SourceMessage,
|
||||
parsed: ParsedCaption,
|
||||
*,
|
||||
state: State = State.PENDING,
|
||||
skip_reason: str | None = None,
|
||||
sequence_no: int | None = None,
|
||||
task_id: int | None = None,
|
||||
) -> bool:
|
||||
"""Insert a newly discovered video. Returns False if it was already known."""
|
||||
media: MediaInfo = msg.media # type: ignore[assignment]
|
||||
now = iso(utcnow())
|
||||
dl, ul = _STAGE_STATUS[state]
|
||||
with self._tx() as db:
|
||||
cur = db.execute(
|
||||
"""INSERT INTO videos (
|
||||
channel_id, message_id, state, caption, indexes, title, topic, batch, extracted_by,
|
||||
warnings, media_kind, media_id, mime_type, original_filename, expected_size,
|
||||
message_date, sequence_no, skip_reason, download_status, upload_status,
|
||||
discovered_at, updated_at, task_id)
|
||||
VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)
|
||||
ON CONFLICT (channel_id, message_id) DO NOTHING""",
|
||||
(
|
||||
msg.channel_id,
|
||||
msg.message_id,
|
||||
state.value,
|
||||
msg.caption,
|
||||
json.dumps(parsed.indexes, ensure_ascii=False),
|
||||
parsed.title,
|
||||
parsed.topic,
|
||||
parsed.batch,
|
||||
parsed.extracted_by,
|
||||
json.dumps(parsed.warnings, ensure_ascii=False),
|
||||
media.kind,
|
||||
media.media_id,
|
||||
media.mime_type,
|
||||
media.file_name,
|
||||
media.size,
|
||||
iso(msg.date) if msg.date else None,
|
||||
sequence_no,
|
||||
skip_reason,
|
||||
dl,
|
||||
ul,
|
||||
now,
|
||||
now,
|
||||
task_id,
|
||||
),
|
||||
)
|
||||
inserted = cur.rowcount == 1
|
||||
if inserted:
|
||||
db.execute(
|
||||
"INSERT INTO transitions (channel_id, message_id, from_state, to_state, at, note) VALUES (?,?,?,?,?,?)",
|
||||
(msg.channel_id, msg.message_id, None, state.value, now, skip_reason or "discovered"),
|
||||
)
|
||||
return inserted
|
||||
|
||||
def exists(self, channel_id: int, message_id: int) -> bool:
|
||||
return bool(self._query("SELECT 1 FROM videos WHERE channel_id=? AND message_id=?", (channel_id, message_id)))
|
||||
|
||||
def get(self, key: tuple[int, int]) -> VideoRecord | None:
|
||||
rows = self._query("SELECT * FROM videos WHERE channel_id=? AND message_id=?", tuple(key))
|
||||
return VideoRecord.from_row(rows[0]) if rows else None
|
||||
|
||||
def find_by_message_id(self, message_id: int) -> list[VideoRecord]:
|
||||
rows = self._query("SELECT * FROM videos WHERE message_id=? ORDER BY channel_id", (message_id,))
|
||||
return [VideoRecord.from_row(r) for r in rows]
|
||||
|
||||
def find_by_media_id(self, media_id: int) -> list[VideoRecord]:
|
||||
rows = self._query("SELECT * FROM videos WHERE media_id=? ORDER BY message_id", (media_id,))
|
||||
return [VideoRecord.from_row(r) for r in rows]
|
||||
|
||||
def transition(self, key: tuple[int, int], new_state: State, *, note: str | None = None, **fields) -> VideoRecord:
|
||||
"""Atomically move an item to ``new_state`` (validated) and update fields."""
|
||||
bad = set(fields) - _MUTABLE
|
||||
if bad:
|
||||
raise ValueError(f"unknown or immutable fields: {sorted(bad)}")
|
||||
now = iso(utcnow())
|
||||
with self._tx() as db:
|
||||
row = db.execute(
|
||||
f"SELECT state FROM videos WHERE channel_id=? AND message_id=?{self._backend.for_update}", tuple(key)
|
||||
).fetchone()
|
||||
if row is None:
|
||||
raise KeyError(f"unknown video {key}")
|
||||
current = State(row["state"])
|
||||
if not can_transition(current, new_state):
|
||||
raise InvalidTransition(f"{key}: {current} -> {new_state} is not allowed")
|
||||
fields.setdefault("updated_at", now)
|
||||
if new_state in _STAGE_STATUS:
|
||||
fields.setdefault("download_status", _STAGE_STATUS[new_state][0])
|
||||
fields.setdefault("upload_status", _STAGE_STATUS[new_state][1])
|
||||
else:
|
||||
failed_stage = "upload_status" if current in (State.DOWNLOADED, State.UPLOADING) else "download_status"
|
||||
fields.setdefault(failed_stage, "failed" if new_state == State.FAILED else "retry_wait")
|
||||
self._update(db, key, state=new_state.value, **fields)
|
||||
db.execute(
|
||||
"INSERT INTO transitions (channel_id, message_id, from_state, to_state, at, note) VALUES (?,?,?,?,?,?)",
|
||||
(key[0], key[1], current.value, new_state.value, now, note),
|
||||
)
|
||||
return self.get(key) # type: ignore[return-value]
|
||||
|
||||
def update_fields(self, key: tuple[int, int], **fields) -> None:
|
||||
bad = set(fields) - _MUTABLE
|
||||
if bad:
|
||||
raise ValueError(f"unknown or immutable fields: {sorted(bad)}")
|
||||
fields.setdefault("updated_at", iso(utcnow()))
|
||||
with self._tx() as db:
|
||||
self._update(db, key, **fields)
|
||||
|
||||
@staticmethod
|
||||
def _update(db, key: tuple[int, int], **fields) -> None:
|
||||
for name in VideoRecord.JSON_FIELDS:
|
||||
if name in fields and not isinstance(fields[name], str):
|
||||
fields[name] = json.dumps(fields[name], ensure_ascii=False)
|
||||
assignments = ", ".join(f"{name}=?" for name in fields)
|
||||
db.execute(
|
||||
f"UPDATE videos SET {assignments} WHERE channel_id=? AND message_id=?",
|
||||
(*[v.value if isinstance(v, State) else v for v in fields.values()], *key),
|
||||
)
|
||||
|
||||
def history(self, key: tuple[int, int]) -> list[tuple[str | None, str, str | None]]:
|
||||
rows = self._query(
|
||||
"SELECT from_state, to_state, note FROM transitions WHERE channel_id=? AND message_id=? ORDER BY id", tuple(key)
|
||||
)
|
||||
return [(r["from_state"], r["to_state"], r["note"]) for r in rows]
|
||||
|
||||
def next_item(self, now: datetime | None = None, *, task_id: int | None = None) -> VideoRecord | None:
|
||||
"""Next item to work on: unfinished in-progress items first, then oldest message.
|
||||
|
||||
With ``task_id`` only that task's items are considered. Without it, items
|
||||
belonging to a paused or queued /upload task are left alone.
|
||||
"""
|
||||
if task_id is not None:
|
||||
scope, params = "task_id = ?", (task_id,)
|
||||
else:
|
||||
scope = (
|
||||
"(task_id IS NULL OR task_id NOT IN "
|
||||
"(SELECT id FROM upload_tasks WHERE status IN ('PAUSED','QUEUED')))"
|
||||
)
|
||||
params = ()
|
||||
rows = self._query(
|
||||
f"""SELECT * FROM videos
|
||||
WHERE {scope}
|
||||
AND (state IN ('UPLOADING','DOWNLOADED','DOWNLOADING','PENDING')
|
||||
OR (state='RETRY_WAIT' AND next_retry_at <= ?))
|
||||
ORDER BY CASE state WHEN 'UPLOADING' THEN 0 WHEN 'DOWNLOADED' THEN 1
|
||||
WHEN 'DOWNLOADING' THEN 2 ELSE 3 END,
|
||||
message_id
|
||||
LIMIT 1""",
|
||||
(*params, iso(now or utcnow())),
|
||||
)
|
||||
return VideoRecord.from_row(rows[0]) if rows else None
|
||||
|
||||
def earliest_retry_at(self) -> datetime | None:
|
||||
rows = self._query("SELECT MIN(next_retry_at) AS t FROM videos WHERE state='RETRY_WAIT'")
|
||||
value = rows[0]["t"] if rows else None
|
||||
return datetime.fromisoformat(value) if value else None
|
||||
|
||||
def recover_stale(self) -> dict[str, int]:
|
||||
"""Called once at startup, after the single-instance lock is held.
|
||||
|
||||
DOWNLOADING items go back to PENDING; a partial file, if it survived, is resumed.
|
||||
UPLOADING items are left as they are: the queue reconciles them against Drive
|
||||
(and resumes the saved upload session) before doing anything else.
|
||||
"""
|
||||
recovered = {}
|
||||
for state in sorted(IN_PROGRESS_STATES):
|
||||
rows = self._query("SELECT channel_id, message_id FROM videos WHERE state=?", (state.value,))
|
||||
recovered[state.value] = len(rows)
|
||||
if state == State.DOWNLOADING:
|
||||
for r in rows:
|
||||
self.transition((r["channel_id"], r["message_id"]), State.PENDING, note="recovered after restart")
|
||||
if any(recovered.values()):
|
||||
log.warning("Recovered interrupted items after restart: %s", recovered)
|
||||
return recovered
|
||||
|
||||
def reset_failed(self, *, has_local) -> int:
|
||||
"""Move FAILED items back into the queue with a fresh attempt budget."""
|
||||
rows = self._query("SELECT * FROM videos WHERE state='FAILED' ORDER BY message_id")
|
||||
for row in rows:
|
||||
rec = VideoRecord.from_row(row)
|
||||
target = State.DOWNLOADED if has_local(rec) else State.PENDING
|
||||
self.transition(rec.key, target, note="manual retry", attempts=0, next_retry_at=None, error_kind=None)
|
||||
return len(rows)
|
||||
|
||||
def counts(self) -> dict[str, int]:
|
||||
counts = {s.value: 0 for s in State}
|
||||
for r in self._query("SELECT state, COUNT(*) AS n FROM videos GROUP BY state"):
|
||||
counts[r["state"]] = int(r["n"])
|
||||
return counts
|
||||
|
||||
def list_by_state(self, state: State, limit: int = 20) -> list[VideoRecord]:
|
||||
rows = self._query("SELECT * FROM videos WHERE state=? ORDER BY message_id LIMIT ?", (state.value, limit))
|
||||
return [VideoRecord.from_row(r) for r in rows]
|
||||
|
||||
def topic_counts(self, limit: int = 30) -> list[tuple[str, int]]:
|
||||
rows = self._query(
|
||||
"""SELECT COALESCE(topic, '') AS topic, COUNT(*) AS n FROM videos WHERE state='COMPLETED'
|
||||
GROUP BY COALESCE(topic, '') ORDER BY n DESC, topic LIMIT ?""",
|
||||
(limit,),
|
||||
)
|
||||
return [(r["topic"], int(r["n"])) for r in rows]
|
||||
|
||||
def videos_for_task(self, task_id: int) -> list[VideoRecord]:
|
||||
rows = self._query("SELECT * FROM videos WHERE task_id=? ORDER BY message_id", (task_id,))
|
||||
return [VideoRecord.from_row(r) for r in rows]
|
||||
|
||||
def task_unfinished_count(self, task_id: int) -> int:
|
||||
rows = self._query(
|
||||
"SELECT COUNT(*) AS n FROM videos WHERE task_id=? AND state NOT IN ('COMPLETED','SKIPPED','FAILED')",
|
||||
(task_id,),
|
||||
)
|
||||
return int(rows[0]["n"])
|
||||
|
||||
# -- upload tasks -----------------------------------------------------------------
|
||||
|
||||
def create_task(self, admin_user_id: int) -> UploadTask:
|
||||
now = iso(utcnow())
|
||||
with self._tx() as db:
|
||||
task_id = db.insert_returning_id(
|
||||
"INSERT INTO upload_tasks (admin_user_id, status, created_at, updated_at) VALUES (?,?,?,?)",
|
||||
(admin_user_id, TaskStatus.AWAITING_START.value, now, now),
|
||||
)
|
||||
return self.get_task(task_id) # type: ignore[return-value]
|
||||
|
||||
def get_task(self, task_id: int) -> UploadTask | None:
|
||||
rows = self._query("SELECT * FROM upload_tasks WHERE id=?", (task_id,))
|
||||
return UploadTask.from_row(rows[0]) if rows else None
|
||||
|
||||
def update_task(self, task_id: int, **fields) -> UploadTask:
|
||||
bad = set(fields) - _TASK_MUTABLE
|
||||
if bad:
|
||||
raise ValueError(f"unknown or immutable task fields: {sorted(bad)}")
|
||||
fields.setdefault("updated_at", iso(utcnow()))
|
||||
values = [v.value if isinstance(v, TaskStatus) else v for v in fields.values()]
|
||||
assignments = ", ".join(f"{name}=?" for name in fields)
|
||||
with self._tx() as db:
|
||||
db.execute(f"UPDATE upload_tasks SET {assignments} WHERE id=?", (*values, task_id))
|
||||
return self.get_task(task_id) # type: ignore[return-value]
|
||||
|
||||
def add_task_counter(self, task_id: int, column: str, amount: int) -> None:
|
||||
if column not in _TASK_COUNTERS:
|
||||
raise ValueError(column)
|
||||
with self._tx() as db:
|
||||
db.execute(f"UPDATE upload_tasks SET {column} = {column} + ? WHERE id=?", (amount, task_id))
|
||||
|
||||
def tasks_with_status(self, *statuses: TaskStatus, admin_user_id: int | None = None) -> list[UploadTask]:
|
||||
marks = ",".join("?" for _ in statuses)
|
||||
sql = f"SELECT * FROM upload_tasks WHERE status IN ({marks})"
|
||||
params: tuple = tuple(s.value for s in statuses)
|
||||
if admin_user_id is not None:
|
||||
sql += " AND admin_user_id=?"
|
||||
params += (admin_user_id,)
|
||||
rows = self._query(sql + " ORDER BY id", params)
|
||||
return [UploadTask.from_row(r) for r in rows]
|
||||
|
||||
def next_runnable_task(self) -> UploadTask | None:
|
||||
"""The task the worker should run: an already RUNNING one, else the oldest QUEUED."""
|
||||
rows = self._query(
|
||||
"""SELECT * FROM upload_tasks WHERE status IN ('RUNNING','QUEUED')
|
||||
ORDER BY CASE status WHEN 'RUNNING' THEN 0 ELSE 1 END, id LIMIT 1"""
|
||||
)
|
||||
return UploadTask.from_row(rows[0]) if rows else None
|
||||
|
||||
def recent_tasks(self, limit: int = 10) -> list[UploadTask]:
|
||||
rows = self._query(
|
||||
"SELECT * FROM upload_tasks WHERE status NOT IN ('AWAITING_START','AWAITING_END','AWAITING_CONFIRMATION') "
|
||||
"ORDER BY id DESC LIMIT ?",
|
||||
(limit,),
|
||||
)
|
||||
return [UploadTask.from_row(r) for r in rows]
|
||||
|
||||
def set_task_status(self, old: TaskStatus, new: TaskStatus) -> int:
|
||||
with self._tx() as db:
|
||||
cur = db.execute(
|
||||
"UPDATE upload_tasks SET status=?, updated_at=? WHERE status=?", (new.value, iso(utcnow()), old.value)
|
||||
)
|
||||
return cur.rowcount
|
||||
|
||||
def add_task_item(self, task_id: int, channel_id: int, message_id: int, result: str) -> None:
|
||||
"""Record a message examined by a task; an existing result is kept (rescans are idempotent)."""
|
||||
with self._tx() as db:
|
||||
db.execute(
|
||||
"""INSERT INTO task_items (task_id, channel_id, message_id, result, updated_at) VALUES (?,?,?,?,?)
|
||||
ON CONFLICT (task_id, channel_id, message_id) DO NOTHING""",
|
||||
(task_id, channel_id, message_id, result, iso(utcnow())),
|
||||
)
|
||||
|
||||
def set_task_item_result(
|
||||
self, task_id: int, channel_id: int, message_id: int, result: str, *, folder: str | None = None
|
||||
) -> None:
|
||||
with self._tx() as db:
|
||||
db.execute(
|
||||
"""UPDATE task_items SET result=?, folder=COALESCE(?, folder), updated_at=?
|
||||
WHERE task_id=? AND channel_id=? AND message_id=?""",
|
||||
(result, folder, iso(utcnow()), task_id, channel_id, message_id),
|
||||
)
|
||||
|
||||
def task_item_results(self, task_id: int) -> dict[str, int]:
|
||||
rows = self._query("SELECT result, COUNT(*) AS n FROM task_items WHERE task_id=? GROUP BY result", (task_id,))
|
||||
return {r["result"]: int(r["n"]) for r in rows}
|
||||
|
||||
def task_items(self, task_id: int) -> list[tuple[int, str]]:
|
||||
rows = self._query("SELECT message_id, result FROM task_items WHERE task_id=? ORDER BY message_id", (task_id,))
|
||||
return [(int(r["message_id"]), r["result"]) for r in rows]
|
||||
|
||||
def task_folders(self, task_id: int) -> list[str]:
|
||||
rows = self._query(
|
||||
"SELECT DISTINCT folder FROM task_items WHERE task_id=? AND folder IS NOT NULL ORDER BY folder", (task_id,)
|
||||
)
|
||||
return [r["folder"] for r in rows]
|
||||
|
||||
# -- folders --------------------------------------------------------------------
|
||||
|
||||
def get_folder(self, parent_id: str, name_key: str) -> str | None:
|
||||
rows = self._query("SELECT folder_id FROM folders WHERE parent_id=? AND name_key=?", (parent_id, name_key))
|
||||
return rows[0]["folder_id"] if rows else None
|
||||
|
||||
def put_folder(self, parent_id: str, name_key: str, name: str, folder_id: str) -> None:
|
||||
with self._tx() as db:
|
||||
db.execute(
|
||||
"""INSERT INTO folders (parent_id, name_key, name, folder_id, created_at) VALUES (?,?,?,?,?)
|
||||
ON CONFLICT (parent_id, name_key) DO UPDATE
|
||||
SET name = excluded.name, folder_id = excluded.folder_id, created_at = excluded.created_at""",
|
||||
(parent_id, name_key, name, folder_id, iso(utcnow())),
|
||||
)
|
||||
|
||||
def delete_folder(self, parent_id: str, name_key: str) -> None:
|
||||
with self._tx() as db:
|
||||
db.execute("DELETE FROM folders WHERE parent_id=? AND name_key=?", (parent_id, name_key))
|
||||
|
||||
def folder_count(self) -> int:
|
||||
return int(self._query("SELECT COUNT(*) AS n FROM folders")[0]["n"])
|
||||
|
||||
# -- key/value --------------------------------------------------------------------
|
||||
|
||||
def get_kv(self, key: str, default: str | None = None) -> str | None:
|
||||
rows = self._query("SELECT value FROM kv WHERE key=?", (key,))
|
||||
return rows[0]["value"] if rows else default
|
||||
|
||||
def set_kv(self, key: str, value: str) -> None:
|
||||
with self._tx() as db:
|
||||
db.execute(_UPSERT_KV, (key, value))
|
||||
|
||||
def get_int(self, key: str, default: int = 0) -> int:
|
||||
value = self.get_kv(key)
|
||||
return int(value) if value is not None else default
|
||||
|
||||
def raise_int(self, key: str, value: int) -> None:
|
||||
"""Store ``value`` only if it is larger than the current one (cursor semantics)."""
|
||||
with self._tx() as db:
|
||||
row = db.execute(f"SELECT value FROM kv WHERE key=?{self._backend.for_update}", (key,)).fetchone()
|
||||
if row is None or int(row["value"]) < value:
|
||||
db.execute(_UPSERT_KV, (key, str(value)))
|
||||
|
||||
def next_sequence(self) -> int:
|
||||
with self._tx() as db:
|
||||
row = db.execute(f"SELECT value FROM kv WHERE key='sequence_counter'{self._backend.for_update}").fetchone()
|
||||
value = (int(row["value"]) if row else 0) + 1
|
||||
db.execute(_UPSERT_KV, ("sequence_counter", str(value)))
|
||||
return value
|
||||
|
||||
|
||||
# -- single-instance lock ----------------------------------------------------------------
|
||||
|
||||
_ADVISORY_LOCK_ID = 0x7464_6272 # arbitrary constant: "tdbr"
|
||||
|
||||
|
||||
class InstanceLock:
|
||||
"""Guarantee that only one worker processes the queue.
|
||||
|
||||
PostgreSQL: a session-level advisory lock held on a dedicated connection, so it
|
||||
also protects against two Heroku dynos (e.g. worker scaled to 2, or an old dyno
|
||||
still shutting down during a deploy). SQLite: an exclusive ``flock`` on a lock
|
||||
file, which only protects processes on the same machine.
|
||||
"""
|
||||
|
||||
def __init__(self, config, *, wait_seconds: float = 0, poll: float = 5.0):
|
||||
self._config = config
|
||||
self._wait = wait_seconds
|
||||
self._poll = poll
|
||||
self._conn = None
|
||||
self._fh = None
|
||||
|
||||
def acquire(self) -> None:
|
||||
deadline = time.monotonic() + self._wait
|
||||
while True:
|
||||
if self._try_acquire():
|
||||
return
|
||||
if time.monotonic() >= deadline:
|
||||
raise InstanceLockError(
|
||||
"Another worker is already processing this queue. Run exactly one worker "
|
||||
"(Heroku: `heroku ps:scale worker=1`)."
|
||||
)
|
||||
log.warning("Another worker holds the queue lock; waiting...")
|
||||
time.sleep(self._poll)
|
||||
|
||||
def _try_acquire(self) -> bool:
|
||||
if self._config.database_url:
|
||||
import psycopg
|
||||
|
||||
if self._conn is None or self._conn.closed:
|
||||
self._conn = psycopg.connect(self._config.database_url, autocommit=True, connect_timeout=15)
|
||||
return bool(self._conn.execute("SELECT pg_try_advisory_lock(%s)", (_ADVISORY_LOCK_ID,)).fetchone()[0])
|
||||
import fcntl
|
||||
|
||||
path = Path(str(self._config.database_path) + ".lock")
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
if self._fh is None:
|
||||
self._fh = open(path, "w")
|
||||
try:
|
||||
fcntl.flock(self._fh, fcntl.LOCK_EX | fcntl.LOCK_NB)
|
||||
return True
|
||||
except BlockingIOError:
|
||||
return False
|
||||
|
||||
def release(self) -> None:
|
||||
if self._conn is not None:
|
||||
try:
|
||||
self._conn.close() # closing the session releases the advisory lock
|
||||
finally:
|
||||
self._conn = None
|
||||
if self._fh is not None:
|
||||
self._fh.close()
|
||||
self._fh = None
|
||||
Reference in New Issue
Block a user