wrolpi/modules/docs/__init__.py

import logging
from typing import Callable

from sqlalchemy import or_
from sqlalchemy.orm import Session

from wrolpi.common import register_modeler, register_refresh_cleanup, \
    truncate_object_bytes, split_lines_by_length
from wrolpi.db import get_db_session
from wrolpi.files.indexers import Indexer
from wrolpi.files.modeler import run_modeler_loop, skip_clause
from wrolpi.files.models import FileGroup
from wrolpi.vars import FILE_MAX_TEXT_SIZE
from .extractors import extract_metadata
from .lib import get_or_create_subject_collection, get_or_create_author_collection, is_valid_author, split_authors, \
    normalize_author, normalize_subject, split_subjects, is_valid_subject, discover_calibre_cover
from .models import Doc, DocSection, DOC_MIMETYPES, COMIC_BOOK_SUFFIXES, EPUB_MIMETYPE, MOBI_MIMETYPE

logger = logging.getLogger(__name__)

DOC_PROCESSING_LIMIT = 10
# Retries for a single doc when it hits SQLite's "database is locked" (SQLITE_BUSY_SNAPSHOT).
DOC_LOCK_RETRIES = 3


def _select_doc_ids(session, skip_ids: set) -> list[int]:
    query = session.query(FileGroup.id) \
        .outerjoin(Doc, Doc.file_group_id == FileGroup.id) \
        .filter(
        or_(
            *[FileGroup.mimetype == mt for mt in DOC_MIMETYPES],
            *[FileGroup.primary_path.like(f'%{suffix}') for suffix in COMIC_BOOK_SUFFIXES],
        ),
        # Model any doc that has no Doc row yet (even if `apply_indexers` already set
        # `indexed=True` before this modeler ran), or any doc explicitly flagged for
        # re-indexing.  Gating solely on `indexed == False` left already-indexed docs
        # permanently unmodeled and invisible in /api/docs.
        or_(
            Doc.id.is_(None),
            FileGroup.indexed != True,
        ),
    )
    query = skip_clause(query, FileGroup.id, skip_ids)
    return [row[0] for row in query.limit(DOC_PROCESSING_LIMIT).all()]


def _apply_doc(session: Session, file_group: FileGroup, prepared):
    doc = _model_doc(file_group, session)
    session.add(doc)
    file_group.model = Doc.__tablename__


@register_modeler
async def doc_modeler(progress_callback: Callable[[int], None] = None):
    """Searches for doc files (epub, mobi, pdf, docx, doc, odt, cbz, cbr) and models them.

    Each doc is modeled in its own short transaction.  Modeling reads a large file (a PDF/epub can
    be tens of MB) and can take seconds; holding one write transaction across a whole batch let
    another connection commit in the meantime, so the deferred transaction failed to upgrade to a
    write with SQLite's `SQLITE_BUSY_SNAPSHOT` ("database is locked", which `busy_timeout` does not
    absorb) -- and the error handler then crashed on the rolled-back session, aborting the entire
    modeler run.  Per-doc transactions keep the write window short and isolate failures so one
    locked/broken doc cannot poison the rest of the run.
    """
    await run_modeler_loop(
        name='doc_modeler',
        batch_size=DOC_PROCESSING_LIMIT,
        select_ids=_select_doc_ids,
        apply=_apply_doc,
        txn='per_item',
        lock_retries=DOC_LOCK_RETRIES,
        progress_callback=progress_callback,
    )


