feat(daemon): grab pipeline for wanted movies (#77)
This commit was merged in pull request #77.
This commit is contained in:
@@ -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());
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user