From 23f06d429f54ac6c4d8dd12551c323134d78290d Mon Sep 17 00:00:00 2001 From: Miguel Palhas Date: Sat, 22 Aug 2026 22:15:05 +0100 Subject: [PATCH] feat(daemon): grab pipeline for wanted movies Closes the "wanted movie, no file" gap: search every indexer, cache each candidate with its verdict and score, send the best eligible one to Transmission with its label and both seeding limits, and record the Grab. Restarting mid-flight cannot double-grab. A title with a live grab is not a gap, selection is a total order over the candidate set, and the add is keyed on the infohash, so a crash between the add and the insert heals on the next tick instead of leaving a second torrent. The policy row to arr-core mapping moves into arr-db, where both the manual-search API and the daemon read it, instead of being copied. --- ...6ccfcb06d257fd4435952c93b6287ebdd78a5.json | 12 + ...905f0091229e0433b038fdb55a49e557efa06.json | 32 + ...eda68b4d031ea128cae2125f7e526660d8b45.json | 12 + ...2fae52ce6394fcdc6c8e42eaeec0d7d6fb663.json | 20 + ...f0d7f4d2742acad26efbc7644193966c39325.json | 20 + ...1351358d4ed583a36fe5d0edc38cc12603e44.json | 38 + ...9e7b21d2906779bfbbe9e375fa3a70a2de555.json | 98 ++ ...624d5d50df897af4a477a9cf3dac196b52ecc.json | 86 -- ...f9caff1712321c6398abed99408a7c6cb31e4.json | 32 + ...a74fe7a4db7bc40477bf963a8f9240d406f13.json | 12 + ...e0cb3697811ab55fe48f224856ed94d59ce92.json | 12 + Cargo.lock | 11 + crates/arr-api/src/search.rs | 162 +-- crates/arr-daemon/Cargo.toml | 7 + crates/arr-daemon/src/config.rs | 96 +- crates/arr-daemon/src/grab.rs | 950 ++++++++++++++++++ crates/arr-daemon/src/main.rs | 31 +- crates/arr-db/Cargo.toml | 4 + crates/arr-db/src/lib.rs | 4 + crates/arr-db/src/policy.rs | 370 +++++++ crates/arr-dl/src/lib.rs | 9 +- 21 files changed, 1775 insertions(+), 243 deletions(-) create mode 100644 .sqlx/query-0744ebfb2d85bbcdba28a21f88e6ccfcb06d257fd4435952c93b6287ebdd78a5.json create mode 100644 .sqlx/query-08313d5c002d1355542a40eb94b905f0091229e0433b038fdb55a49e557efa06.json create mode 100644 .sqlx/query-0a6841e9164a94287732afe2375eda68b4d031ea128cae2125f7e526660d8b45.json create mode 100644 .sqlx/query-0e3b37aa1da32a0a5266c3841012fae52ce6394fcdc6c8e42eaeec0d7d6fb663.json create mode 100644 .sqlx/query-12386f2efee70bd2a8e59ccccb8f0d7f4d2742acad26efbc7644193966c39325.json create mode 100644 .sqlx/query-2c329cab6e494ef1a765a7de05e1351358d4ed583a36fe5d0edc38cc12603e44.json create mode 100644 .sqlx/query-a27c937b6039582e0615a1bec479e7b21d2906779bfbbe9e375fa3a70a2de555.json delete mode 100644 .sqlx/query-b78bef6253289702123c03afcab624d5d50df897af4a477a9cf3dac196b52ecc.json create mode 100644 .sqlx/query-be5b70aed19b1d1e81dd26f84faf9caff1712321c6398abed99408a7c6cb31e4.json create mode 100644 .sqlx/query-e406ad05186e0fee570afea0b24a74fe7a4db7bc40477bf963a8f9240d406f13.json create mode 100644 .sqlx/query-fb28c3b48459572d95cf69439c6e0cb3697811ab55fe48f224856ed94d59ce92.json create mode 100644 crates/arr-daemon/src/grab.rs create mode 100644 crates/arr-db/src/policy.rs 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-a27c937b6039582e0615a1bec479e7b21d2906779bfbbe9e375fa3a70a2de555.json b/.sqlx/query-a27c937b6039582e0615a1bec479e7b21d2906779bfbbe9e375fa3a70a2de555.json new file mode 100644 index 0000000..b66af1f --- /dev/null +++ b/.sqlx/query-a27c937b6039582e0615a1bec479e7b21d2906779bfbbe9e375fa3a70a2de555.json @@ -0,0 +1,98 @@ +{ + "db_name": "SQLite", + "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": "overrides!: String", + "ordinal": 0, + "type_info": "Text" + }, + { + "name": "root_id!: i64", + "ordinal": 1, + "type_info": "Integer" + }, + { + "name": "root_kind!: String", + "ordinal": 2, + "type_info": "Text" + }, + { + "name": "root_audience!: String", + "ordinal": 3, + "type_info": "Text" + }, + { + "name": "root_path!: String", + "ordinal": 4, + "type_info": "Text" + }, + { + "name": "policy_id!: i64", + "ordinal": 5, + "type_info": "Integer" + }, + { + "name": "policy_name!: String", + "ordinal": 6, + "type_info": "Text" + }, + { + "name": "required_audio!: String", + "ordinal": 7, + "type_info": "Text" + }, + { + "name": "dub_blacklist!: String", + "ordinal": 8, + "type_info": "Text" + }, + { + "name": "hdr_rules!: String", + "ordinal": 9, + "type_info": "Text" + }, + { + "name": "size_bands!: String", + "ordinal": 10, + "type_info": "Text" + }, + { + "name": "resolution_pref!: String", + "ordinal": 11, + "type_info": "Text" + }, + { + "name": "source_weights!: String", + "ordinal": 12, + "type_info": "Text" + }, + { + "name": "score_weights!: String", + "ordinal": 13, + "type_info": "Text" + } + ], + "parameters": { + "Right": 1 + }, + "nullable": [ + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false + ] + }, + "hash": "a27c937b6039582e0615a1bec479e7b21d2906779bfbbe9e375fa3a70a2de555" +} diff --git a/.sqlx/query-b78bef6253289702123c03afcab624d5d50df897af4a477a9cf3dac196b52ecc.json b/.sqlx/query-b78bef6253289702123c03afcab624d5d50df897af4a477a9cf3dac196b52ecc.json deleted file mode 100644 index 04ae29f..0000000 --- a/.sqlx/query-b78bef6253289702123c03afcab624d5d50df897af4a477a9cf3dac196b52ecc.json +++ /dev/null @@ -1,86 +0,0 @@ -{ - "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\" FROM movies m JOIN roots r ON r.id = m.root_id JOIN policies p ON p.id = r.policy_id WHERE m.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" - }, - { - "name": "overrides!: serde_json::Value", - "ordinal": 3, - "type_info": "Text" - }, - { - "name": "policy_id!: i64", - "ordinal": 4, - "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", - "ordinal": 7, - "type_info": "Text" - }, - { - "name": "hdr_rules!: String", - "ordinal": 8, - "type_info": "Text" - }, - { - "name": "size_bands!: String", - "ordinal": 9, - "type_info": "Text" - }, - { - "name": "resolution_pref!: String", - "ordinal": 10, - "type_info": "Text" - }, - { - "name": "source_weights!: String", - "ordinal": 11, - "type_info": "Text" - } - ], - "parameters": { - "Right": 1 - }, - "nullable": [ - false, - false, - true, - false, - false, - false, - false, - false, - false, - false, - false, - false - ] - }, - "hash": "b78bef6253289702123c03afcab624d5d50df897af4a477a9cf3dac196b52ecc" -} 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 17849bf..2faf3dc 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 36cf87f..f2414aa 100644 --- a/crates/arr-api/src/search.rs +++ b/crates/arr-api/src/search.rs @@ -1,11 +1,8 @@ -use std::collections::BTreeMap; use std::time::UNIX_EPOCH; use arr_core::policy::{evaluate, Candidate}; -use arr_core::{ - DolbyVisionProfile, HdrRules, Language, MovieOverrides, Policy, PolicyId, RequiredAudio, - Resolution, Rule, ScoreWeights, SizeBand, Source, Verdict, -}; +use arr_core::{Language, MovieOverrides, Policy, Resolution, Rule, Source, Verdict}; +use arr_db::policy::language; use arr_indexer::{ProwlarrClient, SearchRelease, SearchRequest}; use axum::extract::{Query, State}; use axum::Json; @@ -70,46 +67,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, -} - -#[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 OverridesJson { - #[serde(default)] - only_4k: bool, - #[serde(default)] - allow_english_audio: bool, -} - #[utoipa::path( get, path = "/api/search", tag = "search", params(SearchQuery), responses( @@ -254,8 +211,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" 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( @@ -281,22 +243,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, - })?; - let overrides: OverridesJson = serde_json::from_value(movie.overrides) - .map_err(|error| ApiError::Database(error.to_string()))?; - let overrides = MovieOverrides { - 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 @@ -399,64 +347,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)?; - 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(), - // The row's score_weights column is not selected here yet; the seeded - // values match these defaults. Wiring the column through is #68. - score_weights: ScoreWeights::default(), - }) -} - fn classify( release: SearchRelease, policy: &Policy, @@ -561,17 +451,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") { @@ -583,27 +462,6 @@ 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 { 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..a7fd016 --- /dev/null +++ b/crates/arr-daemon/src/grab.rs @@ -0,0 +1,950 @@ +//! 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::{SystemTime, UNIX_EPOCH}; + +use arr_core::policy::{evaluate, Candidate}; +use arr_core::{score::score, Language, MovieOverrides, Policy, Rule, 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; + +/// 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, + }, +} + +/// 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, +} + +impl GrabAction { + #[must_use] + pub fn new( + prowlarr: ProwlarrClient, + transmission: TransmissionClient, + download_dir: PathBuf, + seeding: SeedingLimits, + ) -> Self { + Self { + prowlarr, + transmission, + download_dir, + seeding, + } + } + + 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 indexers = self.prowlarr.indexers().await?; + let searchable: Vec = indexers + .iter() + .filter(|indexer| indexer.capabilities.search.available) + .map(|indexer| indexer.id) + .collect(); + 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) + } + + /// 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: &MovieOverrides, + 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 score = score( + policy, + Candidate::PreGrab(&parsed), + release.size.unwrap_or(0), + release.seeders.unwrap_or(0), + ) + .total; + 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()); + } + + #[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 877beed..3f8525e 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..4d88dcc --- /dev/null +++ b/crates/arr-db/src/policy.rs @@ -0,0 +1,370 @@ +//! 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, MovieOverrides, Policy, PolicyId, RequiredAudio, + Resolution, ScoreWeights, SizeBand, Source, +}; +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: MovieOverrides, + 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: MovieOverrides { + 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 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, MovieOverrides::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); + } + + #[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)); -- 2.47.3