1 files changed, 22 insertions, 15 deletions

diff --git a/src/database/mod.rs b/src/database/mod.rs
index d22a350..67e05c6 100644
--- a/src/database/mod.rs
+++ b/src/database/mod.rs
@@ -1,3 +1,4 @@
+use rusqlite::Connection;
 use tracing::instrument;
 
 pub mod comments;
@@ -7,23 +8,29 @@ pub mod posts;
 pub mod sessions;
 pub mod users;
 
-pub fn connect() -> Result<rusqlite::Connection, rusqlite::Error> {
-    rusqlite::Connection::open("xssbook.db")
+#[derive(Debug)]
+pub struct Database(Connection);
+
+impl Database {
+    pub fn connect() -> Result<Self, rusqlite::Error> {
+        let conn = rusqlite::Connection::open("xssbook.db")?;
+        Ok(Self(conn))
+    }
+
+    #[instrument(skip(self))]
+    pub fn query(&self, query: String) -> Result<usize, rusqlite::Error> {
+        tracing::trace!("Running custom query");
+        self.0.execute(&query, [])
+    }
 }
 
 pub fn init() -> Result<(), rusqlite::Error> {
-    users::init()?;
-    posts::init()?;
-    sessions::init()?;
-    likes::init()?;
-    comments::init()?;
-    friends::init()?;
+    let db = Database::connect()?;
+    db.init_users()?;
+    db.init_posts()?;
+    db.init_sessions()?;
+    db.init_likes()?;
+    db.init_comments()?;
+    db.init_friends()?;
     Ok(())
 }
-
-#[instrument()]
-pub fn query(query: String) -> Result<usize, rusqlite::Error> {
-    tracing::trace!("Running custom query");
-    let conn = connect()?;
-    conn.execute(&query, [])
-}