import asyncio import functools import json import pathlib import re from datetime import datetime from typing import Callable, List, Tuple, Dict import cachetools.func from libzim import Entry from sqlalchemy import or_ from sqlalchemy.orm import Session from sqlalchemy.orm.exc import NoResultFound # noqa from modules.zim import kiwix from modules.zim.errors import UnknownZim, UnknownZimSubscription from modules.zim.kiwix import KIWIX_CATALOG from modules.zim.models import Zim, Zims, TagZimEntry, ZimSubscription from wrolpi import flags from wrolpi.cmd import run_command from wrolpi.common import register_modeler, logger, extract_html_text, extract_headlines, get_media_directory, walk, \ register_refresh_cleanup, background_task, get_wrolpi_config, unique_by_predicate, aiohttp_post from wrolpi.switches import register_switch_handler from wrolpi.dates import Seconds 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.worker import file_worker from wrolpi.files.models import FileGroup from wrolpi.vars import PYTEST, DOCKERIZED logger = logger.getChild(__name__) CONTROLLER_SERVICE = 'http://controller' __all__ = [ 'flag_outdated_zim_files', 'get_all_entries_tags', 'get_entries_tags', 'get_entry', 'get_kiwix_catalog', 'get_unique_paths', 'get_zim', 'get_zims', 'model_zim', 'restart_kiwix_handler', 'zim_download_url_to_name', 'zim_modeler', ] @register_switch_handler('restart_kiwix') async def restart_kiwix_handler(): """Switch handler that restarts the Kiwix service after zim files are added.""" await restart_kiwix() def model_zim(file_group: FileGroup, session: Session) -> Zim: """Create a Zim record for the provided FileGroup.""" # Check if a Zim already exists for this FileGroup zim = session.query(Zim).filter_by(file_group_id=file_group.id).one_or_none() if not zim: # Create new Zim if it doesn't exist zim = Zim(file_group=file_group, path=file_group.primary_path) session.add(zim) # Trigger Kiwix restart so the new zim file is available for browsing. restart_kiwix_handler.activate_switch() 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' return 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) def get_all_entries_tags(): """Returns a dict of dicts which contain the names of all Tags of each Entry's path. Example: { 1: { 'path1': ['tag1', 'tag2'], }, 2: { 'path2': ['tag1'], }, """ stmt = ''' SELECT tz.zim_id, tz.zim_entry, json_group_array(t.name) FROM tag t LEFT JOIN tag_zim tz on t.id = tz.tag_id GROUP BY 1, 2 \ ''' with get_db_curs() as curs: curs.execute(stmt) entries = dict() for zim_id, zim_entry, tag_names in curs.fetchall(): if not zim_id: continue tag_names = json.loads(tag_names) try: entries[zim_id][zim_entry] = tag_names except KeyError: entries[zim_id] = {zim_entry: tag_names} return entries def get_entries_tags(paths: List[str], zim_id: int): all_entries_tags = get_all_entries_tags() entries = dict() for path in paths: zim_tag_entries = all_entries_tags.get(zim_id, dict()) entries[path] = zim_tag_entries.get(path, list()) return entries def get_unique_paths(*paths: str) -> List[str]: """Return a new list which contains unique Zim paths.""" unique_paths = unique_by_predicate( paths, lambda i: i.replace('-', '').replace('_', '').replace('/', '').lower() ) return list(unique_paths) def headline_zim(session: Session, search_str: str, zim_id: int, tag_names: List[str] = None, offset: int = 0, limit: int = 10) -> Dict: zim_id = int(zim_id) zim = get_zim(session, zim_id) if tag_names: # TODO what is the count goes over the limit? entries = zim.entries_with_tags(session, tag_names, offset=offset, limit=1_000) estimate = 0 elif search_str: # Get suggested results (this searches titles). results = zim.suggest(search_str, offset, limit) search_results = zim.search(search_str, offset, limit) # Remove duplicate entries (disambiguation or redirect). results = get_unique_paths(*results, *search_results) results = results[:limit] entries = [zim.find_entry(i) for i in results] estimate = zim.estimate(search_str) else: raise RuntimeError('Must provide search_str or tag_names.') articles = [bytes(i.get_item().content).decode('UTF-8') for i in entries] articles = [extract_html_text(i) for i in articles] entries_tag_names = get_entries_tags([i.path for i in entries], zim_id) # Always get headlines even if there is no search_str. Non-matches get the leading text of each article. headlines = extract_headlines(articles, search_str) search = list() if search_str: # `search_str` allows us to headline the title. titles = [i.title for i in entries] headline_titles = extract_headlines(titles, search_str) for entry in entries: title_headline, title_rank = headline_titles.pop(0) headline, rank = headlines.pop(0) if rank > 0 or title_rank > 0: # The result for each Entry. path = entry.path entry_tag_names = entries_tag_names[path] if path in entries_tag_names else list() search.append(dict( headline=headline, path=path, rank=max(rank, title_rank), tag_names=entry_tag_names, title=title_headline, )) if tag_names: estimate += 1 else: # No `search_str`, we return all entries. for entry in entries: headline, _ = headlines.pop(0) path = entry.path entry_tag_names = entries_tag_names[path] if path in entries_tag_names else list() search.append(dict( headline=headline, path=path, rank=0, tag_names=entry_tag_names, title=entry.title, )) estimate = len(entries) results = dict( path=zim.path, metadata=zim.zim_metadata, search=search, estimate=estimate, ) return results def search_all_zims(session: Session, search_str: str, tag_names: List[str] = None, offset: int = 0, limit: int = 10) -> List[Dict]: """Search all Zim files that have auto_search enabled. Returns a list of per-zim result dicts, each containing metadata and search results. Each search entry includes `zim_id` so callers can fetch the full entry afterward.""" zims = Zims.get_all(session) all_results = [] for zim in zims: if not zim.auto_search: continue try: result = headline_zim(session, search_str, zim.id, tag_names=tag_names, offset=offset, limit=limit) if result.get('search'): for entry in result['search']: entry['zim_id'] = zim.id all_results.append(result) except Exception as e: logger.error(f'Failed to search Zim {zim}', exc_info=e) return all_results @cachetools.func.mru_cache(maxsize=1_000) def get_estimates(search_str: str) -> List[int]: zims = Zims.get_all() estimates = [i.estimate(search_str) for i in zims] return estimates def get_zim(session: Session, zim_id: int) -> Zim: """Return the Zim record of the provided path. @raise UnknownZim: When no record exists with the id. @warning: Results are cached.""" try: zim: Zim = session.query(Zim) \ .join(FileGroup, FileGroup.id == Zim.file_group_id) \ .filter(Zim.id == zim_id) \ .one() return zim except Exception as e: raise UnknownZim() from e @functools.lru_cache(maxsize=1_000) def get_entry(path: str, zim_id: int) -> Entry: with get_db_session() as session: zim = get_zim(session, zim_id) entry = zim.find_entry(path) return entry def get_zims(session: Session, ids: List[int] = None) -> List[Zim]: if ids: zims = session.query(Zim).filter(Zim.id.in_(ids)).all() else: zims = Zims.get_all(session) return zims def delete_zims(session: Session, ids: List[int]): zims = get_zims(session, ids) for zim in zims: zim.delete() async def add_tag(session: Session, tag_name: str, zim_id: int, zim_entry: str) -> TagZimEntry: zim = session.query(Zim).filter_by(id=zim_id).one() tag_zim_entry = zim.tag_entry(session, tag_name, zim_entry) session.commit() return tag_zim_entry async def untag(session: Session, tag_name: str, zim_id: int, zim_entry: str): zim: Zim = session.query(Zim).filter_by(id=zim_id).one() zim.untag_entry(session, tag_name, zim_entry) session.commit() def get_kiwix_catalog(): return KIWIX_CATALOG.copy() def get_kiwix_subscriptions(session: Session) -> Dict[str, ZimSubscription]: subscriptions = session.query(ZimSubscription) subscriptions_by_name = {i.name: i for i in subscriptions} return subscriptions_by_name async def subscribe(session: Session, name: str, language: str, frequency: int = DownloadFrequency.days180) -> ZimSubscription: subscription = ZimSubscription.get_or_create(session, name) subscription.language = language if name == 'Wikipedia (with images)': url = f'https://download.kiwix.org/zim/wikipedia/wikipedia_{language}_all_maxi_' elif name == 'Wikipedia (mini)': url = f'https://download.kiwix.org/zim/wikipedia/wikipedia_{language}_all_mini_' elif name == 'Wikipedia (no images)': url = f'https://download.kiwix.org/zim/wikipedia/wikipedia_{language}_all_nopic_' elif name == 'Wikiversity (with images)': url = f'https://download.kiwix.org/zim/wikiversity/wikiversity_{language}_all_maxi_' elif name == 'Wikiversity (no images)': url = f'https://download.kiwix.org/zim/wikiversity/wikiversity_{language}_all_nopic_' elif name == 'Wikitionary (no images)': url = f'https://download.kiwix.org/zim/wiktionary/wiktionary_{language}_all_nopic_' elif name == 'Wikitionary (with images)': url = f'https://download.kiwix.org/zim/wikiversity/wikiversity_{language}_all_maxi_' elif name == 'Wikibooks (no images)': url = f'https://download.kiwix.org/zim/wikibooks/wikibooks_{language}_all_nopic_' elif name == 'Wikibooks (with images)': url = f'https://download.kiwix.org/zim/wikibooks/wikibooks_{language}_all_maxi_' elif name == 'Wikisource': url = f'https://download.kiwix.org/zim/wikisource/wikisource_{language}_all_maxi_' elif name == 'Raspberry Pi (Stack Exchange)': url = f'https://download.kiwix.org/zim/stack_exchange/raspberrypi.stackexchange.com_{language}_all_' elif name == 'Vikidia': url = f'https://download.kiwix.org/zim/vikidia/vikidia_{language}_all_maxi_' elif name == 'Stackoverflow (Stack Exchange)': if language == 'en': url = f'https://download.kiwix.org/zim/stack_exchange/stackoverflow.com_{language}_all_' else: url = f'https://download.kiwix.org/zim/stack_exchange/{language}.stackoverflow.com_{language}_all_' elif name == 'Superuser (Stack Exchange)': url = f'https://download.kiwix.org/zim/stack_exchange/superuser.com_{language}_all_' elif name == 'iFixit': url = f'https://download.kiwix.org/zim/ifixit/ifixit_{language}_all_' elif name == 'Gutenberg': url = f'https://download.kiwix.org/zim/gutenberg/gutenberg_{language}_all_' elif name == 'Amateur Radio (Stack Exchange)': url = f'https://download.kiwix.org/zim/stack_exchange/ham.stackexchange.com_{language}_all_' elif name == 'DIY (Stack Exchange)': url = f'https://download.kiwix.org/zim/stack_exchange/diy.stackexchange.com_{language}_all_' elif name == 'Electronics (Stack Exchange)': url = f'https://download.kiwix.org/zim/stack_exchange/electronics.stackexchange.com_{language}_all_' elif name == 'Unix (Stack Exchange)': url = f'https://download.kiwix.org/zim/stack_exchange/unix.stackexchange.com_{language}_all_' elif name == 'Ask Ubuntu (Stack Exchange)': url = f'https://download.kiwix.org/zim/stack_exchange/askubuntu.com_{language}_all_' elif name == 'Bicycles (Stack Exchange)': url = f'https://download.kiwix.org/zim/stack_exchange/bicycles.stackexchange.com_{language}_all_' elif name == 'Biology (Stack Exchange)': url = f'https://download.kiwix.org/zim/stack_exchange/biology.stackexchange.com_{language}_all_' elif name == 'Arduino (Stack Exchange)': url = f'https://download.kiwix.org/zim/stack_exchange/arduino.stackexchange.com_{language}_all_' elif name == '3D Printing (Stack Exchange)': url = f'https://download.kiwix.org/zim/stack_exchange/3dprinting.stackexchange.com_{language}_all_' elif name == 'MD Wiki (Wiki Project Med Foundation)': url = f'https://download.kiwix.org/zim/other/mdwiki_{language}_all_' elif name == 'WikiMed Medical Encyclopedia': url = f'https://download.kiwix.org/zim/wikipedia/wikipedia_{language}_medicine_' elif name == 'Post Disaster Resource Library': url = f'https://download.kiwix.org/zim/other/zimgit-post-disaster_{language}_' elif name == 'Military Medicine': url = f'https://download.kiwix.org/zim/zimit/fas-military-medicine_{language}_' elif name == 'WikEm (Global Emergency Medicine Wiki)': url = f'https://download.kiwix.org/zim/other/wikem_{language}_all_maxi_' elif name == 'Health (Stack Exchange)': url = f'https://download.kiwix.org/zim/stack_exchange/health.stackexchange.com_{language}_all_' elif name == 'Khan Academy': url = f'https://download.kiwix.org/zim/other/khanacademy_{language}_all_' else: raise ValueError(f'{name} not a valid Kiwix subscription name!') if language not in kiwix.KIWIX_CATALOG_BY_NAME[name]['languages']: raise ValueError(f'Language {repr(str(language))} is invalid for {name}') subscription.change_download(session, url, frequency) session.add(subscription) session.flush([subscription]) session.commit() return subscription async def unsubscribe(session: Session, subscription_id: int): subscription: ZimSubscription = session.query(ZimSubscription).filter_by(id=subscription_id).one_or_none() if subscription: session.delete(subscription) if subscription.download: session.delete(subscription.download) session.commit() else: raise UnknownZimSubscription(f'{subscription_id=}') KIWIX_URL_PARSER = re.compile(r'https:\/\/download\.kiwix\.org\/zim\/(.+?)\/(.+?)[_-](\w{2,3})([_-].+)') KIWIX_URL_NO_FLAVOR_PARSER = re.compile(r'https:\/\/download\.kiwix\.org\/zim\/(.+?)\/(.+)[_-]([a-zA-Z]{2,3})') def zim_download_url_to_name(url: str) -> Tuple[str, str]: try: _, project, language, flavor = KIWIX_URL_PARSER.match(url).groups() except AttributeError: # May not have a "flavor". _, project, language = KIWIX_URL_NO_FLAVOR_PARSER.match(url).groups() flavor = None if project == 'wikipedia' and flavor == '_all_maxi_': name = 'Wikipedia (with images)' elif project == 'wikipedia' and flavor == '_all_nopic_': name = 'Wikipedia (no images)' elif project == 'wikipedia' and flavor == '_all_mini_': name = 'Wikipedia (mini)' elif project == 'wikiversity' and flavor == '_all_maxi_': name = 'Wikiversity (with images)' elif project == 'wikiversity' and flavor == '_all_nopic_': name = 'Wikiversity (no images)' elif project == 'wiktionary' and flavor == '_all_maxi_': name = 'Wikitionary (with images)' elif project == 'wiktionary' and flavor == '_all_nopic_': name = 'Wikitionary (no images)' elif project == 'wikibooks' and flavor == '_all_maxi_': name = 'Wikibooks (with images)' elif project == 'wikibooks' and flavor == '_all_nopic_': name = 'Wikibooks (no images)' elif project == 'wikisource' and flavor == '_all_maxi_': name = 'Wikisource' elif project == 'raspberrypi.stackexchange.com' and flavor == '_all_': name = 'Raspberry Pi (Stack Exchange)' elif project == 'vikidia' and flavor == '_all_maxi_': name = 'Vikidia' elif 'stackoverflow.com' in project and flavor == '_all_': name = 'Stackoverflow (Stack Exchange)' elif project == 'superuser.com' and flavor == '_all_': name = 'Superuser (Stack Exchange)' elif project == 'ifixit' and flavor == '_all_': name = 'iFixit' elif project == 'gutenberg' and flavor == '_all_': name = 'Gutenberg' elif project == 'ham.stackexchange.com' and flavor == '_all_': name = 'Amateur Radio (Stack Exchange)' elif project == 'diy.stackexchange.com' and flavor == '_all_': name = 'DIY (Stack Exchange)' elif project == 'electronics.stackexchange.com' and flavor == '_all_': name = 'Electronics (Stack Exchange)' elif project == 'unix.stackexchange.com' and flavor == '_all_': name = 'Unix (Stack Exchange)' elif project == 'askubuntu.com' and flavor == '_all_': name = 'Ask Ubuntu (Stack Exchange)' elif project == 'bicycles.stackexchange.com' and flavor == '_all_': name = 'Bicycles (Stack Exchange)' elif project == 'biology.stackexchange.com' and flavor == '_all_': name = 'Biology (Stack Exchange)' elif project == 'arduino.stackexchange.com' and flavor == '_all_': name = 'Arduino (Stack Exchange)' elif project == '3dprinting.stackexchange.com' and flavor == '_all_': name = '3D Printing (Stack Exchange)' elif project == 'mdwiki' and flavor == '_all_': name = 'MD Wiki (Wiki Project Med Foundation)' elif project == 'wikipedia' and flavor == '_medicine_': name = 'WikiMed Medical Encyclopedia' elif project == 'zimgit-post-disaster' and flavor is None: name = 'Post Disaster Resource Library' elif project == 'fas-military-medicine' and flavor is None: name = 'Military Medicine' elif project == 'wikem' and flavor == '_all_maxi_': name = 'WikEm (Global Emergency Medicine Wiki)' elif project == 'health.stackexchange.com' and flavor == '_all_': name = 'Health (Stack Exchange)' elif project == 'khanacademy' and flavor == '_all_': name = 'Khan Academy' else: raise RuntimeError(f'Could not find name for Zim URL: {url} {project=} {flavor=}') if language not in kiwix.KIWIX_CATALOG_BY_NAME[name]['languages']: raise ValueError(f'Language {repr(str(language))} is invalid for {name}') return name, language def get_zim_directory() -> pathlib.Path: zim_directory = get_media_directory() / get_wrolpi_config().zims_destination return zim_directory async def check_zim(path: pathlib.Path) -> int: """ Runs `zimcheck` against the Zim file. Zim is valid only if this returns `0`. @warning: Only performs the checksum check! """ cmd = ('zimcheck', '-C', path.absolute()) cmd_str = " ".join(str(i) for i in cmd) logger.debug(f'check_zim: {cmd_str}') result = await run_command(cmd) logger.debug(f'zimcheck returned {result.return_code}') if result.return_code != 0 and result.stderr: logger.debug(result.stderr) return result.return_code ZIM_NAME_PARSER = re.compile(r'(.+?)_(\d{4}-\d{2}).zim') def parse_name(path: pathlib.Path) -> Tuple[str, datetime]: match = ZIM_NAME_PARSER.match(path.name) if not match: raise ValueError(f'Zim file name does not match expected pattern: {path.name}') name, date = match.groups() date = datetime.strptime(date, '%Y-%m') return name, date def find_outdated_zim_files(path: pathlib.Path = None) -> Tuple[List[pathlib.Path], List[pathlib.Path]]: """Search the Zim directory for outdated Zim files. Returns a list of outdated, and a list of current Zim files.""" path = path or get_zim_directory() if not path.is_dir(): logger.error(f'Cannot search for outdated Zim files because the directory does not exist: {path}') return list(), list() # Find all non-empty Zim files. files = [i for i in walk(path) if i.is_file() and i.suffix == '.zim' and i.stat().st_size] zims = list() for file in files: try: name, date = parse_name(file) except Exception: logger.warning(f'Found Zim {file} but its name does not match.') continue zims.append((name, date, file)) # Order by newest first. zims = sorted(zims, key=lambda i: i[1], reverse=True) outdated_files = list() current_files = list() newest_names = list() for name, date, file in zims: if name in newest_names: # Already saw this name, this date is older. outdated_files.append(file) else: newest_names.append(name) current_files.append(file) return sorted(outdated_files), sorted(current_files) async def remove_outdated_zim_files(path: pathlib.Path = None) -> int: """Deletes all old Zim files.""" outdated, _ = find_outdated_zim_files(path) logger.info(f'Deleting {len(outdated)} outdated Zim files: {outdated}') deleted_count = 0 for file in outdated: file.unlink() deleted_count += 1 if deleted_count: # Refresh synchronously to clean up DB before returning. await file_worker.refresh_sync(list(outdated)) # Trigger Kiwix restart so the deleted zim files are no longer served. restart_kiwix_handler.activate_switch() return deleted_count @register_refresh_cleanup def flag_outdated_zim_files(): """Set the `outdated_zims` Flag if outdated Zim files can be found, otherwise clear it.""" outdated, _ = find_outdated_zim_files() if outdated: logger.info('Outdated Zims were found') flags.outdated_zims.set() else: logger.info('Outdated Zims were not found') flags.outdated_zims.clear() async def restart_kiwix(): if PYTEST: logger.warning('Unable to restart Kiwix serve because we are testing') return if DOCKERIZED: logger.info('Restarting Kiwix via controller API') try: async with aiohttp_post(f'{CONTROLLER_SERVICE}/api/services/zim/restart', None, timeout=Seconds.minute) as response: if response.status == 200: logger.info('Successfully restarted Kiwix container') return 0 else: logger.error(f'Failed to restart Kiwix: {response.status}') Events.send_user_notify( 'Failed to restart Kiwix automatically. Please restart manually: docker compose restart zim') return 1 except Exception as e: logger.error(f'Failed to restart Kiwix: {e}') Events.send_user_notify( 'Failed to restart Kiwix automatically. Please restart manually: docker compose restart zim') return 1 logger.info('Restarting Kiwix serve') cmd = ('sudo', '/usr/bin/systemctl', 'restart', 'wrolpi-kiwix.service') result = await run_command(cmd) logger.debug(f'systemctl returned {result.return_code}') if result.return_code != 0 and result.stderr: logger.debug(result.stderr) return result.return_code def set_zim_auto_search(session: Session, zim_id: int, auto_search: bool): zim = Zim.find_by_id(session, zim_id) zim.auto_search = auto_search session.commit()