From 1933c9ea3d4530b4cb04c7393852163364731fe8 Mon Sep 17 00:00:00 2001
From: Paul Makles <me@insrt.uk>
Date: Sun, 23 Apr 2023 14:59:44 +0100
Subject: [PATCH] feat(core/database): add Reference struct

---
 Cargo.lock                                 |  2 +
 crates/core/database/Cargo.toml            |  5 +++
 crates/core/database/src/lib.rs            |  1 +
 crates/core/database/src/util/mod.rs       |  1 +
 crates/core/database/src/util/reference.rs | 52 ++++++++++++++++++++++
 crates/delta/Cargo.toml                    |  2 +-
 crates/delta/src/routes/bots/fetch.rs      | 10 ++---
 7 files changed, 67 insertions(+), 6 deletions(-)
 create mode 100644 crates/core/database/src/util/mod.rs
 create mode 100644 crates/core/database/src/util/reference.rs

diff --git a/Cargo.lock b/Cargo.lock
index 5129b4cb..bc2cba4c 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2828,6 +2828,8 @@ dependencies = [
  "nanoid",
  "revolt-result",
  "revolt_optional_struct",
+ "rocket",
+ "schemars",
  "serde",
  "serde_json",
 ]
diff --git a/crates/core/database/Cargo.toml b/crates/core/database/Cargo.toml
index 336c01cf..ec998c16 100644
--- a/crates/core/database/Cargo.toml
+++ b/crates/core/database/Cargo.toml
@@ -14,6 +14,7 @@ mongodb = [ "dep:mongodb" ]
 
 # ... Other
 async-std-runtime = [ "async-std" ]
+rocket-impl = [ "rocket", "schemars" ]
 
 # Default Features
 default = [ "mongodb", "async-std-runtime" ]
@@ -42,5 +43,9 @@ async-recursion = "1.0.4"
 # Async
 async-std = { version = "1.8.0", features = ["attributes"], optional = true }
 
+# Rocket Impl
+schemars = { version = "0.8.8", optional = true }
+rocket = { version = "0.5.0-rc.2", default-features = false, features = ["json"], optional = true }
+
 # Authifier
 authifier = { version = "1.0" }
diff --git a/crates/core/database/src/lib.rs b/crates/core/database/src/lib.rs
index 4796b684..5a5c999b 100644
--- a/crates/core/database/src/lib.rs
+++ b/crates/core/database/src/lib.rs
@@ -73,6 +73,7 @@ macro_rules! database_test {
 }
 
 mod models;
+pub mod util;
 pub use models::*;
 
 /// Utility function to check if a boolean value is false
diff --git a/crates/core/database/src/util/mod.rs b/crates/core/database/src/util/mod.rs
new file mode 100644
index 00000000..7b6d3fe1
--- /dev/null
+++ b/crates/core/database/src/util/mod.rs
@@ -0,0 +1 @@
+pub mod reference;
diff --git a/crates/core/database/src/util/reference.rs b/crates/core/database/src/util/reference.rs
new file mode 100644
index 00000000..5a6149a3
--- /dev/null
+++ b/crates/core/database/src/util/reference.rs
@@ -0,0 +1,52 @@
+use revolt_result::Result;
+#[cfg(feature = "rocket-impl")]
+use rocket::request::FromParam;
+#[cfg(feature = "rocket-impl")]
+use schemars::{
+    schema::{InstanceType, Schema, SchemaObject, SingleOrVec},
+    JsonSchema,
+};
+
+use crate::{Bot, Database};
+
+/// Reference to some object in the database
+#[derive(Serialize, Deserialize)]
+pub struct Reference {
+    /// Id of object
+    pub id: String,
+}
+
+impl Reference {
+    /// Create a Ref from an unchecked string
+    pub fn from_unchecked(id: String) -> Reference {
+        Reference { id }
+    }
+
+    /// Fetch bot from Ref
+    pub async fn as_bot(&self, db: &Database) -> Result<Bot> {
+        db.fetch_bot(&self.id).await
+    }
+}
+
+#[cfg(feature = "rocket-impl")]
+impl<'r> FromParam<'r> for Reference {
+    type Error = &'r str;
+
+    fn from_param(param: &'r str) -> Result<Self, Self::Error> {
+        Ok(Reference::from_unchecked(param.into()))
+    }
+}
+
+#[cfg(feature = "rocket-impl")]
+impl JsonSchema for Reference {
+    fn schema_name() -> String {
+        "Id".to_string()
+    }
+
+    fn json_schema(_gen: &mut schemars::gen::SchemaGenerator) -> Schema {
+        Schema::Object(SchemaObject {
+            instance_type: Some(SingleOrVec::Single(Box::new(InstanceType::String))),
+            ..Default::default()
+        })
+    }
+}
diff --git a/crates/delta/Cargo.toml b/crates/delta/Cargo.toml
index b4c7cde6..c85f4a70 100644
--- a/crates/delta/Cargo.toml
+++ b/crates/delta/Cargo.toml
@@ -56,7 +56,7 @@ revolt_rocket_okapi = { version = "0.9.1", features = [ "swagger" ] }
 revolt-quark = { path = "../quark" }
 
 # core
-revolt-database = { path = "../core/database" }
+revolt-database = { path = "../core/database", features = [ "rocket-impl" ] }
 revolt-models = { path = "../core/models", features = [ "schemas" ] }
 
 [build-dependencies]
diff --git a/crates/delta/src/routes/bots/fetch.rs b/crates/delta/src/routes/bots/fetch.rs
index 19a16ee0..ef4cb214 100644
--- a/crates/delta/src/routes/bots/fetch.rs
+++ b/crates/delta/src/routes/bots/fetch.rs
@@ -1,6 +1,6 @@
-use revolt_database::Database;
+use revolt_database::{util::reference::Reference, Database};
 use revolt_models::Bot;
-use revolt_quark::{models::User, Db, Error, Ref, Result};
+use revolt_quark::{models::User, Db, Error, Result};
 use rocket::{serde::json::Json, State};
 use serde::Serialize;
 
@@ -18,18 +18,18 @@ pub struct BotResponse {
 ///
 /// Fetch details of a bot you own by its id.
 #[openapi(tag = "Bots")]
-#[get("/<target>")]
+#[get("/<bot>")]
 pub async fn fetch_bot(
     legacy_db: &Db,
     db: &State<Database>,
     user: User,
-    target: Ref,
+    bot: Reference,
 ) -> Result<Json<BotResponse>> {
     if user.bot.is_some() {
         return Err(Error::IsBot);
     }
 
-    let bot = db.fetch_bot(&target.id).await.map_err(Error::from_core)?;
+    let bot = bot.as_bot(db).await.map_err(Error::from_core)?;
     if bot.owner != user.id {
         return Err(Error::NotFound);
     }