mirror of
https://github.com/lrnselfreliance/wrolpi.git
synced 2026-10-26 19:01:39 -04:00
821c3991be
Addresses review feedback: the new `Doc.id IS NULL` gate keeps matching a doc whose _model_doc fails without persisting a Doc row. With >=DOC_PROCESSING_LIMIT such docs, every batch re-selected the same failing files, `processed` never dropped below the limit, and the loop hung -- starving later docs. Reproduced with a patched _model_doc that always raises. Track attempted FileGroup ids per run and exclude them from the query, so the candidate pool strictly shrinks and the loop always terminates. Failed docs are retried on the next refresh (unchanged behavior). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
228 lines
8.9 KiB
Python
228 lines
8.9 KiB
Python
import asyncio
|
|
import logging
|
|
from typing import Callable, List, Tuple
|
|
|
|
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.models import FileGroup
|
|
from wrolpi.vars import PYTEST, 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
|
|
|
|
|
|
@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."""
|
|
total_processed = 0
|
|
# Track FileGroups already attempted in this run so a doc that fails to model (and therefore
|
|
# keeps matching the `Doc.id IS NULL` gate below) cannot be re-selected forever, which would
|
|
# hang this loop and starve later docs. Failed docs are simply retried on the next refresh.
|
|
attempted_ids: set = set()
|
|
while True:
|
|
with get_db_session(commit=True) as session:
|
|
query = session.query(FileGroup, Doc) \
|
|
.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 == False,
|
|
),
|
|
)
|
|
if attempted_ids:
|
|
query = query.filter(FileGroup.id.notin_(attempted_ids))
|
|
file_groups: List[Tuple[FileGroup, Doc]] = list(query.limit(DOC_PROCESSING_LIMIT))
|
|
|
|
processed = 0
|
|
for file_group, doc in file_groups:
|
|
processed += 1
|
|
attempted_ids.add(file_group.id)
|
|
try:
|
|
if PYTEST:
|
|
session.expire(file_group)
|
|
|
|
doc = _model_doc(file_group, session)
|
|
session.add(doc)
|
|
file_group.model = Doc.__tablename__
|
|
file_group.indexed = True
|
|
except Exception as e:
|
|
logger.error(f'Failed to index doc {file_group}', exc_info=e)
|
|
file_group.indexed = True # Don't retry on failure.
|
|
if PYTEST:
|
|
raise
|
|
|
|
session.commit()
|
|
|
|
total_processed += len(file_groups)
|
|
if progress_callback and len(file_groups) > 0:
|
|
progress_callback(total_processed)
|
|
|
|
if processed < DOC_PROCESSING_LIMIT:
|
|
break
|
|
|
|
await asyncio.sleep(0)
|
|
|
|
|
|
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)
|