mirror of
https://github.com/lrnselfreliance/wrolpi.git
synced 2026-10-26 19:01:39 -04:00
Using aria2c to download files. Removing custom file download functions.
This commit is contained in:
Roland
parent
95a524ba33
commit
220cb8010f
6 changed files with 185 additions and 59 deletions
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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')
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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 = '''<?xml version="1.0" encoding="UTF-8"?>
|
||||
<metalink xmlns="urn:ietf:params:xml:ns:metalink">
|
||||
<file name="{name}">
|
||||
<size>{size}</size>
|
||||
<hash type="sha-256">{sha256}</hash>
|
||||
<url>{url}/{name}</url>
|
||||
</file>
|
||||
</metalink>
|
||||
'''
|
||||
|
||||
|
||||
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()
|
||||
Loading…
Add table
Add a link
Reference in a new issue