arr/crates/arr-db/src/lib.rs

//! arr-db — SQLite persistence and migrations. See DESIGN.md §10.
//!
//! A few thousand rows and a single writer, so the pool exists for readers
//! and the schema does the enforcing: foreign keys on, WAL, `STRICT` tables.

use std::path::Path;

pub mod blacklist;
pub mod policy;

pub use blacklist::Blacklist;
pub use policy::{MoviePolicy, PolicyColumns, PolicyError, TitlePolicy};

use sqlx::sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePoolOptions, SqliteSynchronous};
use sqlx::{migrate::MigrateError, SqlitePool};

/// The migrations embedded in the binary, so a deploy is one file.
pub static MIGRATOR: sqlx::migrate::Migrator = sqlx::migrate!("./migrations");

/// How long a writer waits for the write lock before giving up.
const BUSY_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5);

/// A connection pool with the pragmas this schema assumes already applied.
#[derive(Debug, Clone)]
pub struct Db {
    pool: SqlitePool,
}

impl Db {
    /// Open (creating if absent) the database at `path`.
    ///
    /// Every connection gets `foreign_keys` on — SQLite defaults it *off* per
    /// connection, so setting it anywhere but here is a way to forget it.
    ///
    /// # Errors
    ///
    /// If the file cannot be created or opened.
    pub async fn connect(path: impl AsRef<Path>) -> Result<Self, sqlx::Error> {
        let options = SqliteConnectOptions::new()
            .filename(path)
            .create_if_missing(true)
            .foreign_keys(true)
            .journal_mode(SqliteJournalMode::Wal)
            // WAL makes NORMAL durable against process crashes, which is the
            // failure this actually has. FULL costs an fsync per commit.
            .synchronous(SqliteSynchronous::Normal)
            .busy_timeout(BUSY_TIMEOUT);

        let pool = SqlitePoolOptions::new().connect_with(options).await?;

        Ok(Self { pool })
    }

    /// Apply every migration that has not run yet.
    ///
    /// # Errors
    ///
    /// If a migration fails, or one already applied no longer matches its
    /// recorded checksum.
    pub async fn migrate(&self) -> Result<(), MigrateError> {
        MIGRATOR.run(&self.pool).await
    }

    /// The underlying pool, for crates that own their own queries.
    #[must_use]
    pub fn pool(&self) -> &SqlitePool {
        &self.pool
    }

    /// Every configured root with the policy attached to it (§5.1).
    ///
    /// # Errors
    ///
    /// If the query fails.
    pub async fn list_roots(&self) -> Result<Vec<Root>, sqlx::Error> {
        sqlx::query_as!(
            Root,
            r#"
            SELECT   r.id       AS "id!: i64",
                     r.kind     AS "kind!: String",
                     r.audience AS "audience!: String",
                     r.path     AS "path!: String",
                     r.policy_id,
                     p.name     AS "policy_name!: String"
            FROM     roots r
            JOIN     policies p ON p.id = r.policy_id
            ORDER BY r.kind, r.audience
            "#
        )
        .fetch_all(&self.pool)
        .await
    }
}

/// A library root and the policy it carries (§4, §5.1).
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Root {
    pub id: i64,
    /// `movie` or `tv`.
    pub kind: String,
    /// `main` or `kids`.
    pub audience: String,
    pub path: String,
    pub policy_id: i64,
    pub policy_name: String,
}

#[cfg(test)]
mod tests {
    use super::{Db, MIGRATOR};
    use sqlx::Row;

    /// A migrated database in a directory that lives as long as the guard.
    async fn fresh() -> (tempfile::TempDir, Db) {
        let dir = tempfile::tempdir().expect("tempdir");
        let db = Db::connect(dir.path().join("arr.db"))
            .await
            .expect("connect");
        db.migrate().await.expect("migrate");
        (dir, db)
    }

