From 9d492f6f9195ae80468c668b8dfd873abb48eb66 Mon Sep 17 00:00:00 2001
From: Luke Harding <luke@lukeh990.io>
Date: Fri, 3 May 2024 12:05:46 -0400
Subject: [PATCH] Add database_url to app state

---
 server/src/api/mod.rs | 10 +++++++---
 server/src/db/mod.rs  | 10 ++++------
 server/src/main.rs    | 30 ++++++++++++++++++++++++------
 3 files changed, 35 insertions(+), 15 deletions(-)

diff --git a/server/src/api/mod.rs b/server/src/api/mod.rs
index 83248b0..b47657d 100644
--- a/server/src/api/mod.rs
+++ b/server/src/api/mod.rs
@@ -1,13 +1,17 @@
 // SPDX-License-Identifier: GPL-3.0-Only
 // Copyright (C) 2024 Luke Harding
 
-use actix_web::{get, HttpResponse, Responder};
+use actix_web::{get, HttpResponse, Responder, web};
 
 use crate::db;
 
+pub struct AppState {
+    pub database_url: String,
+}
+
 #[get("/get_tasks")]
-pub async fn get_tasks() -> impl Responder {
-    let tasks = match db::get_tasks().await {
+pub async fn get_tasks(data: web::Data<AppState>) -> impl Responder {
+    let tasks = match db::get_tasks(data.database_url.clone()).await {
         Ok(tasks) => tasks,
         Err(e) => return HttpResponse::InternalServerError().body(e.to_string()),
     };
diff --git a/server/src/db/mod.rs b/server/src/db/mod.rs
index 79e245e..f9804f1 100644
--- a/server/src/db/mod.rs
+++ b/server/src/db/mod.rs
@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: GPL-3.0-Only
 // Copyright (C) 2024 Luke Harding
 
-use std::{env, error};
+use std::error;
 
 use diesel::{Connection, PgConnection, QueryDsl, RunQueryDsl, SelectableHelper};
 use diesel::pg::Pg;
@@ -13,9 +13,7 @@ pub mod models;
 
 pub const MIGRATIONS: EmbeddedMigrations = embed_migrations!("./migrations");
 
-pub fn establish_connection() -> Result<PgConnection, Box<dyn error::Error>> {
-    let database_url = env::var("DATABASE_URL")?;
-
+pub fn establish_connection(database_url: &String) -> Result<PgConnection, Box<dyn error::Error>> {
     Ok(PgConnection::establish(&database_url)?)
 }
 
@@ -27,10 +25,10 @@ pub fn run_migrations(
     Ok(())
 }
 
-pub async fn get_tasks() -> Result<Vec<Task>, Box<dyn error::Error>> {
+pub async fn get_tasks(database_url: String) -> Result<Vec<Task>, Box<dyn error::Error>> {
     use crate::schema::tasks::dsl::*;
 
-    let conn = &mut establish_connection()?;
+    let conn = &mut establish_connection(&database_url)?;
 
     Ok(tasks.select(Task::as_select()).load(conn)?)
 }
diff --git a/server/src/main.rs b/server/src/main.rs
index 1162c84..76c90ca 100644
--- a/server/src/main.rs
+++ b/server/src/main.rs
@@ -3,7 +3,9 @@
 
 use std::{env, io};
 
-use actix_web::{App, HttpServer};
+use actix_web::{App, HttpServer, web};
+
+use crate::api::AppState;
 
 mod api;
 pub mod db;
@@ -17,6 +19,7 @@ async fn main() -> io::Result<()> {
     // Default Server Values
     let mut port = 8000;
     let mut bind_addr = String::from("127.0.0.1");
+    let mut database_url = String::new();
 
     util::notice_println("Loading Environment Variables");
 
@@ -40,8 +43,17 @@ async fn main() -> io::Result<()> {
         }
     }
 
+    if let Ok(env_database) = env::var("DATABASE_URL") {
+        if !env_database.is_empty() {
+            database_url = env_database;
+        } else {
+            util::err_println("No DATABASE_URL found. Aborting,");
+            return Ok(());
+        }
+    }
+
     util::notice_println("Connecting to DB.");
-    let mut conn = match db::establish_connection() {
+    let mut conn = match db::establish_connection(&database_url) {
         Ok(conn) => conn,
         Err(e) => {
             util::err_println("Failed to connect to DB. Aborting.");
@@ -59,8 +71,14 @@ async fn main() -> io::Result<()> {
 
     util::notice_println(format!("Starting server on {}:{}", bind_addr, port));
 
-    HttpServer::new(|| App::new().service(api::get_tasks))
-        .bind((bind_addr, port))?
-        .run()
-        .await
+    HttpServer::new(move || {
+        App::new()
+            .app_data(web::Data::new(AppState {
+                database_url: database_url.clone(),
+            }))
+            .service(api::get_tasks)
+    })
+    .bind((bind_addr, port))?
+    .run()
+    .await
 }