Updated import and export commands to work better

This commit is contained in:
StNicolay 2023-05-06 19:21:40 +03:00
parent fc76d7f65c
commit 8af4d0ab12
Signed by: StNicolay
GPG Key ID: 9693D04DCD962B0D
3 changed files with 41 additions and 34 deletions

View File

@ -96,7 +96,7 @@ impl Entity {
.filter(Column::UserId.eq(user_id)) .filter(Column::UserId.eq(user_id))
.stream(db) .stream(db)
.await?; .await?;
Ok(result.map_err(Into::into)) Ok(result.err_into())
} }
/// Gets a list of account names of a user /// Gets a list of account names of a user
@ -113,7 +113,7 @@ impl Entity {
select = select.order_by_asc(Column::Name); select = select.order_by_asc(Column::Name);
} }
let result = select.into_tuple().stream(db).await?; let result = select.into_tuple().stream(db).await?;
Ok(result.map_err(Into::into)) Ok(result.err_into())
} }
pub async fn exists( pub async fn exists(

View File

@ -8,7 +8,7 @@ use sea_orm::DatabaseConnection;
use serde_json::to_string_pretty; use serde_json::to_string_pretty;
use std::sync::Arc; use std::sync::Arc;
use teloxide::{adaptors::Throttle, prelude::*, types::InputFile}; use teloxide::{adaptors::Throttle, prelude::*, types::InputFile};
use tokio::task::JoinSet; use tokio::{sync::Mutex, task::spawn_blocking};
async fn get_master_pass( async fn get_master_pass(
bot: Throttle<Bot>, bot: Throttle<Bot>,
@ -21,21 +21,27 @@ async fn get_master_pass(
let _ = bot.delete_message(previous.chat.id, previous.id).await; let _ = bot.delete_message(previous.chat.id, previous.id).await;
let master_pass: Arc<str> = master_pass.into(); let master_pass: Arc<str> = master_pass.into();
let user_id = msg.from().unwrap().id.0; let user_id = msg.from().unwrap().id.0;
let mut join_set = JoinSet::new(); let accounts = Arc::new(Mutex::new(Vec::new()));
let mut accounts = Vec::new();
Account::get_all(user_id, &db) Account::get_all(user_id, &db)
.await? .await?
.try_for_each(|account| { .try_for_each_concurrent(None, |account| {
let master_pass = Arc::clone(&master_pass); let master_pass = Arc::clone(&master_pass);
join_set.spawn_blocking(move || DecryptedAccount::from_account(account, &master_pass)); let accounts = Arc::clone(&accounts);
async { crate::Result::Ok(()) } async move {
let account =
spawn_blocking(move || DecryptedAccount::from_account(account, &master_pass))
.await??;
accounts.lock().await.push(account);
Ok(())
}
}) })
.await?; .await?;
drop(master_pass); drop(master_pass);
while let Some(account) = join_set.join_next().await.transpose()?.transpose()? { let mut accounts = match Arc::try_unwrap(accounts) {
accounts.push(account) Ok(account) => account.into_inner(),
} Err(_) => unreachable!(),
accounts.sort_by(|this, other| this.name.cmp(&other.name)); };
accounts.sort_unstable_by(|this, other| this.name.cmp(&other.name));
let json = to_string_pretty(&User { accounts })?; let json = to_string_pretty(&User { accounts })?;
let file = InputFile::memory(json).file_name("accounts.json"); let file = InputFile::memory(json).file_name("accounts.json");
bot.send_document(msg.chat.id, file).await?; bot.send_document(msg.chat.id, file).await?;

View File

@ -2,13 +2,13 @@ use crate::{
handlers::{utils::package_handler, MainDialogue, State}, handlers::{utils::package_handler, MainDialogue, State},
models::{DecryptedAccount, User}, models::{DecryptedAccount, User},
}; };
use futures::TryStreamExt; use futures::{stream, StreamExt, TryStreamExt};
use itertools::Itertools; use itertools::Itertools;
use sea_orm::prelude::*; use sea_orm::prelude::*;
use serde_json::from_slice; use serde_json::from_slice;
use std::sync::Arc; use std::sync::Arc;
use teloxide::{adaptors::Throttle, net::Download, prelude::*}; use teloxide::{adaptors::Throttle, net::Download, prelude::*};
use tokio::task::{spawn_blocking, JoinSet}; use tokio::{sync::Mutex, task::spawn_blocking};
async fn get_master_pass( async fn get_master_pass(
bot: Throttle<Bot>, bot: Throttle<Bot>,
@ -22,28 +22,29 @@ async fn get_master_pass(
let _ = bot.delete_message(previous.chat.id, previous.id).await; let _ = bot.delete_message(previous.chat.id, previous.id).await;
let user_id = msg.from().unwrap().id.0; let user_id = msg.from().unwrap().id.0;
let master_pass: Arc<str> = master_pass.into(); let master_pass: Arc<str> = master_pass.into();
let mut join_set = JoinSet::new(); let failed = Arc::new(Mutex::new(Vec::new()));
let mut failed = Vec::new(); stream::iter(accounts)
for account in accounts { .for_each_concurrent(None, |account| {
let master_pass = Arc::clone(&master_pass); let master_pass = Arc::clone(&master_pass);
let db = db.clone(); let failed = Arc::clone(&failed);
join_set.spawn(async move { let db = db.clone();
let name = account.name.clone(); let name = account.name.clone();
let account = async move {
match spawn_blocking(move || account.into_account(user_id, &master_pass)).await { match spawn_blocking(move || account.into_account(user_id, &master_pass)).await {
Ok(Ok(account)) => account, Ok(Ok(account)) => match account.insert(&db).await {
_ => return Err(name), Ok(_) => (),
}; Err(_) => failed.lock().await.push(name),
account.insert(&db).await.map_err(|_| name)?; },
Ok(()) _ => failed.lock().await.push(name),
}); }
} }
while let Some(result) = join_set.join_next().await.transpose()? { })
match result { .await;
Ok(()) => (), drop(master_pass);
Err(name) => failed.push(name), let failed = match Arc::try_unwrap(failed) {
} Ok(accounts) => accounts.into_inner(),
} Err(_) => unreachable!(),
};
if failed.is_empty() { if failed.is_empty() {
bot.send_message(msg.chat.id, "Success").await?; bot.send_message(msg.chat.id, "Success").await?;
} else { } else {