testsAndMisc/python_pkg/wsg_grabber/_store_queries.py

"""Read-only projections over the file index, plus the atomic claim.

Split out of :mod:`python_pkg.wsg_grabber.store` to keep it under the 250-line
cap. ``store`` owns the writes that move a file through its lifecycle; this
module asks the questions -- what is reviewable, how many of each state there
are, how much work is left -- and owns ``claim_next``, which is the one write
that has to be a single atomic transaction against those same reads.
"""

from __future__ import annotations

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 Mapping, Sequence
    from pathlib import Path
    import sqlite3


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])


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