mirror of
https://github.com/lrnselfreliance/wrolpi.git
synced 2026-10-26 02:48:38 -04:00
The file_group.textsearch tsvector indexes all four text columns;
d_text (captions, document text, zip contents) is ~90% of the data,
making ts_rank searches take seconds on large libraries. Add a
textsearch_abc generated column (title/author/description only, ~8x
smaller) and default all searches to it, with deep=true opting into the
full column.
Database:
- Migration adds the generated column and a GIN index (rewrites the
table; a few minutes on large libraries).
Backend:
- FilesSearchRequest, VideoSearchRequest, ArchiveSearchRequest, and
DocSearchRequest accept `deep`; each search function picks the column.
- /api/search_file_estimates returns both counts in one index-bound
query (COUNT FILTER on the abc subset): `file_groups` (fast) and
`file_groups_deep`.
UI (one coherent story across the count surfaces):
- The typing dropdown always suggests the deep count (what is findable).
- The tab badge shows the active mode's count so it matches the result
list.
- The gap is reconciled by a hint offering deep search: count-aware on
the global search ("No results in fast search. N results available in
deep search." / "N more results available..."), generic on the module
pages (no deep estimate there). The hint resets pagination.
- "Search Depth" toggle in the search filter modals, persisted as
?deep=true.
Tests that relied on d_text matches now pass deep=True, with fast-mode
assertions alongside. Also skip flaky test_download_channel on
CircleCI.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
1749 lines
72 KiB
Python
1749 lines
72 KiB
Python
import base64
|
|
import gzip
|
|
import io
|
|
import json
|
|
import os
|
|
import pathlib
|
|
import re
|
|
import shlex
|
|
import tempfile
|
|
import zipfile
|
|
from dataclasses import dataclass, field
|
|
from datetime import datetime
|
|
from json import JSONDecodeError
|
|
from typing import Optional, Tuple, List, Union
|
|
from urllib.parse import urlparse
|
|
|
|
import pytz
|
|
from sqlalchemy import asc, func, BigInteger
|
|
from sqlalchemy.orm import Session
|
|
|
|
from modules.archive.models import Archive
|
|
from wrolpi import dates
|
|
from wrolpi.cmd import READABILITY_BIN, run_command
|
|
from wrolpi.collections import Collection
|
|
from wrolpi.common import get_media_directory, get_relative_to_media_directory, logger, extract_domain, \
|
|
escape_file_name, aiohttp_post, format_html_string, split_lines_by_length, get_html_soup, get_title_from_html, \
|
|
get_wrolpi_config, html_screenshot, html_file_screenshot, ConfigFile, trim_file_name
|
|
from wrolpi.dates import now, Seconds
|
|
from wrolpi.db import get_db_session, get_db_curs
|
|
from wrolpi.errors import UnknownArchive, InvalidOrderBy, InvalidDatetime
|
|
from wrolpi.events import Events
|
|
from wrolpi.files.lib import handle_file_group_search_results
|
|
from wrolpi.files.models import FileGroup
|
|
from wrolpi.switches import register_switch_handler, ActivateSwitchMethod
|
|
from wrolpi.tags import tag_append_sub_select_where
|
|
from wrolpi.vars import PYTEST, DOCKERIZED
|
|
|
|
logger = logger.getChild(__name__)
|
|
|
|
__all__ = ['DomainsConfig', 'domains_config', 'get_domains_config', 'save_domains_config', 'import_domains_config']
|
|
|
|
ARCHIVE_SERVICE = 'http://archive:8080'
|
|
|
|
|
|
@dataclass
|
|
class DomainsConfigValidator:
|
|
"""Validator for domains config file."""
|
|
version: int = 0
|
|
collections: List[dict] = field(default_factory=list)
|
|
|
|
|
|
class DomainsConfig(ConfigFile):
|
|
"""
|
|
Config file for Domain Collections.
|
|
|
|
This is a specialized config that manages domain collections
|
|
(Collections with kind='domain'). It maintains a domains.yaml file for
|
|
backward compatibility and user convenience.
|
|
|
|
Format:
|
|
collections:
|
|
- name: "example.com"
|
|
kind: "domain"
|
|
description: "Archives from example.com"
|
|
- name: "wikipedia.org"
|
|
kind: "domain"
|
|
|
|
Note: Domain collections can have optional directories for tagging support.
|
|
Items are managed dynamically when Archives are indexed.
|
|
"""
|
|
file_name = 'domains.yaml'
|
|
validator = DomainsConfigValidator
|
|
default_config = dict(
|
|
version=0,
|
|
collections=[],
|
|
)
|
|
# Use wider width to accommodate longer paths
|
|
width = 120
|
|
|
|
def __getitem__(self, item):
|
|
return self._config[item]
|
|
|
|
def __setitem__(self, key, value):
|
|
self._config[key] = value
|
|
|
|
@property
|
|
def collections(self) -> List[dict]:
|
|
"""Get list of collection configs."""
|
|
return self._config.get('collections', [])
|
|
|
|
def import_config(self, file: pathlib.Path = None, send_events=False):
|
|
"""Import domain collections from config file into database using batch operations."""
|
|
file = file or self.get_file()
|
|
file_str = str(get_relative_to_media_directory(file))
|
|
|
|
# If config file doesn't exist, mark import as successful (nothing to import)
|
|
if not file.is_file():
|
|
logger.info('No domains config file, skipping import')
|
|
self.successful_import = True
|
|
return
|
|
|
|
ConfigFile.import_config(self, file, send_events)
|
|
|
|
collections_data = self._config.get('collections', [])
|
|
|
|
# Empty collections list = never delete DB records
|
|
if not collections_data:
|
|
logger.info(f'No domain collections in config, preserving existing DB domain collections')
|
|
self.successful_import = True
|
|
return
|
|
|
|
logger.info(f'Importing {len(collections_data)} domain collections from {file_str}')
|
|
|
|
try:
|
|
with get_db_session(commit=True) as session:
|
|
# Validate and prepare data first
|
|
valid_data_list = []
|
|
for idx, collection_data in enumerate(collections_data):
|
|
try:
|
|
name = collection_data.get('name')
|
|
if not name:
|
|
logger.error(f'Domain collection at index {idx} has no name, skipping')
|
|
continue
|
|
|
|
# Validate domain name format before batching
|
|
if not Collection.is_valid_domain_name(name):
|
|
logger.error(
|
|
f"Domain collection at index {idx} has invalid name '{name}' "
|
|
f"(must contain at least one '.' and not start/end with '.'), skipping"
|
|
)
|
|
continue
|
|
|
|
# Ensure kind is 'domain'
|
|
collection_data = collection_data.copy()
|
|
collection_data['kind'] = 'domain'
|
|
|
|
# Note: Collections can now be tagged even without a directory
|
|
# Tags enable UI search/filtering for directory-less collections
|
|
|
|
valid_data_list.append(collection_data)
|
|
except Exception as e:
|
|
logger.error(f'Failed to validate domain collection at index {idx}', exc_info=e)
|
|
continue
|
|
|
|
# Batch create/update collections
|
|
imported_domains = set()
|
|
if valid_data_list:
|
|
collections = Collection.batch_from_config(session, valid_data_list)
|
|
imported_domains = {c.name for c in collections}
|
|
|
|
# Delete domain collections that are no longer in config
|
|
all_domain_collections = session.query(Collection).filter_by(kind='domain').all()
|
|
for collection in all_domain_collections:
|
|
if collection.name not in imported_domains:
|
|
logger.info(f'Deleting domain collection {repr(collection.name)} (no longer in config)')
|
|
session.delete(collection)
|
|
|
|
logger.info(f'Successfully imported {len(imported_domains)} domain collections from {file_str}')
|
|
self.successful_import = True
|
|
|
|
except Exception as e:
|
|
self.successful_import = False
|
|
message = f'Failed to import {file_str} config!'
|
|
logger.error(message, exc_info=e)
|
|
if send_events:
|
|
Events.send_config_import_failed(message)
|
|
raise
|
|
|
|
def dump_config(self, file: pathlib.Path = None, send_events=False, overwrite=False):
|
|
"""Dump all domain collections from database to config file."""
|
|
logger.info('Dumping domain collections to config')
|
|
|
|
with get_db_session() as session:
|
|
# Get only domain collections, ordered by name
|
|
collections = session.query(Collection).filter_by(kind='domain').order_by(Collection.name).all()
|
|
|
|
# Use to_config to export each collection
|
|
collections_data = [collection.to_config() for collection in collections]
|
|
|
|
if not collections_data and self._config.get('collections'):
|
|
logger.warning('Refusing to save empty domains config because domains exist in the config file!')
|
|
return
|
|
|
|
self._config['collections'] = collections_data
|
|
|
|
logger.info(f'Dumping {len(collections_data)} domain collections to config')
|
|
self.save(file, send_events, overwrite)
|
|
|
|
def preview_backup_import(self, backup_date: str, mode: str) -> dict:
|
|
backup_file = self._get_backup_file(backup_date)
|
|
backup_data = self.read_config_file(backup_file)
|
|
current_data = self.read_config_file() if self.get_file().is_file() else dict(collections=[], version=0)
|
|
|
|
backup_collections = backup_data.get('collections', [])
|
|
current_collections = current_data.get('collections', [])
|
|
|
|
# Identity key for domains is the directory (if present) or name
|
|
def get_key(c):
|
|
return c.get('directory') or c.get('name', '')
|
|
|
|
current_keys = {get_key(c) for c in current_collections}
|
|
backup_keys = {get_key(c) for c in backup_collections}
|
|
|
|
add = []
|
|
remove = []
|
|
unchanged = 0
|
|
|
|
for c in backup_collections:
|
|
key = get_key(c)
|
|
if key not in current_keys:
|
|
add.append(dict(type='domain', name=c.get('name', ''), directory=c.get('directory', '')))
|
|
else:
|
|
unchanged += 1
|
|
if mode == 'overwrite':
|
|
current_key_to_coll = {get_key(c): c for c in current_collections}
|
|
for key in sorted(current_keys - backup_keys):
|
|
c = current_key_to_coll[key]
|
|
remove.append(dict(type='domain', name=c.get('name', ''), directory=c.get('directory', '')))
|
|
|
|
return dict(mode=mode, add=add, remove=remove, unchanged=unchanged)
|
|
|
|
def import_backup(self, backup_date: str, mode: str, send_events: bool = False):
|
|
self._preserve_current_config()
|
|
import shutil
|
|
backup_file = self._get_backup_file(backup_date)
|
|
config_file = self.get_file()
|
|
|
|
if mode == 'overwrite':
|
|
config_file.parent.mkdir(parents=True, exist_ok=True)
|
|
shutil.copy2(backup_file, config_file)
|
|
elif mode == 'merge':
|
|
backup_data = self.read_config_file(backup_file)
|
|
current_data = self.read_config_file() if config_file.is_file() else dict(collections=[], version=0)
|
|
|
|
def get_key(c):
|
|
return c.get('directory') or c.get('name', '')
|
|
|
|
current_keys = {get_key(c) for c in current_data.get('collections', [])}
|
|
merged_collections = list(current_data.get('collections', []))
|
|
for c in backup_data.get('collections', []):
|
|
if get_key(c) not in current_keys:
|
|
merged_collections.append(c)
|
|
|
|
merged_data = dict(
|
|
collections=merged_collections,
|
|
version=current_data.get('version', 0) + 1,
|
|
)
|
|
config_file.parent.mkdir(parents=True, exist_ok=True)
|
|
self.write_config_data(merged_data, config_file)
|
|
|
|
self.import_config(send_events=send_events)
|
|
|
|
|
|
# Global instance
|
|
domains_config = DomainsConfig()
|
|
|
|
|
|
def get_domains_config() -> DomainsConfig:
|
|
"""Get the global domains config instance."""
|
|
return domains_config
|
|
|
|
|
|
# Switch handler for saving domains config
|
|
@register_switch_handler('save_domains_config')
|
|
def save_domains_config():
|
|
"""Save the domains config when the switch is activated."""
|
|
domains_config.background_dump.activate_switch()
|
|
|
|
|
|
# Explicit type for activate_switch helper
|
|
save_domains_config: ActivateSwitchMethod
|
|
|
|
|
|
def import_domains_config():
|
|
"""Import domain collections from config file into database."""
|
|
logger.info('Importing domains config')
|
|
domains_config.import_config()
|
|
|
|
# Link downloads to domain collections after import
|
|
with get_db_session() as session:
|
|
link_domain_and_downloads(session)
|
|
|
|
logger.info('Importing domains config completed')
|
|
|
|
|
|
@dataclass
|
|
class ArchiveDownloaderConfigValidator:
|
|
"""Validator for archives_downloader.yaml config."""
|
|
file_name_format: str
|
|
version: int = 0
|
|
# Browser settings (only used on native deployments, not Docker)
|
|
browser_executable: str = None # None means auto-detect, or absolute path
|
|
browser_args: str = '["--no-sandbox"]' # JSON array of browser arguments
|
|
user_agent: str = None # None means use system default
|
|
|
|
def __post_init__(self):
|
|
# Validate file_name_format ends with .%(ext)s (like video format)
|
|
if not self.file_name_format.endswith('.%(ext)s'):
|
|
raise ValueError('file_name_format must end with .%(ext)s')
|
|
|
|
# Validate browser_args is valid JSON array
|
|
if self.browser_args:
|
|
import json
|
|
try:
|
|
args = json.loads(self.browser_args)
|
|
if not isinstance(args, list):
|
|
raise ValueError('browser_args must be a JSON array')
|
|
except json.JSONDecodeError as e:
|
|
raise ValueError(f'browser_args must be valid JSON: {e}')
|
|
|
|
|
|
class ArchiveDownloaderConfig(ConfigFile):
|
|
"""
|
|
Config file for Archive Downloader settings.
|
|
|
|
This config mirrors videos_downloader.yaml and holds settings for how
|
|
archives are downloaded and named.
|
|
|
|
Format:
|
|
version: 0
|
|
file_name_format: '%(download_date)s_%(title)s.%(ext)s'
|
|
browser_executable: null # null for auto-detect, or absolute path
|
|
browser_args: '["--no-sandbox"]' # JSON array of browser arguments
|
|
user_agent: null # null for system default
|
|
|
|
Variables available in file_name_format:
|
|
- %(title)s - Page title (extracted from HTML)
|
|
- %(download_datetime)s - Full download datetime (YYYY-MM-DD-HH-MM-SS)
|
|
- %(download_date)s - Download date only (YYYY-MM-DD)
|
|
- %(download_year)s - Download year
|
|
- %(download_month)s - Download month (zero-padded)
|
|
- %(download_day)s - Download day (zero-padded)
|
|
- %(domain)s - Domain name
|
|
- %(ext)s - File extension (required, must be at end)
|
|
|
|
Subdirectories can be included in the format, e.g.:
|
|
'%(download_year)s/%(download_datetime)s_%(title)s.%(ext)s'
|
|
|
|
Browser settings (only used on native deployments, not Docker):
|
|
- browser_executable: Path to browser or null for auto-detect
|
|
- browser_args: JSON array of arguments passed to browser
|
|
- user_agent: Custom user agent string or null for default
|
|
"""
|
|
file_name = 'archives_downloader.yaml'
|
|
validator = ArchiveDownloaderConfigValidator
|
|
default_config = dict(
|
|
version=0,
|
|
file_name_format='%(download_datetime)s_%(title)s.%(ext)s',
|
|
browser_executable=None,
|
|
browser_args='["--no-sandbox"]',
|
|
user_agent=None,
|
|
)
|
|
|
|
@property
|
|
def file_name_format(self) -> str:
|
|
return self._config['file_name_format']
|
|
|
|
@property
|
|
def browser_executable(self) -> Optional[str]:
|
|
"""Get the configured browser executable path, or None for auto-detect."""
|
|
return self._config.get('browser_executable')
|
|
|
|
@property
|
|
def browser_args(self) -> str:
|
|
"""Get the browser arguments as a JSON string."""
|
|
return self._config.get('browser_args', '["--no-sandbox"]')
|
|
|
|
@property
|
|
def user_agent(self) -> Optional[str]:
|
|
"""Get the custom user agent, or None for system default."""
|
|
return self._config.get('user_agent')
|
|
|
|
def import_config(self, file: pathlib.Path = None, send_events=False):
|
|
super().import_config(file, send_events)
|
|
self.successful_import = True
|
|
|
|
|
|
# Global instance for archive downloader config
|
|
ARCHIVE_DOWNLOADER_CONFIG: ArchiveDownloaderConfig = ArchiveDownloaderConfig()
|
|
|
|
|
|
def get_archive_downloader_config() -> ArchiveDownloaderConfig:
|
|
"""Get the global archive downloader config instance."""
|
|
return ARCHIVE_DOWNLOADER_CONFIG
|
|
|
|
|
|
def format_archive_filename(
|
|
title: str,
|
|
domain: str = None,
|
|
download_date: datetime = None,
|
|
) -> str:
|
|
"""Format archive filename using config template.
|
|
|
|
This function formats the filename (and optional subdirectory path) for an archive
|
|
using the file_name_format from archives_downloader.yaml.
|
|
|
|
Args:
|
|
title: Page title (extracted from HTML)
|
|
domain: Domain name of the archived URL
|
|
download_date: Date of download (defaults to now)
|
|
|
|
Returns:
|
|
Formatted filename/path (with .html extension included)
|
|
|
|
Example:
|
|
format_archive_filename("My Article", "example.com")
|
|
# With default format: "2025-12-22_My Article.html"
|
|
# With "%(download_year)s/%(title)s.%(ext)s": "2025/My Article.html"
|
|
"""
|
|
config = get_archive_downloader_config()
|
|
template = config.file_name_format
|
|
|
|
download_date = download_date or now()
|
|
|
|
variables = dict(
|
|
title=escape_file_name(title) if title else 'untitled',
|
|
domain=domain or '',
|
|
download_datetime=archive_strftime(download_date),
|
|
download_date=download_date.strftime('%Y-%m-%d'),
|
|
download_year=str(download_date.year),
|
|
download_month=f'{download_date.month:02d}',
|
|
download_day=f'{download_date.day:02d}',
|
|
ext='html',
|
|
)
|
|
|
|
try:
|
|
result = template % variables
|
|
# Trim the filename portion if it's too long (may include subdirectory paths)
|
|
if '/' in result:
|
|
# Has subdirectories - only trim the filename part
|
|
parts = result.rsplit('/', 1)
|
|
subdir = parts[0]
|
|
filename = trim_file_name(parts[1])
|
|
return f'{subdir}/{filename}'
|
|
else:
|
|
return trim_file_name(result)
|
|
except KeyError as e:
|
|
logger.error(f'Invalid variable in archive file_name_format: {e}')
|
|
# Fallback to default format
|
|
return trim_file_name(
|
|
f'{archive_strftime(download_date)}_{escape_file_name(title) if title else "untitled"}.html')
|
|
|
|
|
|
def preview_archive_filename(file_name_format: str) -> str:
|
|
"""Preview the archive filename using sample data.
|
|
|
|
Args:
|
|
file_name_format: The format template to preview
|
|
|
|
Returns:
|
|
A preview string showing what the filename would look like
|
|
|
|
Raises:
|
|
RuntimeError: If the format is invalid or missing required variables
|
|
"""
|
|
if not file_name_format.endswith('.%(ext)s'):
|
|
raise RuntimeError('file_name_format must end with .%(ext)s')
|
|
|
|
sample_date = now()
|
|
variables = dict(
|
|
title='Example Page Title',
|
|
domain='example.com',
|
|
download_datetime=archive_strftime(sample_date),
|
|
download_date=sample_date.strftime('%Y-%m-%d'),
|
|
download_year=str(sample_date.year),
|
|
download_month=f'{sample_date.month:02d}',
|
|
download_day=f'{sample_date.day:02d}',
|
|
ext='html',
|
|
)
|
|
|
|
try:
|
|
return file_name_format % variables
|
|
except KeyError as e:
|
|
raise RuntimeError(f'Invalid variable: {e}')
|
|
except ValueError as e:
|
|
raise RuntimeError(f'Invalid format: {e}')
|
|
|
|
|
|
def format_archive_filename_from_archive(
|
|
archive: Archive,
|
|
file_name_format: str = None,
|
|
) -> str:
|
|
"""Single source of truth for archive filenames when reorganizing.
|
|
|
|
Formats an archive filename using the configured file_name_format template.
|
|
Used by collection reorganization to compute new file paths.
|
|
|
|
Args:
|
|
archive: Archive model with metadata
|
|
file_name_format: Optional override, defaults to config value
|
|
|
|
Returns:
|
|
Formatted filename (may include subdirectory path like "2025/My Article.html")
|
|
"""
|
|
config = get_archive_downloader_config()
|
|
template = file_name_format or config.file_name_format
|
|
|
|
file_group = archive.file_group
|
|
|
|
# Get download date from file_group
|
|
download_date = file_group.download_datetime or file_group.published_datetime or now()
|
|
|
|
# Get title
|
|
title = file_group.title or 'untitled'
|
|
|
|
# Get domain from URL
|
|
domain = ''
|
|
if file_group.url:
|
|
try:
|
|
domain = extract_domain(file_group.url)
|
|
except Exception:
|
|
pass
|
|
|
|
variables = dict(
|
|
title=escape_file_name(title) if title else 'untitled',
|
|
domain=domain,
|
|
download_datetime=archive_strftime(download_date),
|
|
download_date=download_date.strftime('%Y-%m-%d'),
|
|
download_year=str(download_date.year),
|
|
download_month=f'{download_date.month:02d}',
|
|
download_day=f'{download_date.day:02d}',
|
|
ext='html',
|
|
)
|
|
|
|
try:
|
|
result = template % variables
|
|
# Trim the filename portion if it's too long (may include subdirectory paths)
|
|
if '/' in result:
|
|
# Has subdirectories - only trim the filename part
|
|
parts = result.rsplit('/', 1)
|
|
subdir = parts[0]
|
|
filename = trim_file_name(parts[1])
|
|
return f'{subdir}/{filename}'
|
|
else:
|
|
return trim_file_name(result)
|
|
except KeyError as e:
|
|
logger.error(f'Invalid variable in archive file_name_format: {e}')
|
|
# Fallback to default format
|
|
return trim_file_name(
|
|
f'{archive_strftime(download_date)}_{escape_file_name(title) if title else "untitled"}.html')
|
|
|
|
|
|
@dataclass
|
|
class ArchiveFiles:
|
|
"""Every Archive will have some of these files."""
|
|
singlefile: pathlib.Path = None
|
|
readability: pathlib.Path = None
|
|
readability_txt: pathlib.Path = None
|
|
readability_json: pathlib.Path = None
|
|
screenshot: pathlib.Path = None
|
|
|
|
def __repr__(self):
|
|
singlefile = str(self.singlefile.relative_to(get_archive_directory())) if self.singlefile else None
|
|
readability = str(self.readability.relative_to(get_archive_directory())) if self.readability else None
|
|
readability_txt = str(
|
|
self.readability_txt.relative_to(get_archive_directory())) if self.readability_txt else None
|
|
readability_json = str(
|
|
self.readability_json.relative_to(get_archive_directory())) if self.readability_json else None
|
|
screenshot = str(self.screenshot.relative_to(get_archive_directory())) if self.screenshot else None
|
|
return f'<ArchiveFiles {singlefile=} {readability=} {readability_txt=} {readability_json=} {screenshot=}>'
|
|
|
|
|
|
def get_archive_directory() -> pathlib.Path:
|
|
archive_destination = get_wrolpi_config().archive_destination
|
|
variables = dict(domain='', domain_tag='', year='', month='', day='')
|
|
archive_destination = archive_destination % variables
|
|
archive_directory = get_media_directory() / archive_destination
|
|
return archive_directory
|
|
|
|
|
|
# File names include domain and datetime.
|
|
MAXIMUM_ARCHIVE_FILE_CHARACTER_LENGTH = 200
|
|
|
|
|
|
def get_new_archive_files(url: str, title: Optional[str], destination: pathlib.Path = None) -> ArchiveFiles:
|
|
"""Create a list of archive files using a shared name schema. Raise an error if any of them exist.
|
|
|
|
Args:
|
|
url: The URL being archived (used to derive domain if no destination)
|
|
title: Title for the archive files
|
|
destination: Optional directory to save files to. If provided, files go here instead of archive/<domain>
|
|
"""
|
|
domain = extract_domain(url)
|
|
if destination:
|
|
directory = destination
|
|
else:
|
|
# Fallback to archive/<domain>/ when no destination provided
|
|
directory = get_archive_directory() / domain
|
|
directory.mkdir(parents=True, exist_ok=True)
|
|
|
|
# Use the configured file format template (returns full filename with .html extension)
|
|
title = title or 'NA'
|
|
title = title[:MAXIMUM_ARCHIVE_FILE_CHARACTER_LENGTH]
|
|
formatted_filename = format_archive_filename(title, domain=domain)
|
|
|
|
# Handle subdirectories in the format (e.g., "%(download_year)s/%(title)s.%(ext)s")
|
|
if '/' in formatted_filename:
|
|
# Filename includes subdirectories, create them relative to base directory
|
|
full_path = directory / formatted_filename
|
|
full_path.parent.mkdir(parents=True, exist_ok=True)
|
|
singlefile_path = full_path
|
|
stem = full_path.stem
|
|
parent_dir = full_path.parent
|
|
readability_path = trim_file_name(parent_dir / f'{stem}.readability.html')
|
|
readability_txt_path = trim_file_name(parent_dir / f'{stem}.readability.txt')
|
|
readability_json_path = trim_file_name(parent_dir / f'{stem}.readability.json')
|
|
screenshot_path = trim_file_name(parent_dir / f'{stem}.png')
|
|
else:
|
|
singlefile_path = directory / formatted_filename
|
|
stem = singlefile_path.stem
|
|
readability_path = trim_file_name(directory / f'{stem}.readability.html')
|
|
readability_txt_path = trim_file_name(directory / f'{stem}.readability.txt')
|
|
readability_json_path = trim_file_name(directory / f'{stem}.readability.json')
|
|
screenshot_path = trim_file_name(directory / f'{stem}.png')
|
|
|
|
paths = (singlefile_path, readability_path, readability_txt_path, readability_json_path, screenshot_path)
|
|
|
|
for path in paths:
|
|
if path.exists():
|
|
raise FileExistsError(f'New archive file already exists: {path}')
|
|
|
|
archive_files = ArchiveFiles(
|
|
singlefile_path,
|
|
readability_path,
|
|
readability_txt_path,
|
|
readability_json_path,
|
|
screenshot_path,
|
|
)
|
|
return archive_files
|
|
|
|
|
|
ARCHIVE_TIMEOUT = Seconds.minute * 10 # Wait at most 10 minutes for response.
|
|
|
|
|
|
async def request_archive(url: str, singlefile: str = None, compress: bool = False) \
|
|
-> Tuple[str | bytes, Optional[dict], Optional[str]]:
|
|
"""Send a request to the archive service to archive the URL."""
|
|
logger.info(f'Sending archive request to archive service: {url}')
|
|
|
|
data = dict(url=url, singlefile=singlefile, compress=compress)
|
|
try:
|
|
async with aiohttp_post(f'{ARCHIVE_SERVICE}/json', json_=data, timeout=ARCHIVE_TIMEOUT) as response:
|
|
status = response.status
|
|
contents = await response.json()
|
|
if contents and (error := contents.get('error')):
|
|
# Report the error from the archive service.
|
|
raise Exception(f'Received error from archive service: {error}')
|
|
|
|
readability = contents.get('readability')
|
|
# Compressed base64
|
|
singlefile = contents['singlefile']
|
|
screenshot = contents['screenshot']
|
|
|
|
logger.debug(f'archive request status code {status}')
|
|
except Exception as e:
|
|
logger.error('Error when requesting archive', exc_info=e)
|
|
raise
|
|
|
|
if not (screenshot or singlefile or readability):
|
|
raise Exception('singlefile response was empty!')
|
|
|
|
if not readability:
|
|
logger.info(f'Failed to get readability for {url=}')
|
|
|
|
# Decode and decompress.
|
|
singlefile = base64.b64decode(singlefile)
|
|
singlefile = gzip.decompress(singlefile)
|
|
# A compressed (SingleFileZ) singlefile is binary and must stay bytes; only decode regular singlefiles.
|
|
if not is_compressed_singlefile(singlefile):
|
|
singlefile = singlefile.decode()
|
|
|
|
if not screenshot:
|
|
logger.info(f'Failed to get screenshot for {url=}')
|
|
else:
|
|
# Decode and decompress.
|
|
screenshot = base64.b64decode(screenshot)
|
|
screenshot = gzip.decompress(screenshot)
|
|
|
|
return singlefile, readability, screenshot
|
|
|
|
|
|
async def request_screenshot(url: str, singlefile_path: pathlib.Path) -> Optional[bytes]:
|
|
"""Send a request to the archive service to generate a screenshot from the singlefile."""
|
|
logger.info(f'Sending screenshot request to archive service: {url}')
|
|
|
|
# Read, compress, and encode the singlefile
|
|
singlefile_contents = singlefile_path.read_bytes()
|
|
singlefile_compressed = gzip.compress(singlefile_contents)
|
|
singlefile_b64 = base64.b64encode(singlefile_compressed).decode()
|
|
|
|
data = dict(url=url, singlefile=singlefile_b64)
|
|
try:
|
|
async with aiohttp_post(f'{ARCHIVE_SERVICE}/screenshot', json_=data, timeout=ARCHIVE_TIMEOUT) as response:
|
|
status = response.status
|
|
contents = await response.json()
|
|
if contents and (error := contents.get('error')):
|
|
# Report the error from the archive service.
|
|
raise Exception(f'Received error from archive service: {error}')
|
|
|
|
# Compressed base64
|
|
screenshot = contents.get('screenshot')
|
|
if not screenshot:
|
|
logger.warning(f'Failed to get screenshot for {url=}')
|
|
return None
|
|
|
|
logger.debug(f'screenshot request status code {status}')
|
|
except Exception as e:
|
|
logger.error('Error when requesting screenshot', exc_info=e)
|
|
raise
|
|
|
|
# Decode and decompress.
|
|
screenshot = base64.b64decode(screenshot)
|
|
screenshot = gzip.decompress(screenshot)
|
|
|
|
return screenshot
|
|
|
|
|
|
@dataclass
|
|
class WrittenArchiveFiles:
|
|
"""The on-disk artifacts produced by write_archive_files."""
|
|
url: str
|
|
paths: List[pathlib.Path]
|
|
|
|
|
|
def write_archive_files(url: str, singlefile: str | bytes, readability: Optional[dict],
|
|
screenshot: Optional[bytes],
|
|
destination: Optional[pathlib.Path] = None) -> WrittenArchiveFiles:
|
|
"""Pure I/O: write the singlefile/readability/screenshot/json files to disk.
|
|
|
|
Returns the URL and the list of paths to register, leaving DB work to register_archive.
|
|
"""
|
|
compressed = isinstance(singlefile, bytes) and is_compressed_singlefile(singlefile)
|
|
|
|
readability = readability.copy() if readability else readability
|
|
# First try to get the title from Readability.
|
|
title = readability.get('title') if readability else None
|
|
|
|
if not title and singlefile:
|
|
# Try to get the title ourselves from the HTML. A compressed singlefile contains the
|
|
# uncompressed page as index.html.
|
|
title = get_title_from_html(get_compressed_singlefile_index(singlefile) if compressed else singlefile, url=url)
|
|
if readability:
|
|
# Readability could not find title, lets use ours.
|
|
readability['title'] = title
|
|
|
|
archive_files = get_new_archive_files(url, title, destination=destination)
|
|
|
|
if readability:
|
|
# Write the readability parts to their own files. Write what is left after pops to the JSON file.
|
|
with archive_files.readability.open('wt') as fp:
|
|
content = readability.pop('content')
|
|
if compressed:
|
|
# Readability extracted from the ZIP's index.html references images by relative
|
|
# ZIP paths; inline them so the readability file is self-contained.
|
|
content = inline_compressed_singlefile_resources(content, singlefile, url)
|
|
content = format_html_string(content)
|
|
fp.write(content)
|
|
with archive_files.readability_txt.open('wt') as fp:
|
|
readability_txt = readability.pop('textContent')
|
|
readability_txt = split_lines_by_length(readability_txt)
|
|
fp.write(readability_txt)
|
|
else:
|
|
# No readability was returned, so there are no files.
|
|
archive_files.readability_txt = archive_files.readability = None
|
|
|
|
# Store the single-file HTML in its own file.
|
|
if compressed:
|
|
# A compressed (SingleFileZ) singlefile is a binary ZIP; prettifying would corrupt it.
|
|
archive_files.singlefile.write_bytes(singlefile)
|
|
elif isinstance(singlefile, bytes) and b'<html data-sfz' in singlefile[:200]:
|
|
# Marked as compressed but the ZIP is unreadable (e.g. truncated); prettifying would
|
|
# destroy whatever remains, so store it verbatim.
|
|
logger.warning(f'Singlefile of {url} is marked compressed but is not a valid ZIP!')
|
|
archive_files.singlefile.write_bytes(singlefile)
|
|
else:
|
|
singlefile = format_html_string(singlefile)
|
|
archive_files.singlefile.write_text(singlefile)
|
|
|
|
if screenshot:
|
|
archive_files.screenshot.write_bytes(screenshot)
|
|
else:
|
|
archive_files.screenshot = None
|
|
|
|
# Always write a JSON file that contains at least the URL.
|
|
readability = readability or {}
|
|
readability['url'] = url
|
|
with archive_files.readability_json.open('wt') as fp:
|
|
json.dump(readability, fp, indent=2, sort_keys=True)
|
|
|
|
# Create any File models that we can, index them all.
|
|
paths = (
|
|
archive_files.singlefile, archive_files.readability, archive_files.readability_json,
|
|
archive_files.readability_txt,
|
|
archive_files.screenshot)
|
|
paths = list(filter(None, paths))
|
|
|
|
return WrittenArchiveFiles(url=url, paths=paths)
|
|
|
|
|
|
def register_archive(session: Session, written: WrittenArchiveFiles) -> Archive:
|
|
"""Pure DB: create the Archive record from already-written files. Caller commits."""
|
|
archive = Archive.from_paths(session, *written.paths)
|
|
archive.file_group.download_datetime = now()
|
|
archive.url = written.url
|
|
archive.collection = get_or_create_domain_collection(session, written.url)
|
|
|
|
# Update collection's file_format to track the format used for this file
|
|
config = get_archive_downloader_config()
|
|
if archive.collection and config.file_name_format:
|
|
if not archive.collection.file_format:
|
|
# Set initial file_format on collection
|
|
archive.collection.file_format = config.file_name_format
|
|
save_domains_config.activate_switch()
|
|
|
|
archive.flush()
|
|
return archive
|
|
|
|
|
|
async def model_archive_result(url: str, singlefile: str, readability: dict, screenshot: bytes,
|
|
destination: pathlib.Path = None) -> Archive:
|
|
"""
|
|
Convert results from ArchiveDownloader into real files. Create Archive record.
|
|
|
|
Thin wrapper around write_archive_files + register_archive that opens its own
|
|
transaction. New code (notably the phase-split ArchiveDownloader) should call
|
|
the helpers directly so the caller controls the session.
|
|
|
|
Args:
|
|
url: The URL that was archived
|
|
singlefile: The HTML content from SingleFile
|
|
readability: The readability extraction results
|
|
screenshot: The screenshot bytes
|
|
destination: Optional directory to save files to. If provided, files go here instead of archive/<domain>
|
|
"""
|
|
written = write_archive_files(url, singlefile, readability, screenshot, destination=destination)
|
|
with get_db_session(commit=True) as session:
|
|
archive = register_archive(session, written)
|
|
return archive
|
|
|
|
|
|
def detect_domain_directory(collection: Collection) -> Optional[pathlib.Path]:
|
|
"""
|
|
Detect directory for domain collection using format_destination template.
|
|
|
|
Computes the expected domain directory using the collection's format_destination()
|
|
method and checks if that directory exists on the filesystem.
|
|
|
|
Args:
|
|
collection: The domain collection to analyze
|
|
|
|
Returns:
|
|
Path (relative to media directory) if the expected directory exists.
|
|
None if collection is not a domain or directory doesn't exist.
|
|
"""
|
|
if collection.kind != 'domain':
|
|
return None
|
|
|
|
# Get the tag name if collection has a tag
|
|
tag_name = collection.tag.name if collection.tag else None
|
|
|
|
# Compute expected directory using archive_destination template
|
|
expected_directory = collection.format_destination(tag_name)
|
|
|
|
# Simply check if the directory exists
|
|
if expected_directory.is_dir():
|
|
relative_path = get_relative_to_media_directory(expected_directory)
|
|
logger.debug(f'Detected directory for domain {collection.name}: {relative_path}')
|
|
return relative_path
|
|
|
|
return None
|
|
|
|
|
|
def update_domain_directories(session: Session = None) -> int:
|
|
"""
|
|
One-time update to detect and set directories for existing domain collections.
|
|
|
|
This should be run once to fix domain collections that were created without directories
|
|
but have all their archives in a common location.
|
|
|
|
Args:
|
|
session: Database session
|
|
|
|
Returns:
|
|
Number of domain collections updated
|
|
"""
|
|
if session is None:
|
|
with get_db_session(commit=True) as session:
|
|
return update_domain_directories(session)
|
|
|
|
# Find all domain collections without directories
|
|
collections = session.query(Collection).filter_by(kind='domain', directory=None).all()
|
|
|
|
updated_count = 0
|
|
for collection in collections:
|
|
detected_dir = detect_domain_directory(collection)
|
|
if detected_dir:
|
|
collection.directory = detected_dir
|
|
session.flush([collection])
|
|
updated_count += 1
|
|
logger.info(f'Updated domain {collection.name} with directory: {detected_dir}')
|
|
|
|
session.commit()
|
|
logger.info(f'Updated {updated_count} domain collection(s) with auto-detected directories')
|
|
return updated_count
|
|
|
|
|
|
def get_or_create_domain_collection(session: Session, url, directory: pathlib.Path = None) -> Collection:
|
|
"""
|
|
Get or create the domain Collection for this archive.
|
|
|
|
Args:
|
|
session: Database session
|
|
url: URL of the archive
|
|
directory: Optional directory to restrict this domain collection.
|
|
If None, creates unrestricted collection (default).
|
|
If provided, archives will be placed in this directory and collection can be tagged.
|
|
|
|
Returns:
|
|
Collection with kind='domain' for the domain extracted from the URL
|
|
"""
|
|
domain_name = extract_domain(url)
|
|
|
|
# Try to find existing domain collection
|
|
collection = session.query(Collection).filter_by(
|
|
name=domain_name,
|
|
kind='domain'
|
|
).one_or_none()
|
|
|
|
if not collection:
|
|
# Create new domain collection
|
|
collection = Collection(
|
|
name=domain_name,
|
|
kind='domain',
|
|
directory=directory, # Can be None (unrestricted) or a Path (restricted)
|
|
)
|
|
session.add(collection)
|
|
session.flush()
|
|
# Trigger domain config save for new domain
|
|
save_domains_config.activate_switch()
|
|
if directory:
|
|
logger.info(f'Created domain collection with directory: {domain_name} -> {directory}')
|
|
else:
|
|
logger.info(f'Created unrestricted domain collection: {domain_name}')
|
|
|
|
# Auto-detect directory if not explicitly set and collection doesn't have one
|
|
if not directory and not collection.directory:
|
|
detected_dir = detect_domain_directory(collection)
|
|
if detected_dir:
|
|
collection.directory = detected_dir
|
|
session.flush()
|
|
# Trigger domain config save when directory is auto-detected
|
|
save_domains_config.activate_switch()
|
|
logger.info(f'Auto-detected directory for domain {domain_name}: {detected_dir}')
|
|
|
|
return collection
|
|
|
|
|
|
def link_domain_and_downloads(session: Session):
|
|
"""Associate any Download related to a Domain Collection.
|
|
|
|
Downloads are linked to Collections (via collection_id).
|
|
This function finds Domain Collections and links their Downloads.
|
|
|
|
Matching criteria:
|
|
- Recurring downloads whose destination is within the domain collection's directory (including subdirectories)
|
|
- RSS downloads with sub_downloader='archive' (matched by URL domain)
|
|
"""
|
|
# Local import to avoid circular import: downloader -> archive -> downloader
|
|
from wrolpi.downloader import Download
|
|
|
|
# Only Downloads with a frequency can be a Collection Download.
|
|
downloads = list(session.query(Download).filter(Download.frequency.isnot(None)).all())
|
|
|
|
# Get domain collections that have a directory
|
|
domain_collections = session.query(Collection).filter(
|
|
Collection.kind == 'domain',
|
|
Collection.directory.isnot(None)
|
|
).all()
|
|
|
|
need_commit = False
|
|
|
|
# Match downloads by destination directory (including subdirectories)
|
|
downloads_with_destination = [d for d in downloads if (d.settings or {}).get('destination')]
|
|
for collection in domain_collections:
|
|
# Ensure directory ends with / for proper prefix matching
|
|
directory = str(collection.directory)
|
|
directory_prefix = directory if directory.endswith('/') else directory + '/'
|
|
for download in downloads_with_destination:
|
|
dest = download.settings['destination']
|
|
# Match if destination equals directory or is a subdirectory
|
|
if not download.collection_id and (dest == directory or dest.startswith(directory_prefix)):
|
|
download.collection_id = collection.id
|
|
need_commit = True
|
|
|
|
# Match RSS downloads with archive sub_downloader by URL domain
|
|
rss_archive_downloads = [
|
|
d for d in downloads
|
|
if d.downloader == 'rss' and d.sub_downloader == 'archive' and not d.collection_id
|
|
]
|
|
for download in rss_archive_downloads:
|
|
# Extract domain from the RSS URL and find matching domain collection
|
|
domain_name = extract_domain(download.url)
|
|
collection = session.query(Collection).filter_by(
|
|
name=domain_name,
|
|
kind='domain'
|
|
).one_or_none()
|
|
if collection:
|
|
download.collection_id = collection.id
|
|
need_commit = True
|
|
|
|
if need_commit:
|
|
session.commit()
|
|
|
|
|
|
SINGLEFILE_URL_EXTRACTOR = re.compile(r'Page saved with SingleFile \s+url: (http.+?)\n')
|
|
|
|
COMPRESSED_SINGLEFILE_INDEX = 'index.html'
|
|
|
|
|
|
def is_compressed_singlefile(singlefile: bytes) -> bool:
|
|
"""A compressed (SingleFileZ) singlefile is a self-extracting HTML file which is also a valid ZIP
|
|
containing the uncompressed page as index.html."""
|
|
buf = io.BytesIO(singlefile)
|
|
if not zipfile.is_zipfile(buf):
|
|
return False
|
|
try:
|
|
with zipfile.ZipFile(buf) as zip_:
|
|
return COMPRESSED_SINGLEFILE_INDEX in zip_.namelist()
|
|
except zipfile.BadZipFile:
|
|
return False
|
|
|
|
|
|
def get_compressed_singlefile_index(singlefile: bytes) -> bytes:
|
|
"""Extract the uncompressed page (index.html) from a compressed (SingleFileZ) singlefile."""
|
|
with zipfile.ZipFile(io.BytesIO(singlefile)) as zip_:
|
|
return zip_.read(COMPRESSED_SINGLEFILE_INDEX)
|
|
|
|
|
|
def read_singlefile_html(path: pathlib.Path) -> bytes:
|
|
"""Read the page HTML of the singlefile at `path`; extract the uncompressed page when the
|
|
singlefile is compressed (SingleFileZ)."""
|
|
contents = path.read_bytes()
|
|
if is_compressed_singlefile(contents):
|
|
return get_compressed_singlefile_index(contents)
|
|
return contents
|
|
|
|
|
|
def inline_compressed_singlefile_resources(html: str, singlefile: bytes, url: str) -> str:
|
|
"""Inline the images of a compressed (SingleFileZ) singlefile into the provided HTML.
|
|
|
|
The index.html inside a compressed singlefile references its resources by relative ZIP paths
|
|
(e.g. images/0.jpg); HTML extracted from it (readability) dangles those references — either
|
|
still relative, or absolutized against the page URL by readability-extractor
|
|
(https://example.com/images/0.jpg). Replace them with data URIs read from the ZIP, like an
|
|
uncompressed singlefile would contain."""
|
|
import mimetypes
|
|
from urllib.parse import urljoin, unquote
|
|
|
|
# readability-extractor resolves relative srcs against the page URL.
|
|
url_base = urljoin(url, '.')
|
|
|
|
soup = get_html_soup(html)
|
|
with zipfile.ZipFile(io.BytesIO(singlefile)) as zip_:
|
|
names = set(zip_.namelist())
|
|
changed = False
|
|
for img in soup.find_all('img'):
|
|
src = img.get('src') or ''
|
|
name = src.removeprefix('./')
|
|
if src.startswith(url_base):
|
|
name = unquote(src[len(url_base):])
|
|
if not name or name not in names:
|
|
continue
|
|
mimetype = mimetypes.guess_type(name)[0] or 'application/octet-stream'
|
|
img['src'] = f'data:{mimetype};base64,{base64.b64encode(zip_.read(name)).decode()}'
|
|
# The ZIP-path srcset entries are dangling too; the inlined src replaces them.
|
|
if img.get('srcset'):
|
|
del img['srcset']
|
|
changed = True
|
|
return str(soup) if changed else html
|
|
|
|
|
|
def get_url_from_singlefile(html: bytes | str) -> str:
|
|
"""Extract URL from SingleFile contents."""
|
|
# Ignore decode errors; a compressed singlefile has a plain-text prelude but a binary ZIP tail.
|
|
html = html.decode(errors='ignore') if isinstance(html, bytes) else html
|
|
if SINGLEFILE_HEADER not in html:
|
|
raise RuntimeError(f'Not a singlefile!')
|
|
|
|
url, = SINGLEFILE_URL_EXTRACTOR.findall(html)
|
|
url = url.strip()
|
|
|
|
if not url.startswith('http'):
|
|
raise RuntimeError('URL did not start with http')
|
|
|
|
# Verify that the URL can be parsed.
|
|
urlparse(url)
|
|
return url
|
|
|
|
|
|
@dataclass
|
|
class ArticleMetadata:
|
|
title: str = None
|
|
published_datetime: datetime = None
|
|
modified_datetime: datetime = None
|
|
description: str = None
|
|
author: str = None
|
|
|
|
|
|
def parse_article_html_metadata(html: Union[bytes, str], assume_utc: bool = True) -> ArticleMetadata:
|
|
"""
|
|
Read the data from the <meta> tags, extract any data relevant to the article. This function also reads the
|
|
<script type="application/ld+json"> data.
|
|
"""
|
|
metadata = ArticleMetadata()
|
|
|
|
soup = get_html_soup(html)
|
|
|
|
def get_meta_by_property(prop: str):
|
|
return soup.find('meta', attrs={'property': prop})
|
|
|
|
# <meta content="2023-10-18T04:52:23+00:00" property="article:published_time"/>
|
|
if meta_published_time := get_meta_by_property('article:published_time'):
|
|
metadata.published_datetime = dates.strpdate(meta_published_time.attrs['content'])
|
|
# <meta name="article.published" content="2023-04-04T21:52:00.000Z">
|
|
if meta_article_published := soup.find('meta', attrs={'name': 'article.published'}):
|
|
metadata.published_datetime = metadata.published_datetime or dates.strpdate(
|
|
meta_article_published.attrs['content'])
|
|
# <meta itemprop="datePublished" content="2023-04-04T21:52:00.000Z">
|
|
if meta_item_published := soup.find('meta', attrs={'itemprop': 'datePublished'}):
|
|
metadata.published_datetime = metadata.published_datetime or dates.strpdate(
|
|
meta_item_published.attrs['content'])
|
|
# <time itemprop="datePublished" datetime="2023-08-25">
|
|
if time := soup.find('time', attrs={'itemprop': 'datePublished'}):
|
|
metadata.published_datetime = metadata.published_datetime or dates.strpdate(time.attrs['datetime'])
|
|
# <abbr class="published" itemprop="datePublished" title="2022-03-17T03:00:00-07:00">March 17, 2022</abbr>
|
|
if abbr := soup.find('abbr', attrs={'itemprop': 'datePublished'}):
|
|
metadata.published_datetime = metadata.published_datetime or dates.strpdate(abbr.attrs['title'])
|
|
|
|
# <meta content="2023-10-19T05:53:24+00:00" property="article:modified_time"/>
|
|
if meta_modified_time := get_meta_by_property('article:modified_time'):
|
|
metadata.modified_datetime = dates.strpdate(meta_modified_time.attrs['content'])
|
|
# <meta name="article.updated" content="2023-04-04T21:52:00.000Z">
|
|
if meta_article_updated := soup.find('meta', attrs={'name': 'article.updated'}):
|
|
metadata.modified_datetime = metadata.modified_datetime or dates.strpdate(meta_article_updated.attrs['content'])
|
|
|
|
# <meta content="The Title" property="og:title"/>
|
|
if meta_title := get_meta_by_property('og:title'):
|
|
metadata.title = meta_title.attrs['content']
|
|
|
|
# <meta name="author" content="Author Name"/>
|
|
if meta_author := soup.find('meta', attrs={'name': 'author'}):
|
|
metadata.author = meta_author.attrs['content']
|
|
# <meta content="Billy" property="article:author"/>
|
|
if meta_property_author := get_meta_by_property('article:author'):
|
|
metadata.author = metadata.author or meta_property_author.attrs['content']
|
|
# <a href="https://example.com" rel="author">
|
|
if link_author := soup.find('a', attrs={'rel': 'author'}):
|
|
metadata.author = metadata.author or link_author.text.strip()
|
|
|
|
# <script class="sf-hidden" type="application/ld+json">
|
|
if (ld_script := soup.find('script', attrs={'type': 'application/ld+json'})) and ld_script.contents:
|
|
try:
|
|
schema = json.loads(ld_script.text)
|
|
except json.decoder.JSONDecodeError:
|
|
# Was not valid JSON.
|
|
schema = None
|
|
if isinstance(schema, dict) and (context := schema.get('@context')) and '://schema.org' in context:
|
|
# Found https://schema.org/
|
|
if headline := schema.get('headline'):
|
|
metadata.title = metadata.title or headline
|
|
if datePublished := schema.get('datePublished'):
|
|
try:
|
|
metadata.published_datetime = metadata.published_datetime or dates.strpdate(datePublished)
|
|
except InvalidDatetime as e:
|
|
# Invalid date, ignore.
|
|
logger.error('Invalid datetime', exc_info=e)
|
|
if PYTEST:
|
|
raise
|
|
if dateModified := schema.get('dateModified'):
|
|
try:
|
|
metadata.modified_datetime = metadata.modified_datetime or dates.strpdate(dateModified)
|
|
except InvalidDatetime as e:
|
|
# Invalid date, ignore.
|
|
logger.error('Invalid datetime', exc_info=e)
|
|
if PYTEST:
|
|
raise
|
|
if description := schema.get('description'):
|
|
metadata.description = description
|
|
if author := schema.get('author'):
|
|
if isinstance(author, list) and len(author) >= 1:
|
|
# Use the first Author.
|
|
author = author[0]
|
|
|
|
if isinstance(author, dict):
|
|
author = author.get('name') or author
|
|
|
|
if isinstance(author, str):
|
|
metadata.author = author
|
|
else:
|
|
logger.warning(f'Unable to parse author schema: {author}')
|
|
elif authors := schema.get('authors'):
|
|
# Use the first Author.
|
|
author = authors[0]
|
|
|
|
if isinstance(author, dict):
|
|
author = author.get('name') or author
|
|
|
|
if isinstance(author, str):
|
|
metadata.author = author
|
|
else:
|
|
logger.warning(f'Unable to parse author schema: {author}')
|
|
|
|
# Assume UTC if no timezone.
|
|
if metadata.published_datetime and not metadata.published_datetime.tzinfo and assume_utc:
|
|
logger.debug(f'Assuming UTC for {metadata.published_datetime=}')
|
|
metadata.published_datetime = metadata.published_datetime.replace(tzinfo=pytz.UTC)
|
|
if metadata.modified_datetime and not metadata.modified_datetime.tzinfo and assume_utc:
|
|
logger.debug(f'Assuming UTC for {metadata.modified_datetime=}')
|
|
metadata.modified_datetime = metadata.modified_datetime.replace(tzinfo=pytz.UTC)
|
|
|
|
return metadata
|
|
|
|
|
|
def get_archive(session: Session, archive_id: int, skip_viewed: bool = False) -> Archive:
|
|
"""Get an Archive."""
|
|
archive = Archive.find_by_id(session, archive_id)
|
|
if not skip_viewed:
|
|
archive.file_group.set_viewed()
|
|
return archive
|
|
|
|
|
|
def get_archive_by_file_group_id(session: Session, file_group_id: int, skip_viewed: bool = False) -> Archive:
|
|
"""Get an Archive by its FileGroup ID."""
|
|
archive = session.query(Archive).filter_by(file_group_id=file_group_id).one_or_none()
|
|
if not archive:
|
|
raise UnknownArchive(f'Unknown Archive with FileGroup ID: {file_group_id}')
|
|
if not skip_viewed:
|
|
archive.file_group.set_viewed()
|
|
return archive
|
|
|
|
|
|
def delete_archives_by_file_group_ids(*file_group_ids: int, force: bool = False) -> None:
|
|
"""Delete Archives by their FileGroup IDs and all of their files. Pass force=True to delete
|
|
even tagged FileGroups (raises FileGroupIsTagged otherwise via FileGroup.delete)."""
|
|
with get_db_session(commit=True) as session:
|
|
archives: List[Archive] = list(session.query(Archive).filter(Archive.file_group_id.in_(file_group_ids)))
|
|
if not archives:
|
|
raise UnknownArchive(f'Unknown Archives with FileGroup IDs: {", ".join(map(str, file_group_ids))}')
|
|
for archive in archives:
|
|
archive.delete(force=force)
|
|
|
|
|
|
def archive_strftime(dt: datetime) -> str:
|
|
return dt.strftime('%Y-%m-%d-%H-%M-%S')
|
|
|
|
|
|
ARCHIVE_MATCHER = re.compile(r'\d{4}(-\d\d){5}_.*$')
|
|
SINGLEFILE_HEADER = '''<!--
|
|
Page saved with SingleFile'''
|
|
|
|
|
|
def is_singlefile_file(path: pathlib.Path) -> bool:
|
|
"""
|
|
Archive singlefile files are expected to start with the following: %Y-%m-%d-%H-%M-%S
|
|
they must end with .html, but never with .readability.html.
|
|
"""
|
|
if not path.is_file() or \
|
|
path.suffix.lower() != '.html' or \
|
|
path.name.lower().endswith('.readability.html') or \
|
|
path.stat().st_size == 0:
|
|
return False
|
|
if ARCHIVE_MATCHER.match(path.name):
|
|
# A file name matching exactly the Archive format from WROLPi should always be an Archive.
|
|
return True
|
|
with path.open('rb') as fh:
|
|
# Lastly, try to read the header of this HTML file. Read a large part of the head because URLs can be long.
|
|
# Read bytes because a compressed (SingleFileZ) singlefile has a binary ZIP tail.
|
|
header = fh.read(1000).decode(errors='ignore')
|
|
if SINGLEFILE_HEADER in header:
|
|
return True
|
|
return False
|
|
|
|
|
|
def get_domains():
|
|
"""
|
|
Get all domain collections with their archive statistics.
|
|
|
|
This is a thin wrapper around Collection queries that adds archive-specific statistics.
|
|
Returns a list of dicts with collection id, domain name, url_count, and total size.
|
|
"""
|
|
with get_db_session() as session:
|
|
# Query all domain collections with archive statistics in a single query
|
|
# This uses ORM for better maintainability while keeping performance
|
|
query = (
|
|
session.query(
|
|
Collection.id,
|
|
Collection.name.label('domain'),
|
|
func.count(Archive.id).label('url_count'),
|
|
func.sum(FileGroup.size).cast(BigInteger).label('size')
|
|
)
|
|
.outerjoin(Archive, Collection.id == Archive.collection_id)
|
|
.outerjoin(FileGroup, FileGroup.id == Archive.file_group_id)
|
|
.filter(Collection.kind == 'domain')
|
|
.group_by(Collection.id, Collection.name)
|
|
.order_by(Collection.name)
|
|
)
|
|
|
|
domains = [
|
|
{
|
|
'id': row.id,
|
|
'domain': row.domain,
|
|
'url_count': row.url_count or 0,
|
|
'size': row.size or 0,
|
|
}
|
|
for row in query.all()
|
|
]
|
|
|
|
return domains
|
|
|
|
|
|
def get_domain(domain_id: int) -> dict:
|
|
"""
|
|
Get a single domain collection by ID with its archive statistics.
|
|
|
|
Returns a dict with collection details including id, domain name, url_count, size,
|
|
tag_name, directory (relative to media directory), and description.
|
|
|
|
Raises UnknownArchive if domain not found.
|
|
"""
|
|
# Get the collection using find_by_id
|
|
with get_db_session() as session:
|
|
try:
|
|
collection = Collection.find_by_id(session, domain_id)
|
|
except Exception:
|
|
# Collection.find_by_id raises UnknownCollection, but we want UnknownArchive
|
|
raise UnknownArchive(f"Domain collection with ID {domain_id} not found")
|
|
|
|
if collection.kind != 'domain':
|
|
raise UnknownArchive(f"Collection {domain_id} is not a domain")
|
|
|
|
# Get base domain data from __json__()
|
|
domain_data = collection.__json__()
|
|
|
|
# Get archive statistics for this domain
|
|
with get_db_curs() as curs:
|
|
stmt = '''
|
|
SELECT COUNT(a.id) AS url_count, SUM(fg.size)::BIGINT AS size
|
|
FROM archive a
|
|
LEFT JOIN file_group fg on fg.id = a.file_group_id
|
|
WHERE a.collection_id = %(domain_id)s
|
|
'''
|
|
curs.execute(stmt, {'domain_id': domain_id})
|
|
stats = dict(curs.fetchone())
|
|
|
|
# Enhance with archive stats
|
|
domain_data['url_count'] = stats['url_count'] or 0
|
|
domain_data['size'] = stats['size'] or 0
|
|
|
|
# Use 'domain' key for domain collections instead of 'name'
|
|
domain_data['domain'] = domain_data.pop('name')
|
|
|
|
return domain_data
|
|
|
|
|
|
ARCHIVE_ORDERS = {
|
|
# Sometimes we don't have a published_datetime. This is equivalent to COALESCE(fg.published_datetime, fg.download_datetime)
|
|
'published_datetime': (date := 'fg.effective_datetime ASC NULLS LAST'),
|
|
'-published_datetime': (_date := 'fg.effective_datetime DESC NULLS LAST'),
|
|
'published_modified_datetime': f'fg.published_modified_datetime ASC NULLS LAST, {date}',
|
|
'-published_modified_datetime': f'fg.published_modified_datetime DESC NULLS LAST, {_date}',
|
|
'download_datetime': 'fg.download_datetime ASC NULLS LAST',
|
|
'-download_datetime': 'fg.download_datetime DESC NULLS LAST',
|
|
'rank': f'2 DESC, {_date}',
|
|
'-rank': f'2 ASC, {date}',
|
|
'size': 'fg.size ASC, LOWER(fg.primary_path) ASC',
|
|
'-size': 'fg.size DESC NULLS LAST, LOWER(fg.primary_path) DESC',
|
|
'viewed': 'fg.viewed ASC',
|
|
'-viewed': 'fg.viewed DESC',
|
|
}
|
|
ORDER_GROUP_BYS = {
|
|
'published_datetime': 'fg.effective_datetime',
|
|
'-published_datetime': 'fg.effective_datetime',
|
|
'published_modified_datetime': 'fg.published_modified_datetime, fg.published_datetime, fg.download_datetime',
|
|
'-published_modified_datetime': 'fg.published_modified_datetime, fg.published_datetime, fg.download_datetime',
|
|
'download_datetime': 'fg.download_datetime',
|
|
'-download_datetime': 'fg.download_datetime',
|
|
'rank': 'fg.published_datetime, fg.download_datetime',
|
|
'-rank': 'fg.published_datetime, fg.download_datetime',
|
|
'size': 'fg.size, fg.published_datetime, fg.primary_path',
|
|
'-size': 'fg.size, fg.published_datetime, fg.primary_path',
|
|
'viewed': 'fg.viewed',
|
|
'-viewed': 'fg.viewed',
|
|
}
|
|
NO_NULL_ORDERS = {
|
|
'viewed': 'fg.viewed IS NOT NULL',
|
|
'-viewed': 'fg.viewed IS NOT NULL',
|
|
'published_modified_datetime': 'fg.published_modified_datetime IS NOT NULL',
|
|
'-published_modified_datetime': 'fg.published_modified_datetime IS NOT NULL',
|
|
}
|
|
|
|
|
|
def search_archives(search_str: str, domain: str, limit: int, offset: int, order: str, tag_names: List[str],
|
|
headline: bool = False, deep: bool = False) \
|
|
-> Tuple[List[dict], int]:
|
|
# Always filter FileGroups to Archives.
|
|
wheres = []
|
|
group_by = 'fg.effective_datetime, a.file_group_id'
|
|
|
|
params = dict(search_str=search_str, offset=int(offset), limit=int(limit))
|
|
order_by = ARCHIVE_ORDERS['-published_datetime']
|
|
|
|
select_columns = ''
|
|
if search_str:
|
|
# A search_str was provided by the user, modify the query to filter by it.
|
|
# `textsearch` includes d_text (page contents); `textsearch_abc` is much smaller and faster.
|
|
ts_col = 'textsearch' if deep else 'textsearch_abc'
|
|
select_columns = f'ts_rank(fg.{ts_col}, websearch_to_tsquery(%(search_str)s)) AS rank,' \
|
|
' COUNT(*) OVER() AS total'
|
|
wheres.append(f'fg.{ts_col} @@ websearch_to_tsquery(%(search_str)s)')
|
|
params['search_str'] = search_str
|
|
group_by = f'{group_by}, rank'
|
|
|
|
if order:
|
|
try:
|
|
order_by = ARCHIVE_ORDERS[order]
|
|
group_by = f'{group_by}, {ORDER_GROUP_BYS[order]}'
|
|
except KeyError:
|
|
raise InvalidOrderBy(f'Invalid order byy {order}')
|
|
if order in NO_NULL_ORDERS:
|
|
wheres.append(NO_NULL_ORDERS[order])
|
|
|
|
wheres, params = tag_append_sub_select_where(wheres, params, tag_names)
|
|
|
|
if search_str and headline:
|
|
headline = ''',
|
|
ts_headline(fg.title, websearch_to_tsquery(%(search_str)s)) AS "title_headline",
|
|
ts_headline(fg.b_text, websearch_to_tsquery(%(search_str)s)) AS "b_headline",
|
|
ts_headline(fg.c_text, websearch_to_tsquery(%(search_str)s)) AS "c_headline",
|
|
ts_headline(fg.d_text, websearch_to_tsquery(%(search_str)s)) AS "d_headline"'''
|
|
group_by = f'{group_by}, title_headline, b_headline, c_headline, d_headline'
|
|
else:
|
|
headline = ''
|
|
|
|
if domain:
|
|
params['domain'] = domain
|
|
# Use LIMIT 1 to handle potential duplicate domain collections with the same name
|
|
wheres.append(
|
|
"a.collection_id = (select id from collection where collection.name = %(domain)s and collection.kind = 'domain' LIMIT 1)")
|
|
|
|
select_columns = f", {select_columns}" if select_columns else ""
|
|
wheres = '\n AND '.join(wheres)
|
|
where = f'WHERE\n{wheres}' if wheres else ''
|
|
stmt = f'''
|
|
SELECT
|
|
a.file_group_id AS id, -- always get `file_group.id` for `handle_file_group_search_results`
|
|
COUNT(*) OVER() AS total
|
|
{select_columns}
|
|
{headline}
|
|
FROM archive a
|
|
LEFT JOIN file_group fg ON fg.id = a.file_group_id
|
|
{where}
|
|
GROUP BY {group_by}
|
|
ORDER BY {order_by}
|
|
OFFSET %(offset)s LIMIT %(limit)s
|
|
'''.strip()
|
|
logger.debug(stmt, params)
|
|
|
|
results, total = handle_file_group_search_results(stmt, params)
|
|
return results, total
|
|
|
|
|
|
async def search_domains_by_name(session: Session, name: str, limit: int = 5) -> List[dict]:
|
|
"""
|
|
Search for domain collections by name.
|
|
|
|
Args:
|
|
session: Database session
|
|
name: Search string to match against collection names
|
|
limit: Maximum number of results to return
|
|
|
|
Returns:
|
|
List of domain dicts matching the search (in old Domain format for backward compatibility)
|
|
"""
|
|
collections = session.query(Collection) \
|
|
.filter(Collection.kind == 'domain') \
|
|
.filter(Collection.name.ilike(f'%{name}%')) \
|
|
.order_by(asc(Collection.name)) \
|
|
.limit(limit) \
|
|
.all()
|
|
|
|
# Convert to old Domain format for backward compatibility
|
|
archive_dir = get_archive_directory()
|
|
return [
|
|
{
|
|
'id': c.id,
|
|
'domain': c.name,
|
|
# Domain collections can have explicit directories or use default archive path
|
|
'directory': get_relative_to_media_directory(c.directory) if c.directory else str(
|
|
(archive_dir / c.name).relative_to(get_media_directory())),
|
|
}
|
|
for c in collections
|
|
]
|
|
|
|
|
|
async def html_to_readability(html: str | bytes, url: str, timeout: int = 120):
|
|
"""Extract the readability dict from the provided HTML."""
|
|
with tempfile.NamedTemporaryFile('wb', suffix='.html') as singlefile_file:
|
|
singlefile_file.write(html.encode() if isinstance(html, str) else html)
|
|
singlefile_file.flush()
|
|
os.fsync(singlefile_file.fileno())
|
|
|
|
cmd = ('/usr/bin/nice', '-n15', # Nice above map importing, but very low priority.
|
|
READABILITY_BIN, singlefile_file.name, shlex.quote(url))
|
|
# Docker containers may not have this directory. But, this directory is necessary on RPi.
|
|
cwd = '/home/wrolpi' if os.path.isdir('/home/wrolpi') else None
|
|
result = await run_command(cmd, cwd=cwd, timeout=timeout)
|
|
|
|
logger.debug(f'readability for {url} exited with {result.return_code}')
|
|
stdout, stderr = result.stdout.decode(), result.stderr.decode()
|
|
if result.return_code == 0:
|
|
if not stdout:
|
|
raise RuntimeError('readability stdout was empty')
|
|
try:
|
|
readability = json.loads(stdout)
|
|
except TypeError or JSONDecodeError:
|
|
# JSON was invalid.
|
|
e = ChildProcessError(stdout if stdout else 'No stdout')
|
|
if stderr:
|
|
e = ChildProcessError(stderr)
|
|
raise RuntimeError(f'Failed to extract readability from {url}') from e
|
|
logger.debug(f'done readability for {url}')
|
|
return readability
|
|
else:
|
|
logger.error(f'Failed to extract readability for {url}')
|
|
e = ChildProcessError(stdout if stdout else 'No stdout')
|
|
if stderr:
|
|
e = ChildProcessError(stderr)
|
|
raise RuntimeError(f'Failed to extract readability for {url} got {result.return_code}') from e
|
|
|
|
|
|
@dataclass
|
|
class SinglefileArtifacts:
|
|
"""Pure-data result of running a singlefile through readability + screenshot extraction."""
|
|
url: str
|
|
singlefile: bytes
|
|
readability: dict
|
|
screenshot: Optional[bytes]
|
|
|
|
|
|
async def extract_singlefile_artifacts(singlefile: bytes) -> SinglefileArtifacts:
|
|
"""Pure async I/O: extract readability and screenshot from a singlefile. No DB."""
|
|
# Get URL first because it does some simple checking of `singlefile`.
|
|
url = get_url_from_singlefile(singlefile)
|
|
singlefile = singlefile.encode() if isinstance(singlefile, str) else singlefile
|
|
|
|
# A compressed (SingleFileZ) singlefile contains the uncompressed page as index.html; use that
|
|
# for readability and screenshot, but keep the compressed bytes as the singlefile.
|
|
compressed = is_compressed_singlefile(singlefile)
|
|
extractable = get_compressed_singlefile_index(singlefile) if compressed else singlefile
|
|
|
|
if DOCKERIZED:
|
|
# Perform the archive in the Archive docker container. (Typically in the development environment).
|
|
encoded = base64.b64encode(singlefile).decode()
|
|
logger.debug(f'extract_singlefile_artifacts sending to archive service: {url}')
|
|
singlefile, readability, screenshot = await request_archive(url, singlefile=encoded, compress=compressed)
|
|
singlefile = singlefile.encode() if isinstance(singlefile, str) else singlefile
|
|
else:
|
|
# JSON from readability-extractor
|
|
readability = dict()
|
|
try:
|
|
logger.debug(f'extract_singlefile_artifacts extracting readability: {url}')
|
|
readability = await html_to_readability(extractable, url)
|
|
except RuntimeError as e:
|
|
logger.error(f'Failed to extract readability from: {url}', exc_info=e)
|
|
|
|
screenshot = None
|
|
try:
|
|
logger.debug(f'extract_singlefile_artifacts creating screenshot: {url}')
|
|
if compressed:
|
|
# Extract the whole ZIP so the page's relative resources (images, fonts) render.
|
|
with tempfile.TemporaryDirectory(prefix='wrolpi-singlefile-screenshot-') as tmp_dir:
|
|
with zipfile.ZipFile(io.BytesIO(singlefile)) as zip_:
|
|
zip_.extractall(tmp_dir)
|
|
screenshot = html_file_screenshot(pathlib.Path(tmp_dir) / COMPRESSED_SINGLEFILE_INDEX)
|
|
else:
|
|
screenshot = html_screenshot(extractable)
|
|
except Exception as e:
|
|
logger.error(f'Failed to extract screenshot from: {url}', exc_info=e)
|
|
|
|
return SinglefileArtifacts(url=url, singlefile=singlefile, readability=readability, screenshot=screenshot)
|
|
|
|
|
|
async def singlefile_to_archive(singlefile: bytes, destination: pathlib.Path = None) -> Archive:
|
|
"""
|
|
Convert a SingleFile to an Archive.
|
|
|
|
Thin wrapper around extract_singlefile_artifacts + write_archive_files + register_archive
|
|
that opens its own transaction. New code should call the helpers directly so the caller
|
|
controls the session.
|
|
|
|
Args:
|
|
singlefile: The SingleFile HTML bytes
|
|
destination: Optional directory to save files to. If provided, files go here instead of archive/<domain>
|
|
"""
|
|
artifacts = await extract_singlefile_artifacts(singlefile)
|
|
logger.trace(f'singlefile_to_archive modeling: {artifacts.url}')
|
|
written = write_archive_files(artifacts.url, artifacts.singlefile, artifacts.readability,
|
|
artifacts.screenshot, destination=destination)
|
|
with get_db_session(commit=True) as session:
|
|
archive = register_archive(session, written)
|
|
return archive
|
|
|
|
|
|
async def get_statistics():
|
|
"""Get statistics about Archives and Domain Collections."""
|
|
with get_db_curs() as curs:
|
|
curs.execute('''
|
|
SELECT
|
|
-- total archives
|
|
COUNT(a.id) AS "archives",
|
|
-- total archives downloaded over the past week/month/year
|
|
COUNT(a.id) FILTER (WHERE fg.download_datetime >= current_date - interval '1 week') AS "week",
|
|
COUNT(a.id)
|
|
FILTER (WHERE fg.download_datetime >= current_date - interval '1 month') AS "month",
|
|
COUNT(a.id) FILTER (WHERE fg.download_datetime >= current_date - interval '1 year') AS "year",
|
|
-- sum of all archive file sizes
|
|
COALESCE(SUM(fg.size), 0)::BIGINT AS "sum_size",
|
|
-- largest archive
|
|
COALESCE(MAX(fg.size), 0) AS "max_size"
|
|
FROM archive a
|
|
LEFT JOIN file_group fg on a.file_group_id = fg.id
|
|
''')
|
|
archive_stats = dict(curs.fetchone())
|
|
|
|
# Get the total archives downloaded every month for the past two years.
|
|
curs.execute('''
|
|
SELECT DATE_TRUNC('month', months.a),
|
|
COUNT(archive.id)::BIGINT,
|
|
SUM(size)::BIGINT AS "size"
|
|
FROM generate_series(
|
|
date_trunc('month', current_date) - interval '2 years',
|
|
date_trunc('month', current_date) - interval '1 month',
|
|
'1 month'::interval) AS months(a),
|
|
archive
|
|
LEFT JOIN file_group fg on archive.file_group_id = fg.id
|
|
WHERE fg.download_datetime >= date_trunc('month', months.a)
|
|
AND fg.download_datetime < date_trunc('month', months.a) + interval '1 month'
|
|
AND fg.download_datetime IS NOT NULL
|
|
GROUP BY 1
|
|
ORDER BY 1
|
|
''')
|
|
monthly_archives = [dict(i) for i in curs.fetchall()]
|
|
|
|
historical_stats = dict(monthly_archives=monthly_archives)
|
|
historical_stats['average_count'] = (sum(i['count'] for i in monthly_archives) // len(monthly_archives)) \
|
|
if monthly_archives else 0
|
|
historical_stats['average_size'] = (sum(i['size'] for i in monthly_archives) // len(monthly_archives)) \
|
|
if monthly_archives else 0
|
|
|
|
curs.execute('''
|
|
SELECT COUNT(c.id) AS "domains",
|
|
COUNT(c.id) FILTER ( WHERE col.tag_id IS NOT NULL ) AS "tagged_domains"
|
|
FROM collection c
|
|
LEFT JOIN collection col on col.id = c.id
|
|
WHERE c.kind = 'domain'
|
|
''')
|
|
domain_stats = dict(curs.fetchone())
|
|
ret = dict(statistics=dict(
|
|
archives=archive_stats,
|
|
domains=domain_stats,
|
|
historical=historical_stats,
|
|
))
|
|
return ret
|
|
|
|
|
|
async def generate_archive_screenshot(archive_id: int) -> pathlib.Path:
|
|
"""
|
|
Generate a screenshot for an existing Archive that doesn't have one.
|
|
If the Archive already has a screenshot, verify it exists and ensure it's tracked in the FileGroup.
|
|
|
|
Returns the path to the generated screenshot.
|
|
|
|
Raises:
|
|
ValueError: If Archive has no singlefile
|
|
RuntimeError: If screenshot generation fails
|
|
"""
|
|
with get_db_session() as session:
|
|
archive = Archive.find_by_id(session, archive_id)
|
|
|
|
if not archive.singlefile_path:
|
|
raise ValueError(f'Cannot generate screenshot for Archive {archive_id}: no singlefile')
|
|
|
|
# Check if screenshot already exists
|
|
if archive.screenshot_path:
|
|
# Verify the screenshot file actually exists on disk
|
|
if archive.screenshot_path.is_file():
|
|
logger.info(f'Archive {archive_id} already has a screenshot, ensuring it is tracked')
|
|
# Ensure the screenshot is tracked in FileGroup.files and FileGroup.data
|
|
with get_db_session(commit=True) as tracking_session:
|
|
archive = Archive.find_by_id(tracking_session, archive_id)
|
|
file_group = archive.file_group
|
|
# append_files uses unique_by_predicate, so this is safe even if already tracked
|
|
file_group.append_files(archive.screenshot_path)
|
|
|
|
# Also update FileGroup.data (same pattern as set_screenshot)
|
|
data = dict(file_group.data) if file_group.data else {}
|
|
data['screenshot_path'] = str(archive.screenshot_path)
|
|
file_group.data = data
|
|
|
|
archive.validate()
|
|
tracking_session.flush()
|
|
return archive.screenshot_path
|
|
else:
|
|
logger.warning(f'Archive {archive_id} has screenshot_path but file does not exist, regenerating')
|
|
|
|
singlefile_path = archive.singlefile_path
|
|
url = archive.file_group.url
|
|
|
|
# Request screenshot from Archive docker service or generate locally
|
|
if DOCKERIZED:
|
|
logger.debug(f'Requesting screenshot from archive service for Archive {archive_id}')
|
|
screenshot_bytes = await request_screenshot(url, singlefile_path)
|
|
else:
|
|
logger.debug(f'Generating screenshot locally for Archive {archive_id}')
|
|
singlefile_contents = singlefile_path.read_bytes()
|
|
if is_compressed_singlefile(singlefile_contents):
|
|
# Extract the whole ZIP so the page's relative resources (images, fonts) render.
|
|
with tempfile.TemporaryDirectory(prefix='wrolpi-singlefile-screenshot-') as tmp_dir:
|
|
with zipfile.ZipFile(io.BytesIO(singlefile_contents)) as zip_:
|
|
zip_.extractall(tmp_dir)
|
|
screenshot_bytes = html_file_screenshot(pathlib.Path(tmp_dir) / COMPRESSED_SINGLEFILE_INDEX)
|
|
else:
|
|
screenshot_bytes = html_screenshot(singlefile_contents)
|
|
|
|
if not screenshot_bytes:
|
|
raise RuntimeError(f'Failed to generate screenshot for Archive {archive_id}')
|
|
|
|
# Save screenshot next to singlefile with same naming pattern
|
|
screenshot_path = singlefile_path.with_suffix('.png')
|
|
screenshot_path.write_bytes(screenshot_bytes)
|
|
logger.info(f'Generated screenshot for Archive {archive_id}: {screenshot_path}')
|
|
|
|
# Update the Archive to include the new screenshot file
|
|
with get_db_session(commit=True) as session:
|
|
archive = Archive.find_by_id(session, archive_id)
|
|
archive.set_screenshot(screenshot_path)
|
|
session.flush()
|
|
|
|
return screenshot_path
|