def _model_doc(file_group: FileGroup, session: Session) -> Doc:
    """Creates a Doc model from a FileGroup. Extracts metadata and creates collections."""
    doc = session.query(Doc).filter_by(file_group_id=file_group.id).one_or_none()
    if not doc:
        doc = Doc(file_group=file_group)
        session.add(doc)

    metadata = extract_metadata(file_group)

    if metadata.title:
        file_group.title = file_group.a_text = metadata.title
    if metadata.author:
        file_group.author = file_group.b_text = metadata.author
    if metadata.text:
        text = truncate_object_bytes(metadata.text, FILE_MAX_TEXT_SIZE)
        text = split_lines_by_length(text)
        file_group.d_text = text

    # Set c_text to filename for search.
    file_title = Indexer.get_title(file_group.primary_path)
    file_group.c_text = file_title

    if not file_group.title:
        file_group.title = file_title

    if metadata.published_date:
        file_group.published_datetime = metadata.published_date

    # Update Doc-specific fields.
    doc.publisher = metadata.publisher
    doc.language = metadata.language
    doc.page_count = metadata.page_count
    doc.subject = metadata.subject
    doc.description = metadata.description
    doc.size = file_group.primary_path.stat().st_size if file_group.primary_path.exists() else None

    # Handle cover.
    _handle_doc_cover(doc, file_group, metadata)

    # Store data dict.
    data = file_group.data or {}
    if metadata.title:
        data['title'] = metadata.title
    if metadata.author:
        data['author'] = metadata.author
    # Store primary doc path as relative filename.
    data['doc_path'] = file_group.primary_path.name
    file_group.data = data

    # Auto-create collections.
    _auto_create_collections(session, file_group, metadata)

    # Ensure the Doc has an id before inserting DocSection rows that reference it.
    doc.flush()

    # Replace any existing sections with the freshly-extracted ones.
    _replace_doc_sections(session, doc, metadata.sections)

    return doc


def _replace_doc_sections(session: Session, doc: Doc, sections):
    """Delete any DocSection rows for this Doc and insert the new ones.

    Using a bulk delete + bulk insert keeps this O(N) without SQLAlchemy tracking
    each old row. Callers must have flushed `doc` so it has an id.
    """
    session.query(DocSection).filter(DocSection.doc_id == doc.id).delete(synchronize_session=False)
    if not sections:
        return
    session.bulk_save_objects([
        DocSection(
            doc_id=doc.id,
            kind=s.kind,
            ordinal=s.ordinal,
            label=s.label,
            content=s.content,
        ) for s in sections
    ])


def _handle_doc_cover(doc: Doc, file_group: FileGroup, metadata):
    """Handle cover extraction and storage for a doc."""

    # Check for Calibre cover first (epub only) — must happen before poster check
    # since Calibre cover has a different stem ("cover.jpg" not matching epub stem).
    if file_group.mimetype and file_group.mimetype.startswith('application/epub'):
        if calibre_cover := discover_calibre_cover(file_group.primary_path):
            doc._hide_calibre_files(file_group)
            file_group.append_files(calibre_cover)
            file_group.data = file_group.data or {}
            file_group.data['cover_path'] = calibre_cover.name
            return

    # Use existing poster files if present.
    poster_files = file_group.my_poster_files()
    if poster_files:
        file_group.data = file_group.data or {}
        file_group.data['cover_path'] = poster_files[0]['path'].name
        return

    # Use cover from metadata extraction (epub internal cover, PDF first page, etc.)
    cover_bytes = metadata.cover_bytes
    if cover_bytes:
        cover_path = file_group.primary_path.with_suffix('.jpeg')
        cover_path.write_bytes(cover_bytes)
        file_group.append_files(cover_path)
        file_group.data = file_group.data or {}
        file_group.data['cover_path'] = cover_path.name


def _auto_create_collections(session: Session, file_group: FileGroup, metadata):
    """Auto-create author/subject collections and link the file_group."""
    if metadata.author:
        for author in split_authors(metadata.author):
            author = normalize_author(author)
            if author and is_valid_author(author):
                col = get_or_create_author_collection(session, author)
                if col:
                    col.add_file_group(file_group, session=session)

    if metadata.subject:
        for subject in split_subjects(metadata.subject):
            subject = normalize_subject(subject)
            if subject and is_valid_subject(subject):
                col = get_or_create_subject_collection(session, subject)
                if col:
                    col.add_file_group(file_group, session=session)


@register_refresh_cleanup
def doc_cleanup(test_directory=None):
    """Remove empty author/subject collections after refresh."""
    from sqlalchemy import func
    from wrolpi.collections.models import Collection, CollectionItem

    with get_db_session(commit=True) as session:
        empty_collections = session.query(Collection).filter(
            Collection.kind.in_(['author', 'subject']),
        ).outerjoin(CollectionItem).group_by(Collection.id).having(
            func.count(CollectionItem.id) == 0,
        ).all()
        for collection in empty_collections:
            logger.info(f'Removing empty {collection.kind} collection: {collection.name}')
            session.delete(collection)