diff --git a/modules/archive/__init__.py b/modules/archive/__init__.py index 75980466..4a445e21 100644 --- a/modules/archive/__init__.py +++ b/modules/archive/__init__.py @@ -2,7 +2,8 @@ import asyncio import json import pathlib from abc import ABC -from typing import Callable, List, Tuple, Iterable +from dataclasses import dataclass +from typing import Callable, List, Optional, Tuple, Iterable from sqlalchemy import not_ from sqlalchemy.orm import Session @@ -12,14 +13,15 @@ from wrolpi.collections import Collection from wrolpi.common import logger, register_modeler, register_refresh_cleanup, limit_concurrent, split_lines_by_length, \ slow_logger, get_title_from_html, TRACE_LEVEL from wrolpi.db import get_db_session -from wrolpi.downloader import Downloader, Download, DownloadResult +from wrolpi.downloader import Downloader, Download, DownloadContext, DownloadResult from wrolpi.errors import UnrecoverableDownloadError from wrolpi.files.models import FileGroup from wrolpi.vars import PYTEST, DOCKERIZED, DOWNLOAD_USER_AGENT from . import lib from .api import archive_bp # noqa from .errors import InvalidArchive -from .lib import is_singlefile_file, request_archive, SINGLEFILE_HEADER, get_url_from_singlefile +from .lib import is_singlefile_file, request_archive, SINGLEFILE_HEADER, get_url_from_singlefile, \ + SinglefileArtifacts, WrittenArchiveFiles, extract_singlefile_artifacts, register_archive, write_archive_files from .models import Archive PRETTY_NAME = 'Archive' @@ -29,6 +31,24 @@ logger = logger.getChild(__name__) __all__ = ['ArchiveDownloader', 'archive_downloader', 'model_archive'] +@dataclass +class PreparedArchive: + """Plan produced by ArchiveDownloader.prepare_download. + + Holds everything execute_download needs (URL, destination, settings) so the I/O + phase has no DB session. + """ + url: str + destination: Optional[pathlib.Path] + settings: dict + + +@dataclass +class ExecutedArchive: + """Output of ArchiveDownloader.execute_download: files are on disk and ready to register.""" + written: WrittenArchiveFiles + + class ArchiveDownloader(Downloader, ABC): name = 'archive' pretty_name = 'Archive' @@ -40,59 +60,57 @@ class ArchiveDownloader(Downloader, ABC): file_groups = list(session.query(FileGroup).filter(FileGroup.url.in_(urls), FileGroup.model == 'archive')) return file_groups - async def do_download(self, download: Download) -> DownloadResult: - if download.attempts > 3: + def prepare_download(self, session: Session, download: Download) -> PreparedArchive: + """Resolve destination directory. Caller commits (or rolls back on raise).""" + if (download.attempts or 0) > 3: raise UnrecoverableDownloadError(f'Max download attempts reached for {download.url}') - # Get destination from settings (passed by RSS downloader from download.destination column) - destination = None - if download.settings and download.settings.get('destination'): - destination = pathlib.Path(download.settings['destination']) + settings = download.settings or dict() + # Destination from settings (e.g. set by the RSS sub-downloader) wins over collection lookup. + destination = pathlib.Path(settings['destination']) if settings.get('destination') else None - # If no destination provided, use the domain collection's directory - # get_or_set_directory will format from template and save on first use if not destination: - with get_db_session(commit=True) as session: - collection = lib.get_or_create_domain_collection(session, download.url) - if collection: - destination = collection.get_or_set_directory(session) + collection = lib.get_or_create_domain_collection(session, download.url) + if collection: + destination = collection.get_or_set_directory(session) + return PreparedArchive(url=download.url, destination=destination, settings=settings) + + async def execute_download(self, prepared: PreparedArchive, ctx: DownloadContext) -> ExecutedArchive: + """Run the archive container or local singlefile + readability/screenshot extraction, + then write the artifacts to disk. No DB.""" 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) - archive: Archive = await lib.model_archive_result( - download.url, singlefile, readability, screenshot, destination=destination) - archive_id = archive.id + singlefile, readability, screenshot = await request_archive(prepared.url) + artifacts = SinglefileArtifacts(url=prepared.url, singlefile=singlefile, + readability=readability, screenshot=screenshot) else: - # Perform the archive using locally installed executables. - singlefile = await self.do_singlefile(download) - archive = await lib.singlefile_to_archive(singlefile, destination=destination) - archive_id = archive.id + singlefile_bytes = await self.do_singlefile(prepared, ctx) + artifacts = await extract_singlefile_artifacts(singlefile_bytes) + + written = write_archive_files(artifacts.url, artifacts.singlefile, artifacts.readability, + artifacts.screenshot, destination=prepared.destination) + return ExecutedArchive(written=written) + + def finalize_download(self, session: Session, download: Download, + executed: ExecutedArchive) -> DownloadResult: + """Register the Archive record, apply tags, persist parent_download_url. Caller commits.""" + archive = register_archive(session, executed.written) + + if tag_names := download.tag_names: + for name in tag_names: + archive.add_tag(session, name) settings = download.settings or dict() + if parent_download_url := settings.get('parent_download_url'): + try: + archive.file_group.update_wrolpi_json({'parent_download_url': parent_download_url}) + except ValueError as e: + logger.warning(f'Failed to update wrolpi json for {archive}: {e}') - with get_db_session() as session: - archive = Archive.find_by_id(session, archive_id) - need_commit = False - if tag_names := download.tag_names: - for name in tag_names: - archive.add_tag(session, name) - need_commit = True + logger.info(f'Successfully downloaded Archive {download.url} {archive}') + return DownloadResult(success=True, location=archive.location) - if need_commit: - session.commit() - - if parent_download_url := settings.get('parent_download_url'): - try: - archive.file_group.update_wrolpi_json({'parent_download_url': parent_download_url}) - except ValueError as e: - logger.warning(f'Failed to update wrolpi json for {archive}: {e}') - - logger.info(f'Successfully downloaded Archive {download.url} {archive}') - - return DownloadResult(success=True, location=archive.location) - - async def do_singlefile(self, download: Download) -> bytes: + async def do_singlefile(self, prepared: PreparedArchive, ctx: DownloadContext) -> bytes: """Create a Singlefile from the archive's URL.""" # Get browser settings from config config = lib.get_archive_downloader_config() @@ -119,11 +137,12 @@ class ArchiveDownloader(Downloader, ABC): '--user-agent', user_agent, '--dump-content', '--load-deferred-images-dispatch-scroll-event', - download.url, + prepared.url, ) cwd = pathlib.Path('/home/wrolpi') cwd = cwd if cwd.is_dir() else None - result = await self.process_runner(download, cmd, cwd) + # cancel_wrapper reads cancellation from ctx; the Download is only for log messages. + result = await self.process_runner(Download(url=prepared.url), cmd, cwd, ctx=ctx) stderr = result.stderr.decode() log_output = stderr or result.stdout.decode() or 'No stderr or stdout!' @@ -134,11 +153,11 @@ class ArchiveDownloader(Downloader, ABC): if not result.stdout: e = ChildProcessError(log_output[:1000]) - raise RuntimeError(f'Singlefile created was empty: {download.url}') from e + raise RuntimeError(f'Singlefile created was empty: {prepared.url}') from e if SINGLEFILE_HEADER.encode() not in result.stdout: e = ChildProcessError(log_output[:1000]) - raise RuntimeError(f'Singlefile created was invalid: {download.url}') from e + raise RuntimeError(f'Singlefile created was invalid: {prepared.url}') from e return result.stdout diff --git a/modules/archive/conftest.py b/modules/archive/conftest.py index 9fc6648c..07490064 100644 --- a/modules/archive/conftest.py +++ b/modules/archive/conftest.py @@ -9,7 +9,26 @@ import pytz from modules.archive.lib import archive_strftime from modules.archive.models import Archive +from wrolpi.cmd import run_command from wrolpi.collections import Collection +from wrolpi.downloader import DownloadContext + + +def make_test_ctx(*, is_cancelled=None, can_download=None, + outside_download_window=None, run_command_=None, + download_timeout=None) -> DownloadContext: + """Build a DownloadContext suitable for downloader unit tests. + + Defaults to a never-cancelled context that uses the real run_command. Tests can + override any callable to simulate cancellation, throttling, or fake subprocess output. + """ + return DownloadContext( + is_cancelled=is_cancelled or (lambda: False), + can_download=can_download or (lambda: True), + outside_download_window=outside_download_window or (lambda: False), + run_command=run_command_ or run_command, + download_timeout=download_timeout, + ) @pytest.fixture diff --git a/modules/archive/lib.py b/modules/archive/lib.py index b07f4cd6..526d3af3 100644 --- a/modules/archive/lib.py +++ b/modules/archive/lib.py @@ -710,17 +710,19 @@ async def request_screenshot(url: str, singlefile_path: pathlib.Path) -> Optiona return screenshot -async def model_archive_result(url: str, singlefile: str, readability: dict, screenshot: bytes, - destination: pathlib.Path = None) -> Archive: - """ - Convert results from ArchiveDownloader into real files. Create Archive record. +@dataclass +class WrittenArchiveFiles: + """The on-disk artifacts produced by write_archive_files.""" + url: str + paths: List[pathlib.Path] - Args: - url: The URL that was archived - singlefile: The HTML content from SingleFile - readability: The readability extraction results - screenshot: The screenshot bytes - destination: Optional directory to save files to. If provided, files go here instead of archive/ + +def write_archive_files(url: str, singlefile: str, readability: Optional[dict], + screenshot: Optional[bytes], + destination: Optional[pathlib.Path] = None) -> WrittenArchiveFiles: + """Pure I/O: write the singlefile/readability/screenshot/json files to disk. + + Returns the URL and the list of paths to register, leaving DB work to register_archive. """ readability = readability.copy() if readability else readability # First try to get the title from Readability. @@ -770,22 +772,47 @@ async def model_archive_result(url: str, singlefile: str, readability: dict, scr archive_files.screenshot) paths = list(filter(None, paths)) + return WrittenArchiveFiles(url=url, paths=paths) + + +def register_archive(session: Session, written: WrittenArchiveFiles) -> Archive: + """Pure DB: create the Archive record from already-written files. Caller commits.""" + archive = Archive.from_paths(session, *written.paths) + archive.file_group.download_datetime = now() + archive.url = written.url + archive.collection = get_or_create_domain_collection(session, written.url) + + # Update collection's file_format to track the format used for this file + config = get_archive_downloader_config() + if archive.collection and config.file_name_format: + if not archive.collection.file_format: + # Set initial file_format on collection + archive.collection.file_format = config.file_name_format + save_domains_config.activate_switch() + + archive.flush() + return archive + + +async def model_archive_result(url: str, singlefile: str, readability: dict, screenshot: bytes, + destination: pathlib.Path = None) -> Archive: + """ + Convert results from ArchiveDownloader into real files. Create Archive record. + + Thin wrapper around write_archive_files + register_archive that opens its own + transaction. New code (notably the phase-split ArchiveDownloader) should call + the helpers directly so the caller controls the session. + + Args: + url: The URL that was archived + singlefile: The HTML content from SingleFile + readability: The readability extraction results + screenshot: The screenshot bytes + destination: Optional directory to save files to. If provided, files go here instead of archive/ + """ + written = write_archive_files(url, singlefile, readability, screenshot, destination=destination) with get_db_session(commit=True) as session: - archive = Archive.from_paths(session, *paths) - archive.file_group.download_datetime = now() - archive.url = url - archive.collection = get_or_create_domain_collection(session, url) - - # Update collection's file_format to track the format used for this file - config = get_archive_downloader_config() - if archive.collection and config.file_name_format: - if not archive.collection.file_format: - # Set initial file_format on collection - archive.collection.file_format = config.file_name_format - save_domains_config.activate_switch() - - archive.flush() - + archive = register_archive(session, written) return archive @@ -1421,43 +1448,63 @@ async def html_to_readability(html: str | bytes, url: str, timeout: int = 120): raise RuntimeError(f'Failed to extract readability for {url} got {result.return_code}') from e -async def singlefile_to_archive(singlefile: bytes, destination: pathlib.Path = None) -> Archive: - """ - Convert a SingleFile to an Archive. +@dataclass +class SinglefileArtifacts: + """Pure-data result of running a singlefile through readability + screenshot extraction.""" + url: str + singlefile: bytes + readability: dict + screenshot: Optional[bytes] - This is done by extracting readability, creating a screenshot, then attaching them to an Archive/FileGroup. - Args: - singlefile: The SingleFile HTML bytes - destination: Optional directory to save files to. If provided, files go here instead of archive/ - """ - # Get URL first because it does some simple checking of `singlefile` +async def extract_singlefile_artifacts(singlefile: bytes) -> SinglefileArtifacts: + """Pure async I/O: extract readability and screenshot from a singlefile. No DB.""" + # Get URL first because it does some simple checking of `singlefile`. url = get_url_from_singlefile(singlefile) singlefile = singlefile.encode() if isinstance(singlefile, str) else singlefile if DOCKERIZED: # Perform the archive in the Archive docker container. (Typically in the development environment). - singlefile = base64.b64encode(singlefile).decode() - logger.debug(f'singlefile_to_archive sending to archive service: {url}') - singlefile, readability, screenshot = await request_archive(url, singlefile=singlefile) + encoded = base64.b64encode(singlefile).decode() + logger.debug(f'extract_singlefile_artifacts sending to archive service: {url}') + singlefile, readability, screenshot = await request_archive(url, singlefile=encoded) else: # JSON from readability-extractor readability = dict() try: - logger.debug(f'singlefile_to_archive extracting readability: {url}') + logger.debug(f'extract_singlefile_artifacts extracting readability: {url}') readability = await html_to_readability(singlefile, url) except RuntimeError as e: logger.error(f'Failed to extract readability from: {url}', exc_info=e) screenshot = None try: - logger.debug(f'singlefile_to_archive creating screenshot: {url}') + logger.debug(f'extract_singlefile_artifacts creating screenshot: {url}') screenshot = html_screenshot(singlefile) except Exception as e: logger.error(f'Failed to extract screenshot from: {url}', exc_info=e) - logger.trace(f'singlefile_to_archive modeling: {url}') - archive: Archive = await model_archive_result(url, singlefile, readability, screenshot, destination=destination) + return SinglefileArtifacts(url=url, singlefile=singlefile, readability=readability, screenshot=screenshot) + + +async def singlefile_to_archive(singlefile: bytes, destination: pathlib.Path = None) -> Archive: + """ + Convert a SingleFile to an Archive. + + Thin wrapper around extract_singlefile_artifacts + write_archive_files + register_archive + that opens its own transaction. New code should call the helpers directly so the caller + controls the session. + + Args: + singlefile: The SingleFile HTML bytes + destination: Optional directory to save files to. If provided, files go here instead of archive/ + """ + artifacts = await extract_singlefile_artifacts(singlefile) + logger.trace(f'singlefile_to_archive modeling: {artifacts.url}') + written = write_archive_files(artifacts.url, artifacts.singlefile, artifacts.readability, + artifacts.screenshot, destination=destination) + with get_db_session(commit=True) as session: + archive = register_archive(session, written) return archive diff --git a/modules/archive/test/test_archive_downloader.py b/modules/archive/test/test_archive_downloader.py new file mode 100644 index 00000000..2fc24c92 --- /dev/null +++ b/modules/archive/test/test_archive_downloader.py @@ -0,0 +1,130 @@ +"""Phase-split ArchiveDownloader tests. + +These tests exercise prepare_download / execute_download / finalize_download in +isolation. None of them use async_client or test_download_manager — that's the +whole point of the refactor. +""" +import asyncio + +import pytest + +from modules.archive import ( + ExecutedArchive, + PreparedArchive, + archive_downloader, +) +from modules.archive.conftest import make_test_ctx +from modules.archive.lib import write_archive_files +from modules.archive.models import Archive +from modules.archive.test.test_lib import make_fake_archive_result +from wrolpi.collections import Collection +from wrolpi.downloader import Download, DownloadResult +from wrolpi.errors import UnrecoverableDownloadError + + +def test_prepare_max_attempts_raises(test_session): + """prepare_download enforces the attempt cap before any I/O.""" + download = Download(url='https://example.com/x', downloader='archive', attempts=4) + + with pytest.raises(UnrecoverableDownloadError): + archive_downloader.prepare_download(test_session, download) + + +def test_prepare_uses_collection_directory(test_session, test_directory): + """When a domain Collection has a directory, prepare_download routes the archive there.""" + custom_dir = test_directory / 'archive/News/example.com' + custom_dir.mkdir(parents=True, exist_ok=True) + test_session.add(Collection(name='example.com', kind='domain', directory=custom_dir)) + test_session.commit() + + download = Download(url='https://example.com/some-article', downloader='archive', settings={}) + + prepared = archive_downloader.prepare_download(test_session, download) + + assert isinstance(prepared, PreparedArchive) + assert prepared.url == 'https://example.com/some-article' + assert prepared.destination == custom_dir + + +def test_prepare_explicit_destination_overrides_collection(test_session, test_directory): + """An explicit settings['destination'] beats the domain collection's directory.""" + collection_dir = test_directory / 'archive/News/example.com' + collection_dir.mkdir(parents=True, exist_ok=True) + test_session.add(Collection(name='example.com', kind='domain', directory=collection_dir)) + test_session.commit() + + explicit_dir = test_directory / 'archive/Special/my-archives' + explicit_dir.mkdir(parents=True, exist_ok=True) + download = Download( + url='https://example.com/some-article', + downloader='archive', + settings={'destination': str(explicit_dir)}, + ) + + prepared = archive_downloader.prepare_download(test_session, download) + + assert prepared.destination == explicit_dir + + +@pytest.mark.asyncio +async def test_execute_cancels_via_ctx(test_session, test_directory, monkeypatch): + """When ctx reports cancelled, the cancel_wrapper that the dispatch puts around + execute_download returns a cancel DownloadResult — no global download_manager, + no Sanic shared_ctx.""" + # Force the local-singlefile path so the run_command stub is what blocks. + import modules.archive + monkeypatch.setattr(modules.archive, 'DOCKERIZED', False) + monkeypatch.setattr(modules.archive, 'PYTEST', False) + + async def never_completes(*args, **kwargs): + await asyncio.sleep(60) + raise AssertionError('should have been cancelled') + + # Pre-cancelled: the very first poll inside cancel_wrapper aborts the run. + ctx = make_test_ctx(is_cancelled=lambda: True, run_command_=never_completes) + prepared = PreparedArchive( + url='https://example.com/x', + destination=test_directory / 'archive/example.com', + settings={}, + ) + prepared.destination.mkdir(parents=True, exist_ok=True) + + # Mirror the dispatch layer: wrap execute_download in cancel_wrapper. + coro = archive_downloader.execute_download(prepared, ctx) + result = await archive_downloader.cancel_wrapper(coro, Download(url=prepared.url), ctx=ctx) + + assert isinstance(result, DownloadResult) + assert result.success is False + assert 'cancel' in (result.error or '').lower() + + +@pytest.mark.asyncio +async def test_finalize_adds_tags(async_client, test_session, test_directory, tag_factory): + """finalize_download applies download.tag_names to the registered Archive. + + Uses async_client because tag_factory's upsert_tag activates a Sanic switch + (shared_ctx.switches_lock). That's an orthogonal coupling — the rest of + finalize_download is Sanic-free. + """ + singlefile, readability, screenshot = make_fake_archive_result() + destination = test_directory / 'archive/example.com' + destination.mkdir(parents=True, exist_ok=True) + written = write_archive_files('https://example.com/x', singlefile, readability, screenshot, + destination=destination) + executed = ExecutedArchive(written=written) + + tag = await tag_factory() + download = Download( + url='https://example.com/x', + downloader='archive', + tag_names=[tag.name], + ) + + result = archive_downloader.finalize_download(test_session, download, executed) + test_session.commit() + + assert result.success is True + archives = test_session.query(Archive).all() + assert len(archives) == 1 + archive_tag_names = [tf.tag.name for tf in archives[0].file_group.tag_files] + assert tag.name in archive_tag_names diff --git a/modules/archive/test/test_lib.py b/modules/archive/test/test_lib.py index 1fc4659a..46537a00 100644 --- a/modules/archive/test/test_lib.py +++ b/modules/archive/test/test_lib.py @@ -1185,8 +1185,12 @@ async def test_archive_download_uses_domain_collection_directory(async_client, t test_session.add(download) test_session.commit() - # Call the archive downloader - result = await archive_downloader.do_download(download) + # Drive the phase-split API directly: prepare → execute → finalize. + from modules.archive.conftest import make_test_ctx + prepared = archive_downloader.prepare_download(test_session, download) + executed = await archive_downloader.execute_download(prepared, make_test_ctx()) + result = archive_downloader.finalize_download(test_session, download, executed) + test_session.commit() # Verify the download was successful assert result.success is True @@ -1247,8 +1251,12 @@ async def test_archive_download_explicit_destination_overrides_domain_collection test_session.add(download) test_session.commit() - # Call the archive downloader - result = await archive_downloader.do_download(download) + # Drive the phase-split API directly: prepare → execute → finalize. + from modules.archive.conftest import make_test_ctx + prepared = archive_downloader.prepare_download(test_session, download) + executed = await archive_downloader.execute_download(prepared, make_test_ctx()) + result = archive_downloader.finalize_download(test_session, download, executed) + test_session.commit() # Verify the download was successful assert result.success is True diff --git a/wrolpi/downloader.py b/wrolpi/downloader.py index ba25fb5e..227d6a11 100644 --- a/wrolpi/downloader.py +++ b/wrolpi/downloader.py @@ -17,7 +17,7 @@ from enum import Enum from http import HTTPStatus from itertools import filterfalse from typing import List, Dict, Generator, Iterable, Coroutine, Any -from typing import Tuple, Optional +from typing import Tuple, Optional, Callable, Awaitable from urllib.parse import urlparse from zoneinfo import ZoneInfo @@ -173,6 +173,38 @@ class DownloadResult: settings: dict = field(default_factory=dict) +@dataclass +class DownloadContext: + """Per-download seam between DownloadManager and a Downloader. + + Carries cancellation queries, the subprocess runner, and the effective timeout. + Production builds these as closures over api_app.shared_ctx and the singleton + DownloadManager (see DownloadContext.production); tests construct them directly + so a Downloader can run without Sanic's mp stack. + """ + is_cancelled: Callable[[], bool] + can_download: Callable[[], bool] + outside_download_window: Callable[[], bool] + run_command: Callable[..., Awaitable[CommandResult]] + download_timeout: Optional[int] + + @classmethod + def production(cls, manager: 'DownloadManager', download_id: int) -> 'DownloadContext': + """Build a context whose closures read shared_ctx + the manager. + + This is the only place that knows about Sanic's multiprocessing model. + Cross-worker cancellation still flows through + api_app.shared_ctx.download_manager_data['killed_downloads']; this just + names the seam.""" + return cls( + is_cancelled=lambda: download_id in api_app.shared_ctx.download_manager_data['killed_downloads'], + can_download=lambda: manager.can_download, + outside_download_window=lambda: manager.outside_download_window, + run_command=run_command, + download_timeout=get_wrolpi_config().download_timeout or None, + ) + + class DownloadStatus(str, Enum): new = 'new' pending = 'pending' @@ -387,7 +419,30 @@ class Downloader: def __repr__(self): return f'' - async def do_download(self, download: Download) -> DownloadResult: + async def do_download(self, download: Download, ctx: 'DownloadContext' = None) -> DownloadResult: + """Legacy single-phase entry point. + + Subclasses opting in to the phase-split API override prepare_download / + execute_download / finalize_download instead and leave this unimplemented. + """ + raise NotImplementedError() + + def prepare_download(self, session: Session, download: Download) -> Optional[Any]: + """Synchronous DB-bound prep. Caller (DownloadManager) opens and commits the session. + + Returning None signals the subclass has not opted in to the phase-split API and + DownloadManager should fall back to do_download. Returning any other value + (typically a dataclass) is passed unchanged to execute_download. + """ + return None + + async def execute_download(self, prepared: Any, ctx: 'DownloadContext') -> Any: + """Async I/O-bound work. No DB session — does network and subprocess work, + returns plain data for finalize_download to persist.""" + raise NotImplementedError() + + def finalize_download(self, session: Session, download: Download, executed: Any) -> DownloadResult: + """Synchronous DB-bound persistence. Caller opens and commits the session.""" raise NotImplementedError() def already_downloaded(self, session: Session, *urls: List[str]): @@ -397,24 +452,33 @@ class Downloader: async def process_runner(self, download: Download, cmd: Tuple[str | pathlib.Path, ...], cwd: pathlib.Path, timeout: int = None, debug: bool = True, - stdout_callback: callable = None, env: dict = None) -> CommandResult: + stdout_callback: callable = None, env: dict = None, + ctx: 'DownloadContext' = None) -> CommandResult: """ 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) + + When ctx is provided, ctx.run_command and ctx.download_timeout are used and cancellation + is polled through ctx. When ctx is None (legacy callers), behavior is unchanged. """ if cwd and not cwd.is_dir(): raise RuntimeError('cwd directory does not exist') - config_timeout = get_wrolpi_config().download_timeout + if ctx is not None: + config_timeout = ctx.download_timeout + runner = ctx.run_command + else: + config_timeout = get_wrolpi_config().download_timeout + runner = run_command # 0 means no timeout. Prefer config timeout, then argument, then class default. if config_timeout: timeout = config_timeout elif timeout is None: timeout = self.timeout - coro = run_command(cmd, cwd=cwd, timeout=timeout, log_command=debug, - stdout_callback=stdout_callback, env=env) - result = await self.cancel_wrapper(coro, download) + coro = runner(cmd, cwd=cwd, timeout=timeout, log_command=debug, + stdout_callback=stdout_callback, env=env) + result = await self.cancel_wrapper(coro, download, ctx=ctx) if __debug__ and logger.isEnabledFor(TRACE_LEVEL): for line in result.stdout.decode().splitlines(): @@ -425,17 +489,26 @@ class Downloader: return result @staticmethod - async def cancel_wrapper(coro: Coroutine, download: Download): + async def cancel_wrapper(coro: Coroutine, download: Download, ctx: 'DownloadContext' = None): """ Converts an async coroutine to a task. If DownloadManager receives a kill request, this method will cancel the task. + + When ctx is provided, cancellation is polled through ctx.is_cancelled / ctx.can_download / + ctx.outside_download_window. When ctx is None, the global download_manager singleton is + used (legacy path for unmigrated subclasses). """ download_id = download.id task = asyncio.create_task(coro) try: while not task.done(): - if download_manager.download_is_killed(download_id) or not download_manager.can_download: + if ctx is not None: + cancelled = ctx.is_cancelled() or not ctx.can_download() + outside_window = ctx.outside_download_window() + else: + cancelled = download_manager.download_is_killed(download_id) or not download_manager.can_download outside_window = download_manager.outside_download_window + if cancelled: logger.warning(f'Cancel download of {download.url}') task.cancel() try: @@ -449,7 +522,16 @@ class Downloader: error=error, ) finally: - download_manager.unkill_download(download_id) + if ctx is None: + download_manager.unkill_download(download_id) + else: + # Manager still owns the killed-downloads list; clear via singleton. + # The ctx is read-only with respect to cancellation state. + try: + download_manager.unkill_download(download_id) + except Exception: + # In tests without a manager, silently ignore. + pass else: # Wait for the download to complete. Cancel if requested. await asyncio.sleep(0.1) @@ -1449,11 +1531,30 @@ async def signal_download_download(download_id: int, download_url: str): try_again = True try: - # Create download coroutine. Wrap it, so it can be canceled. - if not inspect.iscoroutinefunction(downloader.do_download): - raise RuntimeError(f'Coroutine expected from {downloader} do_download method.') - coro = downloader.do_download(download) - result = await downloader.cancel_wrapper(coro, download) + # Phase-split path: subclass opts in by overriding prepare_download to return non-None. + # Legacy path: do_download remains the single entry point. + ctx = DownloadContext.production(download_manager, download_id) + with get_db_session(commit=True) as s: + prepared = downloader.prepare_download(s, download) + + if prepared is not None: + # New phase-split path: caller (us) owns transaction boundaries. + coro = downloader.execute_download(prepared, ctx) + executed = await downloader.cancel_wrapper(coro, download, ctx=ctx) + if isinstance(executed, DownloadResult): + # cancel_wrapper short-circuited with a cancel result. + result = executed + else: + with get_db_session(commit=True) as s: + result = downloader.finalize_download(s, download, executed) + else: + # Legacy path. Subclasses on this path have signature do_download(self, download) + # without ctx, so we must NOT pass ctx as a positional arg. cancel_wrapper continues + # to use the global download_manager (ctx=None). + if not inspect.iscoroutinefunction(downloader.do_download): + raise RuntimeError(f'Coroutine expected from {downloader} do_download method.') + coro = downloader.do_download(download) + result = await downloader.cancel_wrapper(coro, download) except BotBlockedDownloadError as e: # Bot detection or invalid cookies - stop all pending downloads from this domain. worker_logger.warning(f'BotBlockedDownloadError for {url}', exc_info=e)