diff --git a/alembic/versions/db0a1560d923_downloader_info.py b/alembic/versions/db0a1560d923_downloader_info.py new file mode 100644 index 00000000..184029b3 --- /dev/null +++ b/alembic/versions/db0a1560d923_downloader_info.py @@ -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') diff --git a/modules/archive/__init__.py b/modules/archive/__init__.py index 641e118c..89005692 100644 --- a/modules/archive/__init__.py +++ b/modules/archive/__init__.py @@ -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) diff --git a/modules/videos/conftest.py b/modules/videos/conftest.py index 230b6f7c..55363f85 100644 --- a/modules/videos/conftest.py +++ b/modules/videos/conftest.py @@ -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 diff --git a/modules/videos/downloader.py b/modules/videos/downloader.py index 73a4e29b..72583a07 100755 --- a/modules/videos/downloader.py +++ b/modules/videos/downloader.py @@ -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: diff --git a/modules/videos/test/test_downloader.py b/modules/videos/test/test_downloader.py index a34023e8..d4ace302 100644 --- a/modules/videos/test/test_downloader.py +++ b/modules/videos/test/test_downloader.py @@ -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() diff --git a/wrolpi/conftest.py b/wrolpi/conftest.py index 590ed18d..5eae5e86 100644 --- a/wrolpi/conftest.py +++ b/wrolpi/conftest.py @@ -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 diff --git a/wrolpi/downloader.py b/wrolpi/downloader.py index 9bab06b6..50916281 100644 --- a/wrolpi/downloader.py +++ b/wrolpi/downloader.py @@ -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_: diff --git a/wrolpi/test/test_downloader.py b/wrolpi/test/test_downloader.py index 8f5b53df..bd49cb4b 100644 --- a/wrolpi/test/test_downloader.py +++ b/wrolpi/test/test_downloader.py @@ -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()