use serde::{Deserialize, Serialize}; use sqlx::PgPool; use sqlx::postgres::PgPoolOptions; use uuid::Uuid; /// Generate a short URL-safe slug from a UUID v4. fn generate_slug() -> String { Uuid::new_v4().simple().to_string()[..12].to_owned() } pub async fn connect(database_url: &str) -> Result { PgPoolOptions::new() .max_connections(5) .connect(database_url) .await } pub async fn migrate(pool: &PgPool) -> Result<(), sqlx::migrate::MigrateError> { sqlx::migrate!("./migrations").run(pool).await } // --- Models --- #[derive(Debug, Clone, Serialize, Deserialize, sqlx::FromRow)] pub struct Artist { pub id: i64, pub name: String, } #[derive(Debug, Clone, Serialize, Deserialize, sqlx::FromRow)] pub struct Album { pub id: i64, pub artist_id: i64, pub name: String, pub year: Option, } #[derive(Debug, Clone, Serialize, Deserialize, sqlx::FromRow)] pub struct PendingTrack { pub id: Uuid, pub status: String, pub inbox_path: String, pub file_hash: String, pub file_size: i64, // Raw metadata from file tags pub raw_title: Option, pub raw_artist: Option, pub raw_album: Option, pub raw_year: Option, pub raw_track_number: Option, pub raw_genre: Option, pub duration_secs: Option, // Path-derived hints pub path_artist: Option, pub path_album: Option, pub path_year: Option, pub path_track_number: Option, pub path_title: Option, // Normalized (LLM output) pub norm_title: Option, pub norm_artist: Option, pub norm_album: Option, pub norm_year: Option, pub norm_track_number: Option, pub norm_genre: Option, pub norm_featured_artists: Option, // JSON array pub confidence: Option, pub llm_notes: Option, pub error_message: Option, pub created_at: chrono::DateTime, pub updated_at: chrono::DateTime, } #[derive(Debug, Clone, Serialize, Deserialize, sqlx::FromRow)] pub struct SimilarArtist { pub id: i64, pub name: String, pub similarity: f32, } #[derive(Debug, Clone, Serialize, Deserialize, sqlx::FromRow)] pub struct SimilarAlbum { pub id: i64, pub artist_id: i64, pub name: String, pub year: Option, pub similarity: f32, } #[allow(dead_code)] #[derive(Debug, Clone, Serialize, Deserialize, sqlx::FromRow)] pub struct AlbumImage { pub id: i64, pub album_id: i64, pub image_type: String, pub file_path: String, pub file_hash: String, pub mime_type: String, pub width: Option, pub height: Option, pub file_size: i64, } // --- Queries --- pub async fn file_hash_exists(pool: &PgPool, hash: &str) -> Result { let row: (bool,) = sqlx::query_as( "SELECT EXISTS(SELECT 1 FROM tracks WHERE file_hash = $1) OR EXISTS(SELECT 1 FROM pending_tracks WHERE file_hash = $1 AND status NOT IN ('rejected', 'error'))" ) .bind(hash) .fetch_one(pool) .await?; Ok(row.0) } pub async fn insert_pending( pool: &PgPool, inbox_path: &str, file_hash: &str, file_size: i64, raw: &RawFields, path_hints: &PathHints, duration_secs: Option, ) -> Result { let row: (Uuid,) = sqlx::query_as( r#"INSERT INTO pending_tracks (inbox_path, file_hash, file_size, raw_title, raw_artist, raw_album, raw_year, raw_track_number, raw_genre, path_title, path_artist, path_album, path_year, path_track_number, duration_secs, status) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, 'pending') RETURNING id"#, ) .bind(inbox_path) .bind(file_hash) .bind(file_size) .bind(&raw.title) .bind(&raw.artist) .bind(&raw.album) .bind(raw.year) .bind(raw.track_number) .bind(&raw.genre) .bind(&path_hints.title) .bind(&path_hints.artist) .bind(&path_hints.album) .bind(path_hints.year) .bind(path_hints.track_number) .bind(duration_secs) .fetch_one(pool) .await?; Ok(row.0) } pub async fn update_pending_normalized( pool: &PgPool, id: Uuid, status: &str, norm: &NormalizedFields, error_message: Option<&str>, ) -> Result<(), sqlx::Error> { let featured_json = if norm.featured_artists.is_empty() { None } else { Some(serde_json::to_string(&norm.featured_artists).unwrap_or_default()) }; sqlx::query( r#"UPDATE pending_tracks SET status = $2, norm_title = $3, norm_artist = $4, norm_album = $5, norm_year = $6, norm_track_number = $7, norm_genre = $8, norm_featured_artists = $9, confidence = $10, llm_notes = $11, error_message = $12, updated_at = NOW() WHERE id = $1"#, ) .bind(id) .bind(status) .bind(&norm.title) .bind(&norm.artist) .bind(&norm.album) .bind(norm.year) .bind(norm.track_number) .bind(&norm.genre) .bind(&featured_json) .bind(norm.confidence) .bind(&norm.notes) .bind(error_message) .execute(pool) .await?; Ok(()) } pub async fn update_pending_status( pool: &PgPool, id: Uuid, status: &str, error_message: Option<&str>, ) -> Result<(), sqlx::Error> { sqlx::query("UPDATE pending_tracks SET status = $2, error_message = $3, updated_at = NOW() WHERE id = $1") .bind(id) .bind(status) .bind(error_message) .execute(pool) .await?; Ok(()) } pub async fn find_similar_artists(pool: &PgPool, name: &str, limit: i32) -> Result, sqlx::Error> { // pg_trgm needs at least 3 chars to produce trigrams; for shorter queries use ILIKE prefix if name.chars().count() < 3 { sqlx::query_as::<_, SimilarArtist>( "SELECT id, name, 1.0::real AS similarity FROM artists WHERE name ILIKE $1 || '%' ORDER BY name LIMIT $2" ) .bind(name) .bind(limit) .fetch_all(pool) .await } else { sqlx::query_as::<_, SimilarArtist>( r#"SELECT id, name, MAX(sim) AS similarity FROM ( SELECT id, name, similarity(name, $1) AS sim FROM artists WHERE name % $1 UNION ALL SELECT id, name, 0.01::real AS sim FROM artists WHERE name ILIKE '%' || $1 || '%' ) sub GROUP BY id, name ORDER BY similarity DESC LIMIT $2"# ) .bind(name) .bind(limit) .fetch_all(pool) .await } } pub async fn find_similar_albums(pool: &PgPool, name: &str, limit: i32) -> Result, sqlx::Error> { sqlx::query_as::<_, SimilarAlbum>( "SELECT id, artist_id, name, year, similarity(name, $1) AS similarity FROM albums WHERE name % $1 ORDER BY similarity DESC LIMIT $2" ) .bind(name) .bind(limit) .fetch_all(pool) .await } pub async fn upsert_artist(pool: &PgPool, name: &str) -> Result { let slug = generate_slug(); let row: (i64,) = sqlx::query_as( "INSERT INTO artists (name, slug) VALUES ($1, $2) ON CONFLICT (name) DO UPDATE SET name = EXCLUDED.name RETURNING id" ) .bind(name) .bind(&slug) .fetch_one(pool) .await?; Ok(row.0) } pub async fn upsert_album(pool: &PgPool, artist_id: i64, name: &str, year: Option) -> Result { let slug = generate_slug(); let row: (i64,) = sqlx::query_as( r#"INSERT INTO albums (artist_id, name, year, slug) VALUES ($1, $2, $3, $4) ON CONFLICT (artist_id, name) DO UPDATE SET year = COALESCE(EXCLUDED.year, albums.year) RETURNING id"# ) .bind(artist_id) .bind(name) .bind(year) .bind(&slug) .fetch_one(pool) .await?; Ok(row.0) } pub async fn insert_track( pool: &PgPool, artist_id: i64, album_id: Option, title: &str, track_number: Option, genre: Option<&str>, duration_secs: Option, file_hash: &str, file_size: i64, storage_path: &str, ) -> Result { let slug = generate_slug(); let row: (i64,) = sqlx::query_as( r#"INSERT INTO tracks (artist_id, album_id, title, track_number, genre, duration_secs, file_hash, file_size, storage_path, slug) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10) RETURNING id"# ) .bind(artist_id) .bind(album_id) .bind(title) .bind(track_number) .bind(genre) .bind(duration_secs) .bind(file_hash) .bind(file_size) .bind(storage_path) .bind(&slug) .fetch_one(pool) .await?; Ok(row.0) } pub async fn link_track_artist(pool: &PgPool, track_id: i64, artist_id: i64, role: &str) -> Result<(), sqlx::Error> { sqlx::query( "INSERT INTO track_artists (track_id, artist_id, role) VALUES ($1, $2, $3) ON CONFLICT DO NOTHING" ) .bind(track_id) .bind(artist_id) .bind(role) .execute(pool) .await?; Ok(()) } pub async fn approve_and_finalize( pool: &PgPool, pending_id: Uuid, storage_path: &str, ) -> Result { let pt: PendingTrack = sqlx::query_as("SELECT * FROM pending_tracks WHERE id = $1") .bind(pending_id) .fetch_one(pool) .await?; // Check if track already exists (e.g. previously approved but pending not cleaned up) let existing: Option<(i64,)> = sqlx::query_as("SELECT id FROM tracks WHERE file_hash = $1") .bind(&pt.file_hash) .fetch_optional(pool) .await?; if let Some((track_id,)) = existing { // Already finalized — just mark pending as approved update_pending_status(pool, pending_id, "approved", None).await?; return Ok(track_id); } let artist_name = pt.norm_artist.as_deref().unwrap_or("Unknown Artist"); let artist_id = upsert_artist(pool, artist_name).await?; let album_id = match pt.norm_album.as_deref() { Some(album_name) => Some(upsert_album(pool, artist_id, album_name, pt.norm_year).await?), None => None, }; let title = pt.norm_title.as_deref().unwrap_or("Unknown Title"); let track_id = insert_track( pool, artist_id, album_id, title, pt.norm_track_number, pt.norm_genre.as_deref(), pt.duration_secs, &pt.file_hash, pt.file_size, storage_path, ) .await?; // Link primary artist link_track_artist(pool, track_id, artist_id, "primary").await?; // Link featured artists if let Some(featured_json) = &pt.norm_featured_artists { if let Ok(featured) = serde_json::from_str::>(featured_json) { for feat_name in &featured { let feat_id = upsert_artist(pool, feat_name).await?; link_track_artist(pool, track_id, feat_id, "featured").await?; } } } update_pending_status(pool, pending_id, "approved", None).await?; Ok(track_id) } // --- Album images --- pub async fn image_hash_exists(pool: &PgPool, hash: &str) -> Result { let row: (bool,) = sqlx::query_as("SELECT EXISTS(SELECT 1 FROM album_images WHERE file_hash = $1)") .bind(hash) .fetch_one(pool) .await?; Ok(row.0) } pub async fn insert_album_image( pool: &PgPool, album_id: i64, image_type: &str, file_path: &str, file_hash: &str, mime_type: &str, file_size: i64, ) -> Result { let row: (i64,) = sqlx::query_as( r#"INSERT INTO album_images (album_id, image_type, file_path, file_hash, mime_type, file_size) VALUES ($1, $2, $3, $4, $5, $6) ON CONFLICT (file_hash) DO NOTHING RETURNING id"# ) .bind(album_id) .bind(image_type) .bind(file_path) .bind(file_hash) .bind(mime_type) .bind(file_size) .fetch_one(pool) .await?; Ok(row.0) } #[allow(dead_code)] pub async fn get_album_images(pool: &PgPool, album_id: i64) -> Result, sqlx::Error> { sqlx::query_as::<_, AlbumImage>("SELECT * FROM album_images WHERE album_id = $1 ORDER BY image_type") .bind(album_id) .fetch_all(pool) .await } /// Find album_id by artist+album name (used when linking covers to already-finalized albums) pub async fn find_album_id(pool: &PgPool, artist_name: &str, album_name: &str) -> Result, sqlx::Error> { let row: Option<(i64,)> = sqlx::query_as( r#"SELECT a.id FROM albums a JOIN artists ar ON a.artist_id = ar.id WHERE ar.name = $1 AND a.name = $2"# ) .bind(artist_name) .bind(album_name) .fetch_optional(pool) .await?; Ok(row.map(|r| r.0)) } /// Fetch pending tracks that need (re-)processing by the LLM pipeline. pub async fn list_pending_for_processing(pool: &PgPool, limit: i64) -> Result, sqlx::Error> { sqlx::query_as::<_, PendingTrack>( "SELECT * FROM pending_tracks WHERE status = 'pending' ORDER BY created_at ASC LIMIT $1" ) .bind(limit) .fetch_all(pool) .await } // --- DTOs for insert helpers --- #[derive(Debug, Default)] pub struct RawFields { pub title: Option, pub artist: Option, pub album: Option, pub year: Option, pub track_number: Option, pub genre: Option, } #[derive(Debug, Default)] pub struct PathHints { pub title: Option, pub artist: Option, pub album: Option, pub year: Option, pub track_number: Option, } #[derive(Debug, Default, Serialize, Deserialize)] pub struct NormalizedFields { pub title: Option, pub artist: Option, pub album: Option, pub year: Option, pub track_number: Option, pub genre: Option, #[serde(default)] pub featured_artists: Vec, pub confidence: Option, pub notes: Option, } // --- Admin queries --- pub async fn list_pending(pool: &PgPool, status_filter: Option<&str>, limit: i64, offset: i64) -> Result, sqlx::Error> { match status_filter { Some(status) => { sqlx::query_as::<_, PendingTrack>( "SELECT * FROM pending_tracks WHERE status = $1 ORDER BY created_at DESC LIMIT $2 OFFSET $3" ) .bind(status) .bind(limit) .bind(offset) .fetch_all(pool) .await } None => { sqlx::query_as::<_, PendingTrack>( "SELECT * FROM pending_tracks ORDER BY created_at DESC LIMIT $1 OFFSET $2" ) .bind(limit) .bind(offset) .fetch_all(pool) .await } } } pub async fn get_pending(pool: &PgPool, id: Uuid) -> Result, sqlx::Error> { sqlx::query_as::<_, PendingTrack>("SELECT * FROM pending_tracks WHERE id = $1") .bind(id) .fetch_optional(pool) .await } pub async fn delete_pending(pool: &PgPool, id: Uuid) -> Result { let result = sqlx::query("DELETE FROM pending_tracks WHERE id = $1") .bind(id) .execute(pool) .await?; Ok(result.rows_affected() > 0) } pub async fn list_artists_all(pool: &PgPool) -> Result, sqlx::Error> { sqlx::query_as::<_, Artist>("SELECT id, name FROM artists ORDER BY name") .fetch_all(pool) .await } pub async fn list_albums_by_artist(pool: &PgPool, artist_id: i64) -> Result, sqlx::Error> { sqlx::query_as::<_, Album>("SELECT id, artist_id, name, year FROM albums WHERE artist_id = $1 ORDER BY year, name") .bind(artist_id) .fetch_all(pool) .await } pub async fn update_artist_name(pool: &PgPool, id: i64, name: &str) -> Result { let result = sqlx::query("UPDATE artists SET name = $2 WHERE id = $1") .bind(id) .bind(name) .execute(pool) .await?; Ok(result.rows_affected() > 0) } pub async fn update_album(pool: &PgPool, id: i64, name: &str, year: Option) -> Result { let result = sqlx::query("UPDATE albums SET name = $2, year = $3 WHERE id = $1") .bind(id) .bind(name) .bind(year) .execute(pool) .await?; Ok(result.rows_affected() > 0) } #[derive(Debug, Serialize)] pub struct Stats { pub total_tracks: i64, pub total_artists: i64, pub total_albums: i64, pub pending_count: i64, pub review_count: i64, pub error_count: i64, pub merged_count: i64, } pub async fn get_stats(pool: &PgPool) -> Result { let (total_tracks,): (i64,) = sqlx::query_as("SELECT COUNT(*) FROM tracks").fetch_one(pool).await?; let (total_artists,): (i64,) = sqlx::query_as("SELECT COUNT(*) FROM artists").fetch_one(pool).await?; let (total_albums,): (i64,) = sqlx::query_as("SELECT COUNT(*) FROM albums").fetch_one(pool).await?; let (pending_count,): (i64,) = sqlx::query_as("SELECT COUNT(*) FROM pending_tracks WHERE status = 'pending'").fetch_one(pool).await?; let (review_count,): (i64,) = sqlx::query_as("SELECT COUNT(*) FROM pending_tracks WHERE status = 'review'").fetch_one(pool).await?; let (error_count,): (i64,) = sqlx::query_as("SELECT COUNT(*) FROM pending_tracks WHERE status = 'error'").fetch_one(pool).await?; let (merged_count,): (i64,) = sqlx::query_as("SELECT COUNT(*) FROM pending_tracks WHERE status = 'merged'").fetch_one(pool).await?; Ok(Stats { total_tracks, total_artists, total_albums, pending_count, review_count, error_count, merged_count }) } // =================== Artist Merges =================== #[derive(Debug, Clone, Serialize, Deserialize, sqlx::FromRow)] pub struct ArtistMerge { pub id: Uuid, pub status: String, pub source_artist_ids: String, pub proposal: Option, pub llm_notes: Option, pub error_message: Option, pub created_at: chrono::DateTime, pub updated_at: chrono::DateTime, } #[derive(Debug, Serialize)] pub struct ArtistFullData { pub id: i64, pub name: String, pub albums: Vec, } #[derive(Debug, Serialize)] pub struct AlbumFullData { pub id: i64, pub name: String, pub year: Option, pub tracks: Vec, } #[derive(Debug, Serialize, sqlx::FromRow)] pub struct TrackBasic { pub id: i64, pub title: String, pub track_number: Option, pub storage_path: String, } #[derive(Debug, sqlx::FromRow)] pub struct TrackWithAlbum { pub id: i64, pub storage_path: String, pub album_name: Option, } pub async fn insert_artist_merge(pool: &PgPool, source_artist_ids: &[i64]) -> Result { let ids_json = serde_json::to_string(source_artist_ids).unwrap_or_default(); let row: (Uuid,) = sqlx::query_as( "INSERT INTO artist_merges (source_artist_ids) VALUES ($1) RETURNING id" ).bind(&ids_json).fetch_one(pool).await?; Ok(row.0) } pub async fn list_artist_merges(pool: &PgPool) -> Result, sqlx::Error> { sqlx::query_as::<_, ArtistMerge>("SELECT * FROM artist_merges ORDER BY created_at DESC") .fetch_all(pool).await } pub async fn get_artist_merge(pool: &PgPool, id: Uuid) -> Result, sqlx::Error> { sqlx::query_as::<_, ArtistMerge>("SELECT * FROM artist_merges WHERE id = $1") .bind(id).fetch_optional(pool).await } pub async fn update_merge_status(pool: &PgPool, id: Uuid, status: &str, error: Option<&str>) -> Result<(), sqlx::Error> { sqlx::query("UPDATE artist_merges SET status = $2, error_message = $3, updated_at = NOW() WHERE id = $1") .bind(id).bind(status).bind(error).execute(pool).await?; Ok(()) } pub async fn update_merge_proposal(pool: &PgPool, id: Uuid, proposal_json: &str, notes: &str) -> Result<(), sqlx::Error> { sqlx::query("UPDATE artist_merges SET proposal = $2, llm_notes = $3, status = 'review', error_message = NULL, updated_at = NOW() WHERE id = $1") .bind(id).bind(proposal_json).bind(notes).execute(pool).await?; Ok(()) } pub async fn get_pending_merges_for_processing(pool: &PgPool) -> Result, sqlx::Error> { let rows: Vec<(Uuid,)> = sqlx::query_as( "SELECT id FROM artist_merges WHERE status = 'pending' ORDER BY created_at ASC LIMIT 5" ).fetch_all(pool).await?; Ok(rows.into_iter().map(|(id,)| id).collect()) } pub async fn get_artists_full_data(pool: &PgPool, ids: &[i64]) -> Result, sqlx::Error> { let mut result = Vec::new(); for &id in ids { let artist: Artist = sqlx::query_as("SELECT id, name FROM artists WHERE id = $1") .bind(id).fetch_one(pool).await?; let albums: Vec = sqlx::query_as("SELECT * FROM albums WHERE artist_id = $1 ORDER BY year NULLS LAST, name") .bind(id).fetch_all(pool).await?; let mut album_data = Vec::new(); for album in albums { let tracks: Vec = sqlx::query_as( "SELECT id, title, track_number, storage_path FROM tracks WHERE album_id = $1 ORDER BY track_number NULLS LAST, title" ).bind(album.id).fetch_all(pool).await?; album_data.push(AlbumFullData { id: album.id, name: album.name, year: album.year, tracks }); } result.push(ArtistFullData { id, name: artist.name, albums: album_data }); } Ok(result) } pub async fn get_tracks_with_albums_for_artist(pool: &PgPool, artist_id: i64) -> Result, sqlx::Error> { sqlx::query_as::<_, TrackWithAlbum>( r#"SELECT t.id, t.storage_path, a.name as album_name FROM tracks t LEFT JOIN albums a ON a.id = t.album_id WHERE t.id IN ( SELECT track_id FROM track_artists WHERE artist_id = $1 AND role = 'primary' )"# ).bind(artist_id).fetch_all(pool).await } pub async fn update_track_storage_path(pool: &PgPool, track_id: i64, new_path: &str) -> Result<(), sqlx::Error> { sqlx::query("UPDATE tracks SET storage_path = $2 WHERE id = $1") .bind(track_id).bind(new_path).execute(pool).await?; Ok(()) }