rt/wrolpi
1
0
Fork 0
mirror of https://github.com/lrnselfreliance/wrolpi.git synced 2026-10-26 19:01:39 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

Download info can be saved during the download process and can be used by the Downloader.

Browse source
This commit is contained in:
Roland 2021-12-28 18:07:42 -07:00
parent f830fe3dcc
commit f5ce952da0
No known key found for this signature in database
GPG key ID: 9D6DBD3ACF11AB93
8 changed files with 153 additions and 90 deletions
Download patch file Download diff file Expand all files Collapse all files

33
alembic/versions/db0a1560d923_downloader_info.py Normal file
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')

7
modules/archive/__init__.py
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
@ -13,11 +14,11 @@ logger = logger.getChild(__name__)
class ArchiveDownloader(Downloader, ABC):
@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:
@ -28,4 +29,4 @@ class ArchiveDownloader(Downloader, ABC):
# Archive downloader is the last downloader which should be used.
archive_downloader = ArchiveDownloader(priority=100)
archive_downloader = ArchiveDownloader('archive', priority=100)

20
modules/videos/conftest.py
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_channel')
video_downloader = VideoDownloader('video', 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

32
modules/videos/downloader.py
View file

@ -3,7 +3,7 @@ import json
import pathlib
import re
from abc import ABC
from typing import Tuple, List
from typing import Tuple, List, Optional
from sqlalchemy.orm.exc import NoResultFound
from yt_dlp import YoutubeDL
@ -49,12 +49,12 @@ class ChannelDownloader(Downloader, ABC):
"""
@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):
"""
@ -92,34 +92,30 @@ class VideoDownloader(Downloader, ABC):
"""
@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
channel_name = info.get('channel')
channel_id = info.get('channel_id')
@ -209,9 +205,9 @@ class VideoDownloader(Downloader, ABC):
return final_filename, entry
channel_downloader = ChannelDownloader()
channel_downloader = ChannelDownloader('video_channel')
# Videos may match the ChannelDownloader, give it a higher priority.
video_downloader = VideoDownloader(40)
video_downloader = VideoDownloader('video', 40)
def get_or_create_channel(source_id: str = None, link: str = None, url: str = None, name: str = None) -> Channel:

45
modules/videos/test/test_downloader.py
View file

@ -94,7 +94,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 +103,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',
@ -119,8 +119,8 @@ class TestVideosDownloaders(TestAPI):
def setUp(self) -> None:
super().setUp()
self.vd = VideoDownloader()
self.cd = ChannelDownloader()
self.vd = VideoDownloader('video')
self.cd = ChannelDownloader('channel')
self.mgr = DownloadManager()
self.mgr.register_downloader(self.vd)
self.mgr.register_downloader(self.cd)
@ -130,32 +130,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 +270,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 +361,13 @@ 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()

7
wrolpi/conftest.py
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

58
wrolpi/downloader.py
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 = Column(JSONB)
downloader = Column(Text)
def __repr__(self):
if self.next_download or self.frequency:
@ -103,15 +105,18 @@ class Download(ModelHelper, Base):
class Downloader:
def __init__(self, priority: int = 50):
def __init__(self, name: str, priority: int = 50):
"""
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:
raise ValueError('Downloader must have a name!')
if not 0 <= priority <= 100:
raise ValueError(f'Priority of {priority} of {self} is invalid!')
self.name = name
self.priority = priority
self._kill = multiprocessing.Event()
@ -120,7 +125,7 @@ class Downloader:
download_manager.register_downloader(self)
@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 +153,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 +197,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 +207,22 @@ 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
valid, info = i.valid_url(url)
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]:
@ -251,12 +257,18 @@ class DownloadManager:
Schedule a URL for download. If the URL failed previously, it will be retried.
"""
# Make sure the URL can be downloaded.
self.get_downloader(url)
downloader, info = self.get_downloader(url)
if not isinstance(info, dict) and info is not None:
logger.warning(f'info from {downloader}.is_valid is the wrong type! {type(info)=}')
download = self.get_or_create_download(url, session)
# Download may have failed, try again.
download.renew(reset_attempts=reset_attempts)
download.downloader = downloader.name
if isinstance(info, dict):
download.info = info
session.commit()
if skip_download is True:
@ -316,8 +328,12 @@ 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, _ = self.get_downloader(download.url)
download.downloader = downloader.name
downloader.clear()
download.started()
@ -332,7 +348,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_:

41
wrolpi/test/test_downloader.py
View file

@ -17,7 +17,7 @@ class PermissiveDownloader(Downloader):
"""
def valid_url(self, url: str):
return True
return True, None
class HTTPDownloader(Downloader):
@ -26,7 +26,7 @@ class HTTPDownloader(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):
@ -42,44 +42,25 @@ class TestDownloader(TestAPI):
"""
self.assertEqual(self.mgr.instances, tuple())
http_downloader = HTTPDownloader()
http_downloader = HTTPDownloader('http')
self.mgr.register_downloader(http_downloader)
self.assertEqual(self.mgr.instances, (http_downloader,))
# PermissiveDownloader is first priority.
permissive_downloader = PermissiveDownloader(priority=0)
permissive_downloader = PermissiveDownloader('permissive', priority=0)
self.mgr.register_downloader(permissive_downloader)
self.assertEqual(self.mgr.instances, (permissive_downloader, http_downloader))
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()
self.assertEqual(len(self.mgr.get_downloads(session)), 0)
http_downloader = HTTPDownloader()
http_downloader = HTTPDownloader('http')
http_downloader.do_download.return_value = False
self.mgr.register_downloader(http_downloader)
@ -95,11 +76,11 @@ class TestDownloader(TestAPI):
def test_do_downloads(self):
_, session = get_db_context()
http_downloader = HTTPDownloader()
http_downloader = HTTPDownloader('http')
http_downloader.do_download.return_value = True
self.mgr.register_downloader(http_downloader)
permissive_downloader = PermissiveDownloader(priority=100)
permissive_downloader = PermissiveDownloader('permissive', priority=100)
permissive_downloader.do_download.return_value = False # returns a failure
self.mgr.register_downloader(permissive_downloader)
@ -141,7 +122,7 @@ class TestDownloader(TestAPI):
"""
_, session = get_db_context()
http_downloader = HTTPDownloader()
http_downloader = HTTPDownloader('http')
http_downloader.do_download.return_value = False
self.mgr.register_downloader(http_downloader)
@ -167,7 +148,7 @@ class TestDownloader(TestAPI):
"""
_, session = get_db_context()
http_downloader = HTTPDownloader()
http_downloader = HTTPDownloader('http')
http_downloader.do_download.return_value = True
http_downloader.do_download: MagicMock # noqa
self.mgr.register_downloader(http_downloader)
@ -243,7 +224,7 @@ class TestDownloader(TestAPI):
"""
Once-downloads over a month old should be deleted.
"""
permissive_downloader = PermissiveDownloader(priority=0)
permissive_downloader = PermissiveDownloader('permissive', priority=0)
self.mgr.register_downloader(permissive_downloader)
_, session = get_db_context()
@ -286,7 +267,7 @@ class TestDownloader(TestAPI):
"""
Multiple downloads can be scheduled using DownloadManager.create_downloads.
"""
http_downloader = HTTPDownloader(priority=0)
http_downloader = HTTPDownloader('http', priority=0)
self.mgr.register_downloader(http_downloader)
_, session = get_db_context()