soulsync/core/download_plugins/usenet.py
Broque Thomas c990ce079d feat(downloads): album-bundle flow for torrent/usenet single-source mode
Fixes the core architectural mismatch between indexer-based sources
and the per-track search-and-pick contract every other download
plugin satisfies. Prowlarr returns release-level torrents and NZBs;
searching for "Luther (with SZA)" against the GNX album torrent
scores near-zero on track-title similarity. Per-track candidate
validation rejects every result, every track in the batch flips
to not_found. The album-name fallback added in an earlier commit
papers over it for some cases but doesn't fix the fundamental
behavior: the user wanted the whole album.

New album-bundle flow does what the user actually wanted:
1. Gate fires inside core/downloads/master.py BEFORE the per-track
   analysis loop, strictly when the batch has an album context AND
   download_source.mode is 'torrent' or 'usenet' (single-source —
   hybrid stays per-track to preserve fallback to Soulseek / etc.).
2. Plugin's new download_album_to_staging method searches Prowlarr
   ONCE for the album as a whole ('<artist> <album>'), filters to
   the right protocol, runs results through _pick_best_album_release.
3. Picker prefers seeded FLAC over low-seeded MP3, drops single-
   track torrents that snuck in via the 40 MB size floor (single
   tracks are typically ~10 MB), falls back to most-seeded when
   every candidate is below the floor.
4. Picked release goes to the active adapter (qBit / Transmission /
   Deluge for torrent; SAB / NZBGet for usenet). Polls until
   complete with progress mirrored into the batch state so the
   Downloads page can show meaningful status.
5. On completion the existing archive_pipeline walks the save dir
   (extracting archives if any), every audio file gets copied into
   the staging folder via _unique_staging_path so concurrent batches
   don't collide.
6. Gate exits, master worker continues into the normal per-track
   flow. Each track task hits try_staging_match early in the worker
   and finds its file by fuzzy title match — no Prowlarr search
   ever fires per-track, no candidate rejection, files flow through
   the existing post-processing pipeline (tags, AcoustID, library
   import).

Gate is strictly opt-in. Three orthogonal conditions must all hold:
batch_is_album, mode in ('torrent', 'usenet'), and the plugin must
expose download_album_to_staging. Any other source / hybrid mode /
non-album batch flows through the master worker unchanged. The
existing per-track torrent path still works for basic-search
single-track grabs.

- core/download_plugins/torrent.py: download_album_to_staging plus
  _pick_best_album_release and _unique_staging_path helpers (shared
  with the usenet plugin). _poll_album_download mirrors the existing
  poll loop with progress callback emission.
- core/download_plugins/usenet.py: parallel implementation reusing
  the picker + staging helpers. Different state set ('failed' vs
  'error') from the usenet adapter contract.
- core/downloads/master.py: ~90-line gate right after batch context
  loading. Mirrors plugin lifecycle into batch state under
  ``album_bundle_*`` keys so the Downloads page can render progress
  while the torrent/usenet job runs (per-track tasks don't exist
  yet during this phase). Failed bundle download fails the batch
  with a meaningful error; missing plugin / context falls back to
  the per-track flow with a warning.
- tests/test_torrent_usenet_plugins.py: 5 new tests pinning the
  album picker preferences (FLAC over MP3 with comparable size +
  better seeders, size floor drops singles, fallback when all
  small), staging-path collision suffix, and the not-configured
  short-circuit.
2026-05-20 18:48:48 -07:00

467 lines
18 KiB
Python

