Files
furumi-ng/furumi-agent/src/db.rs

555 lines
16 KiB
Rust

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<PgPool, sqlx::Error> {
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<i32>,
}
#[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<String>,
pub raw_artist: Option<String>,
pub raw_album: Option<String>,
pub raw_year: Option<i32>,
pub raw_track_number: Option<i32>,
pub raw_genre: Option<String>,
pub duration_secs: Option<f64>,
// Path-derived hints
pub path_artist: Option<String>,
pub path_album: Option<String>,
pub path_year: Option<i32>,
pub path_track_number: Option<i32>,
pub path_title: Option<String>,
// Normalized (LLM output)
pub norm_title: Option<String>,
pub norm_artist: Option<String>,
pub norm_album: Option<String>,
pub norm_year: Option<i32>,
pub norm_track_number: Option<i32>,
pub norm_genre: Option<String>,
pub norm_featured_artists: Option<String>, // JSON array
pub confidence: Option<f64>,
pub llm_notes: Option<String>,
pub error_message: Option<String>,
pub created_at: chrono::DateTime<chrono::Utc>,
pub updated_at: chrono::DateTime<chrono::Utc>,
}
#[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<i32>,
pub similarity: f32,
}
#[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<i32>,
pub height: Option<i32>,
pub file_size: i64,
}
// --- Queries ---
pub async fn file_hash_exists(pool: &PgPool, hash: &str) -> Result<bool, sqlx::Error> {
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<f64>,
) -> Result<Uuid, sqlx::Error> {
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<Vec<SimilarArtist>, 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<Vec<SimilarAlbum>, 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<i64, sqlx::Error> {
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<i32>) -> Result<i64, sqlx::Error> {
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<i64>,
title: &str,
track_number: Option<i32>,
genre: Option<&str>,
duration_secs: Option<f64>,
file_hash: &str,
file_size: i64,
storage_path: &str,
) -> Result<i64, sqlx::Error> {
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<i64, sqlx::Error> {
let pt: PendingTrack = sqlx::query_as("SELECT * FROM pending_tracks WHERE id = $1")
.bind(pending_id)
.fetch_one(pool)
.await?;
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::<Vec<String>>(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<bool, sqlx::Error> {
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<i64, sqlx::Error> {
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)
}
pub async fn get_album_images(pool: &PgPool, album_id: i64) -> Result<Vec<AlbumImage>, 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<Option<i64>, 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))
}
// --- DTOs for insert helpers ---
#[derive(Debug, Default)]
pub struct RawFields {
pub title: Option<String>,
pub artist: Option<String>,
pub album: Option<String>,
pub year: Option<i32>,
pub track_number: Option<i32>,
pub genre: Option<String>,
}
#[derive(Debug, Default)]
pub struct PathHints {
pub title: Option<String>,
pub artist: Option<String>,
pub album: Option<String>,
pub year: Option<i32>,
pub track_number: Option<i32>,
}
#[derive(Debug, Default, Serialize, Deserialize)]
pub struct NormalizedFields {
pub title: Option<String>,
pub artist: Option<String>,
pub album: Option<String>,
pub year: Option<i32>,
pub track_number: Option<i32>,
pub genre: Option<String>,
#[serde(default)]
pub featured_artists: Vec<String>,
pub confidence: Option<f64>,
pub notes: Option<String>,
}
// --- Admin queries ---
pub async fn list_pending(pool: &PgPool, status_filter: Option<&str>, limit: i64, offset: i64) -> Result<Vec<PendingTrack>, 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<Option<PendingTrack>, 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<bool, sqlx::Error> {
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<Vec<Artist>, 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<Vec<Album>, 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<bool, sqlx::Error> {
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<i32>) -> Result<bool, sqlx::Error> {
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 async fn get_stats(pool: &PgPool) -> Result<Stats, sqlx::Error> {
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?;
Ok(Stats { total_tracks, total_artists, total_albums, pending_count, review_count, error_count })
}