From 274d54573485f433f21d5a31cf2a627c113b630d Mon Sep 17 00:00:00 2001
From: Syfaro <syfaro@huefox.com>
Date: Fri, 19 Feb 2021 15:38:07 -0500
Subject: [PATCH] Include content rating.

---
 src/handlers.rs |  4 ++++
 src/models.rs   | 12 +++++++++++-
 src/types.rs    | 25 +++++++++++++++++++++++++
 3 files changed, 40 insertions(+), 1 deletion(-)

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::<String, _>("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<Self, Self::Err> {
+        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<SiteInfo>,
 
+    pub rating: Option<Rating>,
+
     #[serde(skip_serializing_if = "Option::is_none")]
     pub hash: Option<i64>,
     #[serde(skip_serializing_if = "Option::is_none")]