diff --git a/.sqlx/query-0744ebfb2d85bbcdba28a21f88e6ccfcb06d257fd4435952c93b6287ebdd78a5.json b/.sqlx/query-0744ebfb2d85bbcdba28a21f88e6ccfcb06d257fd4435952c93b6287ebdd78a5.json new file mode 100644 index 0000000..7c58ae5 --- /dev/null +++ b/.sqlx/query-0744ebfb2d85bbcdba28a21f88e6ccfcb06d257fd4435952c93b6287ebdd78a5.json @@ -0,0 +1,12 @@ +{ + "db_name": "SQLite", + "query": "UPDATE movies SET state = 'grabbed',\n updated_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now')\n WHERE id = ?", + "describe": { + "columns": [], + "parameters": { + "Right": 1 + }, + "nullable": [] + }, + "hash": "0744ebfb2d85bbcdba28a21f88e6ccfcb06d257fd4435952c93b6287ebdd78a5" +} diff --git a/.sqlx/query-08313d5c002d1355542a40eb94b905f0091229e0433b038fdb55a49e557efa06.json b/.sqlx/query-08313d5c002d1355542a40eb94b905f0091229e0433b038fdb55a49e557efa06.json new file mode 100644 index 0000000..4910fcb --- /dev/null +++ b/.sqlx/query-08313d5c002d1355542a40eb94b905f0091229e0433b038fdb55a49e557efa06.json @@ -0,0 +1,32 @@ +{ + "db_name": "SQLite", + "query": "SELECT id AS \"id!: i64\", infohash AS \"infohash!: String\", target_id AS \"target_id!: i64\"\n FROM grabs WHERE state = 'sent'", + "describe": { + "columns": [ + { + "name": "id!: i64", + "ordinal": 0, + "type_info": "Integer" + }, + { + "name": "infohash!: String", + "ordinal": 1, + "type_info": "Text" + }, + { + "name": "target_id!: i64", + "ordinal": 2, + "type_info": "Integer" + } + ], + "parameters": { + "Right": 0 + }, + "nullable": [ + true, + false, + false + ] + }, + "hash": "08313d5c002d1355542a40eb94b905f0091229e0433b038fdb55a49e557efa06" +} diff --git a/.sqlx/query-0a6841e9164a94287732afe2375eda68b4d031ea128cae2125f7e526660d8b45.json b/.sqlx/query-0a6841e9164a94287732afe2375eda68b4d031ea128cae2125f7e526660d8b45.json new file mode 100644 index 0000000..513e301 --- /dev/null +++ b/.sqlx/query-0a6841e9164a94287732afe2375eda68b4d031ea128cae2125f7e526660d8b45.json @@ -0,0 +1,12 @@ +{ + "db_name": "SQLite", + "query": "UPDATE grabs SET state = 'downloaded' WHERE id = ?", + "describe": { + "columns": [], + "parameters": { + "Right": 1 + }, + "nullable": [] + }, + "hash": "0a6841e9164a94287732afe2375eda68b4d031ea128cae2125f7e526660d8b45" +} diff --git a/.sqlx/query-0e3b37aa1da32a0a5266c3841012fae52ce6394fcdc6c8e42eaeec0d7d6fb663.json b/.sqlx/query-0e3b37aa1da32a0a5266c3841012fae52ce6394fcdc6c8e42eaeec0d7d6fb663.json new file mode 100644 index 0000000..450ad73 --- /dev/null +++ b/.sqlx/query-0e3b37aa1da32a0a5266c3841012fae52ce6394fcdc6c8e42eaeec0d7d6fb663.json @@ -0,0 +1,20 @@ +{ + "db_name": "SQLite", + "query": "INSERT INTO grabs (release_id, target_kind, target_id, infohash, state)\n VALUES (?, 'movie', ?, ?, 'sent')\n ON CONFLICT (infohash) DO NOTHING\n RETURNING id AS \"id!: i64\"", + "describe": { + "columns": [ + { + "name": "id!: i64", + "ordinal": 0, + "type_info": "Integer" + } + ], + "parameters": { + "Right": 3 + }, + "nullable": [ + true + ] + }, + "hash": "0e3b37aa1da32a0a5266c3841012fae52ce6394fcdc6c8e42eaeec0d7d6fb663" +} diff --git a/.sqlx/query-12386f2efee70bd2a8e59ccccb8f0d7f4d2742acad26efbc7644193966c39325.json b/.sqlx/query-12386f2efee70bd2a8e59ccccb8f0d7f4d2742acad26efbc7644193966c39325.json new file mode 100644 index 0000000..e918e44 --- /dev/null +++ b/.sqlx/query-12386f2efee70bd2a8e59ccccb8f0d7f4d2742acad26efbc7644193966c39325.json @@ -0,0 +1,20 @@ +{ + "db_name": "SQLite", + "query": "\n INSERT INTO releases\n (indexer_id, guid, name, size, seeders, publish_date, download_url,\n parsed, score, verdict, rejected_rule)\n VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)\n ON CONFLICT (indexer_id, guid) DO UPDATE SET\n name = excluded.name,\n size = excluded.size,\n seeders = excluded.seeders,\n publish_date = excluded.publish_date,\n download_url = excluded.download_url,\n parsed = excluded.parsed,\n score = excluded.score,\n verdict = excluded.verdict,\n rejected_rule = excluded.rejected_rule\n RETURNING id AS \"id!: i64\"\n ", + "describe": { + "columns": [ + { + "name": "id!: i64", + "ordinal": 0, + "type_info": "Integer" + } + ], + "parameters": { + "Right": 11 + }, + "nullable": [ + false + ] + }, + "hash": "12386f2efee70bd2a8e59ccccb8f0d7f4d2742acad26efbc7644193966c39325" +} diff --git a/.sqlx/query-2c329cab6e494ef1a765a7de05e1351358d4ed583a36fe5d0edc38cc12603e44.json b/.sqlx/query-2c329cab6e494ef1a765a7de05e1351358d4ed583a36fe5d0edc38cc12603e44.json new file mode 100644 index 0000000..9e61721 --- /dev/null +++ b/.sqlx/query-2c329cab6e494ef1a765a7de05e1351358d4ed583a36fe5d0edc38cc12603e44.json @@ -0,0 +1,38 @@ +{ + "db_name": "SQLite", + "query": "\n SELECT m.id AS \"id!: i64\",\n m.title AS \"title!: String\",\n m.year,\n m.original_language\n FROM movies m\n WHERE m.wanted = 1\n AND m.blocked = 0\n AND NOT EXISTS (\n SELECT 1 FROM media_files f\n WHERE f.owner_kind = 'movie' AND f.owner_id = m.id\n )\n AND NOT EXISTS (\n SELECT 1 FROM grabs g\n WHERE g.target_kind = 'movie' AND g.target_id = m.id\n AND g.state IN ('sent', 'downloaded', 'imported')\n )\n ORDER BY m.last_searched_at IS NOT NULL, m.last_searched_at, m.id\n LIMIT ?\n ", + "describe": { + "columns": [ + { + "name": "id!: i64", + "ordinal": 0, + "type_info": "Integer" + }, + { + "name": "title!: String", + "ordinal": 1, + "type_info": "Text" + }, + { + "name": "year", + "ordinal": 2, + "type_info": "Integer" + }, + { + "name": "original_language", + "ordinal": 3, + "type_info": "Text" + } + ], + "parameters": { + "Right": 1 + }, + "nullable": [ + true, + false, + true, + true + ] + }, + "hash": "2c329cab6e494ef1a765a7de05e1351358d4ed583a36fe5d0edc38cc12603e44" +} diff --git a/.sqlx/query-6e988f7e8d8ac2167c589a35626b11ab88e1fdcc28c432da4c3debf1202eee7a.json b/.sqlx/query-a27c937b6039582e0615a1bec479e7b21d2906779bfbbe9e375fa3a70a2de555.json similarity index 50% rename from .sqlx/query-6e988f7e8d8ac2167c589a35626b11ab88e1fdcc28c432da4c3debf1202eee7a.json rename to .sqlx/query-a27c937b6039582e0615a1bec479e7b21d2906779bfbbe9e375fa3a70a2de555.json index f613aaf..b66af1f 100644 --- a/.sqlx/query-6e988f7e8d8ac2167c589a35626b11ab88e1fdcc28c432da4c3debf1202eee7a.json +++ b/.sqlx/query-a27c937b6039582e0615a1bec479e7b21d2906779bfbbe9e375fa3a70a2de555.json @@ -1,72 +1,77 @@ { "db_name": "SQLite", - "query": "SELECT m.title AS \"title!: String\", m.tmdb_id AS \"tmdb_id!: i64\", m.original_language, m.overrides AS \"overrides!: serde_json::Value\", 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 = ?", + "query": "\n SELECT m.overrides AS \"overrides!: String\",\n r.id AS \"root_id!: i64\",\n r.kind AS \"root_kind!: String\",\n r.audience AS \"root_audience!: String\",\n r.path AS \"root_path!: String\",\n p.id AS \"policy_id!: i64\",\n p.name AS \"policy_name!: String\",\n p.required_audio AS \"required_audio!: String\",\n p.dub_blacklist AS \"dub_blacklist!: String\",\n p.hdr_rules AS \"hdr_rules!: String\",\n p.size_bands AS \"size_bands!: String\",\n p.resolution_pref AS \"resolution_pref!: String\",\n p.source_weights AS \"source_weights!: String\",\n p.score_weights AS \"score_weights!: String\"\n FROM movies m\n JOIN roots r ON r.id = m.root_id\n JOIN policies p ON p.id = r.policy_id\n WHERE m.id = ?\n ", "describe": { "columns": [ { - "name": "title!: String", + "name": "overrides!: String", "ordinal": 0, "type_info": "Text" }, { - "name": "tmdb_id!: i64", + "name": "root_id!: i64", "ordinal": 1, "type_info": "Integer" }, { - "name": "original_language", + "name": "root_kind!: String", "ordinal": 2, "type_info": "Text" }, { - "name": "overrides!: serde_json::Value", + "name": "root_audience!: String", "ordinal": 3, "type_info": "Text" }, { - "name": "policy_id!: i64", + "name": "root_path!: String", "ordinal": 4, + "type_info": "Text" + }, + { + "name": "policy_id!: i64", + "ordinal": 5, "type_info": "Integer" }, { "name": "policy_name!: String", - "ordinal": 5, - "type_info": "Text" - }, - { - "name": "required_audio!: String", "ordinal": 6, "type_info": "Text" }, { - "name": "dub_blacklist!: String", + "name": "required_audio!: String", "ordinal": 7, "type_info": "Text" }, { - "name": "hdr_rules!: String", + "name": "dub_blacklist!: String", "ordinal": 8, "type_info": "Text" }, { - "name": "size_bands!: String", + "name": "hdr_rules!: String", "ordinal": 9, "type_info": "Text" }, { - "name": "resolution_pref!: String", + "name": "size_bands!: String", "ordinal": 10, "type_info": "Text" }, { - "name": "source_weights!: String", + "name": "resolution_pref!: String", "ordinal": 11, "type_info": "Text" }, { - "name": "score_weights!: String", + "name": "source_weights!: String", "ordinal": 12, "type_info": "Text" + }, + { + "name": "score_weights!: String", + "ordinal": 13, + "type_info": "Text" } ], "parameters": { @@ -75,7 +80,8 @@ "nullable": [ false, false, - true, + false, + false, false, false, false, @@ -88,5 +94,5 @@ false ] }, - "hash": "6e988f7e8d8ac2167c589a35626b11ab88e1fdcc28c432da4c3debf1202eee7a" + "hash": "a27c937b6039582e0615a1bec479e7b21d2906779bfbbe9e375fa3a70a2de555" } diff --git a/.sqlx/query-be5b70aed19b1d1e81dd26f84faf9caff1712321c6398abed99408a7c6cb31e4.json b/.sqlx/query-be5b70aed19b1d1e81dd26f84faf9caff1712321c6398abed99408a7c6cb31e4.json new file mode 100644 index 0000000..a7a631c --- /dev/null +++ b/.sqlx/query-be5b70aed19b1d1e81dd26f84faf9caff1712321c6398abed99408a7c6cb31e4.json @@ -0,0 +1,32 @@ +{ + "db_name": "SQLite", + "query": "SELECT title AS \"title!: String\", tmdb_id AS \"tmdb_id!: i64\", original_language FROM movies WHERE id = ?", + "describe": { + "columns": [ + { + "name": "title!: String", + "ordinal": 0, + "type_info": "Text" + }, + { + "name": "tmdb_id!: i64", + "ordinal": 1, + "type_info": "Integer" + }, + { + "name": "original_language", + "ordinal": 2, + "type_info": "Text" + } + ], + "parameters": { + "Right": 1 + }, + "nullable": [ + false, + false, + true + ] + }, + "hash": "be5b70aed19b1d1e81dd26f84faf9caff1712321c6398abed99408a7c6cb31e4" +} diff --git a/.sqlx/query-e406ad05186e0fee570afea0b24a74fe7a4db7bc40477bf963a8f9240d406f13.json b/.sqlx/query-e406ad05186e0fee570afea0b24a74fe7a4db7bc40477bf963a8f9240d406f13.json new file mode 100644 index 0000000..5e7588d --- /dev/null +++ b/.sqlx/query-e406ad05186e0fee570afea0b24a74fe7a4db7bc40477bf963a8f9240d406f13.json @@ -0,0 +1,12 @@ +{ + "db_name": "SQLite", + "query": "INSERT INTO movie_releases (movie_id, release_id) VALUES (?, ?)\n ON CONFLICT DO NOTHING", + "describe": { + "columns": [], + "parameters": { + "Right": 2 + }, + "nullable": [] + }, + "hash": "e406ad05186e0fee570afea0b24a74fe7a4db7bc40477bf963a8f9240d406f13" +} diff --git a/.sqlx/query-fb28c3b48459572d95cf69439c6e0cb3697811ab55fe48f224856ed94d59ce92.json b/.sqlx/query-fb28c3b48459572d95cf69439c6e0cb3697811ab55fe48f224856ed94d59ce92.json new file mode 100644 index 0000000..65212f8 --- /dev/null +++ b/.sqlx/query-fb28c3b48459572d95cf69439c6e0cb3697811ab55fe48f224856ed94d59ce92.json @@ -0,0 +1,12 @@ +{ + "db_name": "SQLite", + "query": "UPDATE movies\n SET search_attempts = search_attempts + 1,\n last_searched_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now'),\n updated_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now')\n WHERE id = ?", + "describe": { + "columns": [], + "parameters": { + "Right": 1 + }, + "nullable": [] + }, + "hash": "fb28c3b48459572d95cf69439c6e0cb3697811ab55fe48f224856ed94d59ce92" +} diff --git a/Cargo.lock b/Cargo.lock index 71f45be..619249c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -80,13 +80,19 @@ version = "0.1.0" dependencies = [ "arr-api", "arr-compat", + "arr-core", "arr-db", + "arr-dl", + "arr-indexer", "arr-meta", + "arr-parse", "axum", + "chrono", "include_dir", "mime_guess", "reqwest", "serde", + "serde_json", "sqlx", "tempfile", "thiserror", @@ -95,14 +101,19 @@ dependencies = [ "tower-http", "tracing", "tracing-subscriber", + "wiremock", ] [[package]] name = "arr-db" version = "0.1.0" dependencies = [ + "arr-core", + "serde", + "serde_json", "sqlx", "tempfile", + "thiserror", "tokio", ] diff --git a/crates/arr-api/src/search.rs b/crates/arr-api/src/search.rs index db1b4a2..2d6154b 100644 --- a/crates/arr-api/src/search.rs +++ b/crates/arr-api/src/search.rs @@ -1,12 +1,9 @@ -use std::collections::BTreeMap; use std::time::UNIX_EPOCH; use arr_core::policy::{evaluate, Candidate}; use arr_core::score::score; -use arr_core::{ - DolbyVisionProfile, HdrRules, Language, Policy, PolicyId, RequiredAudio, Resolution, Rule, - ScoreWeights, SizeBand, Source, TitleOverrides, Verdict, -}; +use arr_core::{Language, Policy, Rule, TitleOverrides, Verdict}; +use arr_db::policy::language; use arr_indexer::{ProwlarrClient, SearchRelease, SearchRequest}; use axum::extract::{Query, State}; use axum::Json; @@ -71,54 +68,6 @@ pub struct ClassifiedRelease { pub rule: Option, } -#[derive(Debug, Deserialize)] -struct PolicyRow { - policy_id: i64, - policy_name: String, - required_audio: String, - dub_blacklist: String, - hdr_rules: String, - size_bands: String, - resolution_pref: String, - source_weights: String, - score_weights: String, -} - -#[derive(Debug, Deserialize)] -struct RequiredAudioJson { - require: String, - #[serde(default)] - langs: Vec, -} - -#[derive(Debug, Deserialize)] -struct HdrRulesJson { - #[serde(default)] - dv_profile_reject: Vec, -} - -#[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, Deserialize)] -struct OverridesJson { - #[serde(default)] - only_4k: bool, - #[serde(default)] - allow_english_audio: bool, -} - #[utoipa::path( get, path = "/api/search", tag = "search", params(SearchQuery), responses( @@ -263,8 +212,13 @@ pub async fn releases( Query(query): Query, ) -> Result>, ApiError> { let database = state.database().ok_or(ApiError::Unavailable)?; - let movie = sqlx::query!(r#"SELECT m.title AS "title!: String", m.tmdb_id AS "tmdb_id!: i64", m.original_language, m.overrides AS "overrides!: serde_json::Value", 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 = ?"#, query.movie_id) + let movie = sqlx::query!(r#"SELECT title AS "title!: String", tmdb_id AS "tmdb_id!: i64", original_language FROM movies WHERE id = ?"#, query.movie_id) .fetch_optional(database.pool()).await?.ok_or(ApiError::NotFound)?; + let loaded = database + .movie_policy(query.movie_id) + .await + .map_err(|error| ApiError::Database(error.to_string()))? + .ok_or(ApiError::NotFound)?; let tmdb = tmdb_client(&state)? .movie( @@ -290,23 +244,8 @@ pub async fn releases( .indexers() .await .map_err(|_| ApiError::Unavailable)?; - let policy = policy_from_row(PolicyRow { - policy_id: movie.policy_id, - policy_name: movie.policy_name, - required_audio: movie.required_audio, - dub_blacklist: movie.dub_blacklist, - hdr_rules: movie.hdr_rules, - size_bands: movie.size_bands, - resolution_pref: movie.resolution_pref, - source_weights: movie.source_weights, - score_weights: movie.score_weights, - })?; - let overrides: OverridesJson = serde_json::from_value(movie.overrides) - .map_err(|error| ApiError::Database(error.to_string()))?; - let overrides = TitleOverrides { - only_4k: overrides.only_4k, - allow_english_audio: overrides.allow_english_audio, - }; + let policy = loaded.policy; + let overrides = loaded.overrides; let original_language = title_language( movie .original_language @@ -409,67 +348,6 @@ fn upstream_error(error: &arr_meta::Error) -> ApiError { } } -fn policy_from_row(row: PolicyRow) -> Result { - let required: RequiredAudioJson = json(&row.required_audio)?; - let hdr: HdrRulesJson = json(&row.hdr_rules)?; - let bands: BTreeMap = json(&row.size_bands)?; - let resolutions: Vec = json(&row.resolution_pref)?; - let weights: BTreeMap = json(&row.source_weights)?; - let score_weights: ScoreWeightsJson = json(&row.score_weights)?; - Ok(Policy { - id: PolicyId(row.policy_id), - name: row.policy_name, - required_audio: if required.require == "original_language" { - RequiredAudio::OriginalLanguage - } else { - RequiredAudio::AnyOf(required.langs.iter().map(|value| language(value)).collect()) - }, - dub_blacklist: json::>(&row.dub_blacklist)? - .iter() - .map(|value| language(value)) - .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, - }, - }) -} - fn classify( release: SearchRelease, policy: &Policy, @@ -546,17 +424,6 @@ fn bucket(verdict: &str) -> u8 { _ => 2, } } -fn gib(value: u64) -> u64 { - value.saturating_mul(1024 * 1024 * 1024) -} -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()), - } -} fn title_language(value: &str, origin_countries: &[String]) -> Language { if value == "pt" { if origin_countries.iter().any(|country| country == "BR") { @@ -568,30 +435,11 @@ fn title_language(value: &str, origin_countries: &[String]) -> Language { } language(value) } -fn resolution_value(value: &str) -> Option { - match value { - "2160p" => Some(Resolution::R2160p), - "1080p" => Some(Resolution::R1080p), - "720p" => Some(Resolution::R720p), - _ => None, - } -} -fn source_value(value: &str) -> Option { - match value { - "Remux" => Some(Source::Remux), - "BluRay" => Some(Source::BluRay), - "WEB-DL" => Some(Source::WebDl), - "WEBRip" => Some(Source::WebRip), - "HDTV" => Some(Source::Hdtv), - _ => None, - } -} -fn json(value: &str) -> Result { - serde_json::from_str(value).map_err(|error| ApiError::Database(error.to_string())) -} #[cfg(test)] mod tests { + use arr_core::{HdrRules, PolicyId, RequiredAudio, Resolution, ScoreWeights, SizeBand, Source}; + use super::*; use crate::{router, Upstreams}; use wiremock::matchers::{method, path, query_param}; @@ -742,33 +590,6 @@ mod tests { assert_eq!(eligible["score"], 0); } - #[test] - fn policy_row_uses_persisted_score_weights() { - let policy = policy_from_row(PolicyRow { - policy_id: 1, - policy_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(), - }) - .expect("policy row"); - - assert_eq!( - policy.score_weights, - ScoreWeights { - size_at_target: 2000, - source_tier: 7, - seeder_doubling: 11, - } - ); - } - #[test] fn releases_without_sizes_skip_the_size_score() { let policy = Policy { @@ -782,8 +603,8 @@ mod tests { size_bands: std::collections::BTreeMap::from([( Resolution::R2160p, SizeBand { - floor_bytes: gib(8), - target_bytes: gib(22), + floor_bytes: 8 << 30, + target_bytes: 22 << 30, penalty_points_per_gib_over: 60, }, )]), diff --git a/crates/arr-daemon/Cargo.toml b/crates/arr-daemon/Cargo.toml index edda77f..564aaf7 100644 --- a/crates/arr-daemon/Cargo.toml +++ b/crates/arr-daemon/Cargo.toml @@ -13,13 +13,19 @@ path = "src/main.rs" [dependencies] arr-api = { workspace = true } arr-compat = { workspace = true } +arr-core = { workspace = true } arr-db = { workspace = true } +arr-dl = { workspace = true } +arr-indexer = { workspace = true } arr-meta = { workspace = true } +arr-parse = { workspace = true } axum = { workspace = true } +chrono = { workspace = true } include_dir = { workspace = true } mime_guess = { workspace = true } reqwest = { workspace = true } serde = { workspace = true } +serde_json = { workspace = true } sqlx = { workspace = true } thiserror = { workspace = true } tokio = { workspace = true } @@ -30,6 +36,7 @@ tracing-subscriber = { workspace = true } [dev-dependencies] tempfile = { workspace = true } +wiremock = { workspace = true } [lints] workspace = true diff --git a/crates/arr-daemon/src/config.rs b/crates/arr-daemon/src/config.rs index 3b4c1ab..ba96884 100644 --- a/crates/arr-daemon/src/config.rs +++ b/crates/arr-daemon/src/config.rs @@ -20,6 +20,9 @@ pub const ENV_MEDIA_ROOT: &str = "ARR_MEDIA_ROOT"; pub const ENV_PROWLARR_URL: &str = "ARR_PROWLARR_URL"; pub const ENV_PROWLARR_API_KEY: &str = "ARR_PROWLARR_API_KEY"; pub const ENV_TRANSMISSION_URL: &str = "ARR_TRANSMISSION_URL"; +pub const ENV_DOWNLOAD_DIR: &str = "ARR_DOWNLOAD_DIR"; +pub const ENV_SEED_RATIO_LIMIT: &str = "ARR_SEED_RATIO_LIMIT"; +pub const ENV_SEED_IDLE_LIMIT_MINUTES: &str = "ARR_SEED_IDLE_LIMIT_MINUTES"; pub const ENV_TMDB_API_KEY: &str = "ARR_TMDB_API_KEY"; pub const ENV_TMDB_URL: &str = "ARR_TMDB_URL"; pub const ENV_JELLYFIN_URL: &str = "ARR_JELLYFIN_URL"; @@ -31,6 +34,13 @@ pub const DEFAULT_DATABASE_PATH: &str = "arr.db"; pub const DEFAULT_MEDIA_ROOT: &str = "/mnt/media"; pub const DEFAULT_PROWLARR_URL: &str = "http://localhost:9696"; pub const DEFAULT_TRANSMISSION_URL: &str = "http://localhost:9091/transmission/rpc"; +/// Transmission's own view of the download directory (DESIGN.md §3). It +/// shares the media dataset with the library so hardlinks work (§7.2). +pub const DEFAULT_DOWNLOAD_DIR: &str = "/mnt/media/transmission/complete"; +/// Seeding obligation defaults (§7.3), applied to every torrent at add time +/// and enforced by Transmission. Per-tracker rules are issue #25. +pub const DEFAULT_SEED_RATIO_LIMIT: f64 = 1.0; +pub const DEFAULT_SEED_IDLE_LIMIT_MINUTES: u64 = 4320; pub const DEFAULT_JELLYFIN_URL: &str = "http://localhost:8096"; pub const DEFAULT_NTFY_URL: &str = "http://localhost"; @@ -46,6 +56,8 @@ pub enum ConfigError { input: String, source: std::net::AddrParseError, }, + #[error("invalid {env} ({input:?}): expected a number")] + InvalidNumber { env: &'static str, input: String }, } /// On-disk representation. Non-secret fields only — see the module docs. @@ -63,6 +75,12 @@ struct ConfigFile { #[serde(default)] transmission_url: Option, #[serde(default)] + download_dir: Option, + #[serde(default)] + seed_ratio_limit: Option, + #[serde(default)] + seed_idle_limit_minutes: Option, + #[serde(default)] jellyfin_url: Option, #[serde(default)] ntfy_url: Option, @@ -89,6 +107,9 @@ pub struct EnvOverrides { pub prowlarr_url: Option, pub prowlarr_api_key: Option, pub transmission_url: Option, + pub download_dir: Option, + pub seed_ratio_limit: Option, + pub seed_idle_limit_minutes: Option, pub tmdb_api_key: Option, pub tmdb_url: Option, pub jellyfin_url: Option, @@ -106,6 +127,9 @@ impl EnvOverrides { prowlarr_url: std::env::var(ENV_PROWLARR_URL).ok(), prowlarr_api_key: std::env::var(ENV_PROWLARR_API_KEY).ok(), transmission_url: std::env::var(ENV_TRANSMISSION_URL).ok(), + download_dir: std::env::var(ENV_DOWNLOAD_DIR).ok(), + seed_ratio_limit: std::env::var(ENV_SEED_RATIO_LIMIT).ok(), + seed_idle_limit_minutes: std::env::var(ENV_SEED_IDLE_LIMIT_MINUTES).ok(), tmdb_api_key: std::env::var(ENV_TMDB_API_KEY).ok(), tmdb_url: std::env::var(ENV_TMDB_URL).ok(), jellyfin_url: std::env::var(ENV_JELLYFIN_URL).ok(), @@ -116,7 +140,7 @@ impl EnvOverrides { } /// Resolved bootstrap configuration. See DESIGN.md §10. -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq)] pub struct Config { pub bind_addr: SocketAddr, pub database_path: PathBuf, @@ -124,6 +148,11 @@ pub struct Config { pub prowlarr_url: String, pub prowlarr_api_key: Option, pub transmission_url: String, + /// Where Transmission puts completed downloads, in Transmission's own + /// namespace (§7.1). + pub download_dir: PathBuf, + pub seed_ratio_limit: f64, + pub seed_idle_limit_minutes: u64, pub tmdb_api_key: Option, /// E2E seam only, env-only. `None` means the client's built-in TMDB /// address; DESIGN.md §10 keeps the real URL out of configuration. @@ -175,6 +204,21 @@ impl Config { .transmission_url .or(file.transmission_url) .unwrap_or_else(|| DEFAULT_TRANSMISSION_URL.to_string()), + download_dir: env + .download_dir + .map(PathBuf::from) + .or(file.download_dir) + .unwrap_or_else(|| PathBuf::from(DEFAULT_DOWNLOAD_DIR)), + seed_ratio_limit: match &env.seed_ratio_limit { + Some(raw) => parse_number(raw, ENV_SEED_RATIO_LIMIT)?, + None => file.seed_ratio_limit.unwrap_or(DEFAULT_SEED_RATIO_LIMIT), + }, + seed_idle_limit_minutes: match &env.seed_idle_limit_minutes { + Some(raw) => parse_number(raw, ENV_SEED_IDLE_LIMIT_MINUTES)?, + None => file + .seed_idle_limit_minutes + .unwrap_or(DEFAULT_SEED_IDLE_LIMIT_MINUTES), + }, tmdb_api_key: env.tmdb_api_key, tmdb_url: env.tmdb_url, jellyfin_url: env @@ -190,6 +234,13 @@ impl Config { } } +fn parse_number(raw: &str, env: &'static str) -> Result { + raw.parse().map_err(|_| ConfigError::InvalidNumber { + env, + input: raw.to_string(), + }) +} + fn parse_bind_addr(raw: &str, env: &'static str) -> Result { raw.parse().map_err(|source| ConfigError::InvalidBindAddr { env, @@ -211,6 +262,12 @@ mod tests { assert_eq!(config.media_root, PathBuf::from(DEFAULT_MEDIA_ROOT)); assert_eq!(config.prowlarr_url, DEFAULT_PROWLARR_URL); assert_eq!(config.transmission_url, DEFAULT_TRANSMISSION_URL); + assert_eq!(config.download_dir, PathBuf::from(DEFAULT_DOWNLOAD_DIR)); + assert!((config.seed_ratio_limit - DEFAULT_SEED_RATIO_LIMIT).abs() < f64::EPSILON); + assert_eq!( + config.seed_idle_limit_minutes, + DEFAULT_SEED_IDLE_LIMIT_MINUTES + ); assert_eq!(config.jellyfin_url, DEFAULT_JELLYFIN_URL); assert_eq!(config.ntfy_url, DEFAULT_NTFY_URL); assert_eq!(config.prowlarr_api_key, None); @@ -324,6 +381,43 @@ prowlarr_url = "http://prowlarr.internal:9696" )); } + /// §7.3: both seeding limits are set at add time, so both have to be + /// configurable without a rebuild. + #[test] + fn seeding_limits_come_from_file_or_env() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("arr.toml"); + std::fs::write( + &path, + "seed_ratio_limit = 2.5\nseed_idle_limit_minutes = 120\n", + ) + .unwrap(); + let env = EnvOverrides { + config_file: Some(path.to_string_lossy().into_owned()), + ..EnvOverrides::default() + }; + let config = Config::resolve(env.clone()).unwrap(); + assert!((config.seed_ratio_limit - 2.5).abs() < f64::EPSILON); + assert_eq!(config.seed_idle_limit_minutes, 120); + + let config = Config::resolve(EnvOverrides { + seed_ratio_limit: Some("0.5".into()), + ..env.clone() + }) + .unwrap(); + assert!((config.seed_ratio_limit - 0.5).abs() < f64::EPSILON); + + let error = Config::resolve(EnvOverrides { + seed_idle_limit_minutes: Some("soon".into()), + ..env + }) + .unwrap_err(); + assert!(matches!( + error, + ConfigError::InvalidNumber { env, .. } if env == ENV_SEED_IDLE_LIMIT_MINUTES + )); + } + #[test] fn invalid_env_bind_addr_returns_diagnostic() { let env = EnvOverrides { diff --git a/crates/arr-daemon/src/grab.rs b/crates/arr-daemon/src/grab.rs new file mode 100644 index 0000000..63928aa --- /dev/null +++ b/crates/arr-daemon/src/grab.rs @@ -0,0 +1,1073 @@ +//! The grab pipeline: close the "wanted movie, no file" gap by searching, +//! scoring, picking a winner and sending it to Transmission. See DESIGN.md +//! §5.4, §6.2, §7.1, §7.3 and §8. +//! +//! There is no grab delay and there is no job queue. The gap is recomputed +//! from domain rows on every tick, so killing the process mid-grab and +//! restarting converges instead of double-grabbing: +//! +//! - a title with a live `grabs` row is not a gap, so it is never re-searched; +//! - Transmission's `torrent-add` is keyed on the infohash, so re-sending the +//! same release returns the torrent that is already there rather than a +//! second one; +//! - the `grabs` row is written from that response, so a crash between the add +//! and the insert heals on the next tick instead of leaving an orphan. + +use std::collections::HashMap; +use std::path::PathBuf; +use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; + +use arr_core::policy::{evaluate, Candidate}; +use arr_core::{score::score, Language, Policy, Rule, TitleOverrides, Verdict}; +use arr_db::{Db, MoviePolicy}; +use arr_dl::{AddTorrent, TorrentSource, TransmissionClient}; +use arr_indexer::{ProwlarrClient, SearchRelease, SearchRequest}; + +use crate::reconcile::{Action, ActionFuture, Outcome}; + +/// How many titles one tick may search. The reconcile lane has a 25 s budget +/// and a search costs one call per indexer per title (§6.2), so the work is +/// bounded here and the least recently searched titles come first. Backoff +/// and release-date gating are issue #26. +const MOVIES_PER_TICK: i64 = 5; + +/// How long a discovered indexer list is reused. Prowlarr enumerates one +/// indexer per call plus a `t=caps` probe each, all sequential, so paying for +/// it every 30 s would leave the tick no room to search. +const INDEXER_CACHE_TTL: Duration = Duration::from_mins(15); + +/// The share of the reconcile lane's 25 s budget discovery may spend. One +/// unresponsive tracker must not cancel the whole action before a single +/// targeted search has run; past this the last known list is used instead. +const INDEXER_DISCOVERY_TIMEOUT: Duration = Duration::from_secs(8); + +/// Seeding obligations, per tracker in principle (§7.3) and per install in +/// practice until issue #25 gives them a home. Both are set on the torrent at +/// add time and enforced by Transmission. +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct SeedingLimits { + pub ratio: f64, + pub idle_minutes: u64, +} + +/// A failure during one grab tick. +#[derive(Debug, thiserror::Error)] +pub enum GrabError { + #[error("database: {0}")] + Database(#[from] sqlx::Error), + #[error("policy: {0}")] + Policy(#[from] arr_db::PolicyError), + #[error("prowlarr: {0}")] + Prowlarr(#[from] arr_indexer::Error), + #[error("transmission: {0}")] + Transmission(#[from] arr_dl::Error), + #[error("release {name}: {source}")] + Parsed { + name: String, + source: serde_json::Error, + }, + #[error("indexer discovery timed out after {0:?} and none is known yet")] + IndexerDiscoveryTimeout(Duration), +} + +/// The searchable indexer ids, and when they were last discovered. +#[derive(Debug, Default)] +struct IndexerCache { + ids: Vec, + refreshed_at: Option, +} + +/// Sends the best eligible release for every wanted movie that has neither a +/// file nor a grab in flight. +#[derive(Debug)] +pub struct GrabAction { + prowlarr: ProwlarrClient, + transmission: TransmissionClient, + download_dir: PathBuf, + seeding: SeedingLimits, + indexers: tokio::sync::RwLock, + /// [`INDEXER_DISCOVERY_TIMEOUT`], overridden by tests that cannot wait + /// out the real one. Mirrors `ReconcileLoop`'s action timeout override. + discovery_timeout: Duration, +} + +impl GrabAction { + #[must_use] + pub fn new( + prowlarr: ProwlarrClient, + transmission: TransmissionClient, + download_dir: PathBuf, + seeding: SeedingLimits, + ) -> Self { + Self { + prowlarr, + transmission, + download_dir, + seeding, + indexers: tokio::sync::RwLock::new(IndexerCache::default()), + discovery_timeout: INDEXER_DISCOVERY_TIMEOUT, + } + } + + async fn tick(&self, database: &Db) -> Result, GrabError> { + let mut outcomes = self.track_sent_grabs(database).await?; + let gaps = pending_movies(database).await?; + if gaps.is_empty() { + return Ok(outcomes); + } + + let searchable = self.searchable_indexers().await?; + if searchable.is_empty() { + tracing::warn!("no indexer advertises a text search; nothing can be grabbed"); + return Ok(outcomes); + } + + for movie in gaps { + match self.grab_one(database, &movie, &searchable).await { + Ok(Some(outcome)) => outcomes.push(outcome), + Ok(None) => {} + // One title's failure must not cost the rest of the tick. + Err(error) => tracing::error!( + movie_id = movie.id, + title = movie.title, + %error, + "grab failed" + ), + } + } + + Ok(outcomes) + } + + /// Move grabs Transmission reports as complete out of `sent`. + /// + /// Transmission is authoritative and its view is rebuilt on every tick + /// rather than cached (§8), so this is also what reconstructs in-flight + /// state after a restart. + async fn track_sent_grabs(&self, database: &Db) -> Result, GrabError> { + let sent = sqlx::query!( + r#"SELECT id AS "id!: i64", infohash AS "infohash!: String", target_id AS "target_id!: i64" + FROM grabs WHERE state = 'sent'"# + ) + .fetch_all(database.pool()) + .await?; + if sent.is_empty() { + return Ok(Vec::new()); + } + + let torrents: HashMap = self + .transmission + .list_torrents() + .await? + .into_iter() + .map(|torrent| (torrent.hash.to_ascii_lowercase(), torrent.progress)) + .collect(); + + let mut outcomes = Vec::new(); + for grab in sent { + let Some(progress) = torrents.get(&grab.infohash.to_ascii_lowercase()) else { + // Gone from Transmission. Deciding whether that is a failure + // or a manual removal is issue #24's; leaving the row alone + // keeps this tick from re-grabbing behind the operator. + continue; + }; + if *progress < 1.0 { + continue; + } + sqlx::query!( + "UPDATE grabs SET state = 'downloaded' WHERE id = ?", + grab.id + ) + .execute(database.pool()) + .await?; + outcomes.push(Outcome::new( + format!("grab {} downloaded, still marked sent", grab.id), + format!("marked grab {} downloaded", grab.id), + )); + tracing::info!( + grab_id = grab.id, + movie_id = grab.target_id, + "download complete" + ); + } + Ok(outcomes) + } + + /// The indexers that accept a text search, cached across ticks. + /// + /// Discovery is bounded and its result is reused, because Prowlarr probes + /// capabilities one indexer at a time: a single slow tracker would + /// otherwise burn the whole reconcile budget before any search runs. A + /// refresh that fails or times out keeps the previous list rather than + /// stopping the tick, and only an empty cache turns that into an error. + async fn searchable_indexers(&self) -> Result, GrabError> { + { + let cache = self.indexers.read().await; + if cache + .refreshed_at + .is_some_and(|at| at.elapsed() < INDEXER_CACHE_TTL) + { + return Ok(cache.ids.clone()); + } + } + + let discovered = + tokio::time::timeout(self.discovery_timeout, self.prowlarr.indexers()).await; + let mut cache = self.indexers.write().await; + match discovered { + Ok(Ok(indexers)) => { + cache.ids = indexers + .iter() + .filter(|indexer| indexer.capabilities.search.available) + .map(|indexer| indexer.id) + .collect(); + cache.refreshed_at = Some(Instant::now()); + } + Ok(Err(error)) if cache.ids.is_empty() => return Err(error.into()), + Err(_) if cache.ids.is_empty() => { + return Err(GrabError::IndexerDiscoveryTimeout(self.discovery_timeout)) + } + Ok(Err(error)) => { + tracing::warn!(%error, "indexer discovery failed; using the last known list"); + } + Err(_) => tracing::warn!( + timeout_seconds = self.discovery_timeout.as_secs(), + "indexer discovery timed out; using the last known list" + ), + } + Ok(cache.ids.clone()) + } + + /// Search every indexer for one title, cache each candidate with its + /// verdict and score (§9.3), and return the eligible ones best first. + /// + /// The order is total, so the same candidate set picks the same winner + /// after a restart and a re-sent grab is a duplicate rather than a second + /// torrent. Waived releases are a manual, one-click decision (§9.3) and + /// never appear here. + async fn search( + &self, + database: &Db, + movie: &PendingMovie, + indexers: &[i64], + loaded: &MoviePolicy, + original_language: &Language, + ) -> Result, GrabError> { + let request = SearchRequest::Text { + query: search_query(movie), + }; + let mut releases = Vec::new(); + for search in self.prowlarr.search_indexers(indexers, &request).await { + if let Some(error) = search.error { + tracing::warn!(indexer_id = search.indexer_id, %error, "indexer search failed"); + } + releases.extend(search.releases); + } + record_search(database, movie.id).await?; + + let mut candidates = Vec::new(); + for release in releases { + let stored = store_release( + database, + movie.id, + &release, + &loaded.policy, + &loaded.overrides, + original_language, + ) + .await?; + if let Some(candidate) = stored { + candidates.push(candidate); + } + } + candidates.sort_by(|left, right| { + right + .score + .cmp(&left.score) + .then_with(|| left.indexer_id.cmp(&right.indexer_id)) + .then_with(|| left.guid.cmp(&right.guid)) + }); + Ok(candidates) + } + + async fn grab_one( + &self, + database: &Db, + movie: &PendingMovie, + indexers: &[i64], + ) -> Result, GrabError> { + let loaded = database.movie_policy(movie.id).await?; + let Some(loaded) = loaded else { + return Ok(None); + }; + // §5.2: the whole language rule is expressed against the title's own + // original language. Without it there is nothing to evaluate against, + // and guessing gives a child a Brazilian dub or throws away a + // Brazilian film's own soundtrack. + let Some(original_language) = movie.original_language.as_deref() else { + tracing::warn!( + movie_id = movie.id, + title = movie.title, + "no original language yet; not searching" + ); + return Ok(None); + }; + let original_language = arr_db::policy::language(original_language); + + let candidates = self + .search(database, movie, indexers, &loaded, &original_language) + .await?; + let Some(winner) = candidates.into_iter().next() else { + tracing::info!( + movie_id = movie.id, + title = movie.title, + "no eligible release" + ); + return Ok(None); + }; + + let added = self + .transmission + .add_torrent(AddTorrent { + source: torrent_source(&winner.download_url), + label: label(&loaded), + download_dir: self.download_dir.clone(), + seed_ratio_limit: self.seeding.ratio, + seed_idle_limit_minutes: self.seeding.idle_minutes, + }) + .await?; + let infohash = added.hash.to_ascii_lowercase(); + + // A duplicate here is the restart case: the torrent was added before + // the process died. `DO NOTHING` keeps the original row. + let inserted = sqlx::query!( + r#"INSERT INTO grabs (release_id, target_kind, target_id, infohash, state) + VALUES (?, 'movie', ?, ?, 'sent') + ON CONFLICT (infohash) DO NOTHING + RETURNING id AS "id!: i64""#, + winner.id, + movie.id, + infohash + ) + .fetch_optional(database.pool()) + .await?; + sqlx::query!( + "UPDATE movies SET state = 'grabbed', + updated_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now') + WHERE id = ?", + movie.id + ) + .execute(database.pool()) + .await?; + + let Some(inserted) = inserted else { + tracing::info!( + movie_id = movie.id, + infohash, + "grab already recorded for this torrent" + ); + return Ok(None); + }; + + tracing::info!( + movie_id = movie.id, + title = movie.title, + release = winner.name, + score = winner.score, + infohash, + was_duplicate = added.was_duplicate, + "grabbed" + ); + Ok(Some(Outcome::new( + format!("movie {} wanted with no file", movie.id), + format!("grabbed {} as grab {}", winner.name, inserted.id), + ))) + } +} + +impl Action for GrabAction { + fn name(&self) -> &'static str { + "grab" + } + + fn run<'a>(&'a self, database: &'a Db) -> ActionFuture<'a> { + Box::pin(async move { self.tick(database).await.map_err(Into::into) }) + } +} + +/// A wanted movie with neither a file nor a grab in flight. +#[derive(Debug, Clone)] +struct PendingMovie { + id: i64, + title: String, + year: Option, + /// §5.2's language rules are expressed against this, and guessing it is + /// worse than not grabbing. + original_language: Option, +} + +/// The eligible view of a stored release, ranked for selection. +#[derive(Debug, Clone)] +struct Eligible { + id: i64, + indexer_id: i64, + guid: String, + name: String, + download_url: String, + score: i64, +} + +/// The gap, straight out of the domain rows (§8). +/// +/// A title with an unfinished grab is not a gap — that is what stops a +/// restart from grabbing twice. `blocked` stops targeted search only (§6.3), +/// so a blocked title still matches RSS results, which is issue #27. +async fn pending_movies(database: &Db) -> Result, GrabError> { + let rows = sqlx::query!( + r#" + SELECT m.id AS "id!: i64", + m.title AS "title!: String", + m.year, + m.original_language + FROM movies m + WHERE m.wanted = 1 + AND m.blocked = 0 + AND NOT EXISTS ( + SELECT 1 FROM media_files f + WHERE f.owner_kind = 'movie' AND f.owner_id = m.id + ) + AND NOT EXISTS ( + SELECT 1 FROM grabs g + WHERE g.target_kind = 'movie' AND g.target_id = m.id + AND g.state IN ('sent', 'downloaded', 'imported') + ) + ORDER BY m.last_searched_at IS NOT NULL, m.last_searched_at, m.id + LIMIT ? + "#, + MOVIES_PER_TICK + ) + .fetch_all(database.pool()) + .await?; + + Ok(rows + .into_iter() + .map(|row| PendingMovie { + id: row.id, + title: row.title, + year: row.year, + original_language: row.original_language, + }) + .collect()) +} + +/// Cache the classified release and associate it with the title. +/// +/// Returns the candidate only when the release is eligible: automatic +/// selection never takes a waiver (§9.3). +async fn store_release( + database: &Db, + movie_id: i64, + release: &SearchRelease, + policy: &Policy, + overrides: &TitleOverrides, + original_language: &Language, +) -> Result, GrabError> { + let parsed = arr_parse::parse(&release.name); + let evaluation = evaluate( + policy, + overrides, + original_language, + Candidate::PreGrab(&parsed), + release.size, + ); + let scored = score( + policy, + Candidate::PreGrab(&parsed), + release.size.unwrap_or_default(), + release.seeders.unwrap_or_default(), + ); + // A release that did not say its size is not a tiny one: scoring it + // against the band's floor would bury it. Same treatment as the manual + // search view, so the ranking the operator sees is the one that picks. + let score = if release.size.is_some() { + scored.total + } else { + scored.source.saturating_add(scored.seeders) + }; + let (verdict, rule) = verdict_columns(&evaluation.verdict); + let parsed_json = serde_json::to_string(&parsed).map_err(|source| GrabError::Parsed { + name: release.name.clone(), + source, + })?; + let size = i64::try_from(release.size.unwrap_or(0)).unwrap_or(i64::MAX); + let seeders = release.seeders.map(i64::from); + let publish_date = release.publish_date.and_then(rfc3339); + + let id = sqlx::query_scalar!( + r#" + INSERT INTO releases + (indexer_id, guid, name, size, seeders, publish_date, download_url, + parsed, score, verdict, rejected_rule) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT (indexer_id, guid) DO UPDATE SET + name = excluded.name, + size = excluded.size, + seeders = excluded.seeders, + publish_date = excluded.publish_date, + download_url = excluded.download_url, + parsed = excluded.parsed, + score = excluded.score, + verdict = excluded.verdict, + rejected_rule = excluded.rejected_rule + RETURNING id AS "id!: i64" + "#, + release.indexer_id, + release.guid, + release.name, + size, + seeders, + publish_date, + release.download_url, + parsed_json, + score, + verdict, + rule + ) + .fetch_one(database.pool()) + .await?; + + sqlx::query!( + "INSERT INTO movie_releases (movie_id, release_id) VALUES (?, ?) + ON CONFLICT DO NOTHING", + movie_id, + id + ) + .execute(database.pool()) + .await?; + + if verdict != "eligible" { + return Ok(None); + } + Ok(Some(Eligible { + id, + indexer_id: release.indexer_id, + guid: release.guid.clone(), + name: release.name.clone(), + download_url: release.download_url.clone(), + score, + })) +} + +/// Record that the title was searched, so the next tick takes a different one. +async fn record_search(database: &Db, movie_id: i64) -> Result<(), GrabError> { + sqlx::query!( + "UPDATE movies + SET search_attempts = search_attempts + 1, + last_searched_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now'), + updated_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now') + WHERE id = ?", + movie_id + ) + .execute(database.pool()) + .await?; + Ok(()) +} + +/// `movies-main`, `tv-kids` (§7.1). Distinct from Radarr's own labels, so +/// both stacks can run against one Transmission. +fn label(loaded: &MoviePolicy) -> String { + let kind = if loaded.root_kind == "movie" { + "movies" + } else { + &loaded.root_kind + }; + format!("{kind}-{}", loaded.root_audience) +} + +/// Indexers hand out magnets and `.torrent` links interchangeably; +/// Transmission takes either in the same field. +fn torrent_source(download_url: &str) -> TorrentSource { + if download_url.starts_with("magnet:") { + TorrentSource::Magnet(download_url.to_owned()) + } else { + TorrentSource::Url(download_url.to_owned()) + } +} + +fn search_query(movie: &PendingMovie) -> String { + movie.year.map_or_else( + || movie.title.clone(), + |year| format!("{} {year}", movie.title), + ) +} + +fn verdict_columns(verdict: &Verdict) -> (&'static str, Option) { + match verdict { + Verdict::Eligible => ("eligible", None), + Verdict::Waived(_) => ("waived", None), + Verdict::Rejected(rule) => ("rejected", Some(rule_name(rule))), + } +} + +fn rule_name(rule: &Rule) -> String { + match rule { + Rule::RequiredAudio => "required_audio".into(), + Rule::DubBlacklist(_) => "dub_blacklist".into(), + Rule::PortugueseUnverified => "portuguese_unverified".into(), + Rule::DolbyVisionProfile(_) => "dolby_vision_profile".into(), + Rule::Resolution(_) => "resolution".into(), + Rule::Source(_) => "source".into(), + Rule::Size => "size".into(), + Rule::Other(name) => name.clone(), + } +} + +fn rfc3339(time: SystemTime) -> Option { + let seconds = time.duration_since(UNIX_EPOCH).ok()?.as_secs(); + let seconds = i64::try_from(seconds).ok()?; + chrono::DateTime::from_timestamp(seconds, 0).map(|date| date.to_rfc3339()) +} + +#[cfg(test)] +#[allow(clippy::unwrap_used)] +mod tests { + use std::sync::{Arc, Mutex}; + + use serde_json::{json, Value}; + use wiremock::matchers::{method, path, query_param}; + use wiremock::{Mock, MockServer, Request, Respond, ResponseTemplate}; + + use super::*; + + /// A Transmission that dedupes on the infohash, like the real one: the + /// same source added twice is one torrent and a `torrent-duplicate` + /// response. That is the property the restart case leans on. + #[derive(Clone, Debug, Default)] + struct FakeTransmission { + torrents: Arc>>, + } + + #[derive(Clone, Debug)] + struct FakeTorrent { + id: i64, + hash: String, + source: String, + labels: Vec, + progress: f64, + } + + impl FakeTransmission { + fn torrents(&self) -> Vec { + self.torrents.lock().unwrap().clone() + } + + fn complete_all(&self) { + for torrent in self.torrents.lock().unwrap().iter_mut() { + torrent.progress = 1.0; + } + } + + fn add(&self, arguments: &Value) -> ResponseTemplate { + let source = arguments["filename"] + .as_str() + .unwrap_or_default() + .to_owned(); + let labels: Vec = arguments["labels"] + .as_array() + .map(|values| { + values + .iter() + .filter_map(|value| value.as_str().map(ToOwned::to_owned)) + .collect() + }) + .unwrap_or_default(); + let mut torrents = self.torrents.lock().unwrap(); + + if let Some(existing) = torrents.iter().find(|torrent| torrent.source == source) { + return success(&json!({"torrent-duplicate": { + "id": existing.id, "name": existing.source, "hashString": existing.hash + }})); + } + + let id = i64::try_from(torrents.len()).unwrap() + 1; + // Deterministic stand-in for the real infohash, which likewise + // comes out the same for the same torrent. + let hash = format!("{:040x}", id * 7); + torrents.push(FakeTorrent { + id, + hash: hash.clone(), + source: source.clone(), + labels, + progress: 0.0, + }); + success(&json!({"torrent-added": {"id": id, "name": source, "hashString": hash}})) + } + } + + impl Respond for FakeTransmission { + fn respond(&self, request: &Request) -> ResponseTemplate { + let body: Value = serde_json::from_slice(&request.body).unwrap(); + let arguments = &body["arguments"]; + match body["method"].as_str().unwrap_or_default() { + "torrent-add" => self.add(arguments), + "torrent-get" => { + let torrents: Vec = self + .torrents() + .into_iter() + .map(|torrent| { + json!({ + "id": torrent.id, "name": torrent.source, + "hashString": torrent.hash, "status": 4, + "percentDone": torrent.progress, + "downloadDir": "/downloads", "labels": torrent.labels + }) + }) + .collect(); + success(&json!({"torrents": torrents})) + } + _ => success(&json!({})), + } + } + } + + fn success(arguments: &Value) -> ResponseTemplate { + ResponseTemplate::new(200).set_body_json(json!({ + "result": "success", "arguments": arguments + })) + } + + const RSS: &str = r#" + + Dune.Part.Two.2024.2160p.WEB-DL.DDP5.1.Atmos + goodhttps://tracker/good.torrent + 23622320128 + + + + Dune.Part.Two.2024.2160p.REMUX + hugehttps://tracker/huge.torrent + 64424509440 + + + + Dune.Part.Two.2024.1080p.CAM + camhttps://tracker/cam.torrent + 4000000000 + + + "#; + + async fn prowlarr() -> MockServer { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/api/v1/indexer")) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(json!([{"id": 7, "name": "tracker", "enable": true}])), + ) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/7/api")) + .and(query_param("t", "caps")) + .respond_with(ResponseTemplate::new(200).set_body_string( + r#""#, + )) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/7/api")) + .and(query_param("t", "search")) + .respond_with(ResponseTemplate::new(200).set_body_string(RSS)) + .mount(&server) + .await; + server + } + + async fn transmission() -> (MockServer, FakeTransmission) { + let server = MockServer::start().await; + let fake = FakeTransmission::default(); + Mock::given(method("POST")) + .respond_with(fake.clone()) + .mount(&server) + .await; + (server, fake) + } + + async fn wanted_movie() -> (tempfile::TempDir, Db) { + let dir = tempfile::tempdir().unwrap(); + let database = Db::connect(dir.path().join("arr.db")).await.unwrap(); + database.migrate().await.unwrap(); + sqlx::query( + "INSERT INTO movies (tmdb_id, title, year, original_language, root_id) + SELECT 693134, 'Dune Part Two', 2024, 'en', id + FROM roots WHERE kind = 'movie' AND audience = 'main'", + ) + .execute(database.pool()) + .await + .unwrap(); + (dir, database) + } + + fn action(prowlarr: &MockServer, transmission: &MockServer) -> GrabAction { + GrabAction::new( + ProwlarrClient::new(prowlarr.uri(), "key").unwrap(), + TransmissionClient::new(&transmission.uri()).unwrap(), + PathBuf::from("/mnt/media/transmission/complete"), + SeedingLimits { + ratio: 1.5, + idle_minutes: 60, + }, + ) + } + + async fn grabs(database: &Db) -> Vec<(i64, String, String)> { + sqlx::query_as::<_, (i64, String, String)>( + "SELECT target_id, infohash, state FROM grabs ORDER BY id", + ) + .fetch_all(database.pool()) + .await + .unwrap() + } + + /// The acceptance case: one wanted movie, one torrent, one grab row. + #[tokio::test] + async fn a_wanted_movie_ends_with_one_torrent_and_one_grab() { + let (_dir, database) = wanted_movie().await; + let indexer = prowlarr().await; + let (downloader, fake) = transmission().await; + + let outcomes = action(&indexer, &downloader).tick(&database).await.unwrap(); + + assert_eq!(outcomes.len(), 1); + assert_eq!(fake.torrents().len(), 1); + let grabs = grabs(&database).await; + assert_eq!(grabs.len(), 1); + assert_eq!(grabs[0].0, 1); + assert_eq!(grabs[0].2, "sent"); + // §5.5: the 22 GB WEB-DL at target beats the 60 GB remux, and the CAM + // is a hard filter rather than a low score however many seeders it has. + assert!(fake.torrents()[0].source.ends_with("good.torrent")); + let state: String = sqlx::query_scalar("SELECT state FROM movies WHERE id = 1") + .fetch_one(database.pool()) + .await + .unwrap(); + assert_eq!(state, "grabbed"); + } + + /// The failure mode the issue names: killed after the torrent is sent and + /// before the row is written, a restart must not send a second one. + #[tokio::test] + async fn a_restart_mid_flight_does_not_grab_twice() { + let (_dir, database) = wanted_movie().await; + let indexer = prowlarr().await; + let (downloader, fake) = transmission().await; + + action(&indexer, &downloader).tick(&database).await.unwrap(); + // The crash: Transmission has the torrent, the database does not know. + sqlx::query("DELETE FROM grabs") + .execute(database.pool()) + .await + .unwrap(); + sqlx::query("UPDATE movies SET state = 'missing'") + .execute(database.pool()) + .await + .unwrap(); + + // A fresh action, as a restarted process would build. + action(&indexer, &downloader).tick(&database).await.unwrap(); + + assert_eq!(fake.torrents().len(), 1); + assert_eq!(grabs(&database).await.len(), 1); + } + + /// A title with a grab in flight is not a gap, so a settled tick is idle. + #[tokio::test] + async fn a_second_tick_grabs_nothing_new() { + let (_dir, database) = wanted_movie().await; + let indexer = prowlarr().await; + let (downloader, fake) = transmission().await; + let action = action(&indexer, &downloader); + + action.tick(&database).await.unwrap(); + let outcomes = action.tick(&database).await.unwrap(); + + assert!(outcomes.is_empty()); + assert_eq!(fake.torrents().len(), 1); + assert_eq!(grabs(&database).await.len(), 1); + } + + /// §7.1 and §7.3: the label and both seeding limits are set at add time, + /// not patched afterwards. + #[tokio::test] + async fn the_label_and_both_seed_limits_are_set_at_add_time() { + let (_dir, database) = wanted_movie().await; + let indexer = prowlarr().await; + let (downloader, fake) = transmission().await; + + action(&indexer, &downloader).tick(&database).await.unwrap(); + + assert_eq!(fake.torrents()[0].labels, vec!["movies-main".to_owned()]); + let add = downloader + .received_requests() + .await + .unwrap() + .into_iter() + .filter_map(|request| serde_json::from_slice::(&request.body).ok()) + .find(|body| body["method"] == "torrent-add") + .expect("torrent-add"); + assert_eq!(add["arguments"]["seedRatioLimit"], json!(1.5)); + assert_eq!(add["arguments"]["seedIdleLimit"], json!(60)); + assert_eq!(add["arguments"]["seedRatioMode"], json!(1)); + assert_eq!(add["arguments"]["seedIdleMode"], json!(1)); + assert_eq!( + add["arguments"]["download-dir"], + json!("/mnt/media/transmission/complete") + ); + } + + /// Every candidate is cached with its verdict, which is what the manual + /// search view and the attention queues read (§9.3). + #[tokio::test] + async fn every_candidate_is_recorded_with_its_verdict() { + let (_dir, database) = wanted_movie().await; + let indexer = prowlarr().await; + let (downloader, _fake) = transmission().await; + + action(&indexer, &downloader).tick(&database).await.unwrap(); + + let rows = sqlx::query_as::<_, (String, String, Option)>( + "SELECT r.name, r.verdict, r.rejected_rule + FROM releases r JOIN movie_releases mr ON mr.release_id = r.id + WHERE mr.movie_id = 1 ORDER BY r.guid", + ) + .fetch_all(database.pool()) + .await + .unwrap(); + assert_eq!(rows.len(), 3); + let cam = rows.iter().find(|row| row.0.contains("CAM")).unwrap(); + assert_eq!(cam.1, "rejected"); + assert_eq!(cam.2.as_deref(), Some("source")); + let attempts: i64 = sqlx::query_scalar("SELECT search_attempts FROM movies WHERE id = 1") + .fetch_one(database.pool()) + .await + .unwrap(); + assert_eq!(attempts, 1); + } + + /// Transmission is authoritative (§8): a completed torrent moves its grab + /// out of `sent` without the process having watched it happen. + #[tokio::test] + async fn a_completed_torrent_moves_its_grab_to_downloaded() { + let (_dir, database) = wanted_movie().await; + let indexer = prowlarr().await; + let (downloader, fake) = transmission().await; + let action = action(&indexer, &downloader); + + action.tick(&database).await.unwrap(); + fake.complete_all(); + let outcomes = action.tick(&database).await.unwrap(); + + assert_eq!(outcomes.len(), 1); + assert_eq!(grabs(&database).await[0].2, "downloaded"); + } + + /// §5.2: the language rules are expressed against the title's original + /// language, and guessing it is worse than waiting for it. + #[tokio::test] + async fn a_title_with_no_original_language_is_not_searched() { + let (_dir, database) = wanted_movie().await; + sqlx::query("UPDATE movies SET original_language = NULL") + .execute(database.pool()) + .await + .unwrap(); + let indexer = prowlarr().await; + let (downloader, fake) = transmission().await; + + let outcomes = action(&indexer, &downloader).tick(&database).await.unwrap(); + + assert!(outcomes.is_empty()); + assert!(fake.torrents().is_empty()); + } + + /// §6.3: `blocked` stops targeted search for a title. + #[tokio::test] + async fn a_blocked_title_is_not_searched() { + let (_dir, database) = wanted_movie().await; + sqlx::query("UPDATE movies SET blocked = 1") + .execute(database.pool()) + .await + .unwrap(); + let indexer = prowlarr().await; + let (downloader, fake) = transmission().await; + + let outcomes = action(&indexer, &downloader).tick(&database).await.unwrap(); + + assert!(outcomes.is_empty()); + assert!(fake.torrents().is_empty()); + assert!(grabs(&database).await.is_empty()); + } + + /// Prowlarr probes capabilities one indexer at a time, so paying for + /// discovery every 30 s would leave the tick no room to search. + #[tokio::test] + async fn indexer_discovery_is_cached_across_ticks() { + let (_dir, database) = wanted_movie().await; + let indexer = prowlarr().await; + let (downloader, _fake) = transmission().await; + let action = action(&indexer, &downloader); + + action.tick(&database).await.unwrap(); + action.tick(&database).await.unwrap(); + + let enumerations = indexer + .received_requests() + .await + .unwrap() + .into_iter() + .filter(|request| request.url.path() == "/api/v1/indexer") + .count(); + assert_eq!(enumerations, 1); + } + + /// One unresponsive tracker must not cancel the action before a single + /// search has run. + #[tokio::test] + async fn a_stalled_discovery_gives_up_instead_of_hanging() { + let (_dir, database) = wanted_movie().await; + let indexer = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/api/v1/indexer")) + .respond_with( + ResponseTemplate::new(200) + .set_delay(Duration::from_secs(30)) + .set_body_json(json!([])), + ) + .mount(&indexer) + .await; + let (downloader, fake) = transmission().await; + + let mut action = action(&indexer, &downloader); + action.discovery_timeout = Duration::from_millis(50); + + let error = action + .tick(&database) + .await + .expect_err("discovery cannot complete"); + + assert!(matches!(error, GrabError::IndexerDiscoveryTimeout(_))); + assert!(fake.torrents().is_empty()); + } + + #[test] + fn indexer_links_and_magnets_both_reach_transmission() { + assert_eq!( + torrent_source("magnet:?xt=urn:btih:abc"), + TorrentSource::Magnet("magnet:?xt=urn:btih:abc".into()) + ); + assert_eq!( + torrent_source("https://tracker/good.torrent"), + TorrentSource::Url("https://tracker/good.torrent".into()) + ); + } +} diff --git a/crates/arr-daemon/src/main.rs b/crates/arr-daemon/src/main.rs index 4299ccc..4a407a8 100644 --- a/crates/arr-daemon/src/main.rs +++ b/crates/arr-daemon/src/main.rs @@ -1,6 +1,7 @@ //! arr — reconcile loop and process entry point. See DESIGN.md §8. mod config; +mod grab; pub mod reconcile; mod web; @@ -12,7 +13,8 @@ use arr_compat::CompatState; use arr_db::Db; use arr_meta::TmdbClient; use config::Config; -use reconcile::ReconcileLoop; +use grab::{GrabAction, SeedingLimits}; +use reconcile::{ReconcileLoop, Tick}; use tower_http::trace::TraceLayer; /// Dump the `OpenAPI` document and exit, instead of serving. `just gen-client` @@ -67,6 +69,10 @@ enum Error { Migration(#[from] sqlx::migrate::MigrateError), #[error("tmdb client: {0}")] Tmdb(#[from] arr_meta::Error), + #[error("prowlarr client: {0}")] + Prowlarr(#[from] arr_indexer::Error), + #[error("transmission client: {0}")] + Transmission(#[from] arr_dl::Error), #[error("bind {addr}: {source}")] Bind { addr: std::net::SocketAddr, @@ -82,7 +88,28 @@ async fn run() -> Result<(), Error> { let config = Config::load()?; let database = Db::connect(&config.database_path).await?; database.migrate().await?; - let reconcile = ReconcileLoop::new(database.clone()); + + let mut reconcile = ReconcileLoop::new(database.clone()); + // Without a Prowlarr key nothing can be searched, so the grab lane stays + // unregistered rather than failing a tick every 30 seconds. + if let Some(key) = config.prowlarr_api_key.clone() { + let prowlarr = arr_indexer::ProwlarrClient::new(config.prowlarr_url.clone(), key)?; + let transmission = arr_dl::TransmissionClient::new(&config.transmission_url)?; + reconcile = reconcile.register( + Tick::Reconcile, + GrabAction::new( + prowlarr, + transmission, + config.download_dir.clone(), + SeedingLimits { + ratio: config.seed_ratio_limit, + idle_minutes: config.seed_idle_limit_minutes, + }, + ), + ); + } else { + tracing::warn!("no Prowlarr API key: nothing will be grabbed"); + } // Jellyseerr's Radarr shim (DESIGN.md §9.4) reads the same database and // needs its own TMDB client for `movie/lookup`. diff --git a/crates/arr-db/Cargo.toml b/crates/arr-db/Cargo.toml index 9cf2896..bdb3a89 100644 --- a/crates/arr-db/Cargo.toml +++ b/crates/arr-db/Cargo.toml @@ -7,7 +7,11 @@ repository.workspace = true publish = false [dependencies] +arr-core = { workspace = true } +serde = { workspace = true } +serde_json = { workspace = true } sqlx = { workspace = true } +thiserror = { workspace = true } [dev-dependencies] tempfile = { workspace = true } diff --git a/crates/arr-db/src/lib.rs b/crates/arr-db/src/lib.rs index 665fc2d..139f2e2 100644 --- a/crates/arr-db/src/lib.rs +++ b/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}; diff --git a/crates/arr-db/src/policy.rs b/crates/arr-db/src/policy.rs new file mode 100644 index 0000000..9111d0d --- /dev/null +++ b/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 { + let required: RequiredAudioJson = json("required_audio", &self.required_audio)?; + let hdr: HdrRulesJson = json("hdr_rules", &self.hdr_rules)?; + let bands: BTreeMap = json("size_bands", &self.size_bands)?; + let resolutions: Vec = json("resolution_pref", &self.resolution_pref)?; + let weights: BTreeMap = json("source_weights", &self.source_weights)?; + let score_weights: ScoreWeightsJson = json("score_weights", &self.score_weights)?; + let blacklist: Vec = 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, 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 { + 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 { + 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, +} + +#[derive(Debug, Deserialize)] +struct HdrRulesJson { + #[serde(default)] + dv_profile_reject: Vec, +} + +#[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( + column: &'static str, + value: &str, +) -> Result { + 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()); + } +} diff --git a/crates/arr-dl/src/lib.rs b/crates/arr-dl/src/lib.rs index f373614..a2a59fd 100644 --- a/crates/arr-dl/src/lib.rs +++ b/crates/arr-dl/src/lib.rs @@ -15,10 +15,13 @@ const MAX_SESSION_NEGOTIATIONS: usize = 4; const CONNECT_TIMEOUT: Duration = Duration::from_secs(5); const REQUEST_TIMEOUT: Duration = Duration::from_secs(30); -/// A magnet URI or the bytes of a `.torrent` file. +/// Where Transmission is to get the torrent from. #[derive(Clone, Debug, Eq, PartialEq)] pub enum TorrentSource { Magnet(String), + /// An HTTP link to a `.torrent`, fetched by Transmission itself. Indexer + /// download links arrive in this form. + Url(String), Metainfo(Vec), } @@ -138,7 +141,9 @@ impl TransmissionClient { }); match request.source { - TorrentSource::Magnet(uri) => arguments["filename"] = json!(uri), + TorrentSource::Magnet(uri) | TorrentSource::Url(uri) => { + arguments["filename"] = json!(uri); + } TorrentSource::Metainfo(bytes) => { arguments["metainfo"] = json!(base64::engine::general_purpose::STANDARD.encode(bytes));