From e05437b19827aaceebc77f0e3ed8a64bd0d1fda6 Mon Sep 17 00:00:00 2001
From: Knyffen <knyffen@gmail.com>
Date: Tue, 13 Aug 2024 21:50:17 +0200
Subject: [PATCH] Basic error handling

---
 src/main.rs               | 50 ++++++++++++++++++++++++++++++++-------
 src/my_structs.rs         |  2 ++
 src/my_structs/myerror.rs | 25 ++++++++++++++++++++
 3 files changed, 68 insertions(+), 9 deletions(-)
 create mode 100644 src/my_structs/myerror.rs

diff --git a/src/main.rs b/src/main.rs
index 439e73a..bd519af 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -2,7 +2,7 @@
 
 use axum::{
     extract::State,
-    http::header,
+    http::{header, StatusCode, Uri},
     response::{Html, IntoResponse},
     routing::get,
     Router,
@@ -15,7 +15,9 @@ use std::io::{BufWriter, Cursor};
 use std::sync::Arc;
 mod my_structs;
 mod plotting;
+use crate::my_structs::MyError;
 use crate::plotting::Plotter;
+use std::{thread, time};
 
 #[derive(Clone)]
 struct AppState {
@@ -27,6 +29,9 @@ fn get_sql_pool(url: &str) -> Pool {
         let sql_pool_attempt = Pool::new(url);
         if let Ok(sql_pool) = sql_pool_attempt {
             return sql_pool;
+        } else {
+            println!("Cannot connect to SQL database. Retrying");
+            thread::sleep(time::Duration::from_millis(1000));
         }
     }
 }
@@ -45,26 +50,53 @@ async fn main() {
         .route("/image.png", get(png_image))
         .route("/image.svg", get(svg_image))
         .route("/image2.svg", get(svg_image2))
-        .with_state(shared_state);
+        .with_state(shared_state)
+        .fallback(not_found);
 
     // run our app with hyper, listening globally on port 3000
     let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await.unwrap();
     axum::serve(listener, app).await.unwrap();
 }
 
-async fn get_rows(State(state): State<Arc<AppState>>) -> String {
-    let mut conn = state.sql_pool.get_conn().unwrap();
+async fn not_found(uri: Uri) -> (StatusCode, String) {
+    (StatusCode::NOT_FOUND, format!("No route for {uri}"))
+}
+
+async fn get_nutrition_rows(sql_pool: Pool) -> Result<Vec<my_structs::Nutrition>> {
+    let mut conn = sql_pool
+        .get_conn()
+        .expect("Cannot establish database connection");
 
     let (nutrition_query, nutrition_closure) = my_structs::Nutrition::query_map_helper();
-    let nutrition_val: Vec<my_structs::Nutrition> =
-        conn.query_map(nutrition_query, nutrition_closure).unwrap();
+    let nutrition_val: Vec<my_structs::Nutrition> = conn
+        .query_map(nutrition_query, nutrition_closure)
+        .expect("Data in database doesn't match the Nutrition class");
+
+    Ok(nutrition_val)
+}
+
+async fn get_purchases_rows(sql_pool: Pool) -> Result<Vec<my_structs::Purchase>> {
+    let mut conn = sql_pool
+        .get_conn()
+        .expect("Cannot establish database connection");
 
     let (purchases_query, purchases_closure) = my_structs::Purchase::query_map_helper();
-    let purchases_val: Vec<my_structs::Purchase> =
-        conn.query_map(purchases_query, purchases_closure).unwrap();
+    let purchases_val: Vec<my_structs::Purchase> = conn
+        .query_map(purchases_query, purchases_closure)
+        .expect("Data in database doesn't match the Purchase class");
+
+    Ok(purchases_val)
+}
+
+async fn get_rows(State(state): State<Arc<AppState>>) -> Result<String, MyError> {
+    let nutrition_val_promise = get_nutrition_rows(state.sql_pool.clone());
+    let purchases_val_promise = get_purchases_rows(state.sql_pool.clone());
+
+    let nutrition_val = nutrition_val_promise.await?;
+    let purchases_val = purchases_val_promise.await?;
 
     println!("Serving: get_rows");
-    format!("{:#?} {:#?}", nutrition_val, purchases_val)
+    Ok(format!("{:#?} {:#?}", nutrition_val, purchases_val))
 }
 
 async fn html_demo(State(_state): State<Arc<AppState>>) -> Html<String> {
diff --git a/src/my_structs.rs b/src/my_structs.rs
index d82b968..2713bf9 100644
--- a/src/my_structs.rs
+++ b/src/my_structs.rs
@@ -1,4 +1,6 @@
+pub use self::myerror::MyError;
 pub use self::nutrition::Nutrition;
 pub use self::purchases::Purchase;
+mod myerror;
 mod nutrition;
 mod purchases;
diff --git a/src/my_structs/myerror.rs b/src/my_structs/myerror.rs
new file mode 100644
index 0000000..c1da163
--- /dev/null
+++ b/src/my_structs/myerror.rs
@@ -0,0 +1,25 @@
+use axum::{
+    http::StatusCode,
+    response::{IntoResponse, Response},
+};
+
+pub enum MyError {
+    MySQLError(mysql::Error),
+}
+
+impl IntoResponse for MyError {
+    fn into_response(self) -> Response {
+        let (status, message) = match self {
+            // # TODO: Don't expose error messages once this goes into production
+            MyError::MySQLError(e) => (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()),
+        };
+
+        (status, message).into_response()
+    }
+}
+
+impl From<mysql::Error> for MyError {
+    fn from(error: mysql::Error) -> Self {
+        Self::MySQLError(error)
+    }
+}