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

Merge remote-tracking branch 'origin/main' into issue/35-auto-track

Browse Source
This commit is contained in:
Miguel Palhas
2026-08-22 22:16:03 +01:00
parent de0d533a49 b5413eef46
commit e81d5b38bb
10 changed files with 1236 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/arr-api/src/lib.rs
+5 -1
View File
@@ -9,6 +9,7 @@
mod health;
mod movies;
mod owners;
mod roots;
mod search;
mod state;
@@ -22,6 +23,7 @@ use utoipa_scalar::{Scalar, Servable};
pub use health::{Check, Health, HealthReport, Status};
pub use movies::{Accepted, AttentionQueues, CreateMovie, ErrorBody, Movie, Release, UpdateMovie};
pub use owners::{CreateOwner, Owner, UpdateOwner};
pub use roots::Root;
pub use search::{ClassifiedRelease, SearchResponse};
pub use state::{AppState, MovieCommand, Upstreams, DEFAULT_TMDB_URL};
@@ -45,7 +47,8 @@ pub const DOCS_PATH: &str = "/api/docs";
(name = "system", description = "Service health and metadata"),
(name = "movies", description = "Movie library and actions"),
(name = "owners", description = "Owner tags and filtered views (DESIGN.md §4.3)"),
(name = "search", description = "Unified title and release search")
(name = "search", description = "Unified title and release search"),
(name = "roots", description = "Root folders and their policies")
),
)]
struct ApiDoc;
@@ -66,6 +69,7 @@ fn api_router() -> OpenApiRouter<AppState> {
.routes(routes!(owners::get, owners::update, owners::delete))
.routes(routes!(search::search))
.routes(routes!(search::releases))
.routes(routes!(roots::list))
}
/// The generated `OpenAPI` document.
crates/arr-api/src/roots.rs
+79
View File
@@ -0,0 +1,79 @@
use axum::extract::State;
use axum::Json;
use serde::Serialize;
use utoipa::ToSchema;
use crate::movies::{ApiError, ErrorBody};
use crate::state::AppState;
/// A root folder joined with its policy, enough for the add flow to pre-fill
/// root and policy (DESIGN.md §9.2) without a second request.
#[derive(Debug, Clone, Serialize, ToSchema)]
pub struct Root {
pub id: i64,
pub kind: String,
pub audience: String,
pub path: String,
pub policy_id: i64,
pub policy_name: String,
}
#[utoipa::path(
get, path = "/api/roots", tag = "roots",
responses(
(status = 200, body = [Root]),
(status = 500, body = ErrorBody),
(status = 503, body = ErrorBody)
)
)]
pub async fn list(State(state): State<AppState>) -> Result<Json<Vec<Root>>, ApiError> {
let database = state.database().ok_or(ApiError::Unavailable)?;
let roots = sqlx::query_as!(
Root,
r#"SELECT r.id AS "id!: i64", r.kind AS "kind!: String", r.audience AS "audience!: String", r.path AS "path!: String", p.id AS "policy_id!: i64", p.name AS "policy_name!: String" FROM roots r JOIN policies p ON p.id = r.policy_id ORDER BY r.id"#
)
.fetch_all(database.pool())
.await?;
Ok(Json(roots))
}
#[cfg(test)]
mod tests {
use crate::state::Upstreams;
use crate::{router, AppState};
#[tokio::test]
async fn roots_carry_their_policy_name() {
let dir = tempfile::tempdir().expect("tempdir");
let database = arr_db::Db::connect(dir.path().join("arr.db"))
.await
.expect("database");
database.migrate().await.expect("migrate");
let state = AppState::new(Upstreams::new(
"http://127.0.0.1:1".into(),
"http://127.0.0.1:1".into(),
))
.expect("state")
.with_database(database);
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
.await
.expect("bind");
let address = listener.local_addr().expect("address");
let app = router(state);
tokio::spawn(async move { axum::serve(listener, app).await.expect("serve") });
let roots: Vec<serde_json::Value> = reqwest::get(format!("http://{address}/api/roots"))
.await
.expect("roots")
.json()
.await
.expect("json");
assert_eq!(roots.len(), 2, "the two seeded movie roots: {roots:?}");
assert_eq!(roots[0]["audience"], "main");
assert_eq!(roots[0]["policy_name"], "Movies — main");
assert_eq!(roots[1]["audience"], "kids");
assert!(roots[1]["path"]
.as_str()
.is_some_and(|p| p.contains("kids")));
}
}
crates/arr-api/src/search.rs
+115 -44
View File
@@ -2,6 +2,7 @@ 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,
@@ -80,6 +81,7 @@ struct PolicyRow {
size_bands: String,
resolution_pref: String,
source_weights: String,
score_weights: String,
}
#[derive(Debug, Deserialize)]
@@ -102,6 +104,13 @@ struct SizeBandJson {
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)]
@@ -254,7 +263,7 @@ 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" 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 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)
.fetch_optional(database.pool()).await?.ok_or(ApiError::NotFound)?;
let tmdb = tmdb_client(&state)?
@@ -290,6 +299,7 @@ pub async fn releases(
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()))?;
@@ -405,6 +415,7 @@ fn policy_from_row(row: PolicyRow) -> Result<Policy, ApiError> {
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,
@@ -451,9 +462,11 @@ fn policy_from_row(row: PolicyRow) -> Result<Policy, ApiError> {
.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(),
score_weights: ScoreWeights {
size_at_target: score_weights.size_at_target,
source_tier: score_weights.source_tier,
seeder_doubling: score_weights.seeder_doubling,
},
})
}
@@ -472,7 +485,17 @@ fn classify(
release.size,
);
let (verdict, rule) = verdict(&evaluation.verdict);
let score = score(policy, &parsed, release.size, release.seeders);
let score = score(
policy,
Candidate::PreGrab(&parsed),
release.size.unwrap_or_default(),
release.seeders.unwrap_or_default(),
);
let score = if release.size.is_some() {
score.total
} else {
score.source.saturating_add(score.seeders)
};
Ok(ClassifiedRelease {
indexer_id: release.indexer_id,
guid: release.guid,
@@ -495,44 +518,6 @@ fn classify(
})
}
fn score(
policy: &Policy,
parsed: &arr_parse::NameClaims,
size: Option<u64>,
seeders: Option<u32>,
) -> i64 {
let resolution = parsed.resolution.map(Resolution::from);
let resolution_score = resolution
.and_then(|value| {
policy
.resolution_preference
.iter()
.position(|candidate| *candidate == value)
})
.map_or(0, |position| {
10_000 - i64::try_from(position).unwrap_or(0) * 5_000
});
let size_score = resolution
.and_then(|value| policy.size_bands.get(&value))
.zip(size)
.map_or(0, |(band, bytes)| {
if bytes < band.floor_bytes {
-i64::try_from((band.floor_bytes - bytes) / 100_000_000).unwrap_or(i64::MAX)
} else {
-i64::try_from(bytes.abs_diff(band.target_bytes) / 100_000_000).unwrap_or(i64::MAX)
}
});
let source_score = i64::from(
parsed
.source
.map(Source::from)
.and_then(|source| policy.source_weights.get(&source).copied())
.unwrap_or(0),
) * 5;
let seeder_score = seeders.map_or(0, |count| i64::from(count.saturating_add(1).ilog2()) * 2);
resolution_score + size_score + source_score + seeder_score
}
fn verdict(verdict: &Verdict) -> (&'static str, Option<String>) {
match verdict {
Verdict::Eligible => ("eligible", None),
@@ -729,6 +714,12 @@ mod tests {
let (_dir, state, base) = application(&tmdb, &prowlarr).await;
sqlx::query("INSERT INTO movies (tmdb_id, title, year, original_language, root_id) VALUES (693134, 'Dune Part Two', 2024, 'en', 2)")
.execute(state.database().expect("database").pool()).await.expect("movie");
sqlx::query(
"UPDATE policies SET score_weights = '{\"size_at_target\":0,\"source_tier\":0,\"seeder_doubling\":0}'",
)
.execute(state.database().expect("database").pool())
.await
.expect("score weights");
let response = reqwest::get(format!("{base}/api/releases?movie_id=1"))
.await
@@ -744,7 +735,87 @@ mod tests {
.find(|release| release["verdict"] == "rejected")
.expect("rejected");
assert_eq!(rejected["rule"], "source");
assert!(releases[0]["score"].as_i64().is_some());
let eligible = releases
.iter()
.find(|release| release["guid"] == "good")
.expect("eligible");
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 {
id: PolicyId(1),
name: "test".into(),
required_audio: RequiredAudio::OriginalLanguage,
dub_blacklist: Vec::new(),
hdr_rules: HdrRules {
rejected_dolby_vision_profiles: Vec::new(),
},
size_bands: std::collections::BTreeMap::from([(
Resolution::R2160p,
SizeBand {
floor_bytes: gib(8),
target_bytes: gib(22),
penalty_points_per_gib_over: 60,
},
)]),
resolution_preference: vec![Resolution::R2160p],
source_weights: std::collections::BTreeMap::from([(Source::WebDl, 2)]),
score_weights: ScoreWeights::default(),
};
let release = SearchRelease {
indexer_id: 1,
guid: "release".into(),
name: "Dune.Part.Two.2024.2160p.WEB-DL".into(),
size: None,
seeders: Some(8),
publish_date: None,
download_url: "https://tracker/release".into(),
};
let parsed = arr_parse::parse(&release.name);
let core_score = score(&policy, Candidate::PreGrab(&parsed), 0, 8);
let classified = classify(
release,
&policy,
&MovieOverrides::default(),
&Language::Other("en".into()),
)
.expect("classified release");
assert_eq!(
classified.score,
core_score.source.saturating_add(core_score.seeders)
);
assert_ne!(classified.score, core_score.total);
}
#[test]