use migration::{Migrator, MigratorTrait}; use models::prelude::*; use sea_orm::{ ActiveModelTrait, ColumnTrait, Database, DatabaseConnection, DatabaseTransaction, DbErr, EntityTrait, PaginatorTrait, QueryFilter, Set, TransactionTrait, }; use wd_log::{log_error_ln, log_info_ln, log_warn_ln}; const PAGE_SIZE: usize = 25; #[derive(Debug)] pub struct Controller { db: DatabaseConnection, } pub struct PaginatedRecordData { pub items_count: usize, pub pages_count: usize, pub current_data: Vec<(RecordModel, Option)>, } impl Controller { /// Create controller pub async fn new(config: String) -> Result { Ok(Self { db: Database::connect(config).await?, }) } /// Do migrate pub async fn migrate(&self) -> Result<(), DbErr> { if let Err(err) = Migrator::install(&self.db).await { log_warn_ln!("{}", err) } if let Err(err) = Migrator::up(&self.db, None).await { Err(err) } else { log_info_ln!("database initialized."); Ok(()) } } /// register user when `/start` command called. pub async fn register_user(&self, user_id: &i64, username: &String) -> Result<(), DbErr> { let transaction = self.db.begin().await?; self.setup_user(user_id, username, &transaction).await?; transaction.commit().await } /// update user notify when `/mute` or `/unmute` command called. pub async fn set_user_notify(&self, user_id: &i64, notify: bool) -> Result<(), DbErr> { let transaction = self.db.begin().await?; if let Some(user) = self.get_user(user_id, &transaction).await? { let mut user_active: UserActiveModel = user.into(); user_active.notify = Set(notify); user_active.save(&transaction).await?; } transaction.commit().await } pub async fn get_user_notify(&self, user_id: &i64) -> Result { let transaction = self.db.begin().await?; if let Some(user) = self.get_user(&user_id, &transaction).await? { Ok(user.notify) } else { Ok(false) } } async fn setup_user( &self, user_id: &i64, username: &String, transaction: &DatabaseTransaction, ) -> Result { match self.get_user(user_id, &transaction).await? { Some(user) => { let mut user_active: UserActiveModel = user.into(); user_active.username = Set(Some(username.to_string())); user_active.save(transaction).await } None => { UserActiveModel { tg_uid: Set(user_id.to_owned()), username: Set(Some(username.to_string())), notify: Set(true), ..Default::default() } .save(transaction) .await } } } async fn get_user( &self, user_id: &i64, transaction: &DatabaseTransaction, ) -> Result, DbErr> { User::find() .filter(UserColumn::Id.eq(user_id.to_owned())) .one(transaction) .await } pub async fn get_user_by_username(&self, username: &str) -> Result, DbErr> { let transaction = self.db.begin().await?; User::find() .filter(UserColumn::Username.eq(username.to_owned())) .one(&transaction) .await } /// get records when inline query called. pub async fn get_records_by_keywords( &self, key_word: &String, ) -> Result { let pagination = Record::find() .find_also_related(User) .filter(RecordColumn::Message.contains(key_word.as_str())) .paginate(&self.db, PAGE_SIZE * 2); // 50 records seems ok. Ok(PaginatedRecordData { items_count: pagination.num_items().await?, pages_count: pagination.num_pages().await?, current_data: pagination.fetch().await?, }) } /// get records when `/list` command called or inline button request. pub async fn get_records_by_userid_with_pagination( &self, user_id: i64, page: usize, ) -> Result, DbErr> { let transaction = self.db.begin().await?; if let Some(user) = self.get_user(&user_id, &transaction).await? { let pagination = Record::find() .find_also_related(User) .filter(RecordColumn::UserId.eq(user.id)) .paginate(&transaction, PAGE_SIZE); Ok(Some(PaginatedRecordData { current_data: pagination.fetch_page(page).await?, items_count: pagination.num_items().await?, pages_count: pagination.num_pages().await?, })) } else { log_error_ln!("cannot find user tg_uid={}", user_id); Ok(None) } } /// add record forward a message to bot. pub async fn add_record( &self, user_id: i64, username: &String, text: String, ) -> Result<(), DbErr> { let transaction = self.db.begin().await?; let user = self.setup_user(&user_id, &username, &transaction).await?; RecordActiveModel { message: Set(text), user_id: user.id, ..Default::default() } .insert(&transaction) .await?; transaction.commit().await } /// del record when `/delete` command called. pub async fn del_record(&self, id: i64, user_id: i64) -> Result<(), DbErr> { let transaction = self.db.begin().await?; if let Some(user) = self.get_user(&user_id, &transaction).await? { RecordActiveModel { id: Set(id), user_id: Set(user.id), ..Default::default() } .delete(&transaction) .await?; } transaction.commit().await } pub fn err_handler(&self, error: DbErr) { log_error_ln!("{}", error); } }