    #[tokio::test]
    async fn migrations_apply_from_empty() {
        let (_dir, db) = fresh().await;

        let applied = MIGRATOR.iter().count();
        let recorded: i64 = sqlx::query_scalar("SELECT count(*) FROM _sqlx_migrations")
            .fetch_one(db.pool())
            .await
            .expect("count migrations");

        assert_eq!(recorded, i64::try_from(applied).expect("fits"));
    }

    #[tokio::test]
    async fn migrations_are_idempotent() {
        let (_dir, db) = fresh().await;
        db.migrate().await.expect("second run is a no-op");
    }

    #[tokio::test]
    async fn wal_and_foreign_keys_are_on() {
        let (_dir, db) = fresh().await;

        let journal: String = sqlx::query_scalar("PRAGMA journal_mode")
            .fetch_one(db.pool())
            .await
            .expect("journal_mode");
        assert_eq!(journal.to_lowercase(), "wal");

        let foreign_keys: i64 = sqlx::query_scalar("PRAGMA foreign_keys")
            .fetch_one(db.pool())
            .await
            .expect("foreign_keys");
        assert_eq!(foreign_keys, 1);
    }

    #[tokio::test]
    async fn seeds_two_movie_roots_with_distinct_policies() {
        let (_dir, db) = fresh().await;

        let roots = db.list_roots().await.expect("list roots");

        let movie_roots: Vec<_> = roots.iter().filter(|root| root.kind == "movie").collect();
        assert_eq!(movie_roots.len(), 2);
        assert_eq!(movie_roots[0].audience, "kids");
        assert_eq!(movie_roots[0].path, "/mnt/media/movies/kids");
        assert_eq!(movie_roots[1].audience, "main");
        assert_eq!(movie_roots[1].path, "/mnt/media/movies/main");
        assert_ne!(movie_roots[0].policy_id, movie_roots[1].policy_id);
    }

    #[tokio::test]
    async fn kids_policy_blacklists_brazilian_dubs_and_prefers_web_dl() {
        let (_dir, db) = fresh().await;

        let row = sqlx::query(
            "SELECT p.dub_blacklist, p.source_weights
             FROM policies p JOIN roots r ON r.policy_id = p.id
             WHERE r.audience = 'kids'",
        )
        .fetch_one(db.pool())
        .await
        .expect("kids policy");

        let dub_blacklist: String = row.get(0);
        let source_weights: String = row.get(1);

        assert!(dub_blacklist.contains("pt-BR"), "§5.2");
        // §5.2: European Portuguese dubs are a streaming artefact.
        let webdl = source_weights.find("WEB-DL").expect("WEB-DL weight");
        let remux = source_weights.find("Remux").expect("Remux weight");
        assert!(
            webdl < remux,
            "WEB-DL is listed first for kids: {source_weights}"
        );
    }

    #[tokio::test]
    async fn every_policy_carries_its_scoring_numbers() {
        let (_dir, db) = fresh().await;

        let rows = sqlx::query("SELECT size_bands, score_weights FROM policies")
            .fetch_all(db.pool())
            .await
            .expect("policies");

        assert_eq!(rows.len(), 4);
        for row in rows {
            let size_bands: String = row.get(0);
            let score_weights: String = row.get(1);

            // §5.5: floor, target and the penalty above target, per resolution.
            assert!(size_bands.contains("floor_gib"), "{size_bands}");
            assert!(size_bands.contains("target_gib"), "{size_bands}");
            assert!(
                size_bands.contains("penalty_points_per_gib_over"),
                "{size_bands}"
            );
            // §5.5: the weights that keep source tier and seeders small
            // against the size term are tuned by hand too.
            assert!(score_weights.contains("size_at_target"), "{score_weights}");
            assert!(score_weights.contains("source_tier"), "{score_weights}");
            assert!(score_weights.contains("seeder_doubling"), "{score_weights}");
            // §5.5: and the step that ranks resolutions against each other.
            assert!(score_weights.contains("resolution_step"), "{score_weights}");
        }
    }

