make database calls 1 conn

1 files changed, 88 insertions, 85 deletions

diff --git a/src/database/friends.rs b/src/database/friends.rs
index 0b78488..31434d4 100644
--- a/src/database/friends.rs
+++ b/src/database/friends.rs
@@ -1,97 +1,100 @@
 use tracing::instrument;
 
-use crate::{
-    database::{self, users::user_from_row},
-    types::user::{User, FOLLOWED, FOLLOWING, NO_RELATION},
-};
+use crate::types::user::{User, FOLLOWED, FOLLOWING, NO_RELATION};
 
-pub fn init() -> Result<(), rusqlite::Error> {
-    let sql = "
-        CREATE TABLE IF NOT EXISTS friends (
-            follower_id     INTEGER NOT NULL,
-            followee_id     INTEGER NOT NULL,
-            FOREIGN KEY(follower_id) REFERENCES users(user_id),
-            FOREIGN KEY(followee_id) REFERENCES users(user_id),
-            PRIMARY KEY (follower_id, followee_id)
-        );
-    ";
-    let conn = database::connect()?;
-    conn.execute(sql, ())?;
-    Ok(())
-}
+use super::Database;
 
-#[instrument()]
-pub fn get_friend_status(user_id_1: u64, user_id_2: u64) -> Result<u8, rusqlite::Error> {
-    tracing::trace!("Retrieving friend status");
-    let conn = database::connect()?;
-    let mut stmt = conn.prepare("SELECT * FROM friends WHERE (follower_id = ? AND followee_id = ?) OR (follower_id = ? AND followee_id = ?);")?;
-    let mut status = NO_RELATION;
-    let rows: Vec<u64> = stmt
-        .query_map([user_id_1, user_id_2, user_id_2, user_id_1], |row| {
-            let id: u64 = row.get(0)?;
-            Ok(id)
-        })?
-        .into_iter()
-        .flatten()
-        .collect();
+impl Database {
+    pub fn init_friends(&self) -> Result<(), rusqlite::Error> {
+        let sql = "
+            CREATE TABLE IF NOT EXISTS friends (
+                follower_id     INTEGER NOT NULL,
+                followee_id     INTEGER NOT NULL,
+                FOREIGN KEY(follower_id) REFERENCES users(user_id),
+                FOREIGN KEY(followee_id) REFERENCES users(user_id),
+                PRIMARY KEY (follower_id, followee_id)
+            );
+        ";
+        self.0.execute(sql, ())?;
+        Ok(())
+    }
 
-    for follower in rows {
-        if follower == user_id_1 {
-            status |= FOLLOWING;
-        }
+    #[instrument(skip(self))]
+    pub fn get_friend_status(&self, user_id_1: u64, user_id_2: u64) -> Result<u8, rusqlite::Error> {
+        tracing::trace!("Retrieving friend status");
+        let mut stmt = self.0.prepare("SELECT * FROM friends WHERE (follower_id = ? AND followee_id = ?) OR (follower_id = ? AND followee_id = ?);")?;
+        let mut status = NO_RELATION;
+        let rows: Vec<u64> = stmt
+            .query_map([user_id_1, user_id_2, user_id_2, user_id_1], |row| {
+                let id: u64 = row.get(0)?;
+                Ok(id)
+            })?
+            .into_iter()
+            .flatten()
+            .collect();
 
-        if follower == user_id_2 {
-            status |= FOLLOWED;
+        for follower in rows {
+            if follower == user_id_1 {
+                status |= FOLLOWING;
+            }
+
+            if follower == user_id_2 {
+                status |= FOLLOWED;
+            }
         }
-    }
 
-    Ok(status)
-}
+        Ok(status)
+    }
 
-#[instrument()]
-pub fn get_friends(user_id: u64) -> Result<Vec<User>, rusqlite::Error> {
-    tracing::trace!("Retrieving friends");
-    let conn = database::connect()?;
-    let mut stmt = conn.prepare(
-        "
-        SELECT * 
-        FROM users u 
-        WHERE EXISTS (
-            SELECT NULL 
-            FROM friends f 
-            WHERE u.user_id = f.follower_id
-            AND f.followee_id = ?
-        )
-        AND EXISTS (
-            SELECT NULL
-            FROM friends f
-            WHERE u.user_id = f.followee_id
-            AND f.follower_id = ?
-        )
-    ",
-    )?;
-    let row = stmt.query_map([user_id, user_id], |row| {
-        let row = user_from_row(row, true)?;
-        Ok(row)
-    })?;
-    Ok(row.into_iter().flatten().collect())
-}
+    #[instrument(skip(self))]
+    pub fn get_friends(&self, user_id: u64) -> Result<Vec<User>, rusqlite::Error> {
+        tracing::trace!("Retrieving friends");
+        let mut stmt = self.0.prepare(
+            "
+            SELECT * 
+            FROM users u 
+            WHERE EXISTS (
+                SELECT NULL 
+                FROM friends f 
+                WHERE u.user_id = f.follower_id
+                AND f.followee_id = ?
+            )
+            AND EXISTS (
+                SELECT NULL
+                FROM friends f
+                WHERE u.user_id = f.followee_id
+                AND f.follower_id = ?
+            )
+        ",
+        )?;
+        let row = stmt.query_map([user_id, user_id], |row| {
+            let row = Self::user_from_row(row, true)?;
+            Ok(row)
+        })?;
+        Ok(row.into_iter().flatten().collect())
+    }
 
-#[instrument()]
-pub fn set_following(user_id_1: u64, user_id_2: u64) -> Result<bool, rusqlite::Error> {
-    tracing::trace!("Setting following");
-    let conn = database::connect()?;
-    let mut stmt =
-        conn.prepare("INSERT OR REPLACE INTO friends (follower_id, followee_id) VALUES (?,?)")?;
-    let changes = stmt.execute([user_id_1, user_id_2])?;
-    Ok(changes == 1)
-}
+    #[instrument(skip(self))]
+    pub fn set_following(&self, user_id_1: u64, user_id_2: u64) -> Result<bool, rusqlite::Error> {
+        tracing::trace!("Setting following");
+        let mut stmt = self
+            .0
+            .prepare("INSERT OR REPLACE INTO friends (follower_id, followee_id) VALUES (?,?)")?;
+        let changes = stmt.execute([user_id_1, user_id_2])?;
+        Ok(changes == 1)
+    }
 
-#[instrument()]
-pub fn remove_following(user_id_1: u64, user_id_2: u64) -> Result<bool, rusqlite::Error> {
-    tracing::trace!("Removing following");
-    let conn = database::connect()?;
-    let mut stmt = conn.prepare("DELETE FROM friends WHERE follower_id = ? AND followee_id = ?")?;
-    let changes = stmt.execute([user_id_1, user_id_2])?;
-    Ok(changes == 1)
+    #[instrument(skip(self))]
+    pub fn remove_following(
+        &self,
+        user_id_1: u64,
+        user_id_2: u64,
+    ) -> Result<bool, rusqlite::Error> {
+        tracing::trace!("Removing following");
+        let mut stmt = self
+            .0
+            .prepare("DELETE FROM friends WHERE follower_id = ? AND followee_id = ?")?;
+        let changes = stmt.execute([user_id_1, user_id_2])?;
+        Ok(changes == 1)
+    }
 }