notearoni/src/db.rs

use serde;
use rocket::{
  self,
  http::Status,
  serde::Serialize,
  request::{
    self,
    FromRequest,
  },
  Request,
  outcome::{
    try_outcome,
    Outcome::{
      Success,
      Failure,
      Forward,
    },
  },
};
use rocket_db_pools::{
  Database,
  Connection,
};
use sqlx::{
  pool::PoolConnection,
  Postgres,
};

#[derive(Database)]
#[database("notes")]
pub struct Notes(sqlx::PgPool);

pub type Result<T, E = rocket::response::Debug<sqlx::Error>> = std::result::Result<T, E>;

#[derive(Serialize)]
pub struct Note {
  pub id: i32,
  pub uuid: String,
  pub content: String,
  pub list_id: i32,
  pub writer_id: i32,
}
impl ToString for Note {
  fn to_string(&self) -> String {
    format!("Note: {}", self.content)
  }
}

#[derive(Serialize)]
pub struct User {
  pub id: i32,
  pub uuid: String,
  pub username: String,
  pub password: String,
  pub email: String,
}
impl ToString for User {
  fn to_string(&self) -> String {
    format!("User: {}", self.username)
  }
}
#[derive(Debug)]
pub enum UserError {
  NoCookie,
  InvalidCookie,
  WrappedDBError(Option<rocket_db_pools::Error<sqlx::Error>>),
  DBError(sqlx::Error),
}

#[rocket::async_trait]
impl<'a> FromRequest<'a> for User {
  type Error = UserError;

  async fn from_request(req: &'a Request<'_>) -> request::Outcome<Self, Self::Error> {
    let cookies = req.cookies();
    let mut db = match req.guard::<Connection<Notes>>().await {
      Success(dbv) => dbv,
      Failure((http_err, err)) => return Failure((http_err, UserError::WrappedDBError(err))),
      Forward(next) => return Forward(next),
    };
    let user_uuid = match cookies.get_private("user_uuid") {
      Some(crumb) => crumb,
      None => return Forward(()), // this will redirect to the login page
    };
    match sqlx::query!("SELECT * FROM users WHERE uuid = $1", user_uuid.value())
      .fetch_optional(&mut *db)
      .await {
      Ok(Some(row)) => {
        Success(User {
          id: row.id,
          uuid: row.uuid,
          username: row.username,
          password: row.password,
          email: row.email,
        })
      },
      Ok(None) => Forward(()), // this will redirect to the login page
      Err(e) => Failure((Status::InternalServerError, UserError::DBError(e))),
    }
  }
}

#[derive(Serialize)]
pub struct List {
  pub id: i32,
  pub uuid: String,
  pub name: String,
  pub owner_id: i32,
}
impl ToString for List {
  fn to_string(&self) -> String {
    format!("{}: {} (owned by {})", self.id, self.name, self.owner_id)
  }
}
impl List {
  pub async fn from_uuid(list_uuid: String, db: &mut PoolConnection<Postgres>) -> Result<Option<List>> {
    match sqlx::query!("
      SELECT * FROM list WHERE uuid = $1", list_uuid)
      .fetch_optional(db)
      .await? {
      Some(list) => Ok(Some(List {
        id: list.id,
        uuid: list.uuid,
        name: list.name.clone(),
        owner_id: list.owner_id,
      })),
      None => Ok(None),
    }
  }
}

pub async fn get_notes_from_list(list_id: i32, db: &mut PoolConnection<Postgres>) -> Result<Vec<Note>> {
  Ok(sqlx::query!("
    SELECT *
    FROM note
    WHERE list_id = $1", list_id)
  .fetch_all(db)
  .await?
  .iter()
  .map(|r| Note {
    id: r.id,
    uuid: r.uuid.clone(),
    list_id: r.list_id,
    content: r.content.clone(),
    writer_id: r.writer_id,
  })
  .collect())
}

pub async fn get_list_notes(db: &mut PoolConnection<Postgres>, lid: i32) -> Result<Vec<Note>> {
  Ok(sqlx::query!("
    SELECT id,content,list_id,uuid,writer_id
    FROM note
    WHERE list_id = $1", lid)
    .fetch_all(db)
    .await?
    .iter()
    .map(|r| Note {
      id: r.id,
      uuid: r.uuid.clone(),
      content: r.content.clone(),
      list_id: r.list_id,
      writer_id: r.writer_id,
    })
    .collect())
}

pub async fn get_user_lists(db: &mut PoolConnection<Postgres>, uid: i32) -> Result<Vec<List>> {
  Ok(sqlx::query!("
    SELECT id,name,owner_id,uuid
    FROM list
    WHERE owner_id = $1", uid)
    .fetch_all(db)
    .await?
    .iter()
    .map(|r| List { 
      id: r.id,
      uuid: r.uuid.clone(),
      owner_id: r.owner_id,
      name: r.name.clone(),
    })
    .collect())
}

pub async fn get_user_lists_from_perms(db: &mut PoolConnection<Postgres>, uid: i32) -> Result<Vec<List>> {
  Ok(sqlx::query!("
    SELECT list.id,list.name,list.owner_id,uuid
    FROM list
    JOIN perms ON list.id=perms.list_id
    WHERE perms.user_id = $1
      AND perms.read = TRUE", uid)
    .fetch_all(db)
    .await?
    .iter()
    .map(|r| List { 
      name: r.name.clone(),
      id: r.id,
      uuid: r.uuid.clone(),
      owner_id: r.owner_id
    })
    .collect())
}

pub async fn add_note_to_list(db: &mut PoolConnection<Postgres>, list_id: i32, content: String, writer_id: i32) -> Result<bool> {
  match sqlx::query!(
    "INSERT INTO note (list_id, writer_id, content) VALUES ($1, $2, $3)", list_id, writer_id, content)
    .execute(db)
    .await {
      Ok(_) => Ok(true),
      Err(e) => Ok(false),
  }
}

pub async fn get_user_from_email(db: &mut PoolConnection<Postgres>, email: String) -> Result<Option<User>> {
  match sqlx::query!(
    "SELECT * FROM users WHERE email = $1",
    email
  ).fetch_optional(&mut *db)
  .await? {
    Some(row) => Ok(Some(User {
      id: row.id,
      uuid: row.uuid,
      username: row.username,
      password: row.password,
      email: row.email,
    })),
    None => Ok(None),
  }
}

pub async fn add_permission(db: &mut PoolConnection<Postgres>, user_id: i32, list_id: i32, perm: i32) -> Result<()> {
  sqlx::query!("
    INSERT INTO perms (user_id, list_id, read, write) VALUES ($1, $2, $3, $4)",
    user_id,
    list_id,
    perm >= 1,
    perm >= 2
  ).execute(db)
  .await?;
  Ok(())
}