From de2a13060c4536b013e4f235f5aecdd1d5592756 Mon Sep 17 00:00:00 2001 From: "naps62-yolo (agent)" Date: Sat, 22 Aug 2026 21:22:12 +0100 Subject: [PATCH] Search API with classified releases (#65) --- ...568e3e48c0b7a8f593063286c8d8c68159105.json | 86 ++ ...413cc21e26c9cace9a0e2180cca42219cfce3.json | 86 ++ ...624d5d50df897af4a477a9cf3dac196b52ecc.json | 86 ++ ...bf336dfd13920d5439a4a77b92220c62f9526.json | 86 ++ Cargo.lock | 5 + crates/arr-api/Cargo.toml | 5 + crates/arr-api/src/lib.rs | 7 +- crates/arr-api/src/search.rs | 757 ++++++++++++++++++ crates/arr-meta/src/client.rs | 14 +- crates/arr-meta/src/model.rs | 6 + 10 files changed, 1136 insertions(+), 2 deletions(-) create mode 100644 .sqlx/query-139c1c790f1fd1c09ae2edcffb5568e3e48c0b7a8f593063286c8d8c68159105.json create mode 100644 .sqlx/query-8617bd44229736609d860935f22413cc21e26c9cace9a0e2180cca42219cfce3.json create mode 100644 .sqlx/query-b78bef6253289702123c03afcab624d5d50df897af4a477a9cf3dac196b52ecc.json create mode 100644 .sqlx/query-cba9bc8a941c3830d5fdfa3c76bbf336dfd13920d5439a4a77b92220c62f9526.json create mode 100644 crates/arr-api/src/search.rs diff --git a/.sqlx/query-139c1c790f1fd1c09ae2edcffb5568e3e48c0b7a8f593063286c8d8c68159105.json b/.sqlx/query-139c1c790f1fd1c09ae2edcffb5568e3e48c0b7a8f593063286c8d8c68159105.json new file mode 100644 index 0000000..a90671c --- /dev/null +++ b/.sqlx/query-139c1c790f1fd1c09ae2edcffb5568e3e48c0b7a8f593063286c8d8c68159105.json @@ -0,0 +1,86 @@ +{ + "db_name": "SQLite", + "query": "SELECT id AS \"id!: i64\", tmdb_id AS \"tmdb_id!: i64\", title AS \"title!: String\", year, original_language, root_id AS \"root_id!: i64\", wanted AS \"wanted!: bool\", overrides AS \"overrides!: serde_json::Value\", state AS \"state!: String\", blocked AS \"blocked!: bool\", search_attempts AS \"search_attempts!: i64\", last_searched_at FROM movies WHERE title LIKE ? ESCAPE '\\' ORDER BY title, year, id", + "describe": { + "columns": [ + { + "name": "id!: i64", + "ordinal": 0, + "type_info": "Integer" + }, + { + "name": "tmdb_id!: i64", + "ordinal": 1, + "type_info": "Integer" + }, + { + "name": "title!: String", + "ordinal": 2, + "type_info": "Text" + }, + { + "name": "year", + "ordinal": 3, + "type_info": "Integer" + }, + { + "name": "original_language", + "ordinal": 4, + "type_info": "Text" + }, + { + "name": "root_id!: i64", + "ordinal": 5, + "type_info": "Integer" + }, + { + "name": "wanted!: bool", + "ordinal": 6, + "type_info": "Integer" + }, + { + "name": "overrides!: serde_json::Value", + "ordinal": 7, + "type_info": "Text" + }, + { + "name": "state!: String", + "ordinal": 8, + "type_info": "Text" + }, + { + "name": "blocked!: bool", + "ordinal": 9, + "type_info": "Integer" + }, + { + "name": "search_attempts!: i64", + "ordinal": 10, + "type_info": "Integer" + }, + { + "name": "last_searched_at", + "ordinal": 11, + "type_info": "Text" + } + ], + "parameters": { + "Right": 1 + }, + "nullable": [ + false, + false, + false, + true, + true, + false, + false, + false, + false, + false, + false, + true + ] + }, + "hash": "139c1c790f1fd1c09ae2edcffb5568e3e48c0b7a8f593063286c8d8c68159105" +} diff --git a/.sqlx/query-8617bd44229736609d860935f22413cc21e26c9cace9a0e2180cca42219cfce3.json b/.sqlx/query-8617bd44229736609d860935f22413cc21e26c9cace9a0e2180cca42219cfce3.json new file mode 100644 index 0000000..14d3224 --- /dev/null +++ b/.sqlx/query-8617bd44229736609d860935f22413cc21e26c9cace9a0e2180cca42219cfce3.json @@ -0,0 +1,86 @@ +{ + "db_name": "SQLite", + "query": "SELECT id AS \"id!: i64\", tmdb_id AS \"tmdb_id!: i64\", title AS \"title!: String\", year, original_language, root_id AS \"root_id!: i64\", wanted AS \"wanted!: bool\", overrides AS \"overrides!: serde_json::Value\", state AS \"state!: String\", blocked AS \"blocked!: bool\", search_attempts AS \"search_attempts!: i64\", last_searched_at FROM movies WHERE tmdb_id = ? OR title LIKE ? ESCAPE '\\' ORDER BY title, year, id", + "describe": { + "columns": [ + { + "name": "id!: i64", + "ordinal": 0, + "type_info": "Integer" + }, + { + "name": "tmdb_id!: i64", + "ordinal": 1, + "type_info": "Integer" + }, + { + "name": "title!: String", + "ordinal": 2, + "type_info": "Text" + }, + { + "name": "year", + "ordinal": 3, + "type_info": "Integer" + }, + { + "name": "original_language", + "ordinal": 4, + "type_info": "Text" + }, + { + "name": "root_id!: i64", + "ordinal": 5, + "type_info": "Integer" + }, + { + "name": "wanted!: bool", + "ordinal": 6, + "type_info": "Integer" + }, + { + "name": "overrides!: serde_json::Value", + "ordinal": 7, + "type_info": "Text" + }, + { + "name": "state!: String", + "ordinal": 8, + "type_info": "Text" + }, + { + "name": "blocked!: bool", + "ordinal": 9, + "type_info": "Integer" + }, + { + "name": "search_attempts!: i64", + "ordinal": 10, + "type_info": "Integer" + }, + { + "name": "last_searched_at", + "ordinal": 11, + "type_info": "Text" + } + ], + "parameters": { + "Right": 2 + }, + "nullable": [ + false, + false, + false, + true, + true, + false, + false, + false, + false, + false, + false, + true + ] + }, + "hash": "8617bd44229736609d860935f22413cc21e26c9cace9a0e2180cca42219cfce3" +} diff --git a/.sqlx/query-b78bef6253289702123c03afcab624d5d50df897af4a477a9cf3dac196b52ecc.json b/.sqlx/query-b78bef6253289702123c03afcab624d5d50df897af4a477a9cf3dac196b52ecc.json new file mode 100644 index 0000000..04ae29f --- /dev/null +++ b/.sqlx/query-b78bef6253289702123c03afcab624d5d50df897af4a477a9cf3dac196b52ecc.json @@ -0,0 +1,86 @@ +{ + "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-cba9bc8a941c3830d5fdfa3c76bbf336dfd13920d5439a4a77b92220c62f9526.json b/.sqlx/query-cba9bc8a941c3830d5fdfa3c76bbf336dfd13920d5439a4a77b92220c62f9526.json new file mode 100644 index 0000000..a693b15 --- /dev/null +++ b/.sqlx/query-cba9bc8a941c3830d5fdfa3c76bbf336dfd13920d5439a4a77b92220c62f9526.json @@ -0,0 +1,86 @@ +{ + "db_name": "SQLite", + "query": "SELECT id AS \"id!: i64\", tmdb_id AS \"tmdb_id!: i64\", title AS \"title!: String\", year, original_language, root_id AS \"root_id!: i64\", wanted AS \"wanted!: bool\", overrides AS \"overrides!: serde_json::Value\", state AS \"state!: String\", blocked AS \"blocked!: bool\", search_attempts AS \"search_attempts!: i64\", last_searched_at FROM movies WHERE tmdb_id = ?", + "describe": { + "columns": [ + { + "name": "id!: i64", + "ordinal": 0, + "type_info": "Integer" + }, + { + "name": "tmdb_id!: i64", + "ordinal": 1, + "type_info": "Integer" + }, + { + "name": "title!: String", + "ordinal": 2, + "type_info": "Text" + }, + { + "name": "year", + "ordinal": 3, + "type_info": "Integer" + }, + { + "name": "original_language", + "ordinal": 4, + "type_info": "Text" + }, + { + "name": "root_id!: i64", + "ordinal": 5, + "type_info": "Integer" + }, + { + "name": "wanted!: bool", + "ordinal": 6, + "type_info": "Integer" + }, + { + "name": "overrides!: serde_json::Value", + "ordinal": 7, + "type_info": "Text" + }, + { + "name": "state!: String", + "ordinal": 8, + "type_info": "Text" + }, + { + "name": "blocked!: bool", + "ordinal": 9, + "type_info": "Integer" + }, + { + "name": "search_attempts!: i64", + "ordinal": 10, + "type_info": "Integer" + }, + { + "name": "last_searched_at", + "ordinal": 11, + "type_info": "Text" + } + ], + "parameters": { + "Right": 1 + }, + "nullable": [ + true, + false, + false, + true, + true, + false, + false, + false, + false, + false, + false, + true + ] + }, + "hash": "cba9bc8a941c3830d5fdfa3c76bbf336dfd13920d5439a4a77b92220c62f9526" +} diff --git a/Cargo.lock b/Cargo.lock index c1b9880..96373d3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -30,8 +30,13 @@ dependencies = [ name = "arr-api" version = "0.1.0" dependencies = [ + "arr-core", "arr-db", + "arr-indexer", + "arr-meta", + "arr-parse", "axum", + "chrono", "reqwest", "serde", "serde_json", diff --git a/crates/arr-api/Cargo.toml b/crates/arr-api/Cargo.toml index da35873..686ed94 100644 --- a/crates/arr-api/Cargo.toml +++ b/crates/arr-api/Cargo.toml @@ -7,8 +7,13 @@ repository.workspace = true publish = false [dependencies] +arr-core = { workspace = true } arr-db = { workspace = true } +arr-indexer = { workspace = true } +arr-meta = { workspace = true } +arr-parse = { workspace = true } axum = { workspace = true } +chrono = { workspace = true } reqwest = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } diff --git a/crates/arr-api/src/lib.rs b/crates/arr-api/src/lib.rs index f8dcbdf..6650e2f 100644 --- a/crates/arr-api/src/lib.rs +++ b/crates/arr-api/src/lib.rs @@ -8,6 +8,7 @@ mod health; mod movies; +mod search; mod state; use axum::routing::get; @@ -19,6 +20,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 search::{ClassifiedRelease, SearchResponse}; pub use state::{AppState, MovieCommand, Upstreams, DEFAULT_TMDB_URL}; /// Where the generated document is served, and where `just gen-client` reads @@ -39,7 +41,8 @@ pub const DOCS_PATH: &str = "/api/docs"; ), tags( (name = "system", description = "Service health and metadata"), - (name = "movies", description = "Movie library and actions") + (name = "movies", description = "Movie library and actions"), + (name = "search", description = "Unified title and release search") ), )] struct ApiDoc; @@ -54,6 +57,8 @@ fn api_router() -> OpenApiRouter { .routes(routes!(movies::releases)) .routes(routes!(movies::grab)) .routes(routes!(movies::attention)) + .routes(routes!(search::search)) + .routes(routes!(search::releases)) } /// The generated `OpenAPI` document. diff --git a/crates/arr-api/src/search.rs b/crates/arr-api/src/search.rs new file mode 100644 index 0000000..0b3ecad --- /dev/null +++ b/crates/arr-api/src/search.rs @@ -0,0 +1,757 @@ +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, SizeBand, Source, Verdict, +}; +use arr_indexer::{ProwlarrClient, SearchRelease, SearchRequest}; +use axum::extract::{Query, State}; +use axum::Json; +use chrono::{DateTime, Utc}; +use serde::{Deserialize, Serialize}; +use utoipa::{IntoParams, ToSchema}; + +use crate::movies::{ApiError, ErrorBody, Movie}; +use crate::state::AppState; + +#[derive(Debug, Deserialize, IntoParams)] +pub struct SearchQuery { + q: String, +} + +#[derive(Debug, Deserialize, IntoParams)] +pub struct ReleasesQuery { + movie_id: i64, +} + +#[derive(Debug, Clone, Serialize, ToSchema)] +pub struct SearchResponse { + pub kind: SearchInputKind, + pub library: Vec, + pub tmdb: Vec, + pub manual: Option, +} + +#[derive(Debug, Clone, Copy, Serialize, ToSchema)] +#[serde(rename_all = "snake_case")] +pub enum SearchInputKind { + Text, + TmdbId, + ImdbId, + Magnet, + TorrentUrl, +} + +#[derive(Debug, Clone, Serialize, ToSchema)] +pub struct TmdbMovie { + pub tmdb_id: u32, + pub title: String, + pub original_title: String, + pub original_language: String, + pub year: Option, + pub overview: Option, + pub poster_path: Option, +} + +#[derive(Debug, Clone, Serialize, ToSchema)] +pub struct ClassifiedRelease { + pub indexer_id: i64, + pub guid: String, + pub name: String, + pub size: Option, + pub seeders: Option, + pub publish_date: Option, + pub download_url: String, + pub parsed: serde_json::Value, + pub score: i64, + pub verdict: String, + 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_gb: u64, + target_gb: u64, +} + +#[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( + (status = 200, body = SearchResponse), + (status = 422, body = ErrorBody), + (status = 500, body = ErrorBody), + (status = 503, body = ErrorBody) + ) +)] +pub async fn search( + State(state): State, + Query(query): Query, +) -> Result, ApiError> { + let input = query.q.trim(); + if input.is_empty() { + return Err(ApiError::Invalid("q must not be empty".into())); + } + + let kind = input_kind(input); + if matches!(kind, SearchInputKind::Magnet | SearchInputKind::TorrentUrl) { + return Ok(Json(SearchResponse { + kind, + library: Vec::new(), + tmdb: Vec::new(), + manual: Some(input.to_owned()), + })); + } + + let database = state.database().ok_or(ApiError::Unavailable)?; + let mut library = if matches!(kind, SearchInputKind::TmdbId) { + let tmdb_id = input + .strip_prefix("tmdb:") + .unwrap_or(input) + .trim() + .parse::() + .map_err(|_| ApiError::Invalid("invalid TMDB id".into()))?; + let pattern = format!("%{}%", escape_like(input)); + sqlx::query_as!(Movie, r#"SELECT id AS "id!: i64", tmdb_id AS "tmdb_id!: i64", title AS "title!: String", year, original_language, root_id AS "root_id!: i64", wanted AS "wanted!: bool", overrides AS "overrides!: serde_json::Value", state AS "state!: String", blocked AS "blocked!: bool", search_attempts AS "search_attempts!: i64", last_searched_at FROM movies WHERE tmdb_id = ? OR title LIKE ? ESCAPE '\' ORDER BY title, year, id"#, tmdb_id, pattern) + .fetch_all(database.pool()).await? + } else { + let pattern = format!("%{}%", escape_like(input)); + sqlx::query_as!(Movie, r#"SELECT id AS "id!: i64", tmdb_id AS "tmdb_id!: i64", title AS "title!: String", year, original_language, root_id AS "root_id!: i64", wanted AS "wanted!: bool", overrides AS "overrides!: serde_json::Value", state AS "state!: String", blocked AS "blocked!: bool", search_attempts AS "search_attempts!: i64", last_searched_at FROM movies WHERE title LIKE ? ESCAPE '\' ORDER BY title, year, id"#, pattern) + .fetch_all(database.pool()).await? + }; + + let tmdb = tmdb_client(&state)?; + let results = search_tmdb(&tmdb, kind, input).await?; + if matches!(kind, SearchInputKind::ImdbId) { + for result in &results { + let tmdb_id = i64::from(result.tmdb_id); + if let Some(movie) = sqlx::query_as!(Movie, r#"SELECT id AS "id!: i64", tmdb_id AS "tmdb_id!: i64", title AS "title!: String", year, original_language, root_id AS "root_id!: i64", wanted AS "wanted!: bool", overrides AS "overrides!: serde_json::Value", state AS "state!: String", blocked AS "blocked!: bool", search_attempts AS "search_attempts!: i64", last_searched_at FROM movies WHERE tmdb_id = ?"#, tmdb_id) + .fetch_optional(database.pool()).await? + { + library.push(movie); + } + } + } + + Ok(Json(SearchResponse { + kind, + library, + tmdb: results, + manual: None, + })) +} + +async fn search_tmdb( + tmdb: &arr_meta::TmdbClient, + kind: SearchInputKind, + input: &str, +) -> Result, ApiError> { + let results = match kind { + SearchInputKind::TmdbId => { + let id = input + .strip_prefix("tmdb:") + .unwrap_or(input) + .trim() + .parse::() + .map_err(|_| ApiError::Invalid("invalid TMDB id".into()))?; + let mut matches = match tmdb.movie(id).await { + Ok(movie) => { + let year = movie.year(); + vec![TmdbMovie { + tmdb_id: movie.tmdb_id, + title: movie.title, + original_title: movie.original_title, + original_language: movie.original_language, + year, + overview: movie.overview, + poster_path: movie.poster_path, + }] + } + Err(arr_meta::Error::NotFound { .. }) => Vec::new(), + Err(error) => return Err(upstream_error(&error)), + }; + if !input.starts_with("tmdb:") { + for movie in tmdb + .search_movies(input, None) + .await + .map_err(|error| upstream_error(&error))? + .into_iter() + .map(TmdbMovie::from) + { + if !matches + .iter() + .any(|candidate| candidate.tmdb_id == movie.tmdb_id) + { + matches.push(movie); + } + } + } + matches + } + SearchInputKind::ImdbId => match tmdb.find_movie_by_imdb(input).await { + Ok(movies) => movies.into_iter().map(TmdbMovie::from).collect(), + Err(arr_meta::Error::NotFound { .. }) => Vec::new(), + Err(error) => return Err(upstream_error(&error)), + }, + SearchInputKind::Text => tmdb + .search_movies(input, None) + .await + .map_err(|error| upstream_error(&error))? + .into_iter() + .map(TmdbMovie::from) + .collect(), + SearchInputKind::Magnet | SearchInputKind::TorrentUrl => Vec::new(), + }; + Ok(results) +} + +#[utoipa::path( + get, path = "/api/releases", tag = "search", params(ReleasesQuery), + responses( + (status = 200, body = [ClassifiedRelease]), + (status = 404, body = ErrorBody), + (status = 500, body = ErrorBody), + (status = 503, body = ErrorBody) + ) +)] +pub async fn releases( + State(state): State, + 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) + .fetch_optional(database.pool()).await?.ok_or(ApiError::NotFound)?; + + let tmdb = tmdb_client(&state)? + .movie( + u32::try_from(movie.tmdb_id) + .map_err(|_| ApiError::Invalid("movie has invalid TMDB id".into()))?, + ) + .await + .map_err(|error| upstream_error(&error))?; + let request = tmdb.imdb_id.map_or_else( + || SearchRequest::Text { + query: movie.title.clone(), + }, + |imdb_id| SearchRequest::Movie { imdb_id }, + ); + let upstreams = state.upstreams(); + let api_key = upstreams + .prowlarr_api_key + .clone() + .ok_or(ApiError::Unavailable)?; + let prowlarr = ProwlarrClient::new(upstreams.prowlarr_url.clone(), api_key) + .map_err(|_| ApiError::Unavailable)?; + let indexers = prowlarr + .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 original_language = title_language( + movie + .original_language + .as_deref() + .unwrap_or(&tmdb.original_language), + &tmdb.origin_countries, + ); + + let mut classified = Vec::new(); + for indexer in indexers { + let indexer_request = match &request { + SearchRequest::Movie { .. } + if indexer.capabilities.movie.available + && indexer.capabilities.movie.supports_parameter("imdbid") => + { + request.clone() + } + SearchRequest::Movie { .. } if indexer.capabilities.search.available => { + SearchRequest::Text { + query: movie.title.clone(), + } + } + SearchRequest::Text { .. } if indexer.capabilities.search.available => request.clone(), + _ => continue, + }; + match prowlarr.search_indexer(indexer.id, &indexer_request).await { + Ok(releases) => { + for release in releases { + classified.push(classify(release, &policy, &overrides, &original_language)?); + } + } + Err(error) => { + tracing::warn!(indexer_id = indexer.id, %error, "manual release search failed"); + } + } + } + classified.sort_by_key(|release| (bucket(&release.verdict), -release.score)); + Ok(Json(classified)) +} + +impl From for TmdbMovie { + fn from(movie: arr_meta::MovieSearchResult) -> Self { + let year = movie.year(); + Self { + tmdb_id: movie.tmdb_id, + title: movie.title, + original_title: movie.original_title, + original_language: movie.original_language, + year, + overview: movie.overview, + poster_path: movie.poster_path, + } + } +} + +fn input_kind(input: &str) -> SearchInputKind { + let lower = input.to_ascii_lowercase(); + if lower.starts_with("magnet:?") { + SearchInputKind::Magnet + } else if lower.ends_with(".torrent") + && (lower.starts_with("http://") || lower.starts_with("https://")) + { + SearchInputKind::TorrentUrl + } else if lower.starts_with("tt") && lower[2..].chars().all(|c| c.is_ascii_digit()) { + SearchInputKind::ImdbId + } else if lower + .strip_prefix("tmdb:") + .is_some_and(|id| id.trim().chars().all(|c| c.is_ascii_digit())) + || lower.chars().all(|c| c.is_ascii_digit()) + { + SearchInputKind::TmdbId + } else { + SearchInputKind::Text + } +} + +fn escape_like(input: &str) -> String { + input + .replace('\\', "\\\\") + .replace('%', "\\%") + .replace('_', "\\_") +} + +fn tmdb_client(state: &AppState) -> Result { + let upstreams = state.upstreams(); + let key = upstreams + .tmdb_api_key + .clone() + .ok_or(ApiError::Unavailable)?; + arr_meta::TmdbClient::builder(key) + .base_url(&upstreams.tmdb_url) + .build() + .map_err(|_| ApiError::Unavailable) +} + +fn upstream_error(error: &arr_meta::Error) -> ApiError { + match error { + arr_meta::Error::NotFound { .. } => ApiError::NotFound, + _ => ApiError::Unavailable, + } +} + +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: gb(band.floor_gb), + target_bytes: gb(band.target_gb), + }, + ) + }) + }) + .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(), + }) +} + +fn classify( + release: SearchRelease, + policy: &Policy, + overrides: &MovieOverrides, + original_language: &Language, +) -> Result { + let parsed = arr_parse::parse(&release.name); + let evaluation = evaluate( + policy, + overrides, + original_language, + Candidate::PreGrab(&parsed), + ); + let (verdict, rule) = verdict(&evaluation.verdict); + let score = score(policy, &parsed, release.size, release.seeders); + Ok(ClassifiedRelease { + indexer_id: release.indexer_id, + guid: release.guid, + name: release.name, + size: release.size, + seeders: release.seeders, + publish_date: release + .publish_date + .and_then(|time| time.duration_since(UNIX_EPOCH).ok()) + .and_then(|duration| { + DateTime::::from_timestamp(i64::try_from(duration.as_secs()).ok()?, 0) + }) + .map(|date| date.to_rfc3339()), + download_url: release.download_url, + parsed: serde_json::to_value(parsed) + .map_err(|error| ApiError::Database(error.to_string()))?, + score, + verdict: verdict.to_owned(), + rule, + }) +} + +fn score( + policy: &Policy, + parsed: &arr_parse::NameClaims, + size: Option, + seeders: Option, +) -> 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) { + match verdict { + Verdict::Eligible => ("eligible", None), + Verdict::Waived(rule) => ("waived", Some(rule_name(rule))), + 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 bucket(verdict: &str) -> u8 { + match verdict { + "eligible" => 0, + "waived" => 1, + _ => 2, + } +} +fn gb(value: u64) -> u64 { + value.saturating_mul(1_000_000_000) +} +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") { + return Language::PortugueseBrazil; + } + if origin_countries.iter().any(|country| country == "PT") { + return Language::PortuguesePortugal; + } + } + 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 { + use super::*; + use crate::{router, Upstreams}; + use wiremock::matchers::{method, path, query_param}; + use wiremock::{Mock, MockServer, ResponseTemplate}; + + async fn application( + tmdb: &MockServer, + prowlarr: &MockServer, + ) -> (tempfile::TempDir, AppState, String) { + 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(prowlarr.uri(), "http://127.0.0.1:1".into()) + .with_prowlarr_api_key(Some("prowlarr-key".into())) + .with_tmdb_url(tmdb.uri()) + .with_tmdb_api_key(Some("tmdb-key".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.clone()); + tokio::spawn(async move { axum::serve(listener, app).await.expect("serve") }); + (dir, state, format!("http://{address}")) + } + + #[tokio::test] + async fn unified_search_groups_library_before_tmdb() { + let tmdb = MockServer::start().await; + let prowlarr = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/search/movie")) + .and(query_param("query", "Dune")) + .respond_with( + ResponseTemplate::new(200).set_body_json(serde_json::json!({"results":[{ + "id": 693_134, "title": "Dune: Part Two", "original_title": "Dune: Part Two", + "original_language": "en", "release_date": "2024-02-27" + }]})), + ) + .mount(&tmdb) + .await; + let (_dir, state, base) = application(&tmdb, &prowlarr).await; + sqlx::query("INSERT INTO movies (tmdb_id, title, year, original_language, root_id) VALUES (438631, 'Dune', 2021, 'en', 2)") + .execute(state.database().expect("database").pool()).await.expect("movie"); + + let response: serde_json::Value = reqwest::get(format!("{base}/api/search?q=Dune")) + .await + .expect("search") + .json() + .await + .expect("json"); + assert_eq!(response["kind"], "text"); + assert_eq!(response["library"][0]["title"], "Dune"); + assert_eq!(response["tmdb"][0]["tmdb_id"], 693_134); + } + + #[tokio::test] + async fn numeric_titles_survive_a_missing_tmdb_id() { + let tmdb = MockServer::start().await; + let prowlarr = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/movie/1917")) + .respond_with(ResponseTemplate::new(404)) + .mount(&tmdb) + .await; + Mock::given(method("GET")) + .and(path("/search/movie")) + .and(query_param("query", "1917")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "results": [{"id": 530_915, "title": "1917", "original_title": "1917", + "original_language": "en", "release_date": "2019-12-25"}] + }))) + .mount(&tmdb) + .await; + let (_dir, state, base) = application(&tmdb, &prowlarr).await; + sqlx::query("INSERT INTO movies (tmdb_id, title, year, original_language, root_id) VALUES (530915, '1917', 2019, 'en', 2)") + .execute(state.database().expect("database").pool()).await.expect("movie"); + + let response: serde_json::Value = reqwest::get(format!("{base}/api/search?q=1917")) + .await + .expect("search") + .json() + .await + .expect("json"); + assert_eq!(response["library"][0]["title"], "1917"); + assert_eq!(response["tmdb"][0]["title"], "1917"); + } + + #[tokio::test] + async fn manual_releases_are_classified_and_name_rejection_rules() { + let tmdb = MockServer::start().await; + let prowlarr = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/movie/693134")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "id":693_134, "imdb_id":"tt15239678", "title":"Dune Part Two", + "original_title":"Dune Part Two", "original_language":"en", "status":"Released" + }))) + .mount(&tmdb) + .await; + Mock::given(method("GET")) + .and(path("/api/v1/indexer")) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(serde_json::json!([{"id":7,"name":"tracker","enable":true}])), + ) + .mount(&prowlarr) + .await; + Mock::given(method("GET")).and(path("/7/api")).and(query_param("t", "caps")) + .respond_with(ResponseTemplate::new(200).set_body_string("")) + .mount(&prowlarr).await; + Mock::given(method("GET")).and(path("/7/api")).and(query_param("t", "movie")) + .respond_with(ResponseTemplate::new(200).set_body_string(r#"Dune.Part.Two.2024.1080p.CAMbadhttps://tracker/bad4000000000Dune.Part.Two.2024.2160p.WEB-DLgoodhttps://tracker/good22000000000"#)) + .mount(&prowlarr).await; + 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"); + + let response = reqwest::get(format!("{base}/api/releases?movie_id=1")) + .await + .expect("releases"); + assert_eq!(response.status(), 200); + let releases: Vec = response.json().await.expect("json"); + assert_eq!(releases.len(), 2); + assert!(releases + .iter() + .all(|release| release["verdict"].is_string())); + let rejected = releases + .iter() + .find(|release| release["verdict"] == "rejected") + .expect("rejected"); + assert_eq!(rejected["rule"], "source"); + assert!(releases[0]["score"].as_i64().is_some()); + } + + #[test] + fn manual_inputs_skip_title_upstreams() { + assert!(matches!( + input_kind("magnet:?xt=urn:btih:abc"), + SearchInputKind::Magnet + )); + assert!(matches!( + input_kind("https://example.test/a.torrent"), + SearchInputKind::TorrentUrl + )); + assert!(matches!(input_kind("tt15239678"), SearchInputKind::ImdbId)); + assert!(matches!(input_kind("tmdb:693134"), SearchInputKind::TmdbId)); + } +} diff --git a/crates/arr-meta/src/client.rs b/crates/arr-meta/src/client.rs index 93550e7..a4544e2 100644 --- a/crates/arr-meta/src/client.rs +++ b/crates/arr-meta/src/client.rs @@ -8,7 +8,7 @@ use serde::de::DeserializeOwned; use crate::cache::Cache; use crate::error::{Error, Result}; -use crate::model::{Movie, MovieSearchResult, RawMovie, RawSearchPage}; +use crate::model::{Movie, MovieSearchResult, RawFindPage, RawMovie, RawSearchPage}; /// TMDB's v3 API root. pub const DEFAULT_BASE_URL: &str = "https://api.themoviedb.org/3/"; @@ -89,6 +89,18 @@ impl TmdbClient { Ok(page.results.into_iter().map(Into::into).collect()) } + /// Resolve an `IMDb` title id through TMDB's external-id index. + /// + /// # Errors + /// + /// Any of [`Error`]; see its variants for what callers should distinguish. + pub async fn find_movie_by_imdb(&self, imdb_id: &str) -> Result> { + let path = format!("find/{}", imdb_id.trim()); + let params = [("external_source", "imdb_id".to_owned())]; + let page: RawFindPage = self.get_json(&path, ¶ms).await?; + Ok(page.movie_results.into_iter().map(Into::into).collect()) + } + /// Full detail for one movie, including the digital release date. /// /// One HTTP call: release dates come back appended to the same response diff --git a/crates/arr-meta/src/model.rs b/crates/arr-meta/src/model.rs index a1d8cd9..1df921c 100644 --- a/crates/arr-meta/src/model.rs +++ b/crates/arr-meta/src/model.rs @@ -107,6 +107,12 @@ pub(crate) struct RawSearchPage { pub(crate) results: Vec, } +#[derive(Debug, Deserialize)] +pub(crate) struct RawFindPage { + #[serde(default)] + pub(crate) movie_results: Vec, +} + #[derive(Debug, Deserialize)] pub(crate) struct RawSearchResult { id: u32,