From ea1e3c6ac5914d31ccee6c7737013fe887df1b80 Mon Sep 17 00:00:00 2001 From: Connor Johnstone Date: Tue, 17 Mar 2026 21:38:27 -0400 Subject: [PATCH] Updates for the "full flow" --- src/lib.rs | 2 +- src/main.rs | 13 +++++ src/queue.rs | 126 ++++++++++++++++++++++++++++++++++++++----- src/ytdlp.rs | 4 +- tests/integration.rs | 2 +- 5 files changed, 129 insertions(+), 18 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index 50ac56f..44e6924 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -11,5 +11,5 @@ pub mod ytdlp; pub use backend::{AudioFormat, BackendConfig, DownloadBackend, DownloadResult, DownloadTarget, SearchResult}; pub use error::{DlError, DlResult}; -pub use queue::{DlStats, download_single, run_queue}; +pub use queue::{DlStats, SyncStats, download_single, run_queue, sync_wanted_to_queue}; pub use ytdlp::{SearchSource, YtDlpBackend}; diff --git a/src/main.rs b/src/main.rs index 6a8b917..862bf64 100644 --- a/src/main.rs +++ b/src/main.rs @@ -105,6 +105,12 @@ enum QueueAction { }, /// Retry all failed downloads. Retry, + /// Sync wanted items from the watchlist to the download queue. + Sync { + /// Preview what would be enqueued without doing it. + #[arg(long)] + dry_run: bool, + }, } fn default_database_url() -> String { @@ -276,6 +282,13 @@ async fn main() -> anyhow::Result<()> { println!("Requeued {} failed downloads.", failed.len()); } } + QueueAction::Sync { dry_run } => { + if dry_run { + println!("DRY RUN — no items will be enqueued"); + } + let stats = shanty_dl::sync_wanted_to_queue(db.conn(), dry_run).await?; + println!("\nSync complete: {stats}"); + } } } } diff --git a/src/queue.rs b/src/queue.rs index 65d4145..4ab7d20 100644 --- a/src/queue.rs +++ b/src/queue.rs @@ -1,7 +1,7 @@ use std::fmt; use std::time::Duration; -use sea_orm::DatabaseConnection; +use sea_orm::{ActiveValue::Set, DatabaseConnection}; use shanty_db::entities::download_queue::DownloadStatus; use shanty_db::entities::wanted_item::WantedStatus; @@ -106,20 +106,45 @@ pub async fn run_queue( ) .await?; - // Update wanted item status if linked + // Update wanted item status and create track record with MBID if let Some(wanted_id) = item.wanted_item_id { - if let Err(e) = queries::wanted::update_status( - conn, - wanted_id, - WantedStatus::Downloaded, - ) - .await - { - tracing::warn!( - wanted_id = wanted_id, - error = %e, - "failed to update wanted item status" - ); + if let Ok(wanted) = queries::wanted::get_by_id(conn, wanted_id).await { + // Create a track record with the MBID so the tagger + // can skip searching and go straight to the right recording + let now = chrono::Utc::now().naive_utc(); + let file_path = result.file_path.to_string_lossy().to_string(); + let file_size = std::fs::metadata(&result.file_path) + .map(|m| m.len() as i64) + .unwrap_or(0); + + let track_active = shanty_db::entities::track::ActiveModel { + file_path: Set(file_path), + title: Set(Some(result.title.clone())), + artist: Set(result.artist.clone()), + file_size: Set(file_size), + musicbrainz_id: Set(wanted.musicbrainz_id.clone()), + artist_id: Set(wanted.artist_id), + added_at: Set(now), + updated_at: Set(now), + ..Default::default() + }; + if let Err(e) = queries::tracks::upsert(conn, track_active).await { + tracing::warn!(error = %e, "failed to create track record after download"); + } + + if let Err(e) = queries::wanted::update_status( + conn, + wanted_id, + WantedStatus::Downloaded, + ) + .await + { + tracing::warn!( + wanted_id = wanted_id, + error = %e, + "failed to update wanted item status" + ); + } } } @@ -170,6 +195,79 @@ pub async fn run_queue( Ok(stats) } +/// Sync stats from a queue sync operation. +#[derive(Debug, Default)] +pub struct SyncStats { + pub found: u64, + pub enqueued: u64, + pub skipped: u64, +} + +impl fmt::Display for SyncStats { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + f, + "found: {}, enqueued: {}, skipped (already queued): {}", + self.found, self.enqueued, self.skipped, + ) + } +} + +/// Sync wanted items to the download queue. +/// Finds all Track-type Wanted items and enqueues them for download, +/// skipping any that already have a queue entry. +pub async fn sync_wanted_to_queue( + conn: &DatabaseConnection, + dry_run: bool, +) -> DlResult { + let wanted = queries::wanted::list(conn, Some(WantedStatus::Wanted)).await?; + let mut stats = SyncStats::default(); + + for item in &wanted { + stats.found += 1; + + // Build search query from the name + artist + let artist_name = if let Some(id) = item.artist_id { + queries::artists::get_by_id(conn, id) + .await + .map(|a| a.name) + .ok() + } else { + None + }; + + let query = match artist_name { + Some(ref artist) if !item.name.is_empty() => format!("{} {}", artist, item.name), + _ if !item.name.is_empty() => item.name.clone(), + Some(ref artist) => artist.clone(), + None => { + tracing::warn!(id = item.id, "cannot build query — no name or artist"); + continue; + } + }; + + // Check if already queued + if let Some(_existing) = queries::downloads::find_by_wanted_item_id(conn, item.id).await? { + tracing::debug!(id = item.id, name = %item.name, "already in queue, skipping"); + stats.skipped += 1; + continue; + } + + if dry_run { + println!("Would enqueue: {query}"); + stats.enqueued += 1; + continue; + } + + queries::downloads::enqueue(conn, &query, Some(item.id), "ytdlp").await?; + tracing::info!(id = item.id, query = %query, "enqueued for download"); + stats.enqueued += 1; + } + + tracing::info!(%stats, "sync complete"); + Ok(stats) +} + /// Download a single item directly (not from queue). pub async fn download_single( backend: &impl DownloadBackend, diff --git a/src/ytdlp.rs b/src/ytdlp.rs index 7b689a6..fd524f3 100644 --- a/src/ytdlp.rs +++ b/src/ytdlp.rs @@ -280,10 +280,10 @@ impl DownloadBackend for YtDlpBackend { } }; - // Build output template + // Build output template — include artist if available for filename-based metadata fallback let output_template = config .output_dir - .join("%(title)s.%(ext)s") + .join("%(artist,uploader,channel)s - %(title)s.%(ext)s") .to_string_lossy() .to_string(); diff --git a/tests/integration.rs b/tests/integration.rs index a1beca0..cfc4114 100644 --- a/tests/integration.rs +++ b/tests/integration.rs @@ -218,7 +218,7 @@ async fn test_wanted_item_status_updated_on_download() { }; // Create a wanted item - let wanted = queries::wanted::add(db.conn(), ItemType::Track, None, None, None) + let wanted = queries::wanted::add(db.conn(), ItemType::Track, "Wanted Song", None, None, None, None) .await .unwrap(); assert_eq!(wanted.status, WantedStatus::Wanted);