From 0915c0f653f19f4dbeec6f2916eac149d0eb43bd Mon Sep 17 00:00:00 2001
From: "Jakob L. Kreuze" <zerodaysfordays@sdf.org>
Date: Tue, 26 May 2020 17:51:12 -0400
Subject: Refactor shared state for handlers using the database.

---
 src/main.rs | 62 ++++++++++++++++++++++++++++++++++++-------------------------
 1 file changed, 37 insertions(+), 25 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 4294e06..f82f07b 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -22,17 +22,17 @@ extern crate anyhow;
 extern crate image;
 #[macro_use]
 extern crate rocket;
-#[macro_use]
 extern crate rocket_contrib;
 #[macro_use]
 extern crate serde_derive;
 
 mod database;
 
+use anyhow::Result;
 use database::{hash_file, Query, TagDatabase};
-use image::{GenericImageView, ImageFormat};
+use image::ImageFormat;
 use rocket::{Data, State};
-use rocket_contrib::json::{Json, JsonValue};
+use rocket_contrib::json::Json;
 use std::env;
 use std::fs;
 use std::fs::File;
@@ -40,8 +40,26 @@ use std::os::unix::fs::symlink;
 use std::path::Path;
 use std::sync::Mutex;
 
+type SiteState<'a> = State<'a, Mutex<DatabaseConnection>>;
+
+struct DatabaseConnection {
+    tb: TagDatabase,
+    image_dir: String,
+}
+
+impl DatabaseConnection {
+    fn new() -> Result<Self> {
+        let bin = env::current_exe()?.canonicalize()?;
+        let bin_dir = bin.parent().unwrap();
+        let image_dir = String::from(bin_dir.join("image/").to_str().unwrap());
+        let tb_path = bin_dir.join("birka.db");
+        let tb = TagDatabase::new(tb_path.to_str().unwrap())?;
+        Ok(DatabaseConnection { tb, image_dir })
+    }
+}
+
 #[get("/")]
-fn index(tb: State<Mutex<TagDatabase>>) -> &'static str {
+fn index(tb: SiteState) -> &'static str {
     "Hello, world!"
 }
 
@@ -52,8 +70,8 @@ struct ImageResult {
 }
 
 #[get("/posts?<tags>")]
-fn get_posts(tb: State<Mutex<TagDatabase>>, tags: String) -> Json<Vec<ImageResult>> {
-    let tb = tb.inner().lock().unwrap();
+fn get_posts(conn: SiteState, tags: String) -> Json<Vec<ImageResult>> {
+    let tb = &conn.inner().lock().unwrap().tb;
 
     let tags = tags.split(",").collect::<Vec<&str>>();
     let results = tb
@@ -76,29 +94,24 @@ fn get_posts(tb: State<Mutex<TagDatabase>>, tags: String) -> Json<Vec<ImageResul
 }
 
 #[post("/posts/<filename>", data = "<data>")]
-fn put_posts(tb: State<Mutex<TagDatabase>>, filename: String, data: Data) -> Json<&str> {
-    let tb = tb.inner().lock().unwrap();
-
-    let bin = env::current_exe().unwrap();
-    let path = bin.parent().unwrap().join("image/").join(filename);
+fn put_posts(conn: SiteState, filename: String, data: Data) -> Json<&str> {
+    let conn = conn.inner().lock().unwrap();
+    let path = Path::new(&conn.image_dir).join(filename);
     data.stream_to_file(&path).unwrap();
-    tb.import_image(&path, &hash_file(&path).unwrap()).unwrap();
+    conn.tb
+        .import_image(&path, &hash_file(&path).unwrap())
+        .unwrap();
 
     Json("uploaded!")
 }
 
 /// Maybe initialize the directory for storing image symlinks and thumbnails.
-fn create_image_directory() -> TagDatabase {
-    let bin = env::current_exe().unwrap();
-    let image_dir = bin.parent().unwrap().join("image/");
-    fs::create_dir(&image_dir).ok();
+fn create_image_directory(conn: &DatabaseConnection) {
+    fs::create_dir(&conn.image_dir).ok();
 
-    let tb_path = bin.parent().unwrap().join("birka.db");
-    let tb = TagDatabase::new(tb_path.to_str().unwrap()).unwrap();
-
-    for image in tb.images_by_tags(&vec![][..]).unwrap() {
+    for image in conn.tb.images_by_tags(&vec![][..]).unwrap() {
         let src = format!("{}/{}", image.orig_dir, image.filename);
-        let dst = format!("{}/{}", image_dir.to_str().unwrap(), image.filename);
+        let dst = format!("{}/{}", conn.image_dir, image.filename);
         symlink(&src, &dst).ok();
 
         let stem = Path::new(&dst).file_stem().unwrap().to_str().unwrap();
@@ -113,14 +126,13 @@ fn create_image_directory() -> TagDatabase {
                 .unwrap();
         }
     }
-
-    tb
 }
 
 fn main() {
-    let tb = create_image_directory();
+    let conn = DatabaseConnection::new().unwrap();
+    create_image_directory(&conn);
     rocket::ignite()
-        .manage(Mutex::new(tb))
+        .manage(Mutex::new(conn))
         .mount("/", routes![index, get_posts])
         .launch();
 }
-- 
cgit v1.3