From 7e92d8c5be6985d497f30cb6df6e5e2e39cf693c Mon Sep 17 00:00:00 2001
From: BurnyLlama <burnyllama@protonmail.com>
Date: Sun, 6 Nov 2022 15:31:31 +0100
Subject: [PATCH] Refactor.

---
 src/{types/mod.rs => database/map.rs} |  2 +-
 src/database/mod.rs                   | 52 ++++++++++++++++++++-------
 src/main.rs                           | 19 ++++++----
 3 files changed, 54 insertions(+), 19 deletions(-)
 rename src/{types/mod.rs => database/map.rs} (94%)

diff --git a/src/types/mod.rs b/src/database/map.rs
similarity index 94%
rename from src/types/mod.rs
rename to src/database/map.rs
index 0633db2..e4df32c 100644
--- a/src/types/mod.rs
+++ b/src/database/map.rs
@@ -9,7 +9,7 @@ pub struct Map {
     pub points: i8,
     pub stars: i8,
     pub mapper: String,
-    pub timestamp: NaiveDateTime,
+    pub timestamp: NaiveDateTime
 }
 
 /*
diff --git a/src/database/mod.rs b/src/database/mod.rs
index 4e2e18e..3563173 100644
--- a/src/database/mod.rs
+++ b/src/database/mod.rs
@@ -4,23 +4,37 @@ use std::error::Error;
 use mysql::*;
 use mysql::prelude::*;
 
-use crate::types;
+pub mod map;
 
-pub fn get_maps() -> Result<Vec<types::Map>, Box<dyn Error>> {
-    let mut conn = create_pool()?.get_conn()?;
+pub struct DatabasePoolStore(pub Pool);
+
+pub fn get_maps(pool: &Pool) -> Result<Vec<map::Map>, Box<dyn Error>> {
+    let mut conn = pool.get_conn()?;
 
     let maps = conn.query_map(
         "SELECT * FROM record_maps", 
-        |(name, server, points, stars, mapper, timestamp)| types::Map {
-            name, server, points, stars, mapper, timestamp
+        |(
+            name,
+            server,
+            points,
+            stars,
+            mapper,
+            timestamp,
+        )| map::Map {
+            name,
+            server,
+            points,
+            stars,
+            mapper,
+            timestamp
         }
     )?;
     
     Ok(maps)
 }
 
-pub fn get_map_by_name(name: &str) -> Result<Option<types::Map>, Box<dyn Error>> {
-    let mut conn = create_pool()?.get_conn()?;
+pub fn get_map_by_name(pool: &Pool, name: &str) -> Result<Option<map::Map>, Box<dyn Error>> {
+    let mut conn = pool.get_conn()?;
     let stmt = conn.prep("SELECT * FROM record_maps WHERE Map = ?")?;
 
     let map = conn.exec_first(
@@ -29,7 +43,7 @@ pub fn get_map_by_name(name: &str) -> Result<Option<types::Map>, Box<dyn Error>>
     )?;
 
     Ok(match map {
-        Some((name, server, points, stars, mapper, timestamp)) => Some(types::Map {
+        Some((name, server, points, stars, mapper, timestamp)) => Some(map::Map {
             name, server, points, stars, mapper, timestamp
         }),
         None => None
@@ -38,8 +52,22 @@ pub fn get_map_by_name(name: &str) -> Result<Option<types::Map>, Box<dyn Error>>
 
 
 
-fn create_pool() -> Result<Pool, Box<dyn Error>> {
-    let url = env::var("DB_URI")?;
-    let pool = Pool::new(url.as_str())?;
-    Ok(pool)
+pub fn create_pool() -> Pool {
+    let url = match env::var("DB_URI") {
+        Ok(uri) => uri,
+        Err(err) => {
+            println!("You must provide an env var: 'DB_URI'!");
+            panic!("{}", err);
+        }
+    };
+
+    let pool = match Pool::new(url.as_str()) {
+        Ok(pool) => pool,
+        Err(err) => {
+            println!("Couldn't connect to the database!");
+            panic!("{}", err);
+        }
+    };
+
+    pool
 }
\ No newline at end of file
diff --git a/src/main.rs b/src/main.rs
index 1c05bed..18d6948 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,15 +1,19 @@
 #[macro_use] extern crate rocket;
+
+use database::create_pool;
 use rocket::serde::json::Json;
+use  rocket::State;
 
 extern crate dotenv;
 use dotenv::dotenv;
 
-mod types;
 mod database;
+use database::DatabasePoolStore;
+use database::map::Map;
 
 #[get("/maps")]
-fn get_all_maps() -> Option<Json<Vec<types::Map>>> {
-    match database::get_maps() {
+fn get_all_maps(db_pool: &State<DatabasePoolStore>) -> Option<Json<Vec<Map>>> {
+    match database::get_maps(&db_pool.0) {
         Ok(maps) => Some(Json(maps)),
         Err(err) => {
             println!("{err}");
@@ -19,8 +23,8 @@ fn get_all_maps() -> Option<Json<Vec<types::Map>>> {
 }
 
 #[get("/maps/<map>")]
-fn get_map_by_name(map: &str) -> Option<Json<types::Map>> {
-    match database::get_map_by_name(map) {
+fn get_map_by_name(db_pool: &State<DatabasePoolStore>, map: &str) -> Option<Json<Map>> {
+    match database::get_map_by_name(&db_pool.0, map) {
         Ok(map) => match map {
             Some(map) => Some(Json(map)),
             None => None
@@ -32,10 +36,13 @@ fn get_map_by_name(map: &str) -> Option<Json<types::Map>> {
     }
 }
 
-
 #[launch]
 fn rocket() -> _ {
     dotenv().ok();
+
+    let db_pool = create_pool();
+
     rocket::build()
+        .manage(DatabasePoolStore(db_pool))
         .mount("/", routes![get_all_maps, get_map_by_name])
 }
\ No newline at end of file