mirror of
https://github.com/lrnselfreliance/wrolpi.git
synced 2026-10-26 19:01:39 -04:00
9310fe801f
A refresh on a large library aborted the whole doc_modeler run on the
first lock error, so the ~20k stranded docs never got modeled.
doc_modeler held one write transaction across a whole batch of 10 docs
while doing slow per-doc extraction (a PDF/epub can be tens of MB). When
another connection committed mid-batch, the deferred transaction failed
to upgrade to a write with SQLITE_BUSY_SNAPSHOT ("database is locked",
which busy_timeout does not absorb -- the same trap that motivated the
short-transaction rewrite of apply_indexers). The per-item error handler
then did `logger.error(f'... {file_group}')`, which lazy-loaded an
attribute on the now-rolled-back session and raised a *second* exception,
escaping the handler and aborting the entire modeler.
Restructure:
- Pick each batch of candidate ids in a short read-only transaction, then
model each doc in its own short write transaction. A lock/error on one
doc rolls back only that doc and cannot poison the rest of the run.
- Retry a doc a few times on "database is locked" -- a fresh transaction
gets a fresh snapshot and succeeds.
- Log the plain file_group id on failure; never touch the ORM object on a
possibly rolled-back session.
- Track only *failed* ids for the termination guard (successes are
excluded by the Doc.id IS NULL gate), keeping the notin_ set small and
well under SQLite's bound-variable limit.
Adds test_doc_modeler_survives_database_locked.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
257 lines
11 KiB
Python
257 lines
11 KiB
Python
import asyncio
|
|
import logging
|
|
from typing import Callable, List
|
|
|
|
from sqlalchemy import or_
|
|
from sqlalchemy.exc import OperationalError
|
|
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
|
|
# Retries for a single doc when it hits SQLite's "database is locked" (SQLITE_BUSY_SNAPSHOT).
|
|
DOC_LOCK_RETRIES = 3
|
|
|
|
|
|
@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.
|
|
"""
|
|
total_processed = 0
|
|
# Ids of docs that failed to model this run. A failed doc keeps matching the `Doc.id IS NULL`
|
|
# gate below, so without excluding it the loop would re-select it forever, hang, and starve
|
|
# later docs. Successfully-modeled docs are excluded automatically (they gain a Doc row), so
|
|
# only failures need tracking -- keeping this set (and the `notin_` clause) small. Failed docs
|
|
# are retried on the next refresh.
|
|
failed_ids: set = set()
|
|
while True:
|
|
# Short read-only transaction to choose the next batch of ids. Do NOT hold it open across
|
|
# the slow per-doc modeling below.
|
|
with get_db_session() as session:
|
|
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 == False,
|
|
),
|
|
)
|
|
if failed_ids:
|
|
query = query.filter(FileGroup.id.notin_(failed_ids))
|
|
fg_ids: List[int] = [row[0] for row in query.limit(DOC_PROCESSING_LIMIT).all()]
|
|
|
|
if not fg_ids:
|
|
break
|
|
|
|
for fg_id in fg_ids:
|
|
for attempt in range(DOC_LOCK_RETRIES):
|
|
try:
|
|
# One short write transaction per doc.
|
|
with get_db_session(commit=True) as session:
|
|
file_group = session.query(FileGroup).get(fg_id)
|
|
if file_group is None:
|
|
# Deleted between discovery and now; nothing to model.
|
|
break
|
|
doc = _model_doc(file_group, session)
|
|
session.add(doc)
|
|
file_group.model = Doc.__tablename__
|
|
file_group.indexed = True
|
|
break
|
|
except OperationalError as e:
|
|
# Another connection committed during our transaction. Retrying with a fresh
|
|
# transaction gets a fresh snapshot and usually succeeds.
|
|
if 'locked' in str(e).lower() and attempt < DOC_LOCK_RETRIES - 1:
|
|
await asyncio.sleep(0.2 * (attempt + 1))
|
|
continue
|
|
logger.error(f'Failed to model doc file_group_id={fg_id}: database error', exc_info=e)
|
|
failed_ids.add(fg_id)
|
|
break
|
|
except Exception as e:
|
|
# Never reference the ORM object here: after a failed flush the session is
|
|
# rolled back, so touching a lazy/expired attribute (e.g. in its repr) raises a
|
|
# second exception that would abort the whole modeler. Log the plain id only.
|
|
logger.error(f'Failed to model doc file_group_id={fg_id}', exc_info=e)
|
|
failed_ids.add(fg_id)
|
|
if PYTEST:
|
|
raise
|
|
break
|
|
# Yield so a cancel/other tasks can run between docs.
|
|
await asyncio.sleep(0)
|
|
|
|
total_processed += len(fg_ids)
|
|
if progress_callback:
|
|
progress_callback(total_processed)
|
|
|
|
|
|
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)
|