diff --git a/modules/zim/conftest.py b/modules/zim/conftest.py index 1ff80a23..7ac88197 100644 --- a/modules/zim/conftest.py +++ b/modules/zim/conftest.py @@ -1,7 +1,6 @@ import asyncio import json import pathlib -from http import HTTPStatus from itertools import zip_longest from typing import List, Dict from uuid import uuid4 @@ -28,33 +27,18 @@ def kiwix_download_manager(test_download_manager) -> DownloadManager: @pytest.fixture -def kiwix_download_zim(test_directory, kiwix_download_manager, test_zim_bytes): +def kiwix_download_zim(test_directory, kiwix_download_manager, test_zim_bytes, mock_downloader_download_file): async def do_download(download_info: DownloadFileInfo = None, expected_url: str = None, download_file_side_effect=None, hrefs: List[str] = None): - async def _download_file(url: str, output_path: pathlib.Path = None, info: DownloadFileInfo = None): - if expected_url: - assert url == expected_url - if download_info: - assert info == download_info - output_path.write_bytes(test_zim_bytes) - - with mock.patch('modules.zim.downloader.fetch_hrefs') as mock_fetch_hrefs, \ - mock.patch('modules.zim.downloader.get_download_info') as mock_get_download_info, \ - mock.patch('modules.zim.downloader.download_file') as mock_download_file: + with mock.patch('modules.zim.downloader.fetch_hrefs') as mock_fetch_hrefs: # Fake the download, but make sure it is called correctly. Also creates a Zim file at the output_path. - mock_download_file.side_effect = download_file_side_effect or _download_file # These are the links in the parent directory of the url. mock_fetch_hrefs.return_value = hrefs or ['?C=N;O=D', '?C=M;O=A', '?C=S;O=A', '?C=D;O=A', '/zim/', 'wikipedia_es_all_maxi_2023-05.zim', 'wikipedia_es_all_maxi_2023-06.zim', 'wikipedia_es_all_nopic_2023-05.zim', 'wikipedia_es_all_nopic_2023-06.zim'] - # This is the DownloadInfo for the latest Zim file that was selected. - mock_get_download_info.return_value = download_info or DownloadFileInfo( - status=HTTPStatus.OK, - name='wikipedia_es_all_maxi_2023-06.zim', - size=len(test_zim_bytes), - ) + mock_downloader_download_file(test_zim_bytes) await kiwix_download_manager.wait_for_all_downloads() # Async download are hard to test, sleep to let tasks finish. diff --git a/modules/zim/downloader.py b/modules/zim/downloader.py index e2c6280c..a70a6a35 100644 --- a/modules/zim/downloader.py +++ b/modules/zim/downloader.py @@ -2,8 +2,7 @@ from http import HTTPStatus from typing import List from modules.zim import lib -from wrolpi.common import logger, aiohttp_get, get_download_info, download_file, background_task, get_html_soup, \ - trim_file_name +from wrolpi.common import logger, aiohttp_get, background_task, get_html_soup from wrolpi.downloader import Downloader, Download, DownloadResult __all__ = ['KiwixCatalogDownloader', 'KiwixZimDownloader', 'kiwix_zim_downloader', 'kiwix_catalog_downloader'] @@ -72,33 +71,14 @@ class KiwixZimDownloader(Downloader): listable = False async def do_download(self, download: Download) -> DownloadResult: - url = download.url - - download_info = await get_download_info(url) - zim_directory = lib.get_zim_directory() zim_directory.mkdir(parents=True, exist_ok=True) - name = trim_file_name(download_info.name) - output_path = zim_directory / name - if output_path.is_file(): - output_size = output_path.stat().st_size - if download_info.size and download_info.size == output_size: - # File is already downloaded. - return DownloadResult(success=True) - - logger.info(f'Downloading Zim {url} to {output_path} of size {download_info.size}') - await download_file(url, output_path=output_path, info=download_info) + output_path = await self.download_file(download.id, download.url, zim_directory) # Notify the maintainer if outdated Zim files are lying around. lib.flag_outdated_zim_files() - output_size = output_path.stat().st_size - if download_info.size and output_size != download_info.size: - return DownloadResult( - success=False, - error=f'Download size does not match: {output_size} != {download_info.size}') - return_code = await lib.check_zim(output_path) if return_code == 127: logger.warning(f'Not validating {output_path} because zimcheck is not installed') diff --git a/wrolpi/downloader.py b/wrolpi/downloader.py index cd1262da..8db4dce9 100644 --- a/wrolpi/downloader.py +++ b/wrolpi/downloader.py @@ -7,9 +7,11 @@ import os import pathlib import tempfile import traceback +import xml.etree.ElementTree as ET from dataclasses import dataclass, field from datetime import timedelta, datetime from enum import Enum +from http import HTTPStatus from itertools import filterfalse from typing import List, Dict, Generator, Iterable, Coroutine from typing import Tuple, Optional @@ -25,11 +27,13 @@ from sqlalchemy.sql import Delete from wrolpi import flags from wrolpi.api_utils import api_app, perpetual_signal +from wrolpi.cmd import which from wrolpi.common import Base, ModelHelper, logger, wrol_mode_check, zig_zag, ConfigFile, \ - wrol_mode_enabled, background_task, get_wrolpi_config, get_absolute_media_path, timer + wrol_mode_enabled, background_task, get_wrolpi_config, get_absolute_media_path, timer, aiohttp_get, \ + get_download_info, trim_file_name from wrolpi.dates import TZDateTime, now, Seconds from wrolpi.db import get_db_session, get_db_curs, optional_session -from wrolpi.errors import InvalidDownload, UnrecoverableDownloadError, UnknownDownload, ValidationError +from wrolpi.errors import InvalidDownload, UnrecoverableDownloadError, UnknownDownload, ValidationError, DownloadError from wrolpi.events import Events from wrolpi.media_path import MediaPathType from wrolpi.switches import register_switch_handler, ActivateSwitchMethod, await_switches @@ -37,6 +41,8 @@ from wrolpi.vars import PYTEST, SIMULTANEOUS_DOWNLOAD_DOMAINS logger = logger.getChild(__name__) +ARIA2C_PATH = which('aria2c', '/usr/bin/aria2c') + @perpetual_signal(sleep=1 if PYTEST else 5) async def perpetual_download_worker(): @@ -377,6 +383,70 @@ class Downloader: # Return the result of the download attempt. return task.result() + @staticmethod + async def get_meta4_contents(url: str) -> bytes | None: + try: + meta4_url = f'{url}.meta4' + async with aiohttp_get(meta4_url, timeout=5) as response: + if response.status == HTTPStatus.OK: + contents = await response.content.read() + try: + root = ET.fromstring(contents) + if root.tag == '{urn:ietf:params:xml:ns:metalink4}metalink': + return contents + logger.debug(f'meta4 file was not a metalink file: {meta4_url}') + except Exception as e: + logger.debug(f'Failed to parse XML {meta4_url}', exc_info=e) + except Exception as e: + logger.debug(f'Failed to fetch meta4 file {url}', exc_info=e) + + return None + + async def download_file(self, download_id: int, url: str, destination: pathlib.Path, check_for_meta4: bool = True) \ + -> pathlib.Path: + meta4_contents = None + if check_for_meta4: + meta4_contents = await self.get_meta4_contents(url) + + info = await get_download_info(url) + + # TODO verify that this output_path is exclusive. + output_path = destination / trim_file_name(info.name) + + with tempfile.NamedTemporaryFile(suffix='.meta4') as meta4_path: + cmd = ( + ARIA2C_PATH, + '-l', '-', + '-j3', # concurrent downloads + '-s3', # split jobs + '-d', destination, + url + ) + meta4_path = pathlib.Path(meta4_path.name) + if meta4_contents: + meta4_path.write_bytes(meta4_contents) + cmd = (*cmd, + '-M', meta4_path, + ) + + return_code, logs, stdout = await self.process_runner( + download_id, + url, # Log the original URL. + cmd, + destination, + ) + + error = logs.get('stderr') or logs.get('stdout') + error = error.decode() if error else error + + if return_code != 0: + raise DownloadError(f'{error}\n\nFileDownloader failed with return code {return_code}') + + if not output_path.is_file(): + raise DownloadError(f'{error}\n\nOutput file not found: {output_path}') + + return output_path + class DownloadManager: """ diff --git a/wrolpi/files/downloader.py b/wrolpi/files/downloader.py index 722cf41e..2a4da67c 100644 --- a/wrolpi/files/downloader.py +++ b/wrolpi/files/downloader.py @@ -1,10 +1,12 @@ import pathlib +import pathlib +import tempfile import traceback from abc import ABC +from wrolpi.cmd import which from wrolpi.common import get_media_directory, logger, get_download_info, \ - download_file, trim_file_name -from wrolpi.db import get_db_session + trim_file_name, TRACE_LEVEL from wrolpi.downloader import Downloader, Download, DownloadResult from wrolpi.errors import UnrecoverableDownloadError from wrolpi.files.lib import upsert_file @@ -14,6 +16,8 @@ __all__ = ['FileDownloader', 'file_downloader'] logger = logger.getChild(__name__) +ARIA2C_PATH = which('aria2c', '/usr/bin/aria2c') + class FileDownloader(Downloader, ABC): """A Downloader which can download an arbitrary HTTP file. @@ -31,8 +35,6 @@ class FileDownloader(Downloader, ABC): if download.attempts > 3: raise UnrecoverableDownloadError(f'Max download attempts reached for {download.url}') - media_directory = get_media_directory() - settings = download.settings or dict() # We don't know where to put the file if we don't have a destination. @@ -40,28 +42,28 @@ class FileDownloader(Downloader, ABC): if not destination: raise UnrecoverableDownloadError(f'Cannot download the file without a destination') + if logger.isEnabledFor(TRACE_LEVEL): + logger.trace(f'FileDownloader: downloading {download.url} to {destination}') + # Create the destination only if it is within the media directory. - destination = pathlib.Path(destination) - if not str(destination.absolute()).startswith(str(media_directory)): + destination = pathlib.Path(destination).resolve() + media_directory = get_media_directory() + if not str(destination).startswith(str(media_directory)): raise UnrecoverableDownloadError(f'Cannot download outside media directory!') destination.mkdir(parents=True, exist_ok=True) - info = await get_download_info(download.url) - - # TODO verify that this output_path is exclusive. - output_path = destination / trim_file_name(info.name) - + error = None try: - await download_file(download.url, output_path, info) - + output_path = await self.download_file(download.id, download.url, destination) fg = await upsert_file(output_path, tag_names=download.tag_names) location = fg.location - return DownloadResult( success=True, location=location, ) except Exception as e: + if error: + logger.error(f'FileDownloader failed with aria2c error:\n{error}') logger.error(f'Failed to download {repr(str(download.url))}', exc_info=e) if PYTEST: raise diff --git a/wrolpi/files/lib.py b/wrolpi/files/lib.py index 5cfc9ec9..69d0d519 100644 --- a/wrolpi/files/lib.py +++ b/wrolpi/files/lib.py @@ -31,7 +31,7 @@ from wrolpi.dates import now, from_timestamp, months_selector_to_where, date_ran from wrolpi.db import get_db_session, get_db_curs, mogrify, optional_session from wrolpi.downloader import download_manager from wrolpi.errors import InvalidFile, UnknownDirectory, UnknownFile, UnknownTag, FileConflict, FileGroupIsTagged, \ - NoPrimaryFile, InvalidDirectory + NoPrimaryFile, InvalidDirectory, IgnoredDirectoryError from wrolpi.events import Events from wrolpi.files.models import FileGroup, Directory from wrolpi.lang import ISO_639_CODES, ISO_3166_CODES @@ -1537,13 +1537,16 @@ def remove_ignored_directory(directory: Union[pathlib.Path, str]): async def upsert_file(file: pathlib.Path | str, tag_names: List[str] = None) -> FileGroup: """Insert/update all files in the provided file's FileGroup.""" + if not file or not file.is_file(): + raise InvalidFile(f'Cannot upsert file that does not exist: {file}') + # Update/Insert all files in the FileGroup. paths = glob_shared_stem(pathlib.Path(file)) # Remove any ignored files. paths = remove_files_in_ignored_directories(paths) if not paths: logger.warning('upsert_file called, but all files are in ignored directories!') - return + raise IgnoredDirectoryError(f'all files are in ignored directories: {file}') for i in range(2): # Try multiple times because uploads happen concurrently and may conflict. diff --git a/wrolpi/files/test/test_downloader.py b/wrolpi/files/test/test_downloader.py index e69de29b..fe856dc2 100644 --- a/wrolpi/files/test/test_downloader.py +++ b/wrolpi/files/test/test_downloader.py @@ -0,0 +1,87 @@ +import hashlib +from http import HTTPStatus + +from wrolpi import common +from wrolpi.downloader import Download +from wrolpi.files.downloader import FileDownloader + + +async def test_file_downloader(test_session, make_files_structure, test_directory, await_switches, + simple_file_server, test_download_manager, video_file, video_bytes): + """The FileDownloader can download a file without a metalink file.""" + file_downloader = FileDownloader() + test_download_manager.register_downloader(file_downloader) + + host, port = simple_file_server.server_address + url = f'http://{host}:{port}' # noqa + + # Create download for the video file. + video_url = f'{url}/{video_file.name}' + test_download_manager.create_download(video_url, file_downloader.name, test_session, + destination=str(test_directory / 'downloads')) + test_session.commit() + + await test_download_manager.wait_for_all_downloads() + + download: Download = test_session.query(Download).one() + assert not download.error, download.error + + # The video was downloaded correctly. + assert (test_directory / f'downloads/{video_file.name}').is_file() + assert (test_directory / f'downloads/{video_file.name}').stat().st_size == 1056318 + assert (test_directory / f'downloads/{video_file.name}').read_bytes() == video_bytes + + +# Minimum Metalink file to download a file. +META4_XML = ''' + + + {size} + {sha256} + {url}/{name} + + +''' + + +async def test_file_downloader_meta4(test_session, make_files_structure, test_directory, await_switches, + simple_file_server, test_download_manager, video_file, video_bytes): + """The FileDownloader can download a file using a Metalink file.""" + file_downloader = FileDownloader() + test_download_manager.register_downloader(file_downloader) + + host, port = simple_file_server.server_address + url = f'http://{host}:{port}' # noqa + + meta4 = META4_XML.format( + url=url, + name=video_file.name, + size=video_file.stat().st_size, + sha256=hashlib.sha256(video_bytes).hexdigest(), + ) + (test_directory / f'{video_file.name}.meta4').write_text(meta4) + + # Ensure files are being served correctly so they can be downloaded. + video_url = f'{url}/{video_file.name}' + video_meta4_url = f'{url}/{video_file.name}.meta4' + async with common.aiohttp_head(video_url, timeout=5) as response: + assert response.status == HTTPStatus.OK, response.content + async with common.aiohttp_head(video_meta4_url, timeout=5) as response: + assert response.status == HTTPStatus.OK, response.content + + # Create download for the video file, this should automatically check for the .meta4 file and use it with aria2c. + test_download_manager.create_download(video_url, file_downloader.name, test_session, + destination=str(test_directory / 'downloads')) + test_session.commit() + + await test_download_manager.wait_for_all_downloads() + + download: Download = test_session.query(Download).one() + assert not download.error, download.error + + # The video was downloaded correctly. + assert (test_directory / f'downloads/{video_file.name}').is_file() + assert (test_directory / f'downloads/{video_file.name}').stat().st_size == 1056318 + assert (test_directory / f'downloads/{video_file.name}').read_bytes() == video_bytes + # The meta4 was saved to a temporary file, but not saved. + assert not (test_directory / 'downloads/foo.txt.meta4').exists()