testsAndMisc/python_pkg/wsg_grabber/store.py

"""The download lifecycle of a file, as typed reads and writes.

One of three query modules, split to stay under this repo's 500-line-per-file
cap: this one owns how a file got downloaded, :mod:`store_threads` owns which
threads were visited, and :mod:`store_verdicts` owns what the user decided.
Between them they hold all the SQL outside :mod:`db`, so the rest of the
package deals in dataclasses. Every function takes an already-open connection
owned by the calling thread.
"""

from __future__ import annotations

from datetime import datetime, timezone
from typing import TYPE_CHECKING

from python_pkg.wsg_grabber.models import (
    RemoteFile,
    ReviewItem,
)
from python_pkg.wsg_grabber.states import CLAIMABLE, FileState

if TYPE_CHECKING:
    from collections.abc import Iterable, Mapping, Sequence
    from pathlib import Path
    import sqlite3


def _now() -> str:
    """Return the current UTC time as an ISO-8601 string.

    Returns:
        str: Timestamp suitable for the index's text date columns.
    """
    return datetime.now(tz=timezone.utc).isoformat()


def local_name(md5_b64: str, tim: int, ext: str) -> str:
    """Return the on-disk filename for a file.

    The poster-supplied name is untrusted text and is never used to build a
    path; it is carried alongside purely for display.

    Args:
        md5_b64: The API's base64 MD5, used only to disambiguate.
        tim: The CDN filename stem.
        ext: Extension including the leading dot.

    Returns:
        str: A filesystem-safe name unique per file.
    """
    suffix = md5_b64.rstrip("=")[-6:].replace("/", "_").replace("+", "-")
    return f"{tim}-{suffix}{ext}"


def record_files(conn: sqlite3.Connection, files: Iterable[RemoteFile]) -> int:
    """Insert any files not already known.

    Dedupe is structural: ``md5`` is the primary key, so a repost of identical
    bytes in another thread collapses onto the existing row and never costs a
    download.

    Args:
        conn: Open connection.
        files: Candidate attachments parsed from a thread.

    Returns:
        int: How many rows were genuinely new.
    """
    rows = [
        (
            item.md5,
            item.tim,
            item.ext,
            item.orig_name,
            item.fsize,
            item.width,
            item.height,
            item.thread_no,
            item.post_no,
            FileState.NEW.value,
            _now(),
        )
        for item in files
    ]
    if not rows:
        return 0
    before = _scalar(conn, "SELECT COUNT(*) FROM files")
    conn.executemany(
        """
        INSERT INTO files (md5, tim, ext, orig_name, fsize, width, height,
                           thread_no, post_no, state, first_seen)
        VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
        ON CONFLICT(md5) DO NOTHING
        """,
        rows,
    )
    return _scalar(conn, "SELECT COUNT(*) FROM files") - before


def known_md5s(conn: sqlite3.Connection) -> set[str]:
    """Return every md5 the index has ever seen.

    Args:
        conn: Open connection.

    Returns:
        set[str]: Base64 md5 values, including terminal ones, so a reviewed or
        written-off file is never offered again.
    """
    return {str(row[0]) for row in conn.execute("SELECT md5 FROM files")}


def claimable_values() -> list[str]:
    """Return the state values a download worker may pick up, in a stable order.

    The queries below spell out one placeholder per member rather than building
    the ``IN`` list at runtime, so the SQL stays a literal. ``test_store``
    pins the member count so adding a state cannot silently desync the two.

    Returns:
        list[str]: Sorted state values.
    """
    return [state.value for state in sorted(CLAIMABLE)]


def claim_next(conn: sqlite3.Connection) -> RemoteFile | None:
    """Atomically take the oldest downloadable file and mark it in flight.

    Args:
        conn: Open connection.

    Returns:
        RemoteFile | None: The claimed file, or None when nothing is pending.
    """
    conn.execute("BEGIN IMMEDIATE")
    try:
        row = conn.execute(
            """
            SELECT md5, tim, ext, orig_name, fsize, width, height,
                   thread_no, post_no
              FROM files
             WHERE state IN (?, ?, ?)
             ORDER BY first_seen, md5
             LIMIT 1
            """,
            claimable_values(),
        ).fetchone()
        if row is None:
            conn.execute("COMMIT")
            return None
        conn.execute(
            "UPDATE files SET state = ?, attempts = attempts + 1 WHERE md5 = ?",
            (FileState.DOWNLOADING.value, row["md5"]),
        )
        conn.execute("COMMIT")
    except Exception:
        conn.execute("ROLLBACK")
        raise
    return RemoteFile(
        md5=row["md5"],
        tim=row["tim"],
        ext=row["ext"],
        orig_name=row["orig_name"],
        fsize=row["fsize"],
        width=row["width"],
        height=row["height"],
        thread_no=row["thread_no"],
        post_no=row["post_no"],
    )


def state_of(conn: sqlite3.Connection, md5_b64: str) -> FileState | None:
    """Return a file's current state.

    Args:
        conn: Open connection.
        md5_b64: File identity.

    Returns:
        FileState | None: The state, or None when the file is unknown.
    """
    row = conn.execute(
        "SELECT state FROM files WHERE md5 = ?",
        (md5_b64,),
    ).fetchone()
    return None if row is None else FileState(row["state"])


