ImageApi/src/database/mod.rs

use bcrypt::{DEFAULT_COST, hash, verify};
use diesel::prelude::*;
use diesel::sqlite::SqliteConnection;
use std::ops::DerefMut;
use std::sync::{Arc, Mutex};

use crate::database::models::{Favorite, InsertFavorite, InsertUser, User};

pub mod models;
pub mod schema;

pub trait UserDao {
    fn create_user(&mut self, user: &str, password: &str) -> Option<User>;
    fn get_user(&mut self, user: &str, password: &str) -> Option<User>;
    fn user_exists(&mut self, user: &str) -> bool;
}

pub struct SqliteUserDao {
    connection: SqliteConnection,
}

impl SqliteUserDao {
    pub fn new() -> Self {
        Self {
            connection: connect(),
        }
    }
}

#[cfg(test)]
pub mod test {
    use diesel::{Connection, SqliteConnection};
    use diesel_migrations::{EmbeddedMigrations, MigrationHarness, embed_migrations};

    const DB_MIGRATIONS: EmbeddedMigrations = embed_migrations!();

    pub fn in_memory_db_connection() -> SqliteConnection {
        let mut connection = SqliteConnection::establish(":memory:")
            .expect("Unable to create in-memory db connection");
        connection
            .run_pending_migrations(DB_MIGRATIONS)
            .expect("Failure running DB migrations");

        connection
    }
}

impl UserDao for SqliteUserDao {
    // TODO: Should probably use Result here
    fn create_user(&mut self, user: &str, pass: &str) -> Option<User> {
        use schema::users::dsl::*;

        let hashed = hash(pass, DEFAULT_COST);
        if let Ok(hash) = hashed {
            diesel::insert_into(users)
                .values(InsertUser {
                    username: user,
                    password: &hash,
                })
                .execute(&mut self.connection)
                .unwrap();

            users
                .filter(username.eq(username))
                .load::<User>(&mut self.connection)
                .unwrap()
                .first()
                .cloned()
        } else {
            None
        }
    }

    fn get_user(&mut self, user: &str, pass: &str) -> Option<User> {
        use schema::users::dsl::*;

        match users
            .filter(username.eq(user))
            .load::<User>(&mut self.connection)
            .unwrap_or_default()
            .first()
        {
            Some(u) if verify(pass, &u.password).unwrap_or(false) => Some(u.clone()),
            _ => None,
        }
    }

    fn user_exists(&mut self, user: &str) -> bool {
        use schema::users::dsl::*;

        users
            .filter(username.eq(user))
            .load::<User>(&mut self.connection)
            .unwrap_or_default()
            .first()
            .is_some()
    }
}

pub fn connect() -> SqliteConnection {
    let db_url = dotenv::var("DATABASE_URL").expect("DATABASE_URL must be set");
    SqliteConnection::establish(&db_url).expect("Error connecting to DB")
}

#[derive(Debug)]
pub struct DbError {
    pub kind: DbErrorKind,
}

impl DbError {
    fn new(kind: DbErrorKind) -> Self {
        DbError { kind }
    }

    fn exists() -> Self {
        DbError::new(DbErrorKind::AlreadyExists)
    }
}

#[derive(Debug, PartialEq)]
pub enum DbErrorKind {
    AlreadyExists,
    InsertError,
    QueryError,
}

pub trait FavoriteDao: Sync + Send {
    fn add_favorite(&mut self, user_id: i32, favorite_path: &str) -> Result<usize, DbError>;
    fn remove_favorite(&mut self, user_id: i32, favorite_path: String);
    fn get_favorites(&mut self, user_id: i32) -> Result<Vec<Favorite>, DbError>;
}

pub struct SqliteFavoriteDao {
    connection: Arc<Mutex<SqliteConnection>>,
}

impl SqliteFavoriteDao {
    pub fn new() -> Self {
        SqliteFavoriteDao {
            connection: Arc::new(Mutex::new(connect())),
        }
    }
}

impl FavoriteDao for SqliteFavoriteDao {
    fn add_favorite(&mut self, user_id: i32, favorite_path: &str) -> Result<usize, DbError> {
        use schema::favorites::dsl::*;

        let mut connection = self.connection.lock().expect("Unable to get FavoriteDao");

        if favorites
            .filter(userid.eq(user_id).and(path.eq(&favorite_path)))
            .first::<Favorite>(connection.deref_mut())
            .is_err()
        {
            diesel::insert_into(favorites)
                .values(InsertFavorite {
                    userid: &user_id,
                    path: favorite_path,
                })
                .execute(connection.deref_mut())
                .map_err(|_| DbError::new(DbErrorKind::InsertError))
        } else {
            Err(DbError::exists())
        }
    }

    fn remove_favorite(&mut self, user_id: i32, favorite_path: String) {
        use schema::favorites::dsl::*;

        diesel::delete(favorites)
            .filter(userid.eq(user_id).and(path.eq(favorite_path)))
            .execute(self.connection.lock().unwrap().deref_mut())
            .unwrap();
    }

    fn get_favorites(&mut self, user_id: i32) -> Result<Vec<Favorite>, DbError> {
        use schema::favorites::dsl::*;

        favorites
            .filter(userid.eq(user_id))
            .load::<Favorite>(self.connection.lock().unwrap().deref_mut())
            .map_err(|_| DbError::new(DbErrorKind::QueryError))
    }
}