Shanty/dl
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Formatting

Browse Source
This commit is contained in:
Connor Johnstone
2026-03-18 15:36:16 -04:00
parent a57df38eb1
commit 2592651c9a
8 changed files with 111 additions and 98 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
src/backend.rs
View File

@@ -80,7 +80,9 @@ impl std::str::FromStr for AudioFormat {
"mp3" => Ok(AudioFormat::Mp3), "mp3" => Ok(AudioFormat::Mp3),
"flac" => Ok(AudioFormat::Flac), "flac" => Ok(AudioFormat::Flac),
"best" => Ok(AudioFormat::Best), "best" => Ok(AudioFormat::Best),
_ => Err(format!("unsupported format: {s} (expected opus, mp3, flac, or best)")), _ => Err(format!(
"unsupported format: {s} (expected opus, mp3, flac, or best)"
)),
} }
} }
} }
@@ -98,7 +100,10 @@ pub trait DownloadBackend: Send + Sync {
fn check_available(&self) -> impl std::future::Future<Output = DlResult<()>> + Send; fn check_available(&self) -> impl std::future::Future<Output = DlResult<()>> + Send;
/// Search for tracks matching a query. /// Search for tracks matching a query.
fn search(&self, query: &str) -> impl std::future::Future<Output = DlResult<Vec<SearchResult>>> + Send; fn search(
&self,
query: &str,
) -> impl std::future::Future<Output = DlResult<Vec<SearchResult>>> + Send;
/// Download a target to the configured output directory. /// Download a target to the configured output directory.
fn download( fn download(

4
src/error.rs
View File

@@ -35,9 +35,7 @@ impl DlError {
pub fn is_transient(&self) -> bool { pub fn is_transient(&self) -> bool {
matches!( matches!(
self, self,
DlError::RateLimited(_) DlError::RateLimited(_) | DlError::Io(_) | DlError::BackendError(_)
| DlError::Io(_)
| DlError::BackendError(_)
) )
} }
} }

9
src/lib.rs
View File

@@ -9,7 +9,12 @@ pub mod queue;
pub mod rate_limit; pub mod rate_limit;
pub mod ytdlp; pub mod ytdlp;
pub use backend::{AudioFormat, BackendConfig, DownloadBackend, DownloadResult, DownloadTarget, SearchResult}; pub use backend::{
AudioFormat, BackendConfig, DownloadBackend, DownloadResult, DownloadTarget, SearchResult,
};
pub use error::{DlError, DlResult}; pub use error::{DlError, DlResult};
pub use queue::{DlStats, ProgressFn, SyncStats, download_single, run_queue, run_queue_with_progress, sync_wanted_to_queue}; pub use queue::{
DlStats, ProgressFn, SyncStats, download_single, run_queue, run_queue_with_progress,
sync_wanted_to_queue,
};
pub use ytdlp::{SearchSource, YtDlpBackend}; pub use ytdlp::{SearchSource, YtDlpBackend};

24
src/main.rs
View File

@@ -156,9 +156,7 @@ fn make_backend_config(
output: &Option<PathBuf>, output: &Option<PathBuf>,
cookies: &Option<PathBuf>, cookies: &Option<PathBuf>,
) -> anyhow::Result<BackendConfig> { ) -> anyhow::Result<BackendConfig> {
let fmt: AudioFormat = format let fmt: AudioFormat = format.parse().map_err(|e: String| anyhow::anyhow!(e))?;
.parse()
.map_err(|e: String| anyhow::anyhow!(e))?;
Ok(BackendConfig { Ok(BackendConfig {
output_dir: output.clone().unwrap_or_else(default_output_dir), output_dir: output.clone().unwrap_or_else(default_output_dir),
format: fmt, format: fmt,
@@ -199,13 +197,12 @@ async fn main() -> anyhow::Result<()> {
let config = make_backend_config(&format, &output, &cookies)?; let config = make_backend_config(&format, &output, &cookies)?;
// Determine if it's a URL or a search query // Determine if it's a URL or a search query
let target = if query_or_url.starts_with("http://") let target =
|| query_or_url.starts_with("https://") if query_or_url.starts_with("http://") || query_or_url.starts_with("https://") {
{ DownloadTarget::Url(query_or_url)
DownloadTarget::Url(query_or_url) } else {
} else { DownloadTarget::Query(query_or_url)
DownloadTarget::Query(query_or_url) };
};
download_single(&backend, target, &config, dry_run).await?; download_single(&backend, target, &config, dry_run).await?;
} }
@@ -235,7 +232,8 @@ async fn main() -> anyhow::Result<()> {
println!("\nQueue processing complete: {stats}"); println!("\nQueue processing complete: {stats}");
} }
QueueAction::Add { query } => { QueueAction::Add { query } => {
let item = queries::downloads::enqueue(db.conn(), &query, None, "ytdlp").await?; let item =
queries::downloads::enqueue(db.conn(), &query, None, "ytdlp").await?;
println!("Added to queue: id={}, query=\"{}\"", item.id, item.query); println!("Added to queue: id={}, query=\"{}\"", item.id, item.query);
} }
QueueAction::List { status } => { QueueAction::List { status } => {
@@ -254,8 +252,8 @@ async fn main() -> anyhow::Result<()> {
println!("Queue is empty."); println!("Queue is empty.");
} else { } else {
println!( println!(
"{:<5} {:<12} {:<6} {:<40} {}", "{:<5} {:<12} {:<6} {:<40} ERROR",
"ID", "STATUS", "RETRY", "QUERY", "ERROR" "ID", "STATUS", "RETRY", "QUERY"
); );
for item in &items { for item in &items {
println!( println!(

92
src/queue.rs
View File

@@ -81,7 +81,11 @@ pub async fn run_queue_with_progress(
stats.downloads_attempted += 1; stats.downloads_attempted += 1;
if let Some(ref cb) = on_progress { if let Some(ref cb) = on_progress {
cb(stats.downloads_attempted, total, &format!("Downloading: {}", item.query)); cb(
stats.downloads_attempted,
total,
&format!("Downloading: {}", item.query),
);
} }
tracing::info!( tracing::info!(
@@ -106,8 +110,7 @@ pub async fn run_queue_with_progress(
} }
// Mark as downloading // Mark as downloading
queries::downloads::update_status(conn, item.id, DownloadStatus::Downloading, None) queries::downloads::update_status(conn, item.id, DownloadStatus::Downloading, None).await?;
.await?;
// Determine download target // Determine download target
let target = if let Some(ref url) = item.source_url { let target = if let Some(ref url) = item.source_url {
@@ -126,53 +129,45 @@ pub async fn run_queue_with_progress(
"download completed" "download completed"
); );
queries::downloads::update_status( queries::downloads::update_status(conn, item.id, DownloadStatus::Completed, None)
conn, .await?;
item.id,
DownloadStatus::Completed,
None,
)
.await?;
// Update wanted item status and create track record with MBID // Update wanted item status and create track record with MBID
if let Some(wanted_id) = item.wanted_item_id { if let Some(wanted_id) = item.wanted_item_id
if let Ok(wanted) = queries::wanted::get_by_id(conn, wanted_id).await { && 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 // Create a track record with the MBID so the tagger
let now = chrono::Utc::now().naive_utc(); // can skip searching and go straight to the right recording
let file_path = result.file_path.to_string_lossy().to_string(); let now = chrono::Utc::now().naive_utc();
let file_size = std::fs::metadata(&result.file_path) let file_path = result.file_path.to_string_lossy().to_string();
.map(|m| m.len() as i64) let file_size = std::fs::metadata(&result.file_path)
.unwrap_or(0); .map(|m| m.len() as i64)
.unwrap_or(0);
let track_active = shanty_db::entities::track::ActiveModel { let track_active = shanty_db::entities::track::ActiveModel {
file_path: Set(file_path), file_path: Set(file_path),
title: Set(Some(result.title.clone())), title: Set(Some(result.title.clone())),
artist: Set(result.artist.clone()), artist: Set(result.artist.clone()),
file_size: Set(file_size), file_size: Set(file_size),
musicbrainz_id: Set(wanted.musicbrainz_id.clone()), musicbrainz_id: Set(wanted.musicbrainz_id.clone()),
artist_id: Set(wanted.artist_id), artist_id: Set(wanted.artist_id),
added_at: Set(now), added_at: Set(now),
updated_at: Set(now), updated_at: Set(now),
..Default::default() ..Default::default()
}; };
if let Err(e) = queries::tracks::upsert(conn, track_active).await { if let Err(e) = queries::tracks::upsert(conn, track_active).await {
tracing::warn!(error = %e, "failed to create track record after download"); tracing::warn!(error = %e, "failed to create track record after download");
} }
if let Err(e) = queries::wanted::update_status( if let Err(e) =
conn, queries::wanted::update_status(conn, wanted_id, WantedStatus::Downloaded)
wanted_id, .await
WantedStatus::Downloaded, {
) tracing::warn!(
.await wanted_id = wanted_id,
{ error = %e,
tracing::warn!( "failed to update wanted item status"
wanted_id = wanted_id, );
error = %e,
"failed to update wanted item status"
);
}
} }
} }
@@ -244,10 +239,7 @@ impl fmt::Display for SyncStats {
/// Sync wanted items to the download queue. /// Sync wanted items to the download queue.
/// Finds all Track-type Wanted items and enqueues them for download, /// Finds all Track-type Wanted items and enqueues them for download,
/// skipping any that already have a queue entry. /// skipping any that already have a queue entry.
pub async fn sync_wanted_to_queue( pub async fn sync_wanted_to_queue(conn: &DatabaseConnection, dry_run: bool) -> DlResult<SyncStats> {
conn: &DatabaseConnection,
dry_run: bool,
) -> DlResult<SyncStats> {
let wanted = queries::wanted::list(conn, Some(WantedStatus::Wanted)).await?; let wanted = queries::wanted::list(conn, Some(WantedStatus::Wanted)).await?;
let mut stats = SyncStats::default(); let mut stats = SyncStats::default();

3
src/rate_limit.rs
View File

@@ -42,8 +42,7 @@ impl RateLimiter {
state.remaining -= 1; state.remaining -= 1;
// Warn when approaching the limit // Warn when approaching the limit
let pct_remaining = let pct_remaining = (state.remaining as f64 / self.max_per_hour as f64) * 100.0;
(state.remaining as f64 / self.max_per_hour as f64) * 100.0;
if pct_remaining < 10.0 && pct_remaining > 0.0 { if pct_remaining < 10.0 && pct_remaining > 0.0 {
tracing::warn!( tracing::warn!(
remaining = state.remaining, remaining = state.remaining,

36
src/ytdlp.rs
View File

@@ -4,7 +4,9 @@ use std::process::Stdio;
use serde::Deserialize; use serde::Deserialize;
use tokio::process::Command; use tokio::process::Command;
use crate::backend::{BackendConfig, DownloadBackend, DownloadResult, DownloadTarget, SearchResult}; use crate::backend::{
BackendConfig, DownloadBackend, DownloadResult, DownloadTarget, SearchResult,
};
use crate::error::{DlError, DlResult}; use crate::error::{DlError, DlResult};
use crate::rate_limit::RateLimiter; use crate::rate_limit::RateLimiter;
@@ -22,7 +24,9 @@ impl std::str::FromStr for SearchSource {
match s.to_lowercase().as_str() { match s.to_lowercase().as_str() {
"ytmusic" | "youtube_music" | "youtubemusic" => Ok(SearchSource::YouTubeMusic), "ytmusic" | "youtube_music" | "youtubemusic" => Ok(SearchSource::YouTubeMusic),
"youtube" | "yt" => Ok(SearchSource::YouTube), "youtube" | "yt" => Ok(SearchSource::YouTube),
_ => Err(format!("unknown search source: {s} (expected ytmusic or youtube)")), _ => Err(format!(
"unknown search source: {s} (expected ytmusic or youtube)"
)),
} }
} }
} }
@@ -63,9 +67,7 @@ impl YtDlpBackend {
self.rate_limiter.acquire().await; self.rate_limiter.acquire().await;
let mut cmd = Command::new("yt-dlp"); let mut cmd = Command::new("yt-dlp");
cmd.args(args) cmd.args(args).stdout(Stdio::piped()).stderr(Stdio::piped());
.stdout(Stdio::piped())
.stderr(Stdio::piped());
// Add cookies if configured // Add cookies if configured
if let Some(ref cookies) = self.cookies_path { if let Some(ref cookies) = self.cookies_path {
@@ -120,7 +122,9 @@ impl YtDlpBackend {
if !output.status.success() { if !output.status.success() {
let stderr = String::from_utf8_lossy(&output.stderr).to_string(); let stderr = String::from_utf8_lossy(&output.stderr).to_string();
tracing::warn!(stderr = %stderr, "ytmusic search failed"); tracing::warn!(stderr = %stderr, "ytmusic search failed");
return Err(DlError::BackendError(format!("ytmusic search failed: {stderr}"))); return Err(DlError::BackendError(format!(
"ytmusic search failed: {stderr}"
)));
} }
let stdout = String::from_utf8_lossy(&output.stdout); let stdout = String::from_utf8_lossy(&output.stdout);
@@ -180,17 +184,17 @@ impl YtDlpBackend {
fn find_ytmusic_script(&self) -> DlResult<PathBuf> { fn find_ytmusic_script(&self) -> DlResult<PathBuf> {
// Check next to the current executable // Check next to the current executable
if let Ok(exe) = std::env::current_exe() { if let Ok(exe) = std::env::current_exe() {
let beside_exe = exe.parent().unwrap_or(std::path::Path::new(".")).join("ytmusic_search.py"); let beside_exe = exe
.parent()
.unwrap_or(std::path::Path::new("."))
.join("ytmusic_search.py");
if beside_exe.exists() { if beside_exe.exists() {
return Ok(beside_exe); return Ok(beside_exe);
} }
} }
// Check common install locations // Check common install locations
for dir in &[ for dir in &["/usr/share/shanty", "/usr/local/share/shanty"] {
"/usr/share/shanty",
"/usr/local/share/shanty",
] {
let path = PathBuf::from(dir).join("ytmusic_search.py"); let path = PathBuf::from(dir).join("ytmusic_search.py");
if path.exists() { if path.exists() {
return Ok(path); return Ok(path);
@@ -303,10 +307,7 @@ impl DownloadBackend for YtDlpBackend {
// Add cookies from backend config or backend's own cookies // Add cookies from backend config or backend's own cookies
let cookies_str; let cookies_str;
let cookies_path = config let cookies_path = config.cookies_path.as_ref().or(self.cookies_path.as_ref());
.cookies_path
.as_ref()
.or(self.cookies_path.as_ref());
if let Some(c) = cookies_path { if let Some(c) = cookies_path {
cookies_str = c.to_string_lossy().to_string(); cookies_str = c.to_string_lossy().to_string();
args.push("--cookies"); args.push("--cookies");
@@ -318,9 +319,8 @@ impl DownloadBackend for YtDlpBackend {
let output = self.run_ytdlp(&args).await?; let output = self.run_ytdlp(&args).await?;
// Parse the JSON output to get the actual file path // Parse the JSON output to get the actual file path
let info: YtDlpDownloadInfo = serde_json::from_str(output.trim()).map_err(|e| { let info: YtDlpDownloadInfo = serde_json::from_str(output.trim())
DlError::BackendError(format!("failed to parse yt-dlp output: {e}")) .map_err(|e| DlError::BackendError(format!("failed to parse yt-dlp output: {e}")))?;
})?;
// --print-json reports the pre-extraction filename (e.g. .webm), // --print-json reports the pre-extraction filename (e.g. .webm),
// but --extract-audio produces a file with the target format extension. // but --extract-audio produces a file with the target format extension.

32
tests/integration.rs
View File

@@ -1,8 +1,10 @@
use shanty_db::entities::download_queue::DownloadStatus; use shanty_db::entities::download_queue::DownloadStatus;
use shanty_db::{Database, queries}; use shanty_db::{Database, queries};
use shanty_dl::backend::{AudioFormat, BackendConfig, DownloadBackend, DownloadResult, DownloadTarget, SearchResult}; use shanty_dl::backend::{
AudioFormat, BackendConfig, DownloadBackend, DownloadResult, DownloadTarget, SearchResult,
};
use shanty_dl::error::DlResult; use shanty_dl::error::DlResult;
use shanty_dl::queue::{run_queue, download_single}; use shanty_dl::queue::{download_single, run_queue};
use tempfile::TempDir; use tempfile::TempDir;
/// Mock backend for testing without yt-dlp. /// Mock backend for testing without yt-dlp.
@@ -84,7 +86,9 @@ async fn test_queue_process_success() {
.unwrap(); .unwrap();
// Process queue // Process queue
let stats = run_queue(db.conn(), &backend, &config, false).await.unwrap(); let stats = run_queue(db.conn(), &backend, &config, false)
.await
.unwrap();
assert_eq!(stats.downloads_attempted, 1); assert_eq!(stats.downloads_attempted, 1);
assert_eq!(stats.downloads_completed, 1); assert_eq!(stats.downloads_completed, 1);
assert_eq!(stats.downloads_failed, 0); assert_eq!(stats.downloads_failed, 0);
@@ -112,7 +116,9 @@ async fn test_queue_process_failure() {
.await .await
.unwrap(); .unwrap();
let stats = run_queue(db.conn(), &backend, &config, false).await.unwrap(); let stats = run_queue(db.conn(), &backend, &config, false)
.await
.unwrap();
assert_eq!(stats.downloads_attempted, 1); assert_eq!(stats.downloads_attempted, 1);
assert_eq!(stats.downloads_failed, 1); assert_eq!(stats.downloads_failed, 1);
@@ -218,9 +224,17 @@ async fn test_wanted_item_status_updated_on_download() {
}; };
// Create a wanted item // Create a wanted item
let wanted = queries::wanted::add(db.conn(), ItemType::Track, "Wanted Song", None, None, None, None) let wanted = queries::wanted::add(
.await db.conn(),
.unwrap(); ItemType::Track,
"Wanted Song",
None,
None,
None,
None,
)
.await
.unwrap();
assert_eq!(wanted.status, WantedStatus::Wanted); assert_eq!(wanted.status, WantedStatus::Wanted);
// Enqueue download linked to the wanted item // Enqueue download linked to the wanted item
@@ -229,7 +243,9 @@ async fn test_wanted_item_status_updated_on_download() {
.unwrap(); .unwrap();
// Process queue // Process queue
run_queue(db.conn(), &backend, &config, false).await.unwrap(); run_queue(db.conn(), &backend, &config, false)
.await
.unwrap();
// Wanted item should now be Downloaded // Wanted item should now be Downloaded
let updated = queries::wanted::get_by_id(db.conn(), wanted.id) let updated = queries::wanted::get_by_id(db.conn(), wanted.id)