yolo/arr
3
0
Fork 0
Code Issues 12 Pull Requests Actions Packages Projects Releases Wiki Activity

feat(daemon): grab pipeline for wanted movies (#77)
ci / web (push) Successful in 53s
Details
ci / rust (push) Successful in 1m47s
Details
e2e / e2e (push) Successful in 2m1s
Details

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