diff --git a/src/handlers.rs b/src/handlers.rs index ca64971..9406694 100644 --- a/src/handlers.rs +++ b/src/handlers.rs @@ -301,6 +301,7 @@ pub async fn search_file( submission.url, submission.filename, submission.file_id, + submission.rating, artist.name, hashes.id hash_id FROM @@ -321,6 +322,7 @@ pub async fn search_file( submission.url, submission.filename, submission.file_id, + submission.rating, artist.name, hashes.id hash_id FROM @@ -341,6 +343,7 @@ pub async fn search_file( submission.url, submission.filename, submission.file_id, + submission.rating, artist.name, hashes.id hash_id FROM @@ -374,6 +377,7 @@ pub async fn search_file( file_id: row.get("file_id"), })), searched_hash: None, + rating: row.get::("rating").parse().ok(), }) .fetch_all(&db) .await; diff --git a/src/models.rs b/src/models.rs index e240c9e..2857ca9 100644 --- a/src/models.rs +++ b/src/models.rs @@ -117,7 +117,16 @@ pub fn image_query_sync( END file_id, CASE WHEN e621_id IS NOT NULL THEN ARRAY(SELECT jsonb_array_elements_text(e.data->'sources')) - END sources + END sources, + CASE + WHEN furaffinity_id IS NOT NULL THEN (f.rating) + WHEN e621_id IS NOT NULL THEN (e.data->>'rating') + WHEN twitter_id IS NOT NULL THEN + CASE + WHEN (tw.data->'possibly_sensitive')::boolean IS true THEN 'adult' + WHEN (tw.data->'possibly_sensitive')::boolean IS false THEN 'general' + END + END rating FROM hashes LEFT JOIN LATERAL ( @@ -169,6 +178,7 @@ pub fn image_query_sync( id: row.id, site_id, site_info, + rating: row.rating.and_then(|rating| rating.parse().ok()), site_id_str: site_id.to_string(), url: row.url.unwrap_or_default(), hash: Some(row.hash), diff --git a/src/types.rs b/src/types.rs index a5bd185..3586539 100644 --- a/src/types.rs +++ b/src/types.rs @@ -23,6 +23,29 @@ pub enum RateLimit { Available((i16, i16)), } +#[derive(Debug, Serialize)] +#[serde(rename_all = "lowercase")] +pub enum Rating { + General, + Mature, + Adult, +} + +impl std::str::FromStr for Rating { + type Err = &'static str; + + fn from_str(s: &str) -> Result { + let rating = match s { + "g" | "s" | "general" => Self::General, + "m" | "q" | "mature" => Self::Mature, + "a" | "e" | "adult" => Self::Adult, + _ => return Err("unknown rating"), + }; + + Ok(rating) + } +} + /// A general type for every file. #[derive(Debug, Default, Serialize)] pub struct File { @@ -39,6 +62,8 @@ pub struct File { #[serde(flatten)] pub site_info: Option, + pub rating: Option, + #[serde(skip_serializing_if = "Option::is_none")] pub hash: Option, #[serde(skip_serializing_if = "Option::is_none")]