From ad4055a18fa18133b0e6151561f23580790a31d8 Mon Sep 17 00:00:00 2001
From: hank <ysh2291939848@outlook.com>
Date: Wed, 31 Jul 2024 02:05:34 +0800
Subject: [PATCH] some db

---
 Cargo.toml    |  4 ++++
 src/config.rs | 11 +++++++++--
 src/db.rs     | 25 +++++++++++++++++++++++++
 src/error.rs  |  0
 src/main.rs   | 16 ++++++++++------
 src/router.rs | 18 ++++++++++++------
 src/state.rs  | 12 ++++++++++++
 7 files changed, 72 insertions(+), 14 deletions(-)
 create mode 100644 src/error.rs
 create mode 100644 src/state.rs

diff --git a/Cargo.toml b/Cargo.toml
index c9dfa85..467fb1b 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -18,6 +18,10 @@ sea-orm = { version = "1.0.0-rc.7", features = [
   "debug-print",
   "runtime-tokio",
   "sqlx-postgres",
+  "macros",
+  "with-uuid",
+  "with-json",
+  "with-chrono",
 ] }
 
 tokio = { version = "1", features = ["full"] }
diff --git a/src/config.rs b/src/config.rs
index 998da36..c88cf27 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -11,17 +11,24 @@ pub struct Config {
     pub port: u32,
     pub workers: u32,
     pub http3: bool,
-    pub database_url: String,
+    pub db_url: String,
+    pub db_user: String,
+    pub db_password: Option<String>,
+    pub db_name: String,
 }
 
 impl Default for Config {
     fn default() -> Self {
+        let current_user = env::var("USER").unwrap_or("hank".to_owned());
         Config {
             host: "127.0.0.1".to_owned(),
             port: 5800,
             workers: 4,
             http3: false,
-            database_url: "".to_owned(), // FIXME: use a default database url
+            db_url: "localhost".to_owned(), // FIXME: use a default database url
+            db_user: current_user.clone(),
+            db_password: None,
+            db_name: current_user,
         }
     }
 }
diff --git a/src/db.rs b/src/db.rs
index e69de29..99999e6 100644
--- a/src/db.rs
+++ b/src/db.rs
@@ -0,0 +1,25 @@
+use sea_orm::{Database, DatabaseConnection, DbErr};
+use anyhow::Result;
+use tracing::info;
+
+use crate::config::Config;
+
+pub async fn init(config: &Config) -> Result<DatabaseConnection, DbErr> {
+    let db_url = if config.db_password.is_some() {
+        format!(
+            "postgresql://{}:{}@{}/{}",
+            config.db_user,
+            config.db_password.as_deref().unwrap(),
+            config.db_url,
+            config.db_name,
+        )
+    } else {
+        format!(
+            "postgresql://{}@{}/{}",
+            config.db_user, config.db_url, config.db_name,
+        )
+    };
+
+    info!("Connecting to database at {}", db_url);
+    Database::connect(&db_url).await
+}
diff --git a/src/error.rs b/src/error.rs
new file mode 100644
index 0000000..e69de29
diff --git a/src/main.rs b/src/main.rs
index 98c538e..d4be580 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,12 +1,14 @@
 pub mod config;
 pub mod db;
+pub mod error;
 pub mod handler;
 pub mod middleware;
 pub mod models;
 pub mod router;
+pub mod state;
 pub mod utils;
 
-use crate::router::route;
+use crate::router::make_route;
 use anyhow::Result;
 use salvo::prelude::*;
 use tokio::signal;
@@ -16,17 +18,20 @@ use tracing_subscriber::EnvFilter;
 #[tokio::main]
 async fn main() -> Result<()> {
     let filter = EnvFilter::from_default_env();
-    tracing_subscriber::fmt().with_env_filter(filter).init();
+    tracing_subscriber::fmt().with_env_filter(filter).with_test_writer().init();
+    // std::env::set_var("RUST_LOG", "trace");
+    // tracing_subscriber::fmt::init();
 
     let config = config::get_config().await.unwrap_or_else(|_| {
         info!("failed to read config file, using default instead");
         config::Config::default()
     });
 
-    let acceptor = TcpListener::new(config.host + ":" + &config.port.to_string())
+    let acceptor = TcpListener::new(config.host.clone() + ":" + &config.port.to_string())
         .bind()
         .await;
-    let router = Router::with_path("api").push(route());
+
+    let router = Router::with_path("api").push(make_route(&config).await);
 
     // TODO: http3
     let server = Server::new(acceptor);
@@ -52,9 +57,8 @@ mod tests {
 
     #[tokio::test]
     async fn test_basic_auth() {
-        let service = Service::new(super::route());
-
         let test_config = Config::default();
+        let service = Service::new(super::make_route(&test_config).await);
 
         let url = format!("http://{}:{}/", test_config.host, test_config.port);
 
diff --git a/src/router.rs b/src/router.rs
index 5bd4046..30e9e0c 100644
--- a/src/router.rs
+++ b/src/router.rs
@@ -1,10 +1,11 @@
-use crate::middleware::basic_auth::{self, Validator};
+use crate::config::Config;
+use crate::db;
+use crate::middleware::basic_auth::Validator;
+use crate::state::AppState;
 use anyhow::Result;
 use salvo::affix;
 use salvo::prelude::*;
-
-#[derive(Debug, Clone)]
-pub struct AppState {}
+use tracing::info;
 
 #[handler]
 async fn hello() -> Result<String> {
@@ -16,10 +17,15 @@ async fn hello_admin() -> Result<String> {
     Ok("Hello, Admin".to_owned())
 }
 
-pub fn route() -> Router {
-    let state = AppState {};
+// Main Router
+pub async fn make_route(config: &Config) -> Router {
     let auth_handler = BasicAuth::new(Validator);
 
+    let db = db::init(config).await.unwrap();
+    info!("Database connection established");
+
+    let state = AppState::new(db);
+
     let router = Router::new()
         .hoop(affix::inject(state))
         .push(Router::with_path("hello").get(hello))
diff --git a/src/state.rs b/src/state.rs
new file mode 100644
index 0000000..a0f0ad5
--- /dev/null
+++ b/src/state.rs
@@ -0,0 +1,12 @@
+use sea_orm::DatabaseConnection;
+
+#[derive(Debug, Clone)]
+pub struct AppState {
+    conn: DatabaseConnection,
+}
+
+impl AppState {
+    pub fn new(conn: DatabaseConnection) -> Self {
+        Self { conn }
+    }
+}