diff --git a/src/entities/bal.rs b/src/entities/bal.rs index 9e7516c..e1ec870 100644 --- a/src/entities/bal.rs +++ b/src/entities/bal.rs @@ -9,6 +9,7 @@ pub struct Model { pub id: u32, pub user_id: u32, pub name: String, + pub ended: bool } #[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)] diff --git a/src/entities/user.rs b/src/entities/user.rs index 25652fc..20670c2 100644 --- a/src/entities/user.rs +++ b/src/entities/user.rs @@ -10,6 +10,7 @@ pub struct Model { #[sea_orm(unique)] pub username: String, pub hashed_password: String, + pub owner_id: Option, pub current_bal_id: Option, } diff --git a/src/main.rs b/src/main.rs index 4229796..d5a8b86 100644 --- a/src/main.rs +++ b/src/main.rs @@ -3,7 +3,7 @@ use std::{net::{Ipv4Addr, SocketAddr}, path::PathBuf, sync::{Arc, LazyLock}}; use axum::{extract::State, http::HeaderMap, middleware, routing::get}; use clap::{Parser, Subcommand}; use reqwest::{header::USER_AGENT}; -use sea_orm::{ConnectionTrait, Database, DatabaseConnection, EntityTrait, PaginatorTrait, Schema}; +use sea_orm::{ConnectionTrait, Database, DatabaseConnection, DbErr, EntityTrait, PaginatorTrait, Schema}; use tokio::{sync::broadcast::{self, Sender}}; use utoipa::{openapi::{security::{HttpAuthScheme, HttpBuilder, SecurityScheme}, ContactBuilder, InfoBuilder, LicenseBuilder}, Modify, OpenApi}; use utoipa_axum::router::OpenApiRouter; @@ -101,28 +101,9 @@ async fn main() { } }); - let builder = db.get_database_backend(); - let schema = Schema::new(builder); - if let Err(err) = db.execute(builder.build(schema.create_table_from_entity(crate::entities::prelude::Book).if_not_exists())).await { - log::error!(target: "database", "Error while creating book table: {err:?}"); + if let Err(_) = create_tables(db.as_ref()).await { return; - } - if let Err(err) = db.execute(builder.build(schema.create_table_from_entity(crate::entities::prelude::BookInstance).if_not_exists())).await { - log::error!(target: "database", "Error while creating book_instance table: {err:?}"); - return; - } - if let Err(err) = db.execute(builder.build(schema.create_table_from_entity(crate::entities::prelude::Owner).if_not_exists())).await { - log::error!(target: "database", "Error while creating owner table: {err:?}"); - return; - } - if let Err(err) = db.execute(builder.build(schema.create_table_from_entity(crate::entities::prelude::User).if_not_exists())).await { - log::error!(target: "database", "Error while creating user table: {err:?}"); - return; - } - if let Err(err) = db.execute(builder.build(schema.create_table_from_entity(crate::entities::prelude::Bal).if_not_exists())).await { - log::error!(target: "database", "Error while creating bal table: {err:?}"); - return; - } + }; match &CLI.command { Commands::Run {port,..} => run_server(db, *port).await, @@ -130,6 +111,33 @@ async fn main() { } } +async fn create_tables(db_conn: &C) -> Result<(), DbErr> + where C: ConnectionTrait,{ + let builder = db_conn.get_database_backend(); + let schema = Schema::new(builder); + if let Err(err) = db_conn.execute(builder.build(schema.create_table_from_entity(crate::entities::prelude::Book).if_not_exists())).await { + log::error!(target: "database", "Error while creating book table: {err:?}"); + return Err(err); + } + if let Err(err) = db_conn.execute(builder.build(schema.create_table_from_entity(crate::entities::prelude::BookInstance).if_not_exists())).await { + log::error!(target: "database", "Error while creating book_instance table: {err:?}"); + return Err(err); + } + if let Err(err) = db_conn.execute(builder.build(schema.create_table_from_entity(crate::entities::prelude::Owner).if_not_exists())).await { + log::error!(target: "database", "Error while creating owner table: {err:?}"); + return Err(err); + } + if let Err(err) = db_conn.execute(builder.build(schema.create_table_from_entity(crate::entities::prelude::User).if_not_exists())).await { + log::error!(target: "database", "Error while creating user table: {err:?}"); + return Err(err); + } + if let Err(err) = db_conn.execute(builder.build(schema.create_table_from_entity(crate::entities::prelude::Bal).if_not_exists())).await { + log::error!(target: "database", "Error while creating bal table: {err:?}"); + return Err(err); + } + Ok(()) +} + async fn run_server(db: Arc, port: u16) { let (event_bus, _) = broadcast::channel(16); diff --git a/src/routes/bal.rs b/src/routes/bal.rs index c59402c..85e6fe9 100644 --- a/src/routes/bal.rs +++ b/src/routes/bal.rs @@ -6,7 +6,7 @@ use sea_orm::{ActiveModelTrait, ActiveValue::{NotSet, Set}, ColumnTrait, EntityT use serde::{Deserialize, Serialize}; use utoipa::IntoParams; -use crate::{entities::{bal, prelude::*, user}, routes::auth::Claims, utils::auth::user_is_bal_owner, AppState}; +use crate::{entities::{bal, book_instance, prelude::*, user}, routes::auth::Claims, AppState}; #[derive(IntoParams)] @@ -72,6 +72,7 @@ pub async fn create_bal( id: NotSet, user_id: Set(claims.user_id), name: Set(instance_payload.name), + ended: Set(false) }; let b = bal.save(state.db_conn.as_ref()).await; @@ -188,16 +189,24 @@ pub async fn get_current_bal( } } +#[derive(Deserialize, PartialEq, utoipa::ToSchema)] +pub enum BookInstanceTransferMode { + All, + None, + OwnedOnly +} + #[derive(Deserialize, utoipa::ToSchema)] -pub struct BalIdParams{ - id: u32, +pub struct SetCurrentBalParams{ + bal_id: u32, + transfer_book_instances: Option } #[axum::debug_handler] #[utoipa::path( post, path = "/bal/current", - request_body = BalIdParams, + request_body = SetCurrentBalParams, security(("jwt" = [])), responses( (status = OK, description = "Successfully set current active BAL"), @@ -210,17 +219,44 @@ pub struct BalIdParams{ pub async fn set_current_bal( State(state): State>, claims: Claims, - Json(payload): Json, + Json(payload): Json, ) -> StatusCode { - if !user_is_bal_owner(claims.user_id, payload.id, state.db_conn.as_ref()).await { - return StatusCode::UNAUTHORIZED; - } - if let Ok(Some(user)) = User::find_by_id(claims.user_id).one(state.db_conn.as_ref()).await { + if let Ok(Some(new_bal)) = Bal::find_by_id(payload.bal_id).one(state.db_conn.as_ref()).await + && let Ok(Some(user)) = User::find_by_id(claims.user_id).one(state.db_conn.as_ref()).await + { + if new_bal.user_id != claims.user_id { + return StatusCode::UNAUTHORIZED; + } + + if let Some(old_bal_id) = user.current_bal_id { + // Optional instance transfer + if let Some(mode) = payload.transfer_book_instances && mode != BookInstanceTransferMode::None { + let mut update_query = BookInstance::update_many() + .set(book_instance::ActiveModel { + bal_id: Set(new_bal.id), + ..Default::default() + }) + .filter(book_instance::Column::BalId.eq(old_bal_id)); + if mode == BookInstanceTransferMode::OwnedOnly { + update_query = update_query.filter(book_instance::Column::OwnerId.eq(user.owner_id)); + } + + let _ = update_query.exec(state.db_conn.as_ref()).await; + } + // Set old bal as ended if it existed + let _ = Bal::update(bal::ActiveModel { + id: Set(old_bal_id), + ended: Set(true), + ..Default::default() + }).exec(state.db_conn.as_ref()).await; + } + + // Set current bal on user let mut user_active_model: user::ActiveModel = user.into_active_model(); - user_active_model.current_bal_id = Set(Some(payload.id)); + user_active_model.current_bal_id = Set(Some(payload.bal_id)); let _ = User::update(user_active_model).exec(state.db_conn.as_ref()).await; StatusCode::OK } else { - StatusCode::INTERNAL_SERVER_ERROR + return StatusCode::NOT_FOUND; } } diff --git a/src/utils/cli.rs b/src/utils/cli.rs index 2cb8b7e..8956350 100644 --- a/src/utils/cli.rs +++ b/src/utils/cli.rs @@ -61,27 +61,29 @@ pub async fn manage_users(db: Arc) { .with_validator(min_length!(10)) .with_display_mode(inquire::PasswordDisplayMode::Masked) .prompt().unwrap(); - let new_user = user::ActiveModel { + let mut new_user = user::ActiveModel { id: NotSet, username: Set(username), hashed_password: Set(hash_password(password)), - current_bal_id: Set(None) + current_bal_id: Set(None), + owner_id: Set(None) }; - let res = new_user.insert(db.as_ref()).await.unwrap(); + let res = new_user.clone().insert(db.as_ref()).await.unwrap(); - if Confirm::new("Add an owner corresponding to this user ?").with_default(true).prompt().is_ok_and(|choice| choice==true) { - let first_name = prompt_text("First Name: ").unwrap(); - let last_name = prompt_text("Last Name: ").unwrap(); - let contact = prompt_text("Contact: ").unwrap(); - let new_owner = owner::ActiveModel { - id: NotSet, - user_id: Set(res.id), - first_name: Set(first_name), - last_name: Set(last_name), - contact: Set(contact) - }; - let _ = new_owner.insert(db.as_ref()).await.unwrap(); - } + println!("Create corresponding owner:"); + let first_name = prompt_text("First Name: ").unwrap(); + let last_name = prompt_text("Last Name: ").unwrap(); + let contact = prompt_text("Contact: ").unwrap(); + let new_owner = owner::ActiveModel { + id: NotSet, + user_id: Set(res.id), + first_name: Set(first_name), + last_name: Set(last_name), + contact: Set(contact) + }; + let owner_res = new_owner.insert(db.as_ref()).await.unwrap(); + new_user.owner_id = Set(Some(owner_res.id)); + let _ = new_user.update(db.as_ref()); } }, Action::Delete => {