    #[tokio::test]
    async fn foreign_keys_are_enforced() {
        let (_dir, db) = fresh().await;

        let err = sqlx::query("INSERT INTO movies (tmdb_id, title, root_id) VALUES (1, 'x', 999)")
            .execute(db.pool())
            .await
            .expect_err("root 999 does not exist");

        assert!(
            err.to_string().to_lowercase().contains("foreign key"),
            "{err}"
        );
    }

    #[tokio::test]
    async fn movie_state_is_constrained() {
        let (_dir, db) = fresh().await;

        sqlx::query(
            "INSERT INTO movies (tmdb_id, title, root_id, state)
             SELECT 693134, 'Dune Part Two', id, 'nonsense' FROM roots WHERE audience = 'main'",
        )
        .execute(db.pool())
        .await
        .expect_err("state is a closed set");
    }

    /// #73: migration 0007 renames the pre-canonical `grabbed`/`imported`
    /// values in place rather than dropping the rows that held them.
    #[tokio::test]
    async fn movie_state_migration_renames_existing_rows_in_place() {
        let dir = tempfile::tempdir().expect("tempdir");
        let db = Db::connect(dir.path().join("arr.db"))
            .await
            .expect("connect");

        let pre_0007 = sqlx::migrate::Migrator {
            migrations: std::borrow::Cow::Owned(
                MIGRATOR.iter().filter(|m| m.version < 7).cloned().collect(),
            ),
            ignore_missing: MIGRATOR.ignore_missing,
            locking: MIGRATOR.locking,
            no_tx: MIGRATOR.no_tx,
        };
        pre_0007
            .run(db.pool())
            .await
            .expect("migrations before #73");

        sqlx::query(
            "INSERT INTO movies (tmdb_id, title, root_id, state)
             SELECT 1, 'Old Vocabulary', id, 'grabbed' FROM roots WHERE kind = 'movie' LIMIT 1",
        )
        .execute(db.pool())
        .await
        .expect("legacy 'grabbed' row, valid under the pre-#73 constraint");
        sqlx::query(
            "INSERT INTO movies (tmdb_id, title, root_id, state)
             SELECT 2, 'Also Old', id, 'imported' FROM roots WHERE kind = 'movie' LIMIT 1",
        )
        .execute(db.pool())
        .await
        .expect("legacy 'imported' row, valid under the pre-#73 constraint");

        db.migrate()
            .await
            .expect("remaining migrations, including #73");

        let renamed_from_grabbed: String =
            sqlx::query_scalar("SELECT state FROM movies WHERE tmdb_id = 1")
                .fetch_one(db.pool())
                .await
                .expect("row survives the migration");
        let renamed_from_imported: String =
            sqlx::query_scalar("SELECT state FROM movies WHERE tmdb_id = 2")
                .fetch_one(db.pool())
                .await
                .expect("row survives the migration");

        assert_eq!(renamed_from_grabbed, "downloading");
        assert_eq!(renamed_from_imported, "available");
    }

    #[tokio::test]
    async fn seeds_two_tv_roots_with_distinct_policies() {
        let (_dir, db) = fresh().await;
        let roots = db.list_roots().await.expect("list roots");
        let tv_roots: Vec<_> = roots.iter().filter(|root| root.kind == "tv").collect();

        assert_eq!(tv_roots.len(), 2);
        assert_eq!(tv_roots[0].audience, "kids");
        assert_eq!(tv_roots[0].path, "/mnt/media/tv/kids");
        assert_eq!(tv_roots[1].audience, "main");
        assert_eq!(tv_roots[1].path, "/mnt/media/tv/main");
        assert_ne!(tv_roots[0].policy_id, tv_roots[1].policy_id);
    }