def attempts_for(conn: sqlite3.Connection, md5_b64: str) -> int:
    """Return how many download attempts a file has had.

    Args:
        conn: Open connection.
        md5_b64: File identity.

    Returns:
        int: Attempt count, zero when the file is unknown.
    """
    row = conn.execute(
        "SELECT attempts FROM files WHERE md5 = ?",
        (md5_b64,),
    ).fetchone()
    return 0 if row is None else int(row["attempts"])


def set_state(
    conn: sqlite3.Connection,
    md5_b64: str,
    state: FileState,
    *,
    error: str | None = None,
) -> None:
    """Move a file to *state*.

    Args:
        conn: Open connection.
        md5_b64: File identity.
        state: New state.
        error: Optional diagnostic kept for the next run's benefit.
    """
    conn.execute(
        "UPDATE files SET state = ?, last_error = ? WHERE md5 = ?",
        (state.value, error, md5_b64),
    )


def mark_downloaded(conn: sqlite3.Connection, md5_b64: str, name: str) -> None:
    """Record that a file finished downloading and is ready to review.

    Args:
        conn: Open connection.
        md5_b64: File identity.
        name: Filename inside the incoming directory.
    """
    conn.execute(
        """
        UPDATE files
           SET state = ?, local_name = ?, downloaded_at = ?, bytes_done = 0,
               last_error = NULL
         WHERE md5 = ?
        """,
        (FileState.READY.value, name, _now(), md5_b64),
    )


def record_progress(conn: sqlite3.Connection, md5_b64: str, done: int) -> None:
    """Persist partial download progress so a restart can resume.

    Args:
        conn: Open connection.
        md5_b64: File identity.
        done: Bytes written to the ``.part`` file so far.
    """
    conn.execute(
        "UPDATE files SET bytes_done = ? WHERE md5 = ?",
        (done, md5_b64),
    )


def ready_items(conn: sqlite3.Connection, directory: Path) -> list[ReviewItem]:
    """Return downloaded files still awaiting a verdict, oldest first.

    Args:
        conn: Open connection.
        directory: Where downloaded files live.

    Returns:
        list[ReviewItem]: Reviewable videos.
    """
    rows = conn.execute(
        """
        SELECT md5, local_name, orig_name, ext, fsize, width, height
          FROM files
         WHERE state = ? AND local_name IS NOT NULL
         ORDER BY downloaded_at, md5
        """,
        (FileState.READY.value,),
    ).fetchall()
    return [
        ReviewItem(
            md5=row["md5"],
            path=directory / row["local_name"],
            orig_name=f"{row['orig_name']}{row['ext']}",
            fsize=row["fsize"],
            width=row["width"],
            height=row["height"],
        )
        for row in rows
    ]


def ready_item(source: RemoteFile, path: Path) -> ReviewItem:
    """Build the reviewer's view of a file that just finished downloading.

    Args:
        source: The catalogue record it came from.
        path: Where the verified bytes now live.

    Returns:
        ReviewItem: Value handed to the UI thread.
    """
    return ReviewItem(
        md5=source.md5,
        path=path,
        orig_name=f"{source.orig_name}{source.ext}",
        fsize=source.fsize,
        width=source.width,
        height=source.height,
    )


def reset_in_flight(conn: sqlite3.Connection) -> int:
    """Return files interrupted mid-download to the pending pool.

    A row left in ``downloading`` means the process died; the partial file on
    disk is still valid and will be resumed.

    Args:
        conn: Open connection.

    Returns:
        int: How many rows were reset.
    """
    cursor = conn.execute(
        "UPDATE files SET state = ? WHERE state = ?",
        (FileState.FAILED.value, FileState.DOWNLOADING.value),
    )
    return int(cursor.rowcount)


def counts(conn: sqlite3.Connection) -> dict[str, int]:
    """Return the number of files in each state.

    Args:
        conn: Open connection.

    Returns:
        dict[str, int]: State value mapped to row count, plus ``total``.
    """
    result = {
        str(row["state"]): int(row["n"])
        for row in conn.execute(
            "SELECT state, COUNT(*) AS n FROM files GROUP BY state",
        )
    }
    result["total"] = _scalar(conn, "SELECT COUNT(*) FROM files")
    return result


def pending_downloads(conn: sqlite3.Connection) -> int:
    """Return how many files are still waiting to be fetched.

    Args:
        conn: Open connection.

    Returns:
        int: Count of claimable rows.
    """
    return _scalar(
        conn,
        "SELECT COUNT(*) FROM files WHERE state IN (?, ?, ?)",
        claimable_values(),
    )


def _scalar(
    conn: sqlite3.Connection,
    sql: str,
    params: Sequence[object] | Mapping[str, object] | None = None,
) -> int:
    """Run a single-value query and return it as an int.

    Args:
        conn: Open connection.
        sql: Query yielding exactly one column.
        params: Optional bound parameters.

    Returns:
        int: The first column of the first row.
    """
    cursor = conn.execute(sql) if params is None else conn.execute(sql, params)
    return int(cursor.fetchone()[0])