Merge pull request #69 from lrnselfreliance/feature/download-info-saving-2

Squashed commit of the following:
This commit is contained in:
lrnselfreliance 2021-12-30 20:47:47 -07:00 • committed by GitHub
commit fbfa8b354f
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
17 changed files with 416 additions and 150 deletions

View file

@ -0,0 +1,33 @@
"""Downloader info
Revision ID: db0a1560d923
Revises: 84b6993a6091
Create Date: 2021-12-27 16:46:02.929119
"""
import os
from alembic import op
from sqlalchemy.orm import Session
# revision identifiers, used by Alembic.
revision = 'db0a1560d923'
down_revision = '84b6993a6091'
branch_labels = None
depends_on = None
DOCKERIZED = True if os.environ.get('DOCKER', '').lower().startswith('t') else False
def upgrade():
bind = op.get_bind()
session = Session(bind=bind)
session.execute('ALTER TABLE download ADD COLUMN IF NOT EXISTS info_json JSONB')
session.execute('ALTER TABLE download ADD COLUMN IF NOT EXISTS downloader TEXT')
def downgrade():
bind = op.get_bind()
session = Session(bind=bind)
session.execute('ALTER TABLE download DROP COLUMN IF EXISTS info_json')
session.execute('ALTER TABLE download DROP COLUMN IF EXISTS downloader')

View file

