diff --git a/main.py b/main.py
index a17c8a05..9bd19c3d 100755
--- a/main.py
+++ b/main.py
@@ -10,7 +10,7 @@ from sanic.signals import Event
from wrolpi import flags, BEFORE_STARTUP_FUNCTIONS
from wrolpi import root_api, admin
from wrolpi.common import logger, get_config, check_media_directory, limit_concurrent, \
- wrol_mode_enabled, cancel_refresh_tasks, set_log_level
+ wrol_mode_enabled, cancel_refresh_tasks, set_log_level, background_task, cancel_background_tasks
from wrolpi.downloader import download_manager, import_downloads_config
from wrolpi.root_api import api_app
from wrolpi.vars import PROJECT_DIR, DOCKERIZED, PYTEST
@@ -170,35 +170,38 @@ def main():
@api_app.before_server_start
@limit_concurrent(1)
async def startup(app: Sanic):
+ from wrolpi.common import LOG_LEVEL
+
# Check database status first. Many functions will reference flags.db_up.
flags.check_db_is_up()
flags.init_flags()
await import_downloads_config()
+ async def periodic_check_db_is_up():
+ while True:
+ flags.check_db_is_up()
+ flags.init_flags()
+ await asyncio.sleep(10)
+
+ background_task(periodic_check_db_is_up())
+
+ async def periodic_check_log_level():
+ while True:
+ log_level = LOG_LEVEL.value
+ if log_level != logger.getEffectiveLevel():
+ set_log_level(log_level)
+ await asyncio.sleep(1)
+
+ background_task(periodic_check_log_level())
+
+ from wrolpi import status
+ background_task(status.bandwidth_worker())
+
from modules.zim.lib import flag_outdated_zim_files
flag_outdated_zim_files()
-@api_app.after_server_start
-@limit_concurrent(1)
-async def periodic_check_db_is_up(app: Sanic):
- while True:
- flags.check_db_is_up()
- flags.init_flags()
- await asyncio.sleep(10)
-
-
-@api_app.after_server_start
-async def periodic_check_log_level(app: Sanic):
- from wrolpi.common import LOG_LEVEL
- while True:
- log_level = LOG_LEVEL.value
- if log_level != logger.getEffectiveLevel():
- set_log_level(log_level)
- await asyncio.sleep(1)
-
-
@api_app.after_server_start
async def periodic_downloads(app: Sanic):
"""
@@ -253,12 +256,6 @@ async def main_import_tags_config(app: Sanic):
tags.import_tags_config()
-@api_app.after_server_start
-async def bandwidth_worker(app: Sanic):
- from wrolpi import status
- app.add_task(status.bandwidth_worker())
-
-
@root_api.api_app.signal(Event.SERVER_SHUTDOWN_BEFORE)
@limit_concurrent(1)
async def handle_server_shutdown(*args, **kwargs):
@@ -266,6 +263,7 @@ async def handle_server_shutdown(*args, **kwargs):
if not PYTEST:
download_manager.stop()
await cancel_refresh_tasks()
+ await cancel_background_tasks()
if __name__ == '__main__':
diff --git a/modules/archive/__init__.py b/modules/archive/__init__.py
index dd20046a..4e49a57c 100644
--- a/modules/archive/__init__.py
+++ b/modules/archive/__init__.py
@@ -1,22 +1,26 @@
import asyncio
import json
+import pathlib
+import tempfile
from abc import ABC
-from typing import List
+from typing import List, Tuple
+from selenium import webdriver
from sqlalchemy import not_
from sqlalchemy.orm import Session
+from wrolpi.cmd import SINGLE_FILE_BIN, CHROMIUM, READABILITY_BIN
from wrolpi.common import logger, register_modeler, register_refresh_cleanup, limit_concurrent, split_lines_by_length, \
- slow_logger
+ slow_logger, html_screenshot
from wrolpi.db import optional_session, get_db_session
from wrolpi.downloader import Downloader, Download, DownloadResult
from wrolpi.errors import UnrecoverableDownloadError
from wrolpi.files.models import FileGroup
-from wrolpi.vars import PYTEST
+from wrolpi.vars import PYTEST, DOCKERIZED
from . import lib
from .api import bp # noqa
from .errors import InvalidArchive
-from .lib import is_singlefile_file, get_title_from_html
+from .lib import is_singlefile_file, get_title_from_html, request_archive, SINGLEFILE_HEADER
from .models import Archive, Domain
PRETTY_NAME = 'Archive'
@@ -37,7 +41,14 @@ class ArchiveDownloader(Downloader, ABC):
if download.attempts > 3:
raise UnrecoverableDownloadError(f'Max download attempts reached for {download.url}')
- archive: Archive = await lib.do_archive(download.url)
+ if DOCKERIZED or PYTEST:
+ # Perform the archive in the Archive docker container. (Typically in the development environment).
+ singlefile, readability, screenshot = await request_archive(download.url)
+ else:
+ # Perform the archive using locally installed executables.
+ singlefile, readability, screenshot = await self.do_archive(download)
+
+ archive: Archive = await lib.model_archive_result(download.url, singlefile, readability, screenshot)
if download.settings and (tag_names := download.settings.get('tag_names')):
for name in tag_names:
@@ -53,6 +64,99 @@ class ArchiveDownloader(Downloader, ABC):
archives = list(session.query(Archive).filter(Archive.url.in_(urls)))
return archives
+ async def do_singlefile(self, download: Download) -> bytes:
+ """Create a Singlefile from the archive's URL."""
+ cmd = (str(SINGLE_FILE_BIN),
+ download.url,
+ '--browser-executable-path', CHROMIUM,
+ '--browser-args', '["--no-sandbox"]',
+ '--dump-content')
+ return_code, _, stdout = await self.process_runner(
+ download.url,
+ cmd,
+ pathlib.Path('/home/wrolpi'),
+ )
+ if return_code != 0:
+ raise RuntimeError(f'Archive singlefile exited with {return_code}')
+
+ return stdout
+
+ async def do_readability(self, download: Download, html: bytes) -> dict:
+ """Extract the readability dict from the provided HTML."""
+ with tempfile.NamedTemporaryFile('wb', suffix='.html') as fh:
+ fh.write(html)
+
+ cmd = (READABILITY_BIN, fh.name, download.url)
+ logger.debug(f'readability cmd: {cmd}')
+ return_code, logs, stdout = await self.process_runner(
+ download.url,
+ cmd,
+ pathlib.Path('/home/wrolpi'),
+ )
+ if return_code == 0:
+ readability = json.loads(stdout)
+ logger.debug(f'done readability for {download.url}')
+ return readability
+ else:
+ logger.error(f'Failed to extract readability for {download.url}')
+ raise RuntimeError(f'Failed to extract readability for {download.url}')
+
+ @staticmethod
+ async def do_screenshot_url(download: Download) -> bytes:
+ """Use Chromium to get a screenshot of the Download's URL."""
+ # Set Chromium to headless. Use a wide window size so that screenshot will be the "desktop" version of
+ # the page.
+ options = webdriver.ChromeOptions()
+ options.add_argument('headless')
+ options.add_argument('disable-gpu')
+ options.add_argument('window-size=1280x720')
+
+ driver = webdriver.Chrome(chrome_options=options)
+ driver.get(download.url)
+ screenshot = driver.get_screenshot_as_png()
+ return screenshot
+
+ async def do_archive(self, download: Download) -> Tuple[bytes, dict, bytes]:
+ """Use locally installed executables to create an Archive.
+
+ Creates a Singlefile, readability file, and screenshot file.
+
+ @warning: Will not raise errors if readability or screenshot cannot be extracted.
+ """
+ singlefile = await self.do_singlefile(download)
+
+ if SINGLEFILE_HEADER.encode() not in singlefile[:1000]:
+ raise RuntimeError(f'Singlefile created was invalid: {download.ur}')
+
+ # Extract Readability from the Singlefile.
+ try:
+ readability = await self.do_readability(download, singlefile)
+ except RuntimeError:
+ # Readability is not required.
+ readability = None
+
+ screenshot = b''
+ try:
+ # Screenshot the Singlefile first.
+ screenshot = html_screenshot(singlefile)
+ except Exception as e:
+ logger.error(f'Failed to screenshot file for {download.url}', exc_info=e)
+
+ if not screenshot:
+ # Screenshot of singlefile failed. Download the URL again.
+ try:
+ screenshot = await self.do_screenshot_url(download)
+ except Exception as e:
+ # Screenshot failed.
+ logger.error(f'Failed to screenshot {download.url}', exc_info=e)
+
+ singlefile_len = len(singlefile) if singlefile else None
+ readability_len = len(readability) if readability else None
+ screenshot_len = len(screenshot) if screenshot else None
+ logger.debug(f'do_archive of {download.url} finished: {singlefile_len=} {readability_len=} {screenshot_len=}')
+
+ return singlefile, readability, screenshot
+
archive_downloader = ArchiveDownloader()
diff --git a/modules/archive/lib.py b/modules/archive/lib.py
index efa4ff4d..676a40a7 100644
--- a/modules/archive/lib.py
+++ b/modules/archive/lib.py
@@ -3,26 +3,21 @@ import gzip
import json
import pathlib
import re
-import subprocess
-import tempfile
from dataclasses import dataclass
from datetime import datetime
from typing import Optional, Tuple, List
from bs4 import BeautifulSoup
-from selenium import webdriver
from sqlalchemy import asc
from sqlalchemy.orm import Session
from modules.archive.models import Domain, Archive
-from wrolpi.cmd import READABILITY_BIN, SINGLE_FILE_BIN, CHROMIUM
-from wrolpi.common import get_media_directory, logger, extract_domain, chdir, escape_file_name, aiohttp_post, \
+from wrolpi.common import get_media_directory, logger, extract_domain, escape_file_name, aiohttp_post, \
format_html_string, split_lines_by_length
from wrolpi.dates import now, Seconds
from wrolpi.db import get_db_session, get_db_curs, optional_session
from wrolpi.errors import UnknownArchive, InvalidOrderBy
from wrolpi.tags import tag_names_to_file_group_sub_select
-from wrolpi.vars import DOCKERIZED, PYTEST
logger = logger.getChild(__name__)
@@ -146,84 +141,11 @@ async def request_archive(url: str) -> Tuple[str, Optional[str], Optional[str]]:
return singlefile, readability, screenshot
-def local_singlefile(url: str):
- """Run the single-file executable to create an HTML file archive."""
- if not SINGLE_FILE_BIN or not SINGLE_FILE_BIN.is_file():
- raise FileNotFoundError(f'single-file not found. Is it installed?')
-
- cmd = (str(SINGLE_FILE_BIN),
- url,
- '--browser-executable-path', CHROMIUM,
- '--browser-args', '["--no-sandbox"]',
- '--dump-content')
- logger.debug(f'archive cmd: {cmd}')
- output = subprocess.check_output(cmd, timeout=60 * 3)
- logger.debug(f'done archiving for {url}')
- return output
-
-
-def local_screenshot(url: str) -> bytes:
- """Take a screenshot of the URL using chromedriver."""
- logger.info(f'Screenshot: {url}')
-
- # Set Chromium to headless. Use a wide window size so that screenshot will be the "desktop" version of the page.
- options = webdriver.ChromeOptions()
- options.add_argument('headless')
- options.add_argument('disable-gpu')
- options.add_argument('window-size=1280x720')
-
- driver = webdriver.Chrome(chrome_options=options)
- try:
- driver.get(url)
- png = driver.get_screenshot_as_png()
- except Exception as e:
- logger.warning(f'Failed to screenshot {url}', exc_info=e)
- return b''
- return png
-
-
-def local_extract_readability(path: str, url: str) -> dict:
- """Extract the readability from an HTML file, typically from single-file."""
- logger.info(f'readability for {url}')
- if not READABILITY_BIN.is_file():
- raise FileNotFoundError(f'Readability extractor not found')
-
- cmd = (READABILITY_BIN, path, url)
- logger.debug(f'readability cmd: {cmd}')
- output = subprocess.check_output(cmd, timeout=60 * 3)
- output = json.loads(output)
- logger.debug(f'done readability for {url}')
- return output
-
-
-def local_archive(url: str):
- """Perform an archive of the provided URL using local resources (without the Archive docker container)."""
- # Archives must be performed in the wrolpi home directory because chrome saves the screenshot there, and will
- # raise an error if global $HOME is not a real user directory. :(
- with chdir('/home/wrolpi', with_home=True):
- singlefile = local_singlefile(url)
- with tempfile.NamedTemporaryFile('wb') as fh:
- fh.write(singlefile)
- readability = local_extract_readability(fh.name, url)
- screenshot = local_screenshot(url)
- singlefile = singlefile.decode()
- return singlefile, readability, screenshot
-
-
-async def do_archive(url: str) -> Archive:
+async def model_archive_result(url: str, singlefile: str, readability: dict, screenshot: bytes) -> Archive:
"""
- Perform the real archive request to the archiving service. Store the resulting data into files. Create an Archive
- record in the DB. Create Domain if missing.
+ Convert results from ArchiveDownloader into real files. Create Archive record.
"""
- logger.info(f'Archiving {url}')
-
- if DOCKERIZED or PYTEST:
- # Perform the archive in the Archive docker container. (Typically in the development environment).
- singlefile, readability, screenshot = await request_archive(url)
- else:
- # Perform the archive using locally installed executables.
- singlefile, readability, screenshot = local_archive(url)
-
+ readability = readability.copy() if readability else readability
# First try to get the title from Readability.
title = readability.get('title') if readability else None
diff --git a/modules/archive/test/test_lib.py b/modules/archive/test/test_lib.py
index 67d57026..ad2b590c 100644
--- a/modules/archive/test/test_lib.py
+++ b/modules/archive/test/test_lib.py
@@ -4,12 +4,12 @@ import tempfile
from datetime import datetime, timedelta
from http import HTTPStatus
-import mock
import pytest
from PIL import Image
from modules.archive import lib
-from modules.archive.lib import get_or_create_domain, get_new_archive_files, delete_archives, do_archive, get_domains
+from modules.archive.lib import get_or_create_domain, get_new_archive_files, delete_archives, model_archive_result, \
+ get_domains
from modules.archive.models import Archive, Domain
from wrolpi.db import get_db_session
from wrolpi.files import lib as files_lib
@@ -18,51 +18,48 @@ from wrolpi.root_api import CustomJSONEncoder
from wrolpi.test.common import skip_circleci
-def make_fake_request_archive(readability=True, screenshot=True, title=True):
- async def fake_request_archive(_):
- with tempfile.NamedTemporaryFile(suffix='.png') as fh:
- Image.new('RGB', (1, 1), color='grey').save(fh)
- fh.seek(0)
- image = fh.read()
- singlefile = '\ntest single-file\nジにてこちら\nsome title'
- r = dict(
- content=f'test readability content',
- textContent='test readability textContent',
- title='ジにてこちら' if title else None,
- ) if readability else None
- s = image if screenshot else None
- return singlefile, r, s
-
- return fake_request_archive
+def make_fake_archive_result(readability=True, screenshot=True, title=True):
+ with tempfile.NamedTemporaryFile(suffix='.png') as fh:
+ Image.new('RGB', (1, 1), color='grey').save(fh)
+ fh.seek(0)
+ image = fh.read()
+ singlefile = '\ntest single-file\nジにてこちら\nsome title'
+ r = dict(
+ content=f'test readability content',
+ textContent='test readability textContent',
+ title='ジにてこちら' if title else None,
+ ) if readability else None
+ s = image if screenshot else None
+ return singlefile, r, s
@pytest.mark.asyncio
async def test_no_screenshot(test_session):
- with mock.patch('modules.archive.lib.request_archive', make_fake_request_archive(screenshot=False)):
- archive = await do_archive('https://example.com')
- assert isinstance(archive.singlefile_path, pathlib.Path)
- assert isinstance(archive.readability_path, pathlib.Path)
- # Screenshot was empty
- assert archive.screenshot_path is None
+ singlefile, readability, screenshot = make_fake_archive_result(screenshot=False)
+ archive = await model_archive_result('https://example.com', singlefile, readability, screenshot)
+ assert isinstance(archive.singlefile_path, pathlib.Path)
+ assert isinstance(archive.readability_path, pathlib.Path)
+ # Screenshot was empty
+ assert archive.screenshot_path is None
@pytest.mark.asyncio
async def test_no_readability(test_session):
- with mock.patch('modules.archive.lib.request_archive', make_fake_request_archive(readability=False)):
- archive = await do_archive('https://example.com')
- assert isinstance(archive.singlefile_path, pathlib.Path)
- assert isinstance(archive.screenshot_path, pathlib.Path)
- # Readability empty
- assert archive.readability_path is None
- assert archive.readability_txt_path is None
+ singlefile, readability, screenshot = make_fake_archive_result(readability=False)
+ archive = await model_archive_result('https://example.com', singlefile, readability, screenshot)
+ assert isinstance(archive.singlefile_path, pathlib.Path)
+ assert isinstance(archive.screenshot_path, pathlib.Path)
+ # Readability empty
+ assert archive.readability_path is None
+ assert archive.readability_txt_path is None
@pytest.mark.asyncio
async def test_dict(test_session):
- with mock.patch('modules.archive.lib.request_archive', make_fake_request_archive()):
- d = (await do_archive('https://example.com')).dict()
- assert isinstance(d, dict)
- json.dumps(d, cls=CustomJSONEncoder)
+ singlefile, readability, screenshot = make_fake_archive_result()
+ d = (await model_archive_result('https://example.com', singlefile, readability, screenshot)).dict()
+ assert isinstance(d, dict)
+ json.dumps(d, cls=CustomJSONEncoder)
@pytest.mark.asyncio
@@ -70,11 +67,9 @@ async def test_relationships(test_session, test_directory, example_singlefile):
with get_db_session(commit=True) as session:
url = 'https://wrolpi.org:443'
domain = get_or_create_domain(session, url)
- archive = Archive(
- file_group=FileGroup.from_paths(test_session, example_singlefile),
- url=url,
- domain_id=domain.id,
- )
+ archive = Archive.from_paths(test_session, example_singlefile)
+ archive.url = url
+ archive.domain_id = domain.id
session.add(archive)
session.flush()
@@ -284,22 +279,22 @@ def test_get_domains(test_session, archive_factory):
@pytest.mark.asycio
async def test_new_archive(test_session, fake_now):
- with mock.patch('modules.archive.lib.request_archive', make_fake_request_archive()):
- fake_now(datetime(2000, 1, 1))
- archive1 = await do_archive('https://example.com')
- # Everything is filled out.
- assert isinstance(archive1, Archive)
- assert archive1.archive_datetime is not None
- assert isinstance(archive1.singlefile_path, pathlib.Path)
- assert isinstance(archive1.readability_path, pathlib.Path)
- assert isinstance(archive1.readability_txt_path, pathlib.Path)
- assert isinstance(archive1.screenshot_path, pathlib.Path)
- assert archive1.file_group.title == 'ジにてこちら'
- assert archive1.url is not None
- assert archive1.domain is not None
+ singlefile, readability, screenshot = make_fake_archive_result()
+ fake_now(datetime(2000, 1, 1))
+ archive1 = await model_archive_result('https://example.com', singlefile, readability, screenshot)
+ # Everything is filled out.
+ assert isinstance(archive1, Archive)
+ assert archive1.archive_datetime is not None
+ assert isinstance(archive1.singlefile_path, pathlib.Path)
+ assert isinstance(archive1.readability_path, pathlib.Path)
+ assert isinstance(archive1.readability_txt_path, pathlib.Path)
+ assert isinstance(archive1.screenshot_path, pathlib.Path)
+ assert archive1.file_group.title == 'ジにてこちら'
+ assert archive1.url is not None
+ assert archive1.domain is not None
- # The actual files were dumped and read correctly. The HTML/JSON files are reformatted.
- assert archive1.singlefile_path.read_text() == '''
+ # The actual files were dumped and read correctly. The HTML/JSON files are reformatted.
+ assert archive1.singlefile_path.read_text() == '''
@@ -313,7 +308,7 @@ async def test_new_archive(test_session, fake_now):
'''
- assert archive1.readability_path.read_text() == '''
+ assert archive1.readability_path.read_text() == '''
test readability content
@@ -321,53 +316,47 @@ async def test_new_archive(test_session, fake_now):
'''
- with open(archive1.readability_txt_path) as fh:
- assert fh.read() == 'test readability textContent'
- assert archive1.readability_json_path.read_text() == '''{
+ with open(archive1.readability_txt_path) as fh:
+ assert fh.read() == 'test readability textContent'
+ assert archive1.readability_json_path.read_text() == '''{
"title": "\\u30b8\\u306b\\u3066\\u3053\\u3061\\u3089",
"url": "https://example.com"
}'''
- fake_now(datetime(2000, 1, 2))
- archive2 = await do_archive('https://example.com')
- # Domain is reused.
- assert archive1.domain == archive2.domain
+ fake_now(datetime(2000, 1, 2))
+ archive2 = await model_archive_result('https://example.com', singlefile, readability, screenshot)
+ # Domain is reused.
+ assert archive1.domain == archive2.domain
@pytest.mark.asyncio
async def test_get_title_from_html(test_session, fake_now):
fake_now(datetime(2000, 1, 1))
- with mock.patch('modules.archive.lib.request_archive', make_fake_request_archive()):
- archive = await do_archive('https://example.com')
- assert archive.file_group.title == 'ジにてこちら'
+ singlefile, readability, screenshot = make_fake_archive_result()
+ archive = await model_archive_result('https://example.com', singlefile, readability, screenshot)
+ assert archive.file_group.title == 'ジにてこちら'
- async def fake_request_archive(_):
- singlefile = '\ntest single-file\nジにてこちら\n
some title'
- r = dict(
- content=f'test readability content',
- textContent='test readability textContent',
- )
- s = b'screenshot data'
- return singlefile, r, s
+ singlefile = '\ntest single-file\nジにてこちら\nsome title'
+ readability = dict(
+ content=f'test readability content',
+ textContent='test readability textContent',
+ )
+ screenshot = b'screenshot data'
fake_now(datetime(2000, 1, 2))
- with mock.patch('modules.archive.lib.request_archive', fake_request_archive):
- archive = await do_archive('https://example.com')
- assert archive.file_group.title == 'some title'
+ archive = await model_archive_result('https://example.com', singlefile, readability, screenshot)
+ assert archive.file_group.title == 'some title'
- async def fake_request_archive(_):
- singlefile = ''
- r = dict(
- content=f'missing a title',
- textContent='',
- )
- s = b'screenshot data'
- return singlefile, r, s
+ singlefile = ''
+ readability = dict(
+ content=f'missing a title',
+ textContent='',
+ )
+ screenshot = b'screenshot data'
fake_now(datetime(2000, 1, 3))
- with mock.patch('modules.archive.lib.request_archive', fake_request_archive):
- archive = await do_archive('https://example.com')
- assert archive.file_group.title is None
+ archive = await model_archive_result('https://example.com', singlefile, readability, screenshot)
+ assert archive.file_group.title is None
@skip_circleci
@@ -411,8 +400,8 @@ async def test_title_in_filename(test_session, fake_now, test_directory, image_b
"""
fake_now(datetime(2000, 1, 1))
- with mock.patch('modules.archive.lib.request_archive', make_fake_request_archive()):
- archive1 = await do_archive('https://example.com')
+ singlefile, readability, screenshot = make_fake_archive_result()
+ archive1 = await model_archive_result('https://example.com', singlefile, readability, screenshot)
assert archive1.file_group.title == 'ジにてこちら'
@@ -427,18 +416,15 @@ async def test_title_in_filename(test_session, fake_now, test_directory, image_b
assert str(archive1.screenshot_path.relative_to(test_directory)) == \
'archive/example.com/2000-01-01-00-00-00_ジにてこちら.png'
- async def fake_request_archive(_):
- singlefile = '\ntest single-file\nジにてこちら\n' # no title in HTML
- r = dict(
- # No Title from Readability.
- content=f'test readability content',
- textContent='test readability textContent',
- )
- s = image_bytes_factory()
- return singlefile, r, s
+ singlefile = '\ntest single-file\nジにてこちら\n' # no title in HTML
+ readability = dict(
+ # No Title from Readability.
+ content=f'test readability content',
+ textContent='test readability textContent',
+ )
+ screenshot = image_bytes_factory()
- with mock.patch('modules.archive.lib.request_archive', fake_request_archive):
- archive2 = await do_archive('https://example.com')
+ archive2 = await model_archive_result('https://example.com', singlefile, readability, screenshot)
assert str(archive2.singlefile_path.relative_to(test_directory)) == \
'archive/example.com/2000-01-01-00-00-00_NA.html'
@@ -451,18 +437,15 @@ async def test_title_in_filename(test_session, fake_now, test_directory, image_b
assert str(archive2.screenshot_path.relative_to(test_directory)) == \
'archive/example.com/2000-01-01-00-00-00_NA.png'
- async def fake_request_archive(_):
- singlefile = '\ntest single-file\nジにてこちら\ndangerous ;\\//_title</html></html>'
- r = dict(
- # No Title from Readability.
- content=f'<html>test readability content</html>',
- textContent='<html>test readability textContent</html>',
- )
- s = image_bytes_factory()
- return singlefile, r, s
+ singlefile = '<html>\ntest single-file\nジにてこちら\n<title>dangerous ;\\//_title</html></html>'
+ readability = dict(
+ # No Title from Readability.
+ content=f'<html>test readability content</html>',
+ textContent='<html>test readability textContent</html>',
+ )
+ screenshot = image_bytes_factory()
- with mock.patch('modules.archive.lib.request_archive', fake_request_archive):
- archive3 = await do_archive('https://example.com')
+ archive3 = await model_archive_result('https://example.com', singlefile, readability, screenshot)
assert str(archive3.singlefile_path.relative_to(test_directory)) == \
'archive/example.com/2000-01-01-00-00-00_dangerous ;_title.html'
@@ -620,16 +603,14 @@ def test_archive_order(test_session, test_directory, archive_factory):
@pytest.mark.asyncio
-async def test_archive_download_index(test_session, test_directory, image_file):
+async def test_archive_download_index(test_session, test_directory, image_bytes_factory):
"""An Archive is indexed when it is downloaded."""
- with mock.patch('modules.archive.lib.request_archive') as mock_request_archive:
- singlefile = '<html><title>the singlefile'
- readability = dict(
- content='the readability',
- textContent='the readability',
- )
- mock_request_archive.return_value = (singlefile, readability, image_file.read_bytes())
- archive = await lib.do_archive('https://example.com')
+ singlefile = 'the singlefile'
+ readability = dict(
+ content='the readability',
+ textContent='the readability',
+ )
+ archive = await lib.model_archive_result('https://example.com', singlefile, readability, image_bytes_factory())
assert isinstance(archive, lib.Archive), 'Did not get an archive'
assert archive.singlefile_path and archive.singlefile_path.exists() and \
diff --git a/modules/videos/conftest.py b/modules/videos/conftest.py
index eb3b13dc..0c620b59 100644
--- a/modules/videos/conftest.py
+++ b/modules/videos/conftest.py
@@ -185,7 +185,7 @@ def mock_video_prepare_filename():
@pytest.fixture
def mock_video_process_runner():
with mock.patch('modules.videos.downloader.VideoDownloader.process_runner') as mock_process_runner:
- mock_process_runner.return_value = (0, {'stdout': None, 'stderr': None})
+ mock_process_runner.return_value = (0, {'stdout': None, 'stderr': None}, None)
yield mock_process_runner
diff --git a/modules/videos/downloader.py b/modules/videos/downloader.py
index e0ebc0b5..79269e97 100755
--- a/modules/videos/downloader.py
+++ b/modules/videos/downloader.py
@@ -230,8 +230,12 @@ class VideoDownloader(Downloader, ABC):
destination = settings.get('destination')
if not channel and destination:
# Destination may override the real Channel.
- channel = get_channel(directory=destination, return_dict=False)
- found_channel = 'download_settings_directory'
+ try:
+ channel = get_channel(directory=destination, return_dict=False)
+ found_channel = 'download_settings_directory'
+ except UnknownChannel:
+ # Destination must not be a channel.
+ pass
local_channel_id = settings.get('channel_id')
channel_url = settings.get('channel_url')
@@ -298,7 +302,7 @@ class VideoDownloader(Downloader, ABC):
'--ppa', 'Merger+ffmpeg_o1:-strict -2',
url,
)
- return_code, logs = await self.process_runner(url, cmd, out_dir)
+ return_code, logs, _ = await self.process_runner(url, cmd, out_dir)
stdout = logs['stdout'].decode() if hasattr(logs['stdout'], 'decode') else logs['stdout']
stderr = logs['stderr'].decode() if hasattr(logs['stderr'], 'decode') else logs['stderr']
diff --git a/wrolpi/common.py b/wrolpi/common.py
index f8e7f9e3..269136ba 100644
--- a/wrolpi/common.py
+++ b/wrolpi/common.py
@@ -30,6 +30,7 @@ import aiohttp
import yaml
from aiohttp import ClientResponse
from bs4 import BeautifulSoup
+from selenium import webdriver
from sqlalchemy import types
from sqlalchemy.ext.compiler import compiles
from sqlalchemy.ext.declarative import declarative_base
@@ -97,6 +98,7 @@ __all__ = [
'ConfigFile',
'DownloadFileInfo',
'DownloadFileInfoLink',
+ 'LOG_LEVEL',
'ModelHelper',
'WROLPI_CONFIG',
'aiohttp_get',
@@ -135,8 +137,6 @@ __all__ = [
'limit_concurrent',
'logger',
'make_media_directory',
- 'split_lines_by_length',
- 'slow_logger',
'match_paths_to_suffixes',
'minimize_dict',
'native_only',
@@ -150,6 +150,8 @@ __all__ = [
'set_log_level',
'set_test_config',
'set_test_media_directory',
+ 'slow_logger',
+ 'split_lines_by_length',
'timer',
'truncate_generator_bytes',
'truncate_object_bytes',
@@ -1063,7 +1065,6 @@ async def download_file(url: str, output_path: pathlib.Path = None, info: Downlo
last_report = datetime.now()
bytes_received = 0
- # TODO this cannot be canceled.
# Sleep to catch cancel.
await asyncio.sleep(0)
elif output_path.is_file():
@@ -1319,6 +1320,11 @@ def truncate_generator_bytes(gen: Generator, maximum_bytes: int) -> Generator:
BACKGROUND_TASKS = set()
+def add_background_task(task: Task):
+ BACKGROUND_TASKS.add(task)
+ task.add_done_callback(BACKGROUND_TASKS.discard)
+
+
def background_task(coro) -> Task:
"""Convenience function which creates an asyncio task for the provided coroutine.
@@ -1332,11 +1338,19 @@ def background_task(coro) -> Task:
logger_.error('Background task had error:', exc_info=e)
task = asyncio.create_task(error_logger())
- BACKGROUND_TASKS.add(task)
- task.add_done_callback(BACKGROUND_TASKS.discard)
+ add_background_task(task)
return task
+async def cancel_background_tasks():
+ """Cancels any background async tasks, if any."""
+ if BACKGROUND_TASKS:
+ logger_.warning(f'Canceling {len(BACKGROUND_TASKS)} background tasks')
+ for task in BACKGROUND_TASKS:
+ task.cancel()
+ await asyncio.gather(*BACKGROUND_TASKS)
+
+
def get_warn_once(message: str, logger__: logging.Logger, level=logging.ERROR):
"""Create a function that will report an error only once.
@@ -1523,7 +1537,7 @@ def format_json_file(file: pathlib.Path, indent: int = 2):
def format_html_string(html: str) -> str:
- soup = BeautifulSoup(html)
+ soup = BeautifulSoup(html, features='lxml')
return soup.prettify()
@@ -1542,3 +1556,20 @@ def format_html_file(file: pathlib.Path):
finally:
if copy.is_file():
copy.unlink()
+
+
+def html_screenshot(html: bytes) -> bytes:
+ """Return a PNG screenshot of the provided HTML."""
+ # Set Chromium to headless. Use a wide window size so that screenshot will be the "desktop" version of the page.
+ options = webdriver.ChromeOptions()
+ options.add_argument('headless')
+ options.add_argument('disable-gpu')
+ options.add_argument('window-size=1280x720')
+
+ with tempfile.NamedTemporaryFile('wb', suffix='.html') as fh:
+ fh.write(html)
+
+ driver = webdriver.Chrome(chrome_options=options)
+ driver.get(f'file://{fh.name}')
+ screenshot = driver.get_screenshot_as_png()
+ return screenshot
diff --git a/wrolpi/conftest.py b/wrolpi/conftest.py
index 39b69a8b..ce39ddd1 100644
--- a/wrolpi/conftest.py
+++ b/wrolpi/conftest.py
@@ -232,6 +232,20 @@ def test_downloader(test_download_manager):
def __repr__(self):
return ''
+ do_download = MagicMock()
+
+ def set_test_success(self):
+ async def _(*a, **kwargs):
+ return DownloadResult(success=True)
+
+ self.do_download.side_effect = _
+
+ def set_test_failure(self):
+ async def _(*a, **kwargs):
+ return DownloadResult(success=False)
+
+ self.do_download.side_effect = _
+
test_downloader = TestDownloader()
test_download_manager.register_downloader(test_downloader)
diff --git a/wrolpi/downloader.py b/wrolpi/downloader.py
index b14b5f57..99aa50ac 100644
--- a/wrolpi/downloader.py
+++ b/wrolpi/downloader.py
@@ -11,7 +11,7 @@ from datetime import timedelta, datetime
from enum import Enum
from itertools import filterfalse
from queue import Empty
-from typing import List, Dict, Generator, Iterable
+from typing import List, Dict, Generator, Iterable, Coroutine
from typing import Tuple, Optional
from urllib.parse import urlparse
@@ -209,13 +209,13 @@ class Downloader:
self._kill.clear()
async def process_runner(self, url: str, cmd: Tuple[str, ...], cwd: pathlib.Path, timeout: int = None,
- **kwargs) -> Tuple[int, dict]:
+ **kwargs) -> Tuple[int, dict, bytes]:
"""
Run a subprocess using the provided arguments. This process can be killed by the Download Manager.
Global timeout takes precedence over the timeout argument, unless it is 0. (Smaller global timeout wins)
"""
- logger.debug(f'{self} launching download process with args: {" ".join(cmd)}')
+ logger.debug(f'{self} launching download process with args: {" ".join(list(map(str, cmd)))}')
start = now()
proc = await asyncio.create_subprocess_exec(*cmd,
stdout=asyncio.subprocess.PIPE,
@@ -266,7 +266,34 @@ class Downloader:
logger.debug(f'Download exited with {proc.returncode}')
logs = {'stdout': stdout, 'stderr': stderr}
- return proc.returncode, logs
+ return proc.returncode, logs, stdout
+
+ async def cancel_wrapper(self, coro: Coroutine, download: Download):
+ """
+ Converts an async coroutine to a task. If DownloadManager receives a kill request, this method will cancel
+ the task.
+ """
+ task = asyncio.create_task(coro)
+ while not task.done():
+ if self._kill.is_set():
+ logger.warning(f'Cancel download of {download.url}')
+ task.cancel()
+ try:
+ await task
+ except asyncio.CancelledError as e:
+ logger.debug(f'Successful cancel of {download.url}', exc_info=e)
+ return DownloadResult(
+ success=False,
+ error='Download was canceled',
+ )
+ finally:
+ self.clear()
+ else:
+ # Wait for the download to complete. Cancel if requested.
+ await asyncio.sleep(0.1)
+
+ # Return the result of the download attempt.
+ return task.result()
class DownloadManager:
@@ -359,10 +386,9 @@ class DownloadManager:
try_again = True
try:
- if asyncio.iscoroutinefunction(downloader.do_download):
- result = await downloader.do_download(download)
- else:
- result = downloader.do_download(download)
+ # Create download coroutine; wrap it, so it can be canceled
+ coro = downloader.do_download(download)
+ result = await downloader.cancel_wrapper(coro, download)
except UnrecoverableDownloadError as e:
# Download failed and should not be retried.
worker_logger.warning(f'UnrecoverableDownloadError for {url}', exc_info=e)
@@ -406,8 +432,8 @@ class DownloadManager:
background_task(self.queue_downloads())
# Save the config now that the Download has finished.
background_task(save_downloads_config())
- except asyncio.CancelledError:
- worker_logger.warning('Canceled!')
+ except asyncio.CancelledError as e:
+ worker_logger.warning('Canceled!', exc_info=e)
self.download_queue.task_done()
return
except Empty:
@@ -654,10 +680,11 @@ class DownloadManager:
next(self.get_new_downloads())
continue
except StopIteration:
- pass
+ # Give any new downloads a chance to start up.
+ await asyncio.sleep(0.1)
if self.download_queue.empty():
- # Queue is empty.
+ # Queue is empty. Wait for background tasks.
break
@staticmethod
@@ -1188,7 +1215,7 @@ class RSSDownloader(Downloader, ABC):
def __repr__(self):
return ''
- def do_download(self, download: Download) -> DownloadResult:
+ async def do_download(self, download: Download) -> DownloadResult:
if isinstance(download.info_json, dict) and download.info_json.get('feed'):
feed = download.info_json['feed']
else:
diff --git a/wrolpi/files/downloader.py b/wrolpi/files/downloader.py
index 367fd04b..8c553e9a 100644
--- a/wrolpi/files/downloader.py
+++ b/wrolpi/files/downloader.py
@@ -1,3 +1,4 @@
+import asyncio
import pathlib
import traceback
from abc import ABC
diff --git a/wrolpi/status.py b/wrolpi/status.py
index 7f428a54..97b09b58 100755
--- a/wrolpi/status.py
+++ b/wrolpi/status.py
@@ -465,6 +465,8 @@ async def bandwidth_worker(count: int = None):
if not psutil:
return
+ logger.info('Bandwidth worker started')
+
nic_names = get_nic_names()
def append_all_stats():
diff --git a/wrolpi/test/test_downloader.py b/wrolpi/test/test_downloader.py
index 9cdfa8be..236eab61 100644
--- a/wrolpi/test/test_downloader.py
+++ b/wrolpi/test/test_downloader.py
@@ -177,8 +177,7 @@ async def test_calculate_next_download(test_session, test_download_manager, fake
@pytest.mark.asyncio
async def test_recurring_downloads(test_session, test_download_manager, fake_now, test_downloader):
"""A recurring Download should be downloaded forever."""
- test_downloader.do_download = MagicMock()
- test_downloader.do_download.return_value = DownloadResult(success=True)
+ test_downloader.set_test_success()
# Download every hour.
test_download_manager.recurring_download('https://example.com', Seconds.hour, test_downloader.name)
@@ -217,7 +216,7 @@ async def test_recurring_downloads(test_session, test_download_manager, fake_now
# Try the download, but it fails.
test_downloader.do_download.reset_mock()
- test_downloader.do_download.return_value = DownloadResult(success=False)
+ test_downloader.set_test_failure()
await test_download_manager.wait_for_all_downloads()
test_downloader.do_download.assert_called_once()
download = test_session.query(Download).one()
@@ -231,7 +230,7 @@ async def test_recurring_downloads(test_session, test_download_manager, fake_now
# Try the download again, it finally succeeds.
test_downloader.do_download.reset_mock()
now_ = fake_now(datetime(2020, 1, 1, 4, 0, 1, tzinfo=pytz.UTC))
- test_downloader.do_download.return_value = DownloadResult(success=True)
+ test_downloader.set_test_success()
test_download_manager.renew_recurring_downloads(test_session)
await test_download_manager.wait_for_all_downloads()
test_downloader.do_download.assert_called_once()
@@ -271,12 +270,11 @@ async def test_max_attempts(test_session, test_download_manager, test_downloader
@pytest.mark.asyncio
async def test_skip_urls(test_session, test_download_manager, test_client, assert_download_urls, test_downloader):
- """The DownloadManager will not create downloads for URLs in it's skip list."""
+ """The DownloadManager will not create downloads for URLs in its skip list."""
_, session = get_db_context()
get_download_manager_config().skip_urls = ['https://example.com/skipme']
- test_downloader.do_download = MagicMock()
- test_downloader.do_download.return_value = DownloadResult(success=True)
+ test_downloader.set_test_success()
test_download_manager.create_downloads([
'https://example.com/1',
diff --git a/wrolpi/test/test_rss.py b/wrolpi/test/test_rss.py
index fe034a1c..8fafcfc0 100644
--- a/wrolpi/test/test_rss.py
+++ b/wrolpi/test/test_rss.py
@@ -5,13 +5,13 @@ import mock
import pytest
from wrolpi.db import optional_session
-from wrolpi.downloader import Download, DownloadResult, RSSDownloader, Downloader
+from wrolpi.downloader import Download, DownloadResult, RSSDownloader, Downloader, DownloadManager
class RSSHTTPDownloader(Downloader):
name = 'rss_http'
- def do_download(self, download: Download) -> DownloadResult:
+ async def do_download(self, download: Download) -> DownloadResult:
return DownloadResult(success=True)
@optional_session
@@ -45,6 +45,7 @@ async def test_rss_download(test_session, test_download_manager):
dict(link='https://example.com/c'),
]
)
+ test_download_manager: DownloadManager
test_download_manager.create_download('https://example.com/feed', rss_downloader.name,
sub_downloader_name='rss_http')
await test_download_manager.wait_for_all_downloads()
@@ -92,4 +93,5 @@ async def test_rss_no_entries(test_session, test_download_manager):
await test_download_manager.wait_for_all_downloads()
(download,) = test_download_manager.get_downloads(test_session)
+ assert download.status == 'deferred'
assert 'entries' in download.error