RUSTSUR

#[macro_use]
extern crate diesel;

use diesel::prelude::*;
use diesel::sqlite::SqliteConnection;
use dotenv::dotenv;
use std::env;

pub mod schema;
pub mod models;

use self::models::{Page, NewPage};
use warp::{Reply, Filter, http::{StatusCode}};

type Conn = r2d2::Pool<diesel::r2d2::ConnectionManager<SqliteConnection>>;

fn establish_connection() -> r2d2::Pool<diesel::r2d2::ConnectionManager<SqliteConnection>> {
    dotenv().ok();

    let database_url = env::var("DATABASE_URL")
        .expect("DATABASE_URL must be set");
    let manager = diesel::r2d2::ConnectionManager::new(database_url);
    r2d2::Pool::builder()
        .max_size(4)
        .build(manager)
        .expect("Failed to create connection pool!")
}

fn main() {
    let connection = establish_connection();
    let connection = warp::any().map(move || connection.clone());

    let page_name = warp::path::param2::<String>();

    let load_page = warp::get2().and(page_name).and(connection.clone()).and_then(load_page);
    let upload_page = warp::put2().and(page_name).and(connection).and(warp::body::json()).and_then(upload_page);
    let api = warp::path("pages").and(load_page.or(upload_page)).recover(error_handler);

    let routes = warp::path("api").and(api);

    warp::serve(routes)
        .run(([127, 0, 0, 1], 3030));
}

#[derive(Debug)]
enum Error {
    DatabaseError(diesel::result::Error),
    NotFound(String)
}

impl std::fmt::Display for Error {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        match self {
            Error::DatabaseError(err) => write!(f, "Database error: {}", err),
            Error::NotFound(thing) => write!(f, "{} not found", thing)
        }
    }
}

impl std::error::Error for Error {}

fn error_handler(r: warp::Rejection) -> Result<impl Reply, warp::Rejection> {
    if let Some(err) = r.find_cause::<Error>() {
        let code = match err {
            Error::DatabaseError(_) => StatusCode::INTERNAL_SERVER_ERROR,
            Error::NotFound(_) => StatusCode::NOT_FOUND
        };
        Ok(warp::reply::with_status(err.to_string(), code))
    } else {
        Err(r)
    }
}

fn load_page(name: String, conn: Conn) -> Result<impl Reply, warp::Rejection> {
    match get_page(&conn.get().unwrap(), &name) {
        Ok(page) => Ok(warp::reply::json(&page)),
        Err(err) => Err(match err {
            diesel::result::Error::NotFound => warp::reject::custom(Error::NotFound(name)), // https://github.com/seanmonstar/warp/issues/77
            _ => warp::reject::custom(Error::DatabaseError(err))
        })
    }
}

fn upload_page(name: String, conn: Conn, body: String) -> Result<impl Reply, warp::Rejection> {
    match set_page(&conn.get().unwrap(), &name, &body) {
        Ok(_) => Ok(warp::reply()),
        Err(err) => Err(warp::reject::custom(Error::DatabaseError(err)))
    }
}

fn get_page(conn: &SqliteConnection, name_: &str) -> Result<Page, diesel::result::Error> {
    use crate::schema::pages::dsl::*;
    pages.filter(name.eq(name_))
        .get_result(conn)
}

fn set_page<'a>(conn: &SqliteConnection, title: &'a str, body: &'a str) -> Result<(), diesel::result::Error> {
    let new_page = NewPage {
        name: title,
        body: body
    };

    diesel::replace_into(schema::pages::table)
        .values(&new_page)
        .execute(conn)
        .map(|_| ())
}