diff --git a/modules/archive/__init__.py b/modules/archive/__init__.py index 3ad5554b..85af6412 100644 --- a/modules/archive/__init__.py +++ b/modules/archive/__init__.py @@ -1,13 +1,11 @@ -import asyncio import json import os import pathlib import tempfile from abc import ABC from dataclasses import dataclass -from typing import Callable, List, Optional, Tuple, Iterable +from typing import Callable, List, Optional, Tuple -from sqlalchemy import not_ from sqlalchemy.orm import Session from wrolpi.cmd import SINGLE_FILE_BIN, CHROMIUM, FIREFOX, DENO_BIN, SINGLE_FILE_DENO_SCRIPT @@ -17,6 +15,7 @@ from wrolpi.common import logger, register_modeler, register_refresh_cleanup, li from wrolpi.db import get_db_session from wrolpi.downloader import Downloader, Download, DownloadContext, DownloadResult from wrolpi.errors import UnrecoverableDownloadError +from wrolpi.files.modeler import run_modeler_loop, skip_clause, SkipModeler from wrolpi.files.models import FileGroup from wrolpi.vars import PYTEST, DOCKERIZED, DOWNLOAD_USER_AGENT, PROJECT_DIR from . import lib @@ -358,85 +357,46 @@ def model_archive(session: Session, file_group: FileGroup) -> Archive: raise InvalidArchive(f'Failed to model Archive {file_group}') from e +ARCHIVE_PROCESSING_LIMIT = 20 + + +def _select_archive_ids(session, skip_ids: set) -> list[int]: + query = session.query(FileGroup.id).filter( + # Archive shares text/html with generic HTML. Unlike video/doc/zim, do NOT + # select Archive.id IS NULL: that would re-probe every leftover HTML file + # every refresh. Non-SingleFile HTML stays indexed=0 so apply_indexers + # can claim it. + FileGroup.indexed != True, + FileGroup.mimetype == 'text/html', + ) + query = skip_clause(query, FileGroup.id, skip_ids) + return [row[0] for row in query.limit(ARCHIVE_PROCESSING_LIMIT).all()] + + +def _apply_archive(session: Session, file_group: FileGroup, prepared): + with slow_logger(1, f'Modeling archive took %(elapsed)s seconds: {file_group}', logger__=logger): + archive = session.query(Archive).filter_by(file_group_id=file_group.id).one_or_none() + if archive: + archive.validate() + return + try: + model_archive(session, file_group) + except InvalidArchive as e: + raise SkipModeler from e + + @register_modeler async def archive_modeler(progress_callback: Callable[[int], None] = None): """Searches DB for FileGroups that contain an HTML file. If the HTML file is a SingleFile, we model it as an Archive.""" - invalid_archives = set() - total_processed = 0 - - while True: - with get_db_session(commit=True) as session: - results = session.query(FileGroup, Archive) \ - .filter( - # Archive shares text/html with generic HTML. Unlike video/doc/zim, do NOT - # select Archive.id IS NULL: that would re-probe every leftover HTML file - # every refresh. Non-SingleFile HTML stays indexed=0 so apply_indexers - # can claim it. - FileGroup.indexed != True, - FileGroup.mimetype == 'text/html', - ).filter(not_(FileGroup.id.in_(list(invalid_archives)))) \ - .outerjoin(Archive, Archive.file_group_id == FileGroup.id) \ - .limit(20) - results: Iterable[Tuple[FileGroup, Archive]] - results_list = list(results) - - processed = 0 - for processed, (file_group, archive) in enumerate(results_list): - with slow_logger(1, f'Modeling archive took %(elapsed)s seconds: {file_group}', - logger__=logger): - if archive: - try: - archive_id = archive.id - archive.validate() - # Successfully validated, mark as indexed - file_group.indexed = True - except Exception as e: - logger.error(f'Unable to validate Archive {archive_id}') - # Don't mark as indexed - will retry later - if PYTEST: - raise - else: - try: - model_archive(session, file_group) - # Successfully modeled, mark as indexed - file_group.indexed = True - except InvalidArchive: - # It was not a real Archive. Many HTML files will not be an Archive. - file_group.indexed = False - invalid_archives.add(file_group.id) - except Exception as e: - # Some other error occurred during modeling - don't mark as indexed so we can retry - logger.error(f'Failed to model Archive for FileGroup {file_group.id}: {e}') - if PYTEST: - raise - - try: - session.commit() - except Exception as e: - # One poisoned row must not kill the modeler: skip this batch for the rest of this - # run and keep modeling. The batch stays unindexed, so the next refresh retries it. - session.rollback() - logger.error(f'archive_modeler failed to commit batch, skipping {len(results_list)} FileGroups', - exc_info=e) - invalid_archives.update(fg.id for fg, _ in results_list) - if PYTEST: - raise - - # Report batch progress - batch_count = len(results_list) - total_processed += batch_count - if progress_callback and batch_count > 0: - progress_callback(total_processed) - - if processed < 19: - # Did not reach limit (enumerate is 0-indexed, so 19 = 20 items), do not query again. - if __debug__ and logger.isEnabledFor(TRACE_LEVEL): - logger.trace(f'archive_modeler: DONE (processed {processed + 1} files)') - break - - # Sleep to catch cancel. - await asyncio.sleep(0) + await run_modeler_loop( + name='archive_modeler', + batch_size=ARCHIVE_PROCESSING_LIMIT, + select_ids=_select_archive_ids, + apply=_apply_archive, + txn='batch', + progress_callback=progress_callback, + ) @register_refresh_cleanup diff --git a/modules/docs/__init__.py b/modules/docs/__init__.py index e7b15bce..b7454156 100644 --- a/modules/docs/__init__.py +++ b/modules/docs/__init__.py @@ -1,17 +1,16 @@ -import asyncio import logging -from typing import Callable, List +from typing import Callable 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.modeler import run_modeler_loop, skip_clause from wrolpi.files.models import FileGroup -from wrolpi.vars import PYTEST, FILE_MAX_TEXT_SIZE +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 @@ -24,6 +23,33 @@ DOC_PROCESSING_LIMIT = 10 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. @@ -36,78 +62,15 @@ async def doc_modeler(progress_callback: Callable[[int], None] = None): 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 != True, - ), - ) - 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) + 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: diff --git a/modules/docs/test/test_doc_modeler.py b/modules/docs/test/test_doc_modeler.py index 315c91df..8fcda4fd 100644 --- a/modules/docs/test/test_doc_modeler.py +++ b/modules/docs/test/test_doc_modeler.py @@ -167,7 +167,7 @@ async def test_doc_modeler_terminates_on_persistent_failure(async_client, test_s # PYTEST re-raises inside the loop's per-item handler; disable that so we exercise the failure # path and the loop's termination guard, not the re-raise. with patch('modules.docs._model_doc', side_effect=RuntimeError('boom')), \ - patch('modules.docs.PYTEST', False): + patch('wrolpi.files.modeler.PYTEST', False): await asyncio.wait_for(doc_modeler(), timeout=10) # No Doc rows were created (every attempt failed), but the loop still terminated. diff --git a/modules/videos/__init__.py b/modules/videos/__init__.py index a2d1f27a..9341cef6 100644 --- a/modules/videos/__init__.py +++ b/modules/videos/__init__.py @@ -10,6 +10,7 @@ from modules.videos.common import get_or_create_ffprobe_json from modules.videos.models import Video, AUDIO_PLAYLIST_MIMETYPES from wrolpi.common import logger, limit_concurrent, register_modeler, register_refresh_cleanup from wrolpi.db import get_db_curs, get_db_session +from wrolpi.files.modeler import run_modeler_loop, skip_clause from wrolpi.files.models import FileGroup from wrolpi.vars import PYTEST from .downloader import video_downloader # Import downloaders so they are registered. @@ -43,76 +44,55 @@ def _model_video(session: Session, file_group: FileGroup, video: Video | None, p return video +def _select_video_ids(session, skip_ids: set) -> list[int]: + query = session.query(FileGroup.id) \ + .outerjoin(Video, Video.file_group_id == FileGroup.id) \ + .filter( + or_(FileGroup.mimetype.like('video/%'), FileGroup.mimetype.like('audio/%')), + FileGroup.mimetype.notin_(AUDIO_PLAYLIST_MIMETYPES), + # Exclusive-mimetype modeler: claim unmodeled rows even if apply_indexers + # already set indexed=True, and re-model when files changed (indexed != True). + or_(Video.id.is_(None), FileGroup.indexed != True), + ) + query = skip_clause(query, FileGroup.id, skip_ids) + return [row[0] for row in query.limit(VIDEO_PROCESSING_LIMIT).all()] + + +async def _prepare_videos(ids: list[int]) -> dict: + """ffprobe with no write transaction open. A subprocess per video can run for seconds.""" + with get_db_session() as session: + rows = session.query(FileGroup.id, FileGroup.primary_path).filter(FileGroup.id.in_(ids)).all() + probed = dict() + for file_group_id, primary_path in rows: + try: + probed[file_group_id] = await get_or_create_ffprobe_json(pathlib.Path(str(primary_path))) + except Exception as e: + if PYTEST: + raise + logger.error(f'Unable to ffprobe Video: {primary_path}', exc_info=e) + await asyncio.sleep(0) + return probed + + +def _apply_video(session: Session, file_group: FileGroup, prepared: dict): + video = session.query(Video).filter_by(file_group_id=file_group.id).one_or_none() + _model_video(session, file_group, video, prepared or {}) + + @register_modeler async def video_modeler(progress_callback: Callable[[int], None] = None): - total_processed = 0 - # Ids that failed to model this run. A failure leaves no Video row, so the `Video.id IS - # NULL` gate below would re-select it forever; failures are retried on the next refresh. - failed_ids: set = set() - while True: - # Read the batch; nothing is claimed yet, so the write lock stays free while ffprobe runs. - with get_db_session() as session: - query = session.query(FileGroup.id, FileGroup.primary_path) \ - .outerjoin(Video, Video.file_group_id == FileGroup.id) \ - .filter( - or_(FileGroup.mimetype.like('video/%'), FileGroup.mimetype.like('audio/%')), - FileGroup.mimetype.notin_(AUDIO_PLAYLIST_MIMETYPES), - # Exclusive-mimetype modeler: claim unmodeled rows even if apply_indexers - # already set indexed=True, and re-model when files changed (indexed != True). - or_(Video.id.is_(None), FileGroup.indexed != True), - ) - if failed_ids: - query = query.filter(FileGroup.id.notin_(failed_ids)) - batch: List[Tuple[int, pathlib.Path]] = list(query.limit(VIDEO_PROCESSING_LIMIT).all()) - - if not batch: - break - - # ffprobe each file with no transaction open. This is a subprocess per video and can run - # for seconds; inside the write transaction below it would hold the write lock for the whole - # batch, and every other writer on the box would wait out `busy_timeout` (30s). - probed = dict() - for file_group_id, primary_path in batch: - try: - probed[file_group_id] = await get_or_create_ffprobe_json(pathlib.Path(str(primary_path))) - except Exception as e: - if PYTEST: - raise - logger.error(f'Unable to ffprobe Video: {primary_path}', exc_info=e) - # Sleep to catch cancel. - await asyncio.sleep(0) - - with get_db_session(commit=True) as session: - file_groups: List[Tuple[FileGroup, Video]] = list( - session.query(FileGroup, Video) - .filter(FileGroup.id.in_([i for i, _ in batch])) - .outerjoin(Video, Video.file_group_id == FileGroup.id)) - - for file_group, video in file_groups: - try: - _model_video(session, file_group, video, probed) - except Exception as e: - # Before the re-raise, so no failure path can re-select this id forever. - failed_ids.add(file_group.id) - if PYTEST: - raise - logger.error(f'Unable to model Video: {file_group.primary_path}', exc_info=e) - - file_group.indexed = True - - # Report batch progress - total_processed += len(batch) - if progress_callback: - progress_callback(total_processed) - - logger.debug(f'Modeled {len(batch)} videos') - - if len(batch) < VIDEO_PROCESSING_LIMIT: - # Did not reach limit, do not query again. - break - - # Sleep to catch cancel. - await asyncio.sleep(0) + await run_modeler_loop( + name='video_modeler', + batch_size=VIDEO_PROCESSING_LIMIT, + select_ids=_select_video_ids, + prepare=_prepare_videos, + apply=_apply_video, + txn='batch', + # A failure leaves no Video row, so Video.id IS NULL would re-select it. Mark indexed + # anyway; the skip set covers this run, the next refresh retries via Video.id IS NULL. + mark_indexed_on_failure=True, + progress_callback=progress_callback, + ) # Rows written per transaction when claiming Videos for their Channels. Small enough that the diff --git a/modules/videos/test/test_audio_modeling.py b/modules/videos/test/test_audio_modeling.py index aaa2a66c..9ec3aae1 100644 --- a/modules/videos/test/test_audio_modeling.py +++ b/modules/videos/test/test_audio_modeling.py @@ -163,7 +163,7 @@ async def test_video_modeler_terminates_on_persistent_failure(async_client, test # loop's per-item handler; disable that to exercise the failure path and the termination # guard, not the re-raise. with patch('modules.videos._model_video', side_effect=RuntimeError('boom')), \ - patch('modules.videos.PYTEST', False): + patch('wrolpi.files.modeler.PYTEST', False): await asyncio.wait_for(video_modeler(), timeout=10) assert test_session.query(Video).count() == 0, 'every attempt failed, yet the loop terminated' diff --git a/modules/zim/lib.py b/modules/zim/lib.py index 84b5297a..95761df2 100644 --- a/modules/zim/lib.py +++ b/modules/zim/lib.py @@ -1,4 +1,3 @@ -import asyncio import functools import json import pathlib @@ -26,6 +25,7 @@ from wrolpi.events import Events from wrolpi.db import get_db_session, get_db_curs from wrolpi.downloader import DownloadFrequency from wrolpi.files.lib import split_file_name_words +from wrolpi.files.modeler import run_modeler_loop, skip_clause from wrolpi.files.worker import file_worker from wrolpi.files.models import FileGroup from wrolpi.vars import PYTEST, DOCKERIZED @@ -73,58 +73,43 @@ def model_zim(file_group: FileGroup, session: Session) -> Zim: return zim +ZIM_PROCESSING_LIMIT = 10 + + +def _select_zim_ids(session, skip_ids: set) -> list[int]: + query = session.query(FileGroup.id) \ + .outerjoin(Zim, Zim.file_group_id == FileGroup.id) \ + .filter(FileGroup.primary_path.ilike('%.zim')) \ + .filter( + # Exclusive-mimetype modeler: claim unmodeled rows even if apply_indexers + # already set indexed=True, and re-model when files changed. + or_(Zim.id.is_(None), FileGroup.indexed != True), + ) + query = skip_clause(query, FileGroup.id, skip_ids) + return [row[0] for row in query.limit(ZIM_PROCESSING_LIMIT).all()] + + +def _apply_zim(session: Session, file_group: FileGroup, prepared): + zim = session.query(Zim).filter_by(file_group_id=file_group.id).one_or_none() + if not zim: + zim = model_zim(file_group, session) + zim.flush() + else: + file_group.title = file_group.primary_path.name + file_group.a_text = split_file_name_words(file_group.primary_path.name) + file_group.model = 'zim' + + @register_modeler async def zim_modeler(progress_callback: Callable[[int], None] = None): - total_processed = 0 - failed_ids: set = set() - while True: - with get_db_session(commit=True) as session: - query = session.query(FileGroup, Zim) \ - .outerjoin(Zim, Zim.file_group_id == FileGroup.id) \ - .filter(FileGroup.primary_path.ilike('%.zim')) \ - .filter( - # Exclusive-mimetype modeler: claim unmodeled rows even if apply_indexers - # already set indexed=True, and re-model when files changed. - or_(Zim.id.is_(None), FileGroup.indexed != True), - ) - if failed_ids: - query = query.filter(FileGroup.id.notin_(failed_ids)) - file_groups: List[Tuple[FileGroup, Zim]] = list(query.limit(10)) - - processed = 0 - for file_group, zim in file_groups: - processed += 1 - - zim_id = None - try: - if not zim: - zim = model_zim(file_group, session) - zim.flush() - else: - zim_id = zim.id - file_group.title = file_group.primary_path.name - file_group.a_text = split_file_name_words(file_group.primary_path.name) - file_group.indexed = True - file_group.model = 'zim' - except Exception as e: - failed_ids.add(file_group.id) - if PYTEST: - raise - logger.error(f'Unable to model Zim {zim_id=} {file_group.primary_path=}', exc_info=e) - - logger.debug(f'Modeled {processed} zim files') - - # Report batch progress - total_processed += len(file_groups) - if progress_callback and len(file_groups) > 0: - progress_callback(total_processed) - - if processed < 10: - # Did not reach limit, do not query again. - break - - # Sleep to catch cancel. - await asyncio.sleep(0) + await run_modeler_loop( + name='zim_modeler', + batch_size=ZIM_PROCESSING_LIMIT, + select_ids=_select_zim_ids, + apply=_apply_zim, + txn='batch', + progress_callback=progress_callback, + ) def get_all_entries_tags(): diff --git a/modules/zim/test/test_zim_modeler_loop.py b/modules/zim/test/test_zim_modeler_loop.py index 5a54de38..0c3baad8 100644 --- a/modules/zim/test/test_zim_modeler_loop.py +++ b/modules/zim/test/test_zim_modeler_loop.py @@ -8,14 +8,11 @@ import shutil import pytest -from modules.zim.lib import zim_modeler +from modules.zim.lib import zim_modeler, ZIM_PROCESSING_LIMIT from modules.zim.models import Zim from wrolpi.files.models import FileGroup from wrolpi.vars import PROJECT_DIR -# The zim_modeler uses a hardcoded limit of 10 -ZIM_PROCESSING_LIMIT = 10 - @pytest.mark.asyncio async def test_zim_modeler_processes_more_than_batch_limit(async_client, test_session, test_directory): diff --git a/wrolpi/files/modeler.py b/wrolpi/files/modeler.py new file mode 100644 index 00000000..ac9b1cc9 --- /dev/null +++ b/wrolpi/files/modeler.py @@ -0,0 +1,171 @@ +"""Shared FileGroup modeler batch loop. + +Each module still owns select / prepare / apply. This loop owns: + +- paging by batch_size +- a per-run skip set so failures cannot re-select forever +- cancel yields +- progress_callback +- two transaction styles: one write per batch (video, zim, archive) or one write + per item with lock retries (docs) + +Video still probes outside the write lock via ``prepare``. Archive uses +``SkipModeler`` for HTML that is not a SingleFile so apply_indexers can claim it. +""" +from __future__ import annotations + +import asyncio +from typing import Awaitable, Callable, Iterable + +from sqlalchemy.exc import InvalidRequestError, OperationalError + +from wrolpi.common import logger +from wrolpi.db import get_db_session +from wrolpi.files.models import FileGroup +from wrolpi.vars import PYTEST + +logger = logger.getChild(__name__) + +SelectIds = Callable # (session, skip_ids: set) -> list[int] +Prepare = Callable # async (ids: list[int]) -> dict +Apply = Callable # (session, file_group, prepared) -> None + + +class SkipModeler(Exception): + """This FileGroup is not this modeler's type. Leave indexed=False so leftovers can claim it.""" + + +async def run_modeler_loop( + *, + name: str, + batch_size: int, + select_ids: SelectIds, + apply: Apply, + prepare: Prepare | None = None, + txn: str = 'batch', + lock_retries: int = 1, + mark_indexed_on_failure: bool = False, + progress_callback: Callable[[int], None] | None = None, +): + """Page FileGroups and apply a modeler until a short batch or empty select. + + ``txn='batch'``: one write transaction per page (video, zim, archive). + ``txn='per_item'``: one write transaction per FileGroup (docs). + """ + if txn not in ('batch', 'per_item'): + raise ValueError(f'Unknown modeler txn {txn!r}') + + skip_ids: set = set() + total_processed = 0 + while True: + with get_db_session() as session: + ids = list(select_ids(session, skip_ids)) + if not ids: + break + + prepared = {} + if prepare is not None: + prepared = await prepare(ids) or {} + + if txn == 'batch': + await _apply_batch( + name, ids, prepared, apply, skip_ids, mark_indexed_on_failure, + ) + else: + await _apply_per_item( + name, ids, prepared, apply, skip_ids, lock_retries, + ) + + total_processed += len(ids) + if progress_callback: + progress_callback(total_processed) + + logger.debug(f'{name}: modeled {len(ids)} FileGroups') + + if len(ids) < batch_size: + break + await asyncio.sleep(0) + + +async def _apply_batch( + name: str, + ids: list[int], + prepared: dict, + apply: Apply, + skip_ids: set, + mark_indexed_on_failure: bool, +): + try: + with get_db_session(commit=True) as session: + file_groups: Iterable[FileGroup] = session.query(FileGroup).filter(FileGroup.id.in_(ids)) + for file_group in file_groups: + try: + apply(session, file_group, prepared) + file_group.indexed = True + except SkipModeler: + file_group.indexed = False + skip_ids.add(file_group.id) + except InvalidRequestError: + # A previous apply already flushed and rolled the session back. + # Remaining items would each log the same error; skip the page. + skip_ids.update(ids) + break + except Exception as e: + skip_ids.add(file_group.id) + if mark_indexed_on_failure: + file_group.indexed = True + logger.error(f'{name}: unable to model file_group_id={file_group.id}', exc_info=e) + if PYTEST: + raise + except Exception as e: + # One poisoned commit must not kill the modeler. Skip this page for the rest of + # the run; the next refresh retries it. + skip_ids.update(ids) + if PYTEST: + raise + logger.error(f'{name}: failed to commit batch of {len(ids)} FileGroups', exc_info=e) + + +async def _apply_per_item( + name: str, + ids: list[int], + prepared: dict, + apply: Apply, + skip_ids: set, + lock_retries: int, +): + for fg_id in ids: + for attempt in range(lock_retries): + try: + with get_db_session(commit=True) as session: + file_group = session.query(FileGroup).get(fg_id) + if file_group is None: + break + apply(session, file_group, prepared) + file_group.indexed = True + break + except SkipModeler: + skip_ids.add(fg_id) + break + except OperationalError as e: + if 'locked' in str(e).lower() and attempt < lock_retries - 1: + await asyncio.sleep(0.2 * (attempt + 1)) + continue + logger.error(f'{name}: failed to model file_group_id={fg_id}: database error', exc_info=e) + skip_ids.add(fg_id) + break + except Exception as e: + # Do not touch the ORM object: after a failed flush the session is rolled back. + logger.error(f'{name}: failed to model file_group_id={fg_id}', exc_info=e) + skip_ids.add(fg_id) + if PYTEST: + raise + break + await asyncio.sleep(0) + + +def skip_clause(query, column, skip_ids: set): + """Exclude ids already failed/rejected this run.""" + if skip_ids: + query = query.filter(column.notin_(skip_ids)) + return query diff --git a/wrolpi/files/test/test_modeler.py b/wrolpi/files/test/test_modeler.py new file mode 100644 index 00000000..c91e110b --- /dev/null +++ b/wrolpi/files/test/test_modeler.py @@ -0,0 +1,108 @@ +"""Tests for the shared FileGroup modeler loop.""" +from unittest.mock import patch + +import pytest + +from wrolpi.files.modeler import run_modeler_loop, SkipModeler +from wrolpi.files.models import FileGroup + + +def _select_unindexed(session, skip_ids): + query = session.query(FileGroup.id).filter(FileGroup.indexed != True) + if skip_ids: + query = query.filter(FileGroup.id.notin_(skip_ids)) + return [row[0] for row in query.limit(2).all()] + + +@pytest.mark.asyncio +async def test_run_modeler_loop_pages_and_skips_failures(test_session, test_directory, make_files_structure): + """The loop pages by batch_size and does not re-select a FileGroup that failed this run.""" + paths = make_files_structure([f'd{i}.txt' for i in range(5)]) + fgs = [FileGroup.from_paths(test_session, p) for p in paths] + test_session.commit() + ids = [fg.id for fg in fgs] + boom_id = ids[1] + applied = [] + + def select_ids(session, skip_ids): + return _select_unindexed(session, skip_ids) + + def apply(session, file_group, prepared): + if file_group.id == boom_id: + raise RuntimeError('boom') + applied.append(file_group.id) + file_group.model = 'file' + + # PYTEST re-raises apply failures; disable that so we exercise skip-and-continue. + with patch('wrolpi.files.modeler.PYTEST', False): + await run_modeler_loop( + name='test_modeler', + batch_size=2, + select_ids=select_ids, + apply=apply, + txn='batch', + ) + + assert boom_id not in applied + assert set(applied) == set(ids) - {boom_id} + test_session.expire_all() + by_id = {fg.id: fg for fg in test_session.query(FileGroup)} + assert by_id[boom_id].indexed is not True + for i in applied: + assert by_id[i].indexed is True + + +@pytest.mark.asyncio +async def test_run_modeler_loop_skip_modeler_leaves_unindexed(test_session, test_directory, make_files_structure): + """SkipModeler leaves indexed=False so a leftover indexer can claim the FileGroup.""" + path, = make_files_structure(['page.html']) + fg = FileGroup.from_paths(test_session, path) + test_session.commit() + + def select_ids(session, skip_ids): + if skip_ids: + return [] + return [fg.id] + + def apply(session, file_group, prepared): + raise SkipModeler + + await run_modeler_loop( + name='test_skip', + batch_size=10, + select_ids=select_ids, + apply=apply, + txn='batch', + ) + + test_session.expire_all() + assert test_session.query(FileGroup).one().indexed is not True + + +@pytest.mark.asyncio +async def test_run_modeler_loop_per_item_lock_retry(test_session, test_directory, make_files_structure): + """per_item txn models each FileGroup in its own write transaction.""" + paths = make_files_structure(['a.pdf', 'b.pdf']) + fgs = [FileGroup.from_paths(test_session, p) for p in paths] + test_session.commit() + ids = [fg.id for fg in fgs] + applied = [] + + def select_ids(session, skip_ids): + return [i for i in ids if i not in skip_ids] + + def apply(session, file_group, prepared): + applied.append(file_group.id) + + await run_modeler_loop( + name='test_per_item', + batch_size=10, + select_ids=select_ids, + apply=apply, + txn='per_item', + lock_retries=3, + ) + + assert applied == ids + test_session.expire_all() + assert all(fg.indexed is True for fg in test_session.query(FileGroup))