music-agregator/src/db.rs

use anyhow::Result;
use sqlx::Row;
use sqlx::postgres::{PgPool, PgPoolOptions};
use uuid::Uuid;

use crate::torrent_parser::{AudioFormat, ParsedTorrent, ReleaseSource, ReleaseType};

const SCHEMA_SQL: &str = include_str!("../db/schema.sql");

pub async fn connect(database_url: &str) -> Result<PgPool> {
    let pool = PgPoolOptions::new()
        .max_connections(5)
        .connect(database_url)
        .await?;
    Ok(pool)
}

pub async fn apply_schema(pool: &PgPool) -> Result<()> {
    sqlx::raw_sql(SCHEMA_SQL).execute(pool).await?;
    Ok(())
}

pub async fn insert_partial(pool: &PgPool, pt: &ParsedTorrent, source_url: &str) -> Result<Uuid> {
    let row = sqlx::query(
        "INSERT INTO parsed_torrents (
            info_hash, raw_title, artist, album, year,
            release_type, genres, label, source, rip_type,
            format, bitrate, bit_depth, sample_rate,
            track_names, track_count, release_count, audio_file_count, total_audio_size,
            has_cover_art, has_cue_sheet, has_rip_log,
            parsed_successfully, parse_errors,
            state, magnet
        ) VALUES (
            $1, $2, $3, $4, $5,
            $6, $7, $8, $9, $10,
            $11, $12, $13, $14,
            $15, $16, $17, $18, $19,
            $20, $21, $22,
            $23, $24,
            'partial', $25
        )
        RETURNING id",
    )
    .bind(&pt.info_hash)
    .bind(&pt.raw_title)
    .bind(&pt.artist)
    .bind(&pt.album)
    .bind(pt.year as i32)
    .bind(pt.release_type.as_str())
    .bind(sqlx::types::Json(&pt.genres))
    .bind(&pt.label)
    .bind(pt.source.as_str())
    .bind(&pt.rip_type)
    .bind(pt.format.as_str())
    .bind(&pt.bitrate)
    .bind(pt.bit_depth as i32)
    .bind(pt.sample_rate as i32)
    .bind(sqlx::types::Json(&pt.track_names))
    .bind(pt.track_count as i32)
    .bind(pt.release_count as i32)
    .bind(pt.audio_file_count as i32)
    .bind(pt.total_audio_size as i64)
    .bind(pt.has_cover_art)
    .bind(pt.has_cue_sheet)
    .bind(pt.has_rip_log)
    .bind(pt.parsed_successfully)
    .bind(sqlx::types::Json(&pt.parse_errors))
    .bind(source_url)
    .fetch_one(pool)
    .await?;

    Ok(row.get("id"))
}

pub async fn update_resolved(pool: &PgPool, id: Uuid, pt: &ParsedTorrent) -> Result<()> {
    sqlx::query(
        "UPDATE parsed_torrents SET
            info_hash = $2,
            raw_title = $3,
            artist = $4,
            album = $5,
            year = $6,
            release_type = $7,
            genres = $8,
            label = $9,
            source = $10,
            rip_type = $11,
            format = $12,
            bitrate = $13,
            bit_depth = $14,
            sample_rate = $15,
            track_names = $16,
            track_count = $17,
            release_count = $18,
            audio_file_count = $19,
            total_audio_size = $20,
            has_cover_art = $21,
            has_cue_sheet = $22,
            has_rip_log = $23,
            parsed_successfully = $24,
            parse_errors = $25,
            state = 'resolved',
            updated_at = now()
        WHERE id = $1",
    )
    .bind(id)
    .bind(&pt.info_hash)
    .bind(&pt.raw_title)
    .bind(&pt.artist)
    .bind(&pt.album)
    .bind(pt.year as i32)
    .bind(pt.release_type.as_str())
    .bind(sqlx::types::Json(&pt.genres))
    .bind(&pt.label)
    .bind(pt.source.as_str())
    .bind(&pt.rip_type)
    .bind(pt.format.as_str())
    .bind(&pt.bitrate)
    .bind(pt.bit_depth as i32)
    .bind(pt.sample_rate as i32)
    .bind(sqlx::types::Json(&pt.track_names))
    .bind(pt.track_count as i32)
    .bind(pt.release_count as i32)
    .bind(pt.audio_file_count as i32)
    .bind(pt.total_audio_size as i64)
    .bind(pt.has_cover_art)
    .bind(pt.has_cue_sheet)
    .bind(pt.has_rip_log)
    .bind(pt.parsed_successfully)
    .bind(sqlx::types::Json(&pt.parse_errors))
    .execute(pool)
    .await?;
    Ok(())
}

pub async fn get(pool: &PgPool, info_hash: &str) -> Result<Option<ParsedTorrent>> {
    let row = sqlx::query(
        "SELECT info_hash, raw_title, artist, album, year,
                release_type, genres, label, source, rip_type,
                format, bitrate, bit_depth, sample_rate,
                track_names, track_count, release_count, audio_file_count, total_audio_size,
                has_cover_art, has_cue_sheet, has_rip_log,
                parsed_successfully, parse_errors
         FROM parsed_torrents
         WHERE info_hash = $1 AND state = 'resolved'",
    )
    .bind(info_hash)
    .fetch_optional(pool)
    .await?;

    let row = match row {
        Some(r) => r,
        None => return Ok(None),
    };

    let genres: sqlx::types::Json<Vec<String>> = row.try_get("genres")?;
    let track_names: sqlx::types::Json<Vec<String>> = row.try_get("track_names")?;
    let parse_errors: sqlx::types::Json<Vec<String>> = row.try_get("parse_errors")?;

    Ok(Some(ParsedTorrent {
        info_hash: row.try_get("info_hash")?,
        raw_title: row.try_get("raw_title")?,
        artist: row.try_get("artist")?,
        album: row.try_get("album")?,
        year: row.try_get::<i32, _>("year")? as u32,
        release_type: ReleaseType::from_str(row.try_get("release_type")?),
        genres: genres.0,
        label: row.try_get("label")?,
        source: ReleaseSource::from_str(row.try_get("source")?),
        rip_type: row.try_get("rip_type")?,
        format: AudioFormat::from_str(row.try_get("format")?),
        bitrate: row.try_get("bitrate")?,
        bit_depth: row.try_get::<i32, _>("bit_depth")? as u32,
        sample_rate: row.try_get::<i32, _>("sample_rate")? as u32,
        track_names: track_names.0,
        track_count: row.try_get::<i32, _>("track_count")? as u32,
        release_count: row.try_get::<i32, _>("release_count")? as u32,
        audio_file_count: row.try_get::<i32, _>("audio_file_count")? as u32,
        total_audio_size: row.try_get::<i64, _>("total_audio_size")? as u64,
        has_cover_art: row.try_get("has_cover_art")?,
        has_cue_sheet: row.try_get("has_cue_sheet")?,
        has_rip_log: row.try_get("has_rip_log")?,
        parsed_successfully: row.try_get("parsed_successfully")?,
        parse_errors: parse_errors.0,
    }))
}