Migrating Kiwix downloaders.

This commit is contained in:
Roland 2026-05-06 19:39:07 -06:00
parent a2be7a3d8b
commit f2dc249e01
No known key found for this signature in database
GPG key ID: 1BA620CE6C9A4680
3 changed files with 272 additions and 23 deletions

View file

@ -1,11 +1,20 @@
import pathlib
from dataclasses import dataclass
from http import HTTPStatus
from typing import List
from typing import List, Optional
from sqlalchemy.orm import Session
from modules.zim import lib
from wrolpi.common import logger, aiohttp_get, get_html_soup
from wrolpi.downloader import Downloader, Download, DownloadResult
from wrolpi.downloader import Downloader, Download, DownloadContext, DownloadResult
__all__ = ['KiwixCatalogDownloader', 'KiwixZimDownloader', 'kiwix_zim_downloader', 'kiwix_catalog_downloader']
__all__ = [
'KiwixCatalogDownloader', 'KiwixZimDownloader',
'PreparedKiwixCatalog', 'ExecutedKiwixCatalog',
'PreparedKiwixZim', 'ExecutedKiwixZim',
'kiwix_zim_downloader', 'kiwix_catalog_downloader',
]
from wrolpi.files.lib import upsert_file
from wrolpi.vars import DOWNLOAD_USER_AGENT
@ -48,6 +57,33 @@ async def fetch_hrefs(url: str) -> List[str]:
return downloads
# ---------------------------------------------------------------------------
# KiwixCatalogDownloader
# ---------------------------------------------------------------------------
@dataclass
class PreparedKiwixCatalog:
"""Plan produced by KiwixCatalogDownloader.prepare_download.
Splits the URL into the parent listing URL and the basename prefix used to match
Zim files in that listing.
"""
url: str
parent_url: str
name: str
@dataclass
class ExecutedKiwixCatalog:
"""Output of KiwixCatalogDownloader.execute_download.
`latest_zim_url` is None when the listing contained no matching files; finalize
converts that into a failure DownloadResult.
"""
latest_zim_url: Optional[str]
class KiwixCatalogDownloader(Downloader):
"""
Downloads the catalog of Zim files from Kiwix. Searches for the latest Zim that matches the provided URL.
@ -58,59 +94,105 @@ class KiwixCatalogDownloader(Downloader):
name = 'kiwix_catalog'
listable = False
async def do_download(self, download: Download) -> DownloadResult:
def prepare_download(self, session: Session, download: Download) -> PreparedKiwixCatalog:
"""Split the download URL into the parent listing URL and the basename prefix."""
url = download.url
# Get the parent directory of the file to be downloaded
*parents, name = url.split('/')
parent_url = '/'.join(parents)
return PreparedKiwixCatalog(url=url, parent_url=parent_url, name=name)
# Find the latest Zim that has the correct name.
downloadable_files = await fetch_hrefs(parent_url)
matching_zims = sorted([i for i in downloadable_files if i.startswith(name) and i.endswith('.zim')])
async def execute_download(self, prepared: PreparedKiwixCatalog, ctx: DownloadContext,
download: Download = None) -> ExecutedKiwixCatalog:
"""Fetch the parent directory listing and pick the latest matching Zim."""
downloadable_files = await fetch_hrefs(prepared.parent_url)
matching_zims = sorted([i for i in downloadable_files
if i.startswith(prepared.name) and i.endswith('.zim')])
if not matching_zims:
return DownloadResult(success=False, error=f'No Zim files match {str(repr(url))}')
return ExecutedKiwixCatalog(latest_zim_url=None)
logger.debug(f'Found matching Zims: {matching_zims}')
latest_zim = f'{parent_url}/{matching_zims[-1]}'
return ExecutedKiwixCatalog(latest_zim_url=f'{prepared.parent_url}/{matching_zims[-1]}')
return DownloadResult(
success=True,
downloads=[latest_zim, ],
)
def finalize_download(self, session: Session, download: Download,
executed: ExecutedKiwixCatalog) -> DownloadResult:
if executed.latest_zim_url is None:
return DownloadResult(success=False, error=f'No Zim files match {str(repr(download.url))}')
return DownloadResult(success=True, downloads=[executed.latest_zim_url])
kiwix_catalog_downloader = KiwixCatalogDownloader()
# ---------------------------------------------------------------------------
# KiwixZimDownloader
# ---------------------------------------------------------------------------
@dataclass
class PreparedKiwixZim:
"""Plan produced by KiwixZimDownloader.prepare_download: zim_directory exists on disk."""
url: str
zim_directory: pathlib.Path
@dataclass
class ExecutedKiwixZim:
"""Output of KiwixZimDownloader.execute_download.
`is_valid` is False when zimcheck flagged the file as corrupt; finalize turns that
into a failure DownloadResult. When zimcheck is not installed (FileNotFoundError),
we treat the file as valid (matching legacy behaviour).
"""
output_path: pathlib.Path
is_valid: bool
class KiwixZimDownloader(Downloader):
"""Downloads a Zim file to the zim directory."""
name = 'kiwix_zim'
listable = False
async def do_download(self, download: Download) -> DownloadResult:
def prepare_download(self, session: Session, download: Download) -> PreparedKiwixZim:
zim_directory = lib.get_zim_directory()
zim_directory.mkdir(parents=True, exist_ok=True)
return PreparedKiwixZim(url=download.url, zim_directory=zim_directory)
output_path = await self.download_file(download, download.url, zim_directory)
async def execute_download(self, prepared: PreparedKiwixZim, ctx: DownloadContext,
download: Download = None) -> ExecutedKiwixZim:
"""Download the Zim file, flag any outdated zims, validate via zimcheck, and
register the new file via upsert_file when valid.
Note: lib.flag_outdated_zim_files (sets a Sanic flag) and upsert_file (opens its
own DB session) are pre-existing side effects we keep on the post-download path.
Splitting upsert_file is a future PR shared with the FileDownloader migration.
"""
download = download if download is not None else Download(url=prepared.url)
output_path = await self.download_file(download, prepared.url, prepared.zim_directory, ctx=ctx)
# Notify the maintainer if outdated Zim files are lying around.
lib.flag_outdated_zim_files()
is_valid = True
try:
return_code = await lib.check_zim(output_path)
if return_code > 0:
# zimcheck ran, but the file is invalid.
return DownloadResult(success=False, error=f'Zim file is invalid')
is_valid = False
except FileNotFoundError:
logger.warning(f'Not validating {output_path} because zimcheck is not installed')
# Add the new Zim to the FileGroups, model it.
# The model_zim() function triggers a Kiwix restart via switch when a new Zim is created.
await upsert_file(output_path)
if is_valid:
# Add the new Zim to the FileGroups, model it.
# The model_zim() function triggers a Kiwix restart via switch when a new Zim is created.
await upsert_file(output_path)
logger.info(f'Successfully downloaded Zim {prepared.url} to {output_path}')
logger.info(f'Successfully downloaded Zim {download.url} to {output_path}')
return ExecutedKiwixZim(output_path=output_path, is_valid=is_valid)
def finalize_download(self, session: Session, download: Download,
executed: ExecutedKiwixZim) -> DownloadResult:
if not executed.is_valid:
return DownloadResult(success=False, error='Zim file is invalid')
# Location is just the generic Kiwix viewer.
return DownloadResult(success=True, location='/zim')

View file

@ -0,0 +1,166 @@
"""Phase-split tests for KiwixCatalogDownloader and KiwixZimDownloader.
These exercise prepare_download / finalize_download (and execute_download for the
catalog, where it's pure HTTP) in isolation — no async_client, no test_download_manager,
no real subprocess. The Zim downloader's execute phase calls into download_file,
flag_outdated_zim_files, check_zim (subprocess), and upsert_file (own DB session); a
unit test of that full pipeline would effectively be an integration test, so it's
deliberately omitted.
"""
from unittest import mock
import pytest
from modules.archive.conftest import make_test_ctx
from modules.zim.downloader import (
ExecutedKiwixCatalog,
ExecutedKiwixZim,
PreparedKiwixCatalog,
PreparedKiwixZim,
kiwix_catalog_downloader,
kiwix_zim_downloader,
)
from wrolpi.downloader import Download
# ---------------------------------------------------------------------------
# KiwixCatalogDownloader
# ---------------------------------------------------------------------------
def test_catalog_prepare_extracts_parent_and_name(test_session):
"""prepare_download splits the URL into the listing parent and the filename prefix."""
download = Download(
url='https://download.kiwix.org/zim/wikipedia/wikipedia_en_all_nopic',
downloader='kiwix_catalog',
)
prepared = kiwix_catalog_downloader.prepare_download(test_session, download)
assert isinstance(prepared, PreparedKiwixCatalog)
assert prepared.url == download.url
assert prepared.parent_url == 'https://download.kiwix.org/zim/wikipedia'
assert prepared.name == 'wikipedia_en_all_nopic'
@pytest.mark.asyncio
async def test_catalog_execute_picks_latest_matching_zim(test_session):
"""execute_download fetches the parent listing and picks the lexicographically last
matching .zim — sorted ascending, so the newest YYYY-MM Kiwix filename wins."""
listing = [
'wikipedia_en_all_nopic_2024-01.zim',
'wikipedia_en_all_nopic_2024-06.zim',
'wikipedia_en_all_nopic_2024-12.zim',
'wikipedia_en_all_nopic_2024-12.zim.meta4', # not a .zim file → skipped
'unrelated_corpus_2024-12.zim', # different prefix → skipped
]
async def fake_fetch_hrefs(url):
assert url == 'https://download.kiwix.org/zim/wikipedia'
return listing
prepared = PreparedKiwixCatalog(
url='https://download.kiwix.org/zim/wikipedia/wikipedia_en_all_nopic',
parent_url='https://download.kiwix.org/zim/wikipedia',
name='wikipedia_en_all_nopic',
)
with mock.patch('modules.zim.downloader.fetch_hrefs', fake_fetch_hrefs):
executed = await kiwix_catalog_downloader.execute_download(prepared, make_test_ctx())
assert isinstance(executed, ExecutedKiwixCatalog)
assert executed.latest_zim_url == \
'https://download.kiwix.org/zim/wikipedia/wikipedia_en_all_nopic_2024-12.zim'
@pytest.mark.asyncio
async def test_catalog_execute_no_matches(test_session):
"""When the listing has no matching files, executed.latest_zim_url is None."""
async def fake_fetch_hrefs(url):
return ['some_other_corpus_2024-01.zim']
prepared = PreparedKiwixCatalog(
url='https://download.kiwix.org/zim/wikipedia/wikipedia_en_all_nopic',
parent_url='https://download.kiwix.org/zim/wikipedia',
name='wikipedia_en_all_nopic',
)
with mock.patch('modules.zim.downloader.fetch_hrefs', fake_fetch_hrefs):
executed = await kiwix_catalog_downloader.execute_download(prepared, make_test_ctx())
assert executed.latest_zim_url is None
def test_catalog_finalize_no_match_returns_failure(test_session):
"""latest_zim_url=None → DownloadResult(success=False) with the URL in the error."""
download = Download(url='https://example.com/zim/foo', downloader='kiwix_catalog')
executed = ExecutedKiwixCatalog(latest_zim_url=None)
result = kiwix_catalog_downloader.finalize_download(test_session, download, executed)
assert result.success is False
assert 'https://example.com/zim/foo' in (result.error or '')
def test_catalog_finalize_with_match_emits_child_download(test_session):
"""The latest matching URL is emitted as a child download for the kiwix_zim downloader."""
download = Download(url='https://example.com/zim/foo', downloader='kiwix_catalog')
executed = ExecutedKiwixCatalog(
latest_zim_url='https://example.com/zim/foo_2024-12.zim',
)
result = kiwix_catalog_downloader.finalize_download(test_session, download, executed)
assert result.success is True
assert result.downloads == ['https://example.com/zim/foo_2024-12.zim']
# ---------------------------------------------------------------------------
# KiwixZimDownloader
# ---------------------------------------------------------------------------
def test_zim_prepare_creates_zim_directory(test_session, test_directory, monkeypatch):
"""prepare_download ensures the zim directory exists and returns it on the plan."""
expected = test_directory / 'zims'
assert not expected.exists()
# Override get_zim_directory so we don't depend on wrolpi config defaults.
monkeypatch.setattr('modules.zim.lib.get_zim_directory', lambda: expected)
download = Download(
url='https://example.com/zim/wikipedia_en_all_nopic_2024-12.zim',
downloader='kiwix_zim',
)
prepared = kiwix_zim_downloader.prepare_download(test_session, download)
assert isinstance(prepared, PreparedKiwixZim)
assert prepared.url == download.url
assert prepared.zim_directory == expected
assert expected.is_dir()
def test_zim_finalize_invalid_zim_returns_failure(test_session, test_directory):
"""is_valid=False → DownloadResult(success=False) with the legacy error message."""
output_path = test_directory / 'zims/wikipedia_en_all_nopic_2024-12.zim'
executed = ExecutedKiwixZim(output_path=output_path, is_valid=False)
download = Download(url='https://example.com/zim/foo', downloader='kiwix_zim')
result = kiwix_zim_downloader.finalize_download(test_session, download, executed)
assert result.success is False
assert result.error == 'Zim file is invalid'
def test_zim_finalize_valid_returns_success(test_session, test_directory):
"""is_valid=True → DownloadResult(success=True) with the generic Kiwix viewer location."""
output_path = test_directory / 'zims/wikipedia_en_all_nopic_2024-12.zim'
executed = ExecutedKiwixZim(output_path=output_path, is_valid=True)
download = Download(url='https://example.com/zim/foo', downloader='kiwix_zim')
result = kiwix_zim_downloader.finalize_download(test_session, download, executed)
assert result.success is True
assert result.location == '/zim'

View file

@ -1059,7 +1059,8 @@ def mock_downloader_download_file():
"""Mock the Downloader.download_file method."""
contents = None
async def _download_file(self, id_, url, destination):
async def _download_file(self, id_, url, destination, **kwargs):
# **kwargs absorbs ctx= and any future optional args added to Downloader.download_file.
name = url.split('/')[-1]
output_path = destination / name
output_path.write_bytes(contents or b'')