mirror of
https://github.com/lrnselfreliance/wrolpi.git
synced 2026-10-26 10:51:36 -04:00
213 lines
8 KiB
Python
213 lines
8 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
|
|
while True:
|
|
with get_db_session(commit=True) as session:
|
|
file_groups = session.query(FileGroup, Doc).filter(
|
|
or_(
|
|
*[FileGroup.mimetype == mt for mt in DOC_MIMETYPES],
|
|
*[FileGroup.primary_path.like(f'%{suffix}') for suffix in COMIC_BOOK_SUFFIXES],
|
|
),
|
|
FileGroup.indexed == False,
|
|
).outerjoin(Doc, Doc.file_group_id == FileGroup.id) \
|
|
.limit(DOC_PROCESSING_LIMIT)
|
|
file_groups: List[Tuple[FileGroup, Doc]] = list(file_groups)
|
|
|
|
processed = 0
|
|
for file_group, doc in file_groups:
|
|
processed += 1
|
|
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)
|