drive/database.py

"""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