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,
},
)]),