From 77243161f7bb05581a22443409e83bb780f024ff Mon Sep 17 00:00:00 2001
From: subcrip <contact@subc.rip>
Date: Fri, 29 Mar 2024 18:07:36 +0800
Subject: [PATCH] feat(db): implement wrapper for database client and unit
 tests

---
 src/db.rs   | 49 ++++++++++++++++++++++++++++++++++++++++++-------
 src/main.rs | 12 +++++++-----
 2 files changed, 49 insertions(+), 12 deletions(-)

diff --git a/src/db.rs b/src/db.rs
index d03b9b6..9569bbb 100644
--- a/src/db.rs
+++ b/src/db.rs
@@ -1,13 +1,48 @@
-use tokio_postgres::{ NoTls, Error, tls::NoTlsStream };
+use std::time::Duration;
 
-pub struct Database {
+pub struct OAuthDatabase {
     client: tokio_postgres::Client,
-    connection: tokio_postgres::Connection<tokio_postgres::Socket, NoTlsStream>,
 }
 
-impl Database {
-    pub async fn new(host: &str, port: u16, user: &str, password: &str, dbname: &str) -> Result<Self, Error> {
-        let (client, connection) = tokio_postgres::connect(format!("host={} port={} user={} password={} dbname={}", host, port, user, password, dbname).as_str(), NoTls).await?;
-        Ok(Self { client, connection })
+impl OAuthDatabase {
+    pub async fn connect(host: &str, port: u16, user: &str, password: &str, db_name: &str) -> Result<Self, tokio_postgres::Error> {
+        let (client, connection) = tokio_postgres::connect(format!("host={} port={} user={} password={} dbname={}", host, port, user, password, db_name).as_str(), tokio_postgres::NoTls).await?;
+        Ok(Self::new(client, connection).await?)
+    }
+    
+    pub async fn new(client: tokio_postgres::Client, connection: tokio_postgres::Connection<tokio_postgres::Socket, tokio_postgres::tls::NoTlsStream>) -> Result<Self, tokio_postgres::Error> {
+        let obj = Self { client };
+        tokio::spawn(connection);
+        Ok(obj)
+    }
+
+    async fn psql_execute<T: ?Sized + tokio_postgres::ToStatement>(&mut self, query: &T, params: &[&(dyn tokio_postgres::types::ToSql + Sync)]) -> Result<u64, tokio_postgres::Error> {
+        self.client.execute(query, params).await
+    }
+
+    async fn psql_query<T: ?Sized + tokio_postgres::ToStatement>(&mut self, query: &T, params: &[&(dyn tokio_postgres::types::ToSql + Sync)]) -> Result<Vec<tokio_postgres::Row>, tokio_postgres::Error> {
+        self.client.query(query, params).await
+    }
+}
+
+#[cfg(test)]
+pub mod db_tests {
+    use super::*;
+
+    pub struct DBTestComponents {
+        d: OAuthDatabase
+    }
+
+    impl DBTestComponents {
+        pub async fn new(mut d: OAuthDatabase) -> Result<Self, tokio_postgres::Error> {
+            d.psql_execute("create table rust_test(name text, age int)", &[]).await?;
+            Ok(Self { d })
+        }
+    }
+
+    impl Drop for DBTestComponents {
+        fn drop(&mut self) {
+            futures::executor::block_on(self.d.psql_execute("drop table rust_test", &[])).unwrap();
+        }
     }
 }
diff --git a/src/main.rs b/src/main.rs
index 7f11fd9..4651124 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,10 +1,6 @@
-use async_std::io::prelude::SeekExt;
 use tide::Request;
 use tide::prelude::*;
 
-
-use futures::executor::block_on;
-
 mod db;
 
 #[derive(Serialize, Deserialize)]
@@ -22,8 +18,14 @@ impl Test {
 #[tokio::main]
 async fn main() {
     let a = Test::new(2, 3);
-    let b = block_on(db::Database::new("localhost", 5432, "postgres", "configjson", "myoauth")).unwrap();
+    let mut b = db::OAuthDatabase::connect("localhost", 5432, "postgres", "configjson", "myoauth").await.unwrap();
     // let s = serde_json::to_string(&a);
     // json!(a);
     println!("{}", json!(a));
 }
+
+#[tokio::test]
+async fn test_table_level_operations() {
+    use db::db_tests;
+    let _b = db_tests::DBTestComponents::new(db::OAuthDatabase::connect("localhost", 5432, "postgres", "configjson", "myoauth").await.unwrap());
+}