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 UTC, datetime
from typing import TYPE_CHECKING

# scalar() is shared: record_files counts rows before and after. The rest are
# re-exported through __all__ so store.claim_next / store.ready_items / etc.
# keep working for downloader.py, cli.py and test_store.py unchanged.
from python_pkg.wsg_grabber._store_queries import (
    attempts_for,
    claim_next,
    claimable_values,
    counts,
    known_md5s,
    pending_downloads,
    ready_item,
    ready_items,
    scalar,
    state_of,
)
from python_pkg.wsg_grabber.states import FileState

__all__ = [
    "attempts_for",
    "claim_next",
    "claimable_values",
    "counts",
    "known_md5s",
    "local_name",
    "mark_downloaded",
    "pending_downloads",
    "ready_item",
    "ready_items",
    "record_files",
    "record_progress",
    "reset_in_flight",
    "set_state",
    "state_of",
]

if TYPE_CHECKING:
    from collections.abc import Iterable
    import sqlite3

    from python_pkg.wsg_grabber.models import (
        RemoteFile,
    )


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