    #[tokio::test]
    async fn episode_wanted_is_independent_of_tracking_rules() {
        let (_dir, db) = fresh().await;
        let root_id: i64 =
            sqlx::query_scalar("SELECT id FROM roots WHERE kind = 'tv' AND audience = 'main'")
                .fetch_one(db.pool())
                .await
                .expect("TV root");
        let series_id = sqlx::query("INSERT INTO series (tmdb_id, title, root_id, auto_track) VALUES (82728, 'Bluey', ?, 1)")
            .bind(root_id)
            .execute(db.pool())
            .await
            .expect("series")
            .last_insert_rowid();
        let season_id =
            sqlx::query("INSERT INTO seasons (series_id, number, tracked) VALUES (?, 1, 1)")
                .bind(series_id)
                .execute(db.pool())
                .await
                .expect("season")
                .last_insert_rowid();
        sqlx::query("INSERT INTO episodes (season_id, number, title) VALUES (?, 2, 'Hospital')")
            .bind(season_id)
            .execute(db.pool())
            .await
            .expect("episode");

        let wanted: bool = sqlx::query_scalar("SELECT wanted FROM episodes")
            .fetch_one(db.pool())
            .await
            .expect("wanted");
        assert!(
            !wanted,
            "auto_track and tracked are rules, not episode intent"
        );

        sqlx::query("INSERT INTO movies (tmdb_id, title, root_id) SELECT 693134, 'Dune Part Two', id FROM roots WHERE kind = 'movie' AND audience = 'main'")
            .execute(db.pool())
            .await
            .expect("existing movie schema remains writable");
    }

    #[tokio::test]
    async fn aggregates_require_roots_of_their_own_kind() {
        let (_dir, db) = fresh().await;
        let movie_root: i64 =
            sqlx::query_scalar("SELECT id FROM roots WHERE kind = 'movie' LIMIT 1")
                .fetch_one(db.pool())
                .await
                .expect("movie root");
        let tv_root: i64 = sqlx::query_scalar("SELECT id FROM roots WHERE kind = 'tv' LIMIT 1")
            .fetch_one(db.pool())
            .await
            .expect("TV root");

        sqlx::query("INSERT INTO movies (tmdb_id, title, root_id) VALUES (1, 'x', ?)")
            .bind(tv_root)
            .execute(db.pool())
            .await
            .expect_err("movie cannot use TV root");
        sqlx::query("INSERT INTO series (tmdb_id, title, root_id) VALUES (2, 'x', ?)")
            .bind(movie_root)
            .execute(db.pool())
            .await
            .expect_err("series cannot use movie root");
    }

    #[tokio::test]
    async fn episode_state_and_upstream_completion_are_constrained() {
        let (_dir, db) = fresh().await;
        let root_id: i64 =
            sqlx::query_scalar("SELECT id FROM roots WHERE kind = 'tv' AND audience = 'main'")
                .fetch_one(db.pool())
                .await
                .expect("TV root");
        let series_id = sqlx::query(
            "INSERT INTO series (tmdb_id, title, root_id, upstream_ended)
             VALUES (82728, 'Bluey', ?, 1)",
        )
        .bind(root_id)
        .execute(db.pool())
        .await
        .expect("ended series")
        .last_insert_rowid();
        let season_id = sqlx::query("INSERT INTO seasons (series_id, number) VALUES (?, 1)")
            .bind(series_id)
            .execute(db.pool())
            .await
            .expect("season")
            .last_insert_rowid();

        sqlx::query(
            "INSERT INTO episodes (season_id, number, title, state)
             VALUES (?, 1, 'Magic Xylophone', 'grabbed')",
        )
        .bind(season_id)
        .execute(db.pool())
        .await
        .expect_err("database mirrors MediaState");
        sqlx::query("UPDATE series SET upstream_ended = 2 WHERE id = ?")
            .bind(series_id)
            .execute(db.pool())
            .await
            .expect_err("upstream completion is boolean");
    }
}