@ -36,7 +36,7 @@ async function apiCall(url, method, body, ms = 10000) {
if (response.status === 403 && (await response.json())['code'] === 17) {
toast({
type: 'error',
header: 'WROL Mode is enabled!',
title: 'WROL Mode is enabled!',
description: 'This functionality is not enabled while WROL Mode is enabled!',
time: 5000,
});
@ -47,14 +47,14 @@ async function apiCall(url, method, body, ms = 10000) {
if (e.message === 'promise timeout') {
toast({
type: 'error',
header: 'Server failed to respond!',
title: 'Server failed to respond!',
description: 'Timeout while waiting for server response. See server logs.',
time: 5000,
});
} else {
toast({
type: 'error',
header: 'Error!',
title: 'Error!',
description: 'See client logs',
time: 5000,
});
@ -138,7 +138,7 @@ export async function deleteVideo(video_id) {
if (data.code === 17) {
toast({
type: 'warning',
header: 'WROL Mode Enabled',
title: 'WROL Mode Enabled',
description: 'This cannot be done while WROL Mode is enabled.',
time: 5000,
});
@ -147,7 +147,7 @@ export async function deleteVideo(video_id) {
if (response.status !== 204) {
toast({
type: 'error',
header: 'Unexpected server response',
title: 'Unexpected server response',
description: 'Failed to delete video. See server logs.',
time: 5000,
});
@ -209,7 +209,16 @@ export async function download() {
export async function downloadChannel(link) {
let url = `${VIDEOS_API}/download/${link}`;
await fetch(url, {method: 'POST'});
let response = await apiPost(url);
let json = await response.json();
if (response.status === 400 && json['code'] === 30) {
toast({
type: 'error',
title: 'Cannot Download!',
description: 'This channel does not have a download record. Create one!',
time: 5000,
});
}
}
export async function refreshChannel(link) {
@ -219,11 +228,11 @@ export async function refreshChannel(link) {
export async function encryptOTP(otp, plaintext) {
let body = {otp, plaintext};
let response = await apiPost(`${OTP_API}:encrypt_otp`, body);
let response = await apiPost(`${OTP_API}/encrypt_otp`, body);
if (response.status !== 200) {
toast({
type: 'error',
header: 'Error!',
title: 'Error!',
description: 'Failed to encrypt OTP',
time: 5000,
});
@ -234,11 +243,11 @@ export async function encryptOTP(otp, plaintext) {
export async function decryptOTP(otp, ciphertext) {
let body = {otp, ciphertext};
let response = await apiPost(`${OTP_API}:decrypt_otp`, body);
let response = await apiPost(`${OTP_API}/decrypt_otp`, body);
if (response.status !== 200) {
toast({
type: 'error',
header: 'Error!',
title: 'Error!',
description: 'Failed to decrypt OTP',
time: 5000,
});
@ -339,8 +348,8 @@ export async function fetchDomains() {
}
}
export async function postDownload(urls) {
let body = {'urls': urls};
export async function postDownload(urls, downloader) {
let body = {urls: urls, downloader: downloader};
let response = await apiPost(`${API_URI}/download`, body);
return response;
}
@ -350,3 +359,8 @@ export async function killDownload(download_id) {
let response = await apiPost(`${API_URI}/download/${download_id}/kill`);
return response;
}
export async function getDownloaders() {
let response = await apiGet(`${API_URI}/downloaders`);
return await response.json();
}

View file

@ -219,7 +219,8 @@ class StoppableRow extends React.Component {
handleStart = async (e) => {
e.preventDefault();
await postDownload(`${this.props.url}`);
let downloader = this.props.downloader || null;
await postDownload(`${this.props.url}`, downloader);
this.closeStart();
await this.props.fetchDownloads();
};

View file

@ -269,6 +269,11 @@ class ChannelPage extends APIForm {
this.setState({deleteOpen: true});
}
handleCancel = (e) => {
e.preventDefault();
this.props.history.goBack();
}
render() {
return (
<Container>
@ -410,7 +415,7 @@ class ChannelPage extends APIForm {
<Button
secondary
floated='right'
onClick={() => this.props.history.goBack()}
onClick={this.handleCancel}
>
Cancel
</Button>

View file

@ -1,6 +1,6 @@
import React from "react";
import {Button, Form, Header, Segment, TextArea} from "semantic-ui-react";
import {postDownload} from "../api";
import {Button, Dropdown, Form, Header, Segment, TextArea} from "semantic-ui-react";
import {getDownloaders, postDownload} from "../api";
const validUrl = /^(http|https):\/\/[^ "]+$/;
@ -11,14 +11,29 @@ class Downloader extends React.Component {
urls: '',
valid: true,
pending: false,
downloaders: [],
downloader: null,
};
}
async componentDidMount() {
await this.fetchDownloaders();
}
fetchDownloaders = async () => {
let {downloaders} = await getDownloaders();
downloaders = downloaders.map((i) => {
return {key: i.name, text: i.pretty_name || i.name, value: i.name}
})
this.setState({downloaders});
}
submitDownload = async () => {
if (this.state.urls) {
let {urls, downloader} = this.state;
if (urls) {
this.setState({pending: true});
try {
let response = await postDownload(this.state.urls);
let response = await postDownload(urls, downloader);
if (response.status === 204) {
this.setState({urls: '', pending: false});
}
@ -52,14 +67,25 @@ class Downloader extends React.Component {
return (
<Form onSubmit={this.submitDownload}>
<Header as='h4'>Enter the URLs you wish to bring offline.</Header>
<TextArea
placeholder={textareaPlaceholder}
name='urls'
onChange={this.handleInputChange}
value={this.state.urls}
error={!this.state.valid}
/>
<Header as='h4'>Enter the URLs you wish to save offline.</Header>
<Form.Field>
<TextArea
placeholder={textareaPlaceholder}
name='urls'
onChange={this.handleInputChange}
value={this.state.urls}
error={!this.state.valid}
/>
</Form.Field>
<Form.Field>
<Dropdown
selection
name='downloader'
options={this.state.downloaders}
placeholder='Automatic'
onChange={this.handleInputChange}
/>
</Form.Field>
<Button primary style={{marginTop: '1em'}} disabled={disabled}>Download</Button>
</Form>
)

View file

@ -1,4 +1,5 @@
from abc import ABC
from typing import Tuple
from wrolpi.common import logger
from wrolpi.downloader import Downloader, Download
@ -12,12 +13,18 @@ logger = logger.getChild(__name__)
class ArchiveDownloader(Downloader, ABC):
name = 'archive'
pretty_name = 'Archive'
def __repr__(self):
return f'<ArchiveDownloader name={self.name}>'
@classmethod
def valid_url(cls, url: str) -> bool:
def valid_url(cls, url: str) -> Tuple[bool, None]:
"""
Archiver will attempt to archive anything, so it should be last!
"""
return True
return True, None
def do_download(self, download: Download):
if download.attempts > 3:

View file

@ -13,7 +13,7 @@ from .schema import EncryptOTPRequest, DecryptOTPRequest
bp = get_blueprint('OTP', '/api/otp')
@bp.post(':encrypt_otp')
@bp.post('/encrypt_otp')
@validate_doc(
summary='Encrypt a message with OTP',
consumes=EncryptOTPRequest,
@ -23,7 +23,7 @@ async def post_encrypt_otp(_: Request, data: dict):
return response.json(data)
@bp.post(':decrypt_otp')
@bp.post('/decrypt_otp')
@validate_doc(
summary='Decrypt a message with OTP',
consumes=DecryptOTPRequest,

View file

@ -1,6 +1,6 @@
from collections import defaultdict
from datetime import timedelta
from typing import List, Dict, Union
from typing import List, Dict, Union, Generator
from sqlalchemy.orm.exc import NoResultFound
@ -8,7 +8,7 @@ from wrolpi.common import sanitize_link, run_after, get_relative_to_media_direct
from wrolpi.dates import today
from wrolpi.db import get_db_session, get_db_curs
from wrolpi.downloader import download_manager, Download
from wrolpi.errors import UnknownChannel, UnknownDirectory, APIError, ValidationError
from wrolpi.errors import UnknownChannel, UnknownDirectory, APIError, ValidationError, InvalidDownload
from ..common import check_for_channel_conflicts
from ..lib import save_channels_config
from ..models import Channel
@ -97,7 +97,7 @@ def get_channel(source_id: str = None, link: str = None, url: str = None, return
return channel
def _spread_by_frequency(channels: List[Channel]):
def _spread_by_frequency(channels: List[Channel]) -> Generator[Dict, None, None]:
channels_by_frequency = defaultdict(lambda: [])
for channel in channels:
channels_by_frequency[channel.download_frequency].append(channel)
@ -129,13 +129,11 @@ def spread_channel_downloads():
url_next_download = _spread_by_frequency(channels)
for info in url_next_download:
download = download_manager.get_download(session, url=info['url'])
if not download:
download = Download(url=info['url'])
session.add(download)
download.frequency = info['frequency']
download.next_download = info['next_download']
url, frequency, next_download = info['url'], info['frequency'], info['next_download']
download = download_manager.get_or_create_download(url, session)
download.downloader = 'video_channel'
download.frequency = frequency
download.next_download = next_download
@run_after(save_channels_config)
@ -243,6 +241,12 @@ def delete_channel(link):
def download_channel(link: str):
"""
Create a Download record for a Channel's entire catalog.
"""
channel = get_channel(link=link, return_dict=False)
with get_db_session(commit=True) as session:
download_manager.create_download(channel.url, session)
if not download_manager.get_download(session, channel.url):
raise InvalidDownload(f'Channel {channel.name} does not have a download!')
download_manager.create_download(channel.url, session, downloader='video_channel')

View file

@ -2,12 +2,14 @@ import pathlib
import shutil
from uuid import uuid4
import mock
import pytest
from modules.videos.channel.lib import spread_channel_downloads
from modules.videos.downloader import VideoDownloader, ChannelDownloader
from modules.videos.models import Channel, Video
from wrolpi.common import sanitize_link
from wrolpi.downloader import DownloadFrequency
from wrolpi.downloader import DownloadFrequency, DownloadManager
from wrolpi.vars import PROJECT_DIR
@ -105,3 +107,19 @@ def video_factory(test_session, test_directory):
return video
return _
@pytest.fixture
def video_download_manager(test_download_manager) -> DownloadManager:
"""
Get a DownloadManager ready to download Videos and Channels.
"""
channel_downloader = ChannelDownloader()
video_downloader = VideoDownloader(priority=40)
test_download_manager.register_downloader(channel_downloader)
test_download_manager.register_downloader(video_downloader)
with mock.patch('modules.videos.downloader.YDL.extract_info') as mock_extract_info: # prevent real downloads
mock_extract_info.side_effect = Exception('You must patch modules.videos.downloader.YDL.extract_info!')
yield test_download_manager

View file

@ -3,8 +3,9 @@ import json
import pathlib
import re
from abc import ABC
from typing import Tuple, List
from typing import Tuple, List, Optional
from sqlalchemy.orm import Session
from sqlalchemy.orm.exc import NoResultFound
from yt_dlp import YoutubeDL
from yt_dlp.extractor import YoutubeTabIE
@ -47,14 +48,20 @@ class ChannelDownloader(Downloader, ABC):
"""
Handling downloading of all Videos in a Channel.
"""
name = 'video_channel'
pretty_name = 'Video Channel'
listable = False
def __repr__(self):
return f'<ChannelDownloader name={self.name}>'
@classmethod
def valid_url(cls, url) -> bool:
def valid_url(cls, url) -> Tuple[bool, None]:
for ie in ChannelIEs:
if ie.suitable(url):
return True
return True, None
logger.debug(f'{cls.__name__} not suitable for {url}')
return False
return False, None
def do_download(self, download: Download):
"""
@ -76,7 +83,8 @@ class ChannelDownloader(Downloader, ABC):
for video_id, source_id, missing_video in missing_videos:
url = video_url_resolver(domain, missing_video)
# Schedule any missing videos for download.
self.manager.get_or_create_download(url, session)
download = self.manager.get_or_create_download(url, session)
download.downloader = 'video'
if PYTEST:
self.manager.do_downloads_sync()
@ -90,36 +98,46 @@ class VideoDownloader(Downloader, ABC):
"""
Download a single video. Store the video in it's channel's directory, otherwise store it in `videos/NO CHANNEL`.
"""
name = 'video'
pretty_name = 'Videos'
def __repr__(self):
return f'<VideoDownloader name={self.name}>'
@classmethod
def valid_url(cls, url) -> bool:
def valid_url(cls, url) -> Tuple[bool, Optional[dict]]:
"""
Match against all Youtube-DL Info Extractors, except those that match a Channel.
"""
for ie in YDL._ies.values():
if ie.suitable(url) and not ChannelDownloader.valid_url(url):
if ie.suitable(url) and not ChannelDownloader.valid_url(url)[0]:
try:
YDL.extract_info(url, download=False, process=False)
return True
info = YDL.extract_info(url, download=False, process=False)
return True, info
except UnsupportedError:
logger.debug(f'Video downloader extract_info failed for {url}')
return False
return False, None
except DownloadError:
logger.debug(f'Video downloader extract_info failed for {url}')
return False
return False, None
logger.debug(f'{cls.__name__} not suitable for {url}')
return False
return False, None
def do_download(self, download: Download):
if download.attempts >= 10:
raise UnrecoverableDownloadError('Max download attempts reached')
url = download.url
try:
info = YDL.extract_info(url, download=False, process=False)
except UnsupportedError as e:
# Video wasn't really valid... probably the GenericIE.
raise UnrecoverableDownloadError() from e
info = download.info_json
if not info:
# Info was not fetched by the DownloadManager, lets get it.
valid, info = self.valid_url(url)
if not valid:
raise UnrecoverableDownloadError(f'{self} cannot download {url}')
session = Session.object_session(download)
download.info_json = info
session.commit()
channel_name = info.get('channel')
channel_id = info.get('channel_id')

View file

@ -255,9 +255,6 @@ class Channel(ModelHelper, Base):
info_json = Column(JSON)
info_date = Column(Date)
# Will be filled later by directory path relative to the media directory.
_directory = None
videos: InstrumentedList = relationship('Video', primaryjoin='Channel.id==Video.channel_id')
def __repr__(self):

View file

@ -3,9 +3,10 @@ from copy import copy
from datetime import datetime
from unittest import mock
import pytest
from yt_dlp.utils import UnsupportedError
from modules.videos.channel.lib import spread_channel_downloads
from modules.videos.channel.lib import spread_channel_downloads, download_channel
from modules.videos.downloader import find_all_missing_videos, VideoDownloader, \
ChannelDownloader, get_or_create_channel
from modules.videos.models import Channel, Video
@ -13,7 +14,9 @@ from modules.videos.test.common import create_channel_structure
from wrolpi.dates import local_timezone
from wrolpi.db import get_db_session, get_db_context
from wrolpi.downloader import DownloadManager, Download, DownloadFrequency
from wrolpi.errors import InvalidDownload
from wrolpi.test.common import wrap_test_db, TestAPI
from wrolpi.test.test_downloader import HTTPDownloader
class FakeYDL:
@ -94,7 +97,7 @@ example_channel_json = {
'ie_key': 'Youtube',
'title': 'video 1 title',
'uploader': None,
'url': 'video 1 url',
'url': 'video_1_url',
'view_count': 58504},
{'_type': 'url',
'description': None,
@ -103,7 +106,7 @@ example_channel_json = {
'ie_key': 'Youtube',
'title': 'video 2 title',
'uploader': None,
'url': 'video 2 url',
'url': 'video_2_url',
'view_count': 1413},
],
'extractor': 'youtube:tab',
@ -130,32 +133,33 @@ class TestVideosDownloaders(TestAPI):
def test_video_valid_url(self):
# A specific video can be downloaded.
self.assertTrue(self.vd.valid_url('https://www.youtube.com/watch?v=31jPEBiAC3c'))
self.assertTrue(self.vd.valid_url('https://www.youtube.com/watch?v=31jPEBiAC3c')[0])
# A channel cannot be downloaded.
self.assertFalse(self.vd.valid_url('https://www.youtube.com/c/LearningSelfReliance/videos'))
self.assertFalse(self.vd.valid_url('https://www.youtube.com/c/LearningSelfReliance'))
self.assertFalse(self.vd.valid_url('https://www.youtube.com/c/LearningSelfReliance/videos')[0])
self.assertFalse(self.vd.valid_url('https://www.youtube.com/c/LearningSelfReliance')[0])
def test_channel_valid_url(self):
# An entire domain cannot be downloaded.
self.assertFalse(self.cd.valid_url('https://example.com'))
self.assertFalse(self.cd.valid_url('https://youtube.com'))
self.assertFalse(self.cd.valid_url('https://example.com')[0])
self.assertFalse(self.cd.valid_url('https://youtube.com')[0])
# Cannot download a single video.
self.assertFalse(self.cd.valid_url('https://www.youtube.com/watch?v=31jPEBiAC3c'))
self.assertFalse(self.cd.valid_url('https://www.youtube.com/watch?v=31jPEBiAC3c')[0])
# Can download entire channels.
self.assertTrue(self.cd.valid_url('https://www.youtube.com/c/LearningSelfReliance/videos'))
self.assertTrue(self.cd.valid_url('https://www.youtube.com/c/LearningSelfReliance'))
self.assertTrue(self.cd.valid_url('https://www.youtube.com/c/LearningSelfReliance/videos')[0])
self.assertTrue(self.cd.valid_url('https://www.youtube.com/c/LearningSelfReliance')[0])
# Can download entire playlists.
self.assertTrue(self.cd.valid_url('https://www.youtube.com/playlist?list=PLCdlMQeP-TbG12nkBCt0E96yr3EfwcH4E'))
self.assertTrue(
self.cd.valid_url('https://www.youtube.com/playlist?list=PLCdlMQeP-TbG12nkBCt0E96yr3EfwcH4E')[0])
def test_get_downloader(self):
# The correct Downloader is gotten.
self.assertEqual(self.mgr.get_downloader('https://www.youtube.com/c/LearningSelfReliance/videos'), self.cd)
self.assertEqual(self.mgr.get_downloader('https://www.youtube.com/c/LearningSelfReliance'), self.cd)
self.assertEqual(self.mgr.get_downloader('https://www.youtube.com/watch?v=31jPEBiAC3c'), self.vd)
self.assertEqual(self.mgr.get_downloader('https://www.youtube.com/c/LearningSelfReliance/videos')[0], self.cd)
self.assertEqual(self.mgr.get_downloader('https://www.youtube.com/c/LearningSelfReliance')[0], self.cd)
self.assertEqual(self.mgr.get_downloader('https://www.youtube.com/watch?v=31jPEBiAC3c')[0], self.vd)
@staticmethod
def _make_video_files(channel_dir):
@ -269,7 +273,7 @@ class TestVideosDownloaders(TestAPI):
_, session = get_db_context()
with mock.patch('modules.videos.downloader.VideoDownloader.valid_url') as mock_valid_url, \
mock.patch('modules.videos.downloader.YDL') as mock_ydl:
mock_valid_url.return_value = True
mock_valid_url.return_value = (True, None)
mock_ydl.extract_info.side_effect = UnsupportedError('oops')
self.mgr.create_download('url', session)
@ -360,3 +364,95 @@ class TestVideosDownloaders(TestAPI):
# New channel can be retrieved.
self.assertEqual(get_or_create_channel(source_id='quux'), channel)
def test_download_info_save(test_session, video_download_manager):
"""
An info dict from Downloader.valid_url is saved to the Downloader and can be used by the Downloader later.
"""
with mock.patch('modules.videos.downloader.YDL.extract_info') as mock_extract_info:
mock_extract_info.return_value = {'some': 'info'}
video_download_manager.create_download('https://www.youtube.com/watch?v=HQ_62YwcA80', test_session)
mock_extract_info.assert_called_once()
def test_selected_downloader(test_session, video_download_manager):
"""
A user can specify which Downloader to use, no other should ever be used (even if the Download fails).
"""
# The test HTTP Downloader will claim to be able to handle any HTTP URL; will pretend to have downloaded anything.
http_downloader = HTTPDownloader()
mock_do_download = http_downloader.do_download = mock.MagicMock()
mock_http_valid_url = http_downloader.valid_url = mock.MagicMock()
mock_http_valid_url.return_value = (True, None)
video_download_manager.register_downloader(http_downloader)
def check_attempts(attempts):
assert test_session.query(Download).one().attempts == attempts
# DownloadManager will attempt to choose the Downloader.
with mock.patch('modules.videos.downloader.YDL.extract_info') as mock_extract_info:
mock_extract_info.side_effect = UnsupportedError('not this one')
video_download_manager.create_downloads(['https://example.com'], test_session)
mock_http_valid_url.assert_called_once()
mock_extract_info.assert_called_once()
mock_do_download.assert_called_once()
check_attempts(1)
mock_http_valid_url.reset_mock()
mock_do_download.reset_mock()
# DownloadManager will only use the HTTPDownloader.
with mock.patch('modules.videos.downloader.YDL.extract_info') as mock_extract_info:
video_download_manager.create_downloads(['https://example.com'], test_session, downloader='http')
# DownloadManager does not try to find the Download, it trusts the "downloader" param.
mock_extract_info.assert_not_called()
mock_http_valid_url.assert_not_called()
# DownloadManager called do_download.
mock_do_download.assert_called_once()
check_attempts(2)
mock_http_valid_url.reset_mock()
mock_do_download.reset_mock()
# DownloadManager will only use the HTTPDownloader, even if the download fails.
with mock.patch('modules.videos.downloader.YDL.extract_info') as mock_extract_info:
mock_do_download.side_effect = Exception('oh no!')
video_download_manager.create_downloads(['https://example.com'], test_session, downloader='http')
# DownloadManager does not try to find the Download, it trusts the "downloader" param.
mock_extract_info.assert_not_called()
mock_http_valid_url.assert_not_called()
# DownloadManager called do_download.
mock_do_download.assert_called_once()
check_attempts(3)
def test_channel_downloader_hidden(video_download_manager):
"""
ChannelDownloader should not be presented to the User.
"""
downloaders = video_download_manager.list_downloaders()
assert [i.__json__() for i in downloaders] == [(dict(name='video', pretty_name='Videos')), ]
def test_bad_downloader(test_session, video_download_manager):
"""
Attempting to use an unknown downloader should raise an error.
"""
with pytest.raises(InvalidDownload):
video_download_manager.create_download('https://example.com', test_session, 'bad downloader')
def test_channel_download_no_download(test_session, video_download_manager, simple_channel):
"""
A Channel without an existing Download cannot be downloaded. This is to avoid the situation where a single
video has been downloaded in a Channel, but the User has not requested downloading of all videos in the channel.
"""
# simple_channel has no download record.
assert video_download_manager.get_downloads(test_session) == []
with pytest.raises(InvalidDownload):
download_channel(simple_channel.link)

View file

@ -16,6 +16,7 @@ from sqlalchemy.orm import Session, sessionmaker
from wrolpi.common import set_test_media_directory, Base, get_example_config
from wrolpi.db import postgres_engine, get_db_args
from wrolpi.downloader import DownloadManager
from wrolpi.root_api import BLUEPRINTS, api_app
os.environ['DB_PORT'] = '54321'
@ -113,3 +114,9 @@ def test_client() -> SanicTestClient:
ROUTES_ATTACHED = True
return api_app.test_client
@pytest.fixture
def test_download_manager():
manager = DownloadManager()
return manager

View file

@ -7,10 +7,10 @@ from datetime import timedelta
from enum import Enum
from functools import partial
from operator import attrgetter
from typing import List, Optional
from typing import List, Optional, Tuple, Dict
from cachetools import cached, TTLCache
from sqlalchemy import Column, Integer, String
from sqlalchemy import Column, Integer, String, Text
from sqlalchemy.dialects.postgresql import JSONB
from sqlalchemy.orm import Session
from wrolpi.common import Base, ModelHelper, logger, iterify
@ -47,6 +47,8 @@ class Download(ModelHelper, Base):
last_successful_download = Column(TZDateTime)
next_download = Column(TZDateTime)
status = Column(String, default='new')
info_json = Column(JSONB)
downloader = Column(Text)
def __repr__(self):
if self.next_download or self.frequency:
@ -102,16 +104,22 @@ class Download(ModelHelper, Base):
class Downloader:
name = None
pretty_name = None
listable = True
def __init__(self, priority: int = 50):
def __init__(self, priority: int = 50, name: str = None):
"""
Lower `priority` means the downloader will be checked first. Valid priority: 0-100
Downloaders of equal priority will be used at random.
"""
if not name and not self.name:
raise NotImplementedError('Downloader must have a name!')
if not 0 <= priority <= 100:
raise ValueError(f'Priority of {priority} of {self} is invalid!')
self.name = self.name or name
self.priority = priority
self._kill = multiprocessing.Event()
@ -119,8 +127,11 @@ class Downloader:
download_manager.register_downloader(self)
def __json__(self):
return dict(name=self.name, pretty_name=self.pretty_name)
@classmethod
def valid_url(cls, url) -> bool:
def valid_url(cls, url) -> Tuple[bool, Optional[dict]]:
raise NotImplementedError()
def do_download(self, download: Download):
@ -148,7 +159,7 @@ class Downloader:
Clear any "kill" request for this Downloader.
"""
if self._kill.is_set():
self.clear()
self._kill.clear()
def process_runner(self, url: str, cmd: List[str], cwd: pathlib.Path, **kwargs) -> int:
"""
@ -192,6 +203,7 @@ class DownloadManager:
def __init__(self):
self.instances = tuple()
self._instances = dict()
def register_downloader(self, instance: Downloader):
if not isinstance(instance, Downloader):
@ -201,22 +213,23 @@ class DownloadManager:
instance.manager = self
instances = (*self.instances, instance)
self.instances = tuple(self.priority_sorter(instances))
i = (*self.instances, instance)
self.instances = tuple(self.priority_sorter(i))
self._instances[instance.name] = instance
@cached(cache=TTLCache(maxsize=100, ttl=600))
def get_downloader(self, url: str):
def get_downloader(self, url: str) -> Tuple[Downloader, Dict]:
for i in self.instances:
if i.valid_url(url):
return i
result = i.valid_url(url)
valid, info = result
if valid:
return i, info
raise InvalidDownload(f'Invalid URL {url=}')
def valid_url(self, url: str) -> bool:
try:
self.get_downloader(url)
return True
except InvalidDownload:
return False
def get_downloader_by_name(self, name: str) -> Downloader:
"""
Attempts to find a registered Downloader by its name. Returns None if it cannot be found.
"""
return self._instances.get(name)
@staticmethod
def get_new_downloads(session: Session) -> List[Download]:
@ -246,33 +259,17 @@ class DownloadManager:
session.flush()
return download
def create_download(self, url: str, session, skip_download: bool = False, reset_attempts: bool = False) -> Download:
def create_download(self, url: str, session, downloader: str = None, skip_download: bool = False,
reset_attempts: bool = False) -> Download:
"""
Schedule a URL for download. If the URL failed previously, it will be retried.
Schedule a URL for download. If the URL failed previously, it may be retried.
"""
# Make sure the URL can be downloaded.
self.get_downloader(url)
downloads = self.create_downloads([url], session, downloader, skip_download=skip_download,
reset_attempts=reset_attempts)
return downloads[0]
download = self.get_or_create_download(url, session)
# Download may have failed, try again.
download.renew(reset_attempts=reset_attempts)
session.commit()
if skip_download is True:
return download
if PYTEST:
# Download now (for testing).
self._do_downloads(session)
else:
# Start downloading ASAP.
self.start_downloads()
return download
def create_downloads(self, urls: List[str], session: Session = None, skip_download: bool = False,
reset_attempts: bool = False) -> List[Download]:
def create_downloads(self, urls: List[str], session: Session = None, downloader: str = None,
skip_download: bool = False, reset_attempts: bool = False) -> List[Download]:
"""
Schedule all URLs for download. If one cannot be downloaded, none will be added.
"""
@ -280,16 +277,24 @@ class DownloadManager:
_, session = get_db_context()
downloads = []
# Verify that all URLs can be downloaded.
for url in urls:
self.get_downloader(url)
forced_downloader = self.get_downloader_by_name(downloader) if downloader else None
if downloader and not forced_downloader:
# Could not find the downloader
raise InvalidDownload(f'Unknown downloader {downloader}')
with session.transaction:
for url in urls:
info_json = None
downloader = forced_downloader
if not forced_downloader:
# User has requested automatic downloader selection, try and find it.
downloader, info_json = self.get_downloader(url)
download = self.get_or_create_download(url, session)
# Download may have failed, try again.
download.renew(reset_attempts=reset_attempts)
download.downloader = downloader.name
download.info_json = info_json
downloads.append(download)
if skip_download is True:
@ -305,6 +310,9 @@ class DownloadManager:
return downloads
def _do_downloads(self, session=None):
"""
This method calls the Downloader's do_download method.
"""
# This is a long-running function, lets get a session that can be used for a long time.
if not session:
_, session = get_db_context()
@ -316,8 +324,13 @@ class DownloadManager:
for download in downloads:
download_count += 1
download_id = download.id
url = download.url
downloader = self.get_downloader(download.url)
downloader = self.get_downloader_by_name(download.downloader)
if not downloader:
downloader, info_json = self.get_downloader(download.url)
download.downloader = downloader.name
download.info_json = info_json
downloader.clear()
download.started()
@ -332,7 +345,7 @@ class DownloadManager:
logger.warning(f'UnrecoverableDownloadError for {download.url}', exc_info=e)
failure = True
except Exception as e:
logger.warning(f'Failed to download {download.url}. Will be tried again later.', exc_info=e)
logger.warning(f'Failed to download {url}. Will be tried again later.', exc_info=e)
# Long-running, get the download again.
with get_db_session(commit=True) as session_:
@ -499,7 +512,10 @@ class DownloadManager:
"""
with get_db_session(commit=True) as session:
download = self.get_download(session, id_=download_id)
downloader = self.get_downloader(download.url)
if not download.downloader:
downloader, _ = self.get_downloader(download.url)
else:
downloader = self.get_downloader_by_name(download.downloader)
logger.warning(f'Killing download {download_id} in {downloader}')
if download.status == 'pending':
downloader.kill()
@ -519,6 +535,12 @@ class DownloadManager:
download.last_successful_download < one_month:
session.delete(download)
def list_downloaders(self) -> List[Downloader]:
"""
Return a list of the Downloaders available on this Download Manager.
"""
return list(filter(lambda i: i.listable, self.instances))
# The global DownloadManager. This should be used everywhere!
download_manager = DownloadManager()

View file

@ -181,7 +181,10 @@ def valid_regex(_: Request, data: dict):
async def post_download(request: Request, data: dict):
# URLs are provided in a textarea, lets split all lines.
urls = [i.strip() for i in str(data['urls']).strip().splitlines()]
download_manager.create_downloads(urls, reset_attempts=True)
downloader = data.get('downloader')
if not downloader or downloader in ('auto', 'None', 'null'):
downloader = None
download_manager.create_downloads(urls, downloader=downloader, reset_attempts=True)
return response.empty()
@ -210,6 +213,14 @@ async def delete_download(request: Request, download_id: int):
return response.empty(HTTPStatus.NO_CONTENT if deleted else HTTPStatus.NOT_FOUND)
@root_api.get('/downloaders')
async def get_downloaders(request: Request):
downloaders = download_manager.list_downloaders()
downloaders.insert(0, dict(name='auto', pretty_name='Automatic'))
ret = dict(downloaders=downloaders)
return json_response(ret)
class CustomJSONEncoder(json.JSONEncoder):
def default(self, obj):

View file

@ -179,7 +179,8 @@ class EventsResponse:
class DownloadRequest:
urls = doc.String()
urls = doc.String(required=True)
downloader = doc.String()
class Trinary(Field):

View file

@ -1,8 +1,11 @@
from abc import ABC
from datetime import datetime, timedelta
from itertools import zip_longest
from unittest import mock
from unittest.mock import MagicMock
import pytest
from wrolpi import downloader
from wrolpi.dates import local_timezone
from wrolpi.db import get_db_context
@ -15,18 +18,26 @@ class PermissiveDownloader(Downloader):
"""
A testing Downloader which always says it's valid.
"""
name = 'permissive'
def __repr__(self):
return '<TESTING Permissive Downloader>'
def valid_url(self, url: str):
return True
return True, None
class HTTPDownloader(Downloader):
"""
A testing Downloader which says its valid when a URL starts with http/https
"""
name = 'http'
def __repr__(self):
return '<TESTING HTTP Downloader>'
def valid_url(self, url: str):
return url.startswith('https://') or url.startswith('http://')
return url.startswith('https://') or url.startswith('http://'), None
class TestDownloader(TestAPI):
@ -54,25 +65,6 @@ class TestDownloader(TestAPI):
self.assertRaises(ValueError, self.mgr.register_downloader, http_downloader)
self.assertRaises(ValueError, self.mgr.register_downloader, permissive_downloader)
def test_valid_url(self):
# No downloaders available.
self.assertFalse(self.mgr.valid_url('foo'))
self.assertFalse(self.mgr.valid_url('https://example.com'))
http_downloader = HTTPDownloader()
self.mgr.register_downloader(http_downloader)
self.assertFalse(self.mgr.valid_url('foo'))
self.assertTrue(self.mgr.valid_url('https://example.com'))
self.assertEqual(self.mgr.get_downloader('https://example.com'), http_downloader)
# Last priority.
permissive_downloader = PermissiveDownloader(priority=100)
self.mgr.register_downloader(permissive_downloader)
self.assertTrue(self.mgr.valid_url('foo'))
self.assertEqual(self.mgr.get_downloader('foo'), permissive_downloader)
self.assertTrue(self.mgr.valid_url('https://example.com'))
self.assertEqual(self.mgr.get_downloader('https://example.com'), http_downloader)
@wrap_test_db
def test_ensure_download(self):
_, session = get_db_context()
@ -301,3 +293,17 @@ class TestDownloader(TestAPI):
skip_download=True)
downloads = self.mgr.get_downloads(session)
self.assertEqual({i.url for i in downloads}, {'https://example.com/1', 'https://example.com/2'})
def test_downloader_must_have_name():
"""
A Downloader class must have a name.
"""
with pytest.raises(NotImplementedError):
Downloader()
class D(Downloader, ABC):
pass
with pytest.raises(NotImplementedError):
D()