feat(daemon): grab pipeline for wanted movies (#77)

crates/arr-db/src/lib.rs

@@ -5,6 +5,10 @@
 
 use std::path::Path;
 
+pub mod policy;
+
+pub use policy::{MoviePolicy, PolicyColumns, PolicyError};
+
 use sqlx::sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePoolOptions, SqliteSynchronous};
 use sqlx::{migrate::MigrateError, SqlitePool};
 

crates/arr-db/src/policy.rs

@@ -0,0 +1,402 @@
+//! Reading a policy row into the pure `arr-core` types. See DESIGN.md §5.1.
+//!
+//! Policy lives in the database because the numbers get tuned by hand (§10),
+//! so every consumer — the API's manual search and the daemon's grab
+//! selection — has to turn the same JSON columns into the same
+//! [`arr_core::Policy`]. That mapping lives here once.
+
+use std::collections::BTreeMap;
+
+use arr_core::{
+    DolbyVisionProfile, HdrRules, Language, Policy, PolicyId, RequiredAudio, Resolution,
+    ScoreWeights, SizeBand, Source, TitleOverrides,
+};
+use serde::Deserialize;
+
+use crate::Db;
+
+/// A failure loading a policy row.
+#[derive(Debug, thiserror::Error)]
+pub enum PolicyError {
+    #[error("database: {0}")]
+    Database(#[from] sqlx::Error),
+    #[error("policy column {column} is not valid JSON: {source}")]
+    Json {
+        column: &'static str,
+        source: serde_json::Error,
+    },
+}
+
+/// The effective policy for one movie, plus the root it is attached to.
+///
+/// The root's `kind` and `audience` are the Transmission label and the
+/// on-disk layout (§7.1, §7.4), and they only exist together with the policy,
+/// so they are returned together.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct MoviePolicy {
+    pub policy: Policy,
+    /// Per-title relaxations and tightenings of the root policy (§5.1).
+    pub overrides: TitleOverrides,
+    pub root_id: i64,
+    /// `movie` or `tv`.
+    pub root_kind: String,
+    /// `main` or `kids`.
+    pub root_audience: String,
+    pub root_path: String,
+}
+
+/// The raw policy columns, as the `policies` table stores them (§5.5, §10).
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct PolicyColumns {
+    pub id: i64,
+    pub name: String,
+    pub required_audio: String,
+    pub dub_blacklist: String,
+    pub hdr_rules: String,
+    pub size_bands: String,
+    pub resolution_pref: String,
+    pub source_weights: String,
+    pub score_weights: String,
+}
+
+impl PolicyColumns {
+    /// Turn the stored JSON into the policy the engine evaluates against.
+    ///
+    /// A resolution or source tier the mapping does not know is dropped
+    /// rather than rejected: an unknown key is a policy with no opinion,
+    /// which is what an empty band or weight already means.
+    ///
+    /// # Errors
+    ///
+    /// If a column does not hold the JSON its migration promises.
+    pub fn to_policy(&self) -> Result<Policy, PolicyError> {
+        let required: RequiredAudioJson = json("required_audio", &self.required_audio)?;
+        let hdr: HdrRulesJson = json("hdr_rules", &self.hdr_rules)?;
+        let bands: BTreeMap<String, SizeBandJson> = json("size_bands", &self.size_bands)?;
+        let resolutions: Vec<String> = json("resolution_pref", &self.resolution_pref)?;
+        let weights: BTreeMap<String, i32> = json("source_weights", &self.source_weights)?;
+        let score_weights: ScoreWeightsJson = json("score_weights", &self.score_weights)?;
+        let blacklist: Vec<String> = json("dub_blacklist", &self.dub_blacklist)?;
+
+        Ok(Policy {
+            id: PolicyId(self.id),
+            name: self.name.clone(),
+            required_audio: if required.require == "original_language" {
+                RequiredAudio::OriginalLanguage
+            } else {
+                RequiredAudio::AnyOf(required.langs.iter().map(|lang| language(lang)).collect())
+            },
+            dub_blacklist: blacklist.iter().map(|lang| language(lang)).collect(),
+            hdr_rules: HdrRules {
+                rejected_dolby_vision_profiles: hdr
+                    .dv_profile_reject
+                    .iter()
+                    .filter_map(|value| value.parse().ok())
+                    .map(|profile| DolbyVisionProfile {
+                        profile,
+                        compatibility_id: None,
+                    })
+                    .collect(),
+            },
+            size_bands: bands
+                .into_iter()
+                .filter_map(|(resolution, band)| {
+                    resolution_value(&resolution).map(|resolution| {
+                        (
+                            resolution,
+                            SizeBand {
+                                floor_bytes: gib(band.floor_gib),
+                                target_bytes: gib(band.target_gib),
+                                penalty_points_per_gib_over: band.penalty_points_per_gib_over,
+                            },
+                        )
+                    })
+                })
+                .collect(),
+            resolution_preference: resolutions
+                .iter()
+                .filter_map(|value| resolution_value(value))
+                .collect(),
+            source_weights: weights
+                .into_iter()
+                .filter_map(|(source, weight)| source_value(&source).map(|source| (source, weight)))
+                .collect(),
+            score_weights: ScoreWeights {
+                size_at_target: score_weights.size_at_target,
+                source_tier: score_weights.source_tier,
+                seeder_doubling: score_weights.seeder_doubling,
+            },
+        })
+    }
+}
+
+impl Db {
+    /// The policy attached to a movie's root, with that movie's overrides.
+    ///
+    /// `None` when the movie does not exist.
+    ///
+    /// # Errors
+    ///
+    /// If the query fails, or a policy column does not hold the JSON its
+    /// migration promises.
+    pub async fn movie_policy(&self, movie_id: i64) -> Result<Option<MoviePolicy>, PolicyError> {
+        let row = sqlx::query!(
+            r#"
+            SELECT   m.overrides       AS "overrides!: String",
+                     r.id              AS "root_id!: i64",
+                     r.kind            AS "root_kind!: String",
+                     r.audience        AS "root_audience!: String",
+                     r.path            AS "root_path!: String",
+                     p.id              AS "policy_id!: i64",
+                     p.name            AS "policy_name!: String",
+                     p.required_audio  AS "required_audio!: String",
+                     p.dub_blacklist   AS "dub_blacklist!: String",
+                     p.hdr_rules       AS "hdr_rules!: String",
+                     p.size_bands      AS "size_bands!: String",
+                     p.resolution_pref AS "resolution_pref!: String",
+                     p.source_weights  AS "source_weights!: String",
+                     p.score_weights   AS "score_weights!: String"
+            FROM     movies m
+            JOIN     roots r    ON r.id = m.root_id
+            JOIN     policies p ON p.id = r.policy_id
+            WHERE    m.id = ?
+            "#,
+            movie_id
+        )
+        .fetch_optional(self.pool())
+        .await?;
+
+        let Some(row) = row else {
+            return Ok(None);
+        };
+
+        let overrides: OverridesJson = json("overrides", &row.overrides)?;
+        let policy = PolicyColumns {
+            id: row.policy_id,
+            name: row.policy_name,
+            required_audio: row.required_audio,
+            dub_blacklist: row.dub_blacklist,
+            hdr_rules: row.hdr_rules,
+            size_bands: row.size_bands,
+            resolution_pref: row.resolution_pref,
+            source_weights: row.source_weights,
+            score_weights: row.score_weights,
+        }
+        .to_policy()?;
+
+        Ok(Some(MoviePolicy {
+            policy,
+            overrides: TitleOverrides {
+                only_4k: overrides.only_4k,
+                allow_english_audio: overrides.allow_english_audio,
+            },
+            root_id: row.root_id,
+            root_kind: row.root_kind,
+            root_audience: row.root_audience,
+            root_path: row.root_path,
+        }))
+    }
+}
+
+/// A BCP-47 tag as the policy columns write it.
+///
+/// `por-unverified` is the tag §5.2 gives a Portuguese track no signal
+/// resolved, and it maps to the same variant as a bare `pt`.
+#[must_use]
+pub fn language(value: &str) -> Language {
+    match value {
+        "pt-PT" => Language::PortuguesePortugal,
+        "pt-BR" => Language::PortugueseBrazil,
+        "pt" | "por-unverified" => Language::PortugueseUnverified,
+        other => Language::Other(other.to_owned()),
+    }
+}
+
+/// A resolution as `resolution_pref` and `size_bands` name it. Anything else
+/// is a resolution this policy has no opinion about.
+#[must_use]
+pub fn resolution_value(value: &str) -> Option<Resolution> {
+    match value {
+        "2160p" => Some(Resolution::R2160p),
+        "1080p" => Some(Resolution::R1080p),
+        "720p" => Some(Resolution::R720p),
+        _ => None,
+    }
+}
+
+/// A source tier as `source_weights` names it.
+#[must_use]
+pub fn source_value(value: &str) -> Option<Source> {
+    match value {
+        "Remux" => Some(Source::Remux),
+        "BluRay" => Some(Source::BluRay),
+        "WEB-DL" => Some(Source::WebDl),
+        "WEBRip" => Some(Source::WebRip),
+        "HDTV" => Some(Source::Hdtv),
+        _ => None,
+    }
+}
+
+#[derive(Debug, Deserialize)]
+struct RequiredAudioJson {
+    require: String,
+    #[serde(default)]
+    langs: Vec<String>,
+}
+
+#[derive(Debug, Deserialize)]
+struct HdrRulesJson {
+    #[serde(default)]
+    dv_profile_reject: Vec<String>,
+}
+
+#[derive(Debug, Deserialize)]
+struct SizeBandJson {
+    floor_gib: u64,
+    target_gib: u64,
+    penalty_points_per_gib_over: i32,
+}
+
+#[derive(Debug, Deserialize)]
+struct ScoreWeightsJson {
+    size_at_target: i32,
+    source_tier: i32,
+    seeder_doubling: i32,
+}
+
+#[derive(Debug, Default, Deserialize)]
+struct OverridesJson {
+    #[serde(default)]
+    only_4k: bool,
+    #[serde(default)]
+    allow_english_audio: bool,
+}
+
+fn gib(value: u64) -> u64 {
+    value.saturating_mul(1 << 30)
+}
+
+fn json<T: serde::de::DeserializeOwned>(
+    column: &'static str,
+    value: &str,
+) -> Result<T, PolicyError> {
+    serde_json::from_str(value).map_err(|source| PolicyError::Json { column, source })
+}
+
+#[cfg(test)]
+#[allow(clippy::unwrap_used)]
+mod tests {
+    use super::*;
+
+    async fn seeded() -> (tempfile::TempDir, Db) {
+        let dir = tempfile::tempdir().unwrap();
+        let db = Db::connect(dir.path().join("arr.db")).await.unwrap();
+        db.migrate().await.unwrap();
+        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
+        .unwrap();
+        (dir, db)
+    }
+
+    #[tokio::test]
+    async fn loads_the_seeded_main_policy() {
+        let (_dir, db) = seeded().await;
+
+        let loaded = db.movie_policy(1).await.unwrap().expect("movie 1");
+
+        assert_eq!(loaded.root_kind, "movie");
+        assert_eq!(loaded.root_audience, "main");
+        assert_eq!(
+            loaded.policy.required_audio,
+            RequiredAudio::OriginalLanguage
+        );
+        assert_eq!(
+            loaded.policy.dub_blacklist,
+            vec![Language::PortugueseBrazil]
+        );
+        assert_eq!(
+            loaded.policy.resolution_preference,
+            vec![Resolution::R2160p, Resolution::R1080p]
+        );
+        assert_eq!(loaded.overrides, TitleOverrides::default());
+    }
+
+    /// §5.5: the scoring numbers are policy rows, so the loader reads them
+    /// rather than falling back to the compiled defaults.
+    #[tokio::test]
+    async fn size_bands_and_score_weights_come_from_the_row() {
+        let (_dir, db) = seeded().await;
+        sqlx::query(
+            r#"UPDATE policies SET score_weights = '{"size_at_target":500,"source_tier":10,"seeder_doubling":1}'"#,
+        )
+        .execute(db.pool())
+        .await
+        .unwrap();
+
+        let loaded = db.movie_policy(1).await.unwrap().expect("movie 1");
+
+        let band = loaded.policy.size_bands[&Resolution::R2160p];
+        assert_eq!(band.floor_bytes, 8 << 30);
+        assert_eq!(band.target_bytes, 22 << 30);
+        assert_eq!(band.penalty_points_per_gib_over, 60);
+        assert_eq!(loaded.policy.score_weights.size_at_target, 500);
+        assert_eq!(loaded.policy.score_weights.source_tier, 10);
+        assert_eq!(loaded.policy.score_weights.seeder_doubling, 1);
+    }
+
+    #[tokio::test]
+    async fn per_title_overrides_ride_along() {
+        let (_dir, db) = seeded().await;
+        sqlx::query(r#"UPDATE movies SET overrides = '{"only_4k":true}'"#)
+            .execute(db.pool())
+            .await
+            .unwrap();
+
+        let loaded = db.movie_policy(1).await.unwrap().expect("movie 1");
+
+        assert!(loaded.overrides.only_4k);
+        assert!(!loaded.overrides.allow_english_audio);
+    }
+
+    /// The mapping is usable without a database, which is how the API's
+    /// classifier and the daemon's selection stay on the same numbers.
+    #[test]
+    fn columns_map_without_a_row() {
+        let policy = PolicyColumns {
+            id: 1,
+            name: "test".into(),
+            required_audio: r#"{"require":"original_language"}"#.into(),
+            dub_blacklist: "[]".into(),
+            hdr_rules: "{}".into(),
+            size_bands:
+                r#"{"2160p":{"floor_gib":8,"target_gib":22,"penalty_points_per_gib_over":60}}"#
+                    .into(),
+            resolution_pref: r#"["2160p"]"#.into(),
+            source_weights: r#"{"WEB-DL":2}"#.into(),
+            score_weights: r#"{"size_at_target":2000,"source_tier":7,"seeder_doubling":11}"#.into(),
+        }
+        .to_policy()
+        .unwrap();
+
+        assert_eq!(
+            policy.score_weights,
+            ScoreWeights {
+                size_at_target: 2000,
+                source_tier: 7,
+                seeder_doubling: 11,
+            }
+        );
+        assert_eq!(policy.source_weights[&Source::WebDl], 2);
+    }
+
+    #[tokio::test]
+    async fn an_unknown_movie_is_not_an_error() {
+        let (_dir, db) = seeded().await;
+        assert!(db.movie_policy(404).await.unwrap().is_none());
+    }
+}