"""UsenetDownloadPlugin — composes Prowlarr search + usenet client
adapter + archive_pipeline into a uniform download source.
Mirrors ``TorrentDownloadPlugin`` in shape and lifecycle (see that
module's docstring for the full pipeline rationale). Differences:
- Search filters Prowlarr results to ``protocol='usenet'``.
- ``add_nzb`` replaces ``add_torrent``; for NZBs we usually have
a direct HTTP URL the indexer exposes via Prowlarr.
- Usenet clients (SABnzbd, NZBGet) typically auto-extract during
post-processing, so ``archive_pipeline.collect_audio_after_extraction``
usually has nothing to extract and just walks loose files.
"""
from __future__ import annotations
import shutil
import threading
import time
import uuid
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple
from core.archive_pipeline import collect_audio_after_extraction
from core.download_plugins.base import DownloadSourcePlugin
from core.download_plugins.torrent import (
_adapter_state_to_display,
_decode_filename,
_guess_quality_from_title,
_parse_indexer_id_filter,
_parse_release_title,
_pick_best_album_release,
_row_to_status,
_unique_staging_path,
_COMPLETE_STATES,
_FILENAME_SEP,
_POLL_INTERVAL_SECONDS,
_POLL_TIMEOUT_SECONDS,
)
from core.download_plugins.types import AlbumResult, DownloadStatus, TrackResult
from core.prowlarr_client import (
DEFAULT_MUSIC_CATEGORIES,
ProwlarrClient,
ProwlarrSearchResult,
)
from core.usenet_clients import get_active_adapter as get_active_usenet_adapter
from utils.async_helpers import run_async
from utils.logging_config import get_logger
logger = get_logger("download_plugins.usenet")
class UsenetDownloadPlugin(DownloadSourcePlugin):
"""Usenet download source backed by Prowlarr + an active usenet
client adapter (SABnzbd or NZBGet)."""
def __init__(self) -> None:
self._prowlarr = ProwlarrClient()
self.active_downloads: Dict[str, Dict[str, Any]] = {}
self._lock = threading.Lock()
self.shutdown_check = None
def set_shutdown_check(self, check_callable):
self.shutdown_check = check_callable
def reload_settings(self) -> None:
self._prowlarr.reload_settings()
def is_configured(self) -> bool:
if not self._prowlarr.is_configured():
return False
adapter = get_active_usenet_adapter()
return bool(adapter and adapter.is_configured())
async def check_connection(self) -> bool:
if not self._prowlarr.is_configured():
return False
adapter = get_active_usenet_adapter()
if not adapter or not adapter.is_configured():
return False
if not await self._prowlarr.check_connection():
return False
return await adapter.check_connection()
# ------------------------------------------------------------------
# Search
# ------------------------------------------------------------------
async def search(
self,
query: str,
timeout: Optional[int] = None,
progress_callback=None,
) -> Tuple[List[TrackResult], List[AlbumResult]]:
if not self._prowlarr.is_configured():
return ([], [])
try:
indexer_ids = _parse_indexer_id_filter()
results = await self._prowlarr.search(
query,
categories=DEFAULT_MUSIC_CATEGORIES,
indexer_ids=indexer_ids,
)
except Exception as e:
logger.error("Usenet plugin search failed: %s", e)
return ([], [])
return self._project_results(results)
def _project_results(
self, results: List[ProwlarrSearchResult]
) -> Tuple[List[TrackResult], List[AlbumResult]]:
tracks: List[TrackResult] = []
albums: List[AlbumResult] = []
for result in results:
if result.protocol != 'usenet':
continue
if not result.download_url:
continue
filename = f"{result.download_url}{_FILENAME_SEP}{result.title}"
quality = _guess_quality_from_title(result.title)
parsed_artist, parsed_title = _parse_release_title(result.title)
tr = TrackResult(
username='usenet',
filename=filename,
size=result.size,
bitrate=None,
duration=None,
quality=quality,
# Usenet doesn't expose per-uploader concurrency the way
# Soulseek does; fill in neutral non-punishing values.
free_upload_slots=1,
upload_speed=0,
queue_length=0,
# Pre-fill artist + title so TrackResult.__post_init__
# doesn't auto-parse the filename — same URL-in-filename
# gotcha as the torrent plugin.
artist=parsed_artist or result.indexer_name or 'Usenet',
title=parsed_title or result.title,
album=parsed_title or None,
track_number=None,
_source_metadata={
'indexer': result.indexer_name,
'indexer_id': result.indexer_id,
'grabs': result.grabs,
'protocol': 'usenet',
},
)
tracks.append(tr)
albums.append(AlbumResult(
username='usenet',
album_path=f"usenet/{result.guid}",
album_title=parsed_title or result.title,
artist=parsed_artist or None,
track_count=1,
total_size=result.size,
tracks=[tr],
dominant_quality=quality,
year=None,
))
return tracks, albums
# ------------------------------------------------------------------
# Download
# ------------------------------------------------------------------
async def download(
self,
username: str,
filename: str,
file_size: int = 0,
) -> Optional[str]:
if not self.is_configured():
return None
nzb_url, display_name = _decode_filename(filename)
if not nzb_url:
logger.error("Usenet download missing URL in filename: %r", filename)
return None
download_id = str(uuid.uuid4())
with self._lock:
self.active_downloads[download_id] = {
'id': download_id,
'filename': filename,
'username': 'usenet',
'display_name': display_name,
'state': 'Initializing',
'progress': 0.0,
'size': file_size,
'transferred': 0,
'speed': 0,
'file_path': None,
'job_id': None,
'error': None,
}
thread = threading.Thread(
target=self._download_thread,
args=(download_id, nzb_url),
daemon=True,
name=f'usenet-dl-{download_id[:8]}',
)
thread.start()
return download_id
def _download_thread(self, download_id: str, nzb_url: str) -> None:
adapter = get_active_usenet_adapter()
if adapter is None or not adapter.is_configured():
self._mark_error(download_id, "No usenet client configured")
return
try:
job_id = run_async(adapter.add_nzb(nzb_url))
except Exception as e:
self._mark_error(download_id, f"add_nzb failed: {e}")
return
if not job_id:
self._mark_error(download_id, "Usenet client refused the NZB")
return
with self._lock:
row = self.active_downloads.get(download_id)
if row is not None:
row['job_id'] = job_id
row['state'] = 'InProgress, Downloading'
deadline = time.monotonic() + _POLL_TIMEOUT_SECONDS
last_save_path: Optional[str] = None
while time.monotonic() < deadline:
if self.shutdown_check and self.shutdown_check():
return
try:
status = run_async(adapter.get_status(job_id))
except Exception as e:
logger.warning("Usenet poll error for %s: %s", job_id, e)
status = None
if status is None:
self._mark_error(download_id, "Usenet job disappeared from client")
return
with self._lock:
row = self.active_downloads.get(download_id)
if row is not None:
row['progress'] = status.progress * 100.0
row['transferred'] = status.downloaded
row['speed'] = status.download_speed
row['size'] = status.size or row.get('size', 0)
row['state'] = _adapter_state_to_display(status.state)
row['error'] = status.error
if status.save_path:
last_save_path = status.save_path
if status.state in _COMPLETE_STATES:
self._finalize_download(download_id, last_save_path)
return
if status.state == 'failed':
self._mark_error(download_id, status.error or "Usenet client reported failure")
return
time.sleep(_POLL_INTERVAL_SECONDS)
self._mark_error(download_id, "Usenet download timed out")
def _finalize_download(self, download_id: str, save_path: Optional[str]) -> None:
if not save_path:
self._mark_error(download_id, "Usenet job completed but no save_path reported")
return
try:
audio_files = collect_audio_after_extraction(Path(save_path))
except Exception as e:
self._mark_error(download_id, f"Post-extract walk failed: {e}")
return
if not audio_files:
self._mark_error(download_id, f"No audio files found in {save_path}")
return
primary = audio_files[0]
with self._lock:
row = self.active_downloads.get(download_id)
if row is not None:
row['state'] = 'Completed, Succeeded'
row['progress'] = 100.0
row['file_path'] = str(primary)
logger.info("Usenet download complete: %s -> %s (%d audio files)",
download_id[:8], primary.name, len(audio_files))
def _mark_error(self, download_id: str, message: str) -> None:
logger.error("Usenet download %s failed: %s", download_id[:8], message)
with self._lock:
row = self.active_downloads.get(download_id)
if row is not None:
row['state'] = 'Completed, Errored'
row['error'] = message
# ------------------------------------------------------------------
# Status / lifecycle
# ------------------------------------------------------------------
async def get_all_downloads(self) -> List[DownloadStatus]:
with self._lock:
rows = list(self.active_downloads.values())
return [_row_to_status(r) for r in rows]
async def get_download_status(self, download_id: str) -> Optional[DownloadStatus]:
with self._lock:
row = self.active_downloads.get(download_id)
if row is None:
return None
return _row_to_status(row)
async def cancel_download(
self,
download_id: str,
username: Optional[str] = None,
remove: bool = False,
) -> bool:
adapter = get_active_usenet_adapter()
with self._lock:
row = self.active_downloads.get(download_id)
job_id = row.get('job_id') if row else None
if adapter and job_id:
try:
await adapter.remove(job_id, delete_files=remove)
except Exception as e:
logger.warning("Usenet cancel via adapter failed: %s", e)
with self._lock:
if remove:
self.active_downloads.pop(download_id, None)
else:
row = self.active_downloads.get(download_id)
if row is not None:
row['state'] = 'Cancelled'
return True
async def clear_all_completed_downloads(self) -> bool:
with self._lock:
for did in list(self.active_downloads.keys()):
state = self.active_downloads[did].get('state', '')
if state.startswith('Completed') or state == 'Cancelled':
self.active_downloads.pop(did, None)
return True
# ------------------------------------------------------------------
# Album-bundle flow
# ------------------------------------------------------------------
def download_album_to_staging(
self,
album_name: str,
artist_name: str,
staging_dir: str,
progress_callback=None,
) -> Dict[str, Any]:
"""Usenet sibling of ``TorrentDownloadPlugin.download_album_to_staging``.
See that method's docstring for the contract."""
result: Dict[str, Any] = {'success': False, 'files': [], 'error': None}
if not self.is_configured():
result['error'] = 'Usenet source not configured'
return result
adapter = get_active_usenet_adapter()
if adapter is None or not adapter.is_configured():
result['error'] = 'No active usenet client'
return result
def _emit(state: str, **extra) -> None:
if progress_callback:
try:
progress_callback({'state': state, **extra})
except Exception as cb_exc:
logger.debug("[Usenet album] progress callback failed: %s", cb_exc)
query = f"{artist_name} {album_name}".strip()
_emit('searching', query=query)
try:
search_results = run_async(self._prowlarr.search(
query, categories=DEFAULT_MUSIC_CATEGORIES,
indexer_ids=_parse_indexer_id_filter(),
))
except Exception as e:
result['error'] = f'Prowlarr search failed: {e}'
return result
candidates = [r for r in search_results
if r.protocol == 'usenet' and r.download_url]
if not candidates:
result['error'] = f'No usenet results found for "{query}"'
return result
picked = _pick_best_album_release(candidates)
if picked is None:
result['error'] = 'No suitable NZB candidate after filtering'
return result
logger.info("[Usenet album] Picked '%s' (size=%.1fMB grabs=%s indexer=%s)",
picked.title, picked.size / 1_048_576, picked.grabs, picked.indexer_name)
_emit('queued', release=picked.title, size=picked.size, grabs=picked.grabs)
try:
job_id = run_async(adapter.add_nzb(picked.download_url))
except Exception as e:
result['error'] = f'Usenet client refused the NZB: {e}'
return result
if not job_id:
result['error'] = 'Usenet client refused the NZB'
return result
_emit('downloading', release=picked.title)
save_path = self._poll_album_download(adapter, job_id, picked.title, _emit)
if save_path is None:
result['error'] = 'Usenet download failed or timed out'
return result
_emit('staging', release=picked.title)
try:
audio_files = collect_audio_after_extraction(Path(save_path))
except Exception as e:
result['error'] = f'Failed to walk audio files: {e}'
return result
if not audio_files:
result['error'] = f'No audio files found in {save_path}'
return result
staging_path = Path(staging_dir)
staging_path.mkdir(parents=True, exist_ok=True)
copied: List[str] = []
for src in audio_files:
dst = _unique_staging_path(staging_path, src)
try:
shutil.copy2(src, dst)
copied.append(str(dst))
except Exception as e:
logger.warning("[Usenet album] Failed to copy %s -> %s: %s", src, dst, e)
if not copied:
result['error'] = 'No audio files copied to staging'
return result
logger.info("[Usenet album] Staged %d audio files for '%s'", len(copied), album_name)
_emit('staged', count=len(copied))
result['success'] = True
result['files'] = copied
return result
def _poll_album_download(self, adapter, job_id, title, emit) -> Optional[str]:
deadline = time.monotonic() + _POLL_TIMEOUT_SECONDS
last_save_path: Optional[str] = None
while time.monotonic() < deadline:
if self.shutdown_check and self.shutdown_check():
return None
try:
status = run_async(adapter.get_status(job_id))
except Exception as e:
logger.warning("[Usenet album] Poll error: %s", e)
status = None
if status is None:
logger.error("[Usenet album] '%s' disappeared from client", title)
return None
emit('downloading', progress=status.progress, downloaded=status.downloaded,
speed=status.download_speed)
if status.save_path:
last_save_path = status.save_path
if status.state in _COMPLETE_STATES:
return last_save_path
if status.state == 'failed':
logger.error("[Usenet album] '%s' failed: %s", title, status.error)
return None
time.sleep(_POLL_INTERVAL_SECONDS)
logger.error("[Usenet album] '%s' timed out", title)
return None