From 17e4f1892cec064b5f43339830648619b3013bc4 Mon Sep 17 00:00:00 2001
From: StNicolay <stnicolay2007@mail.ru>
Date: Thu, 4 May 2023 18:27:19 +0300
Subject: [PATCH] Added more paralization to allow better multi-user perfomace

---
 Cargo.lock                            |  1 +
 Cargo.toml                            |  1 +
 src/entity/account.rs                 | 15 ++++++++++-----
 src/handlers/commands/add_account.rs  | 13 ++++++-------
 src/handlers/commands/get_account.rs  |  6 +++---
 src/handlers/commands/get_accounts.rs | 17 +++++++++--------
 6 files changed, 30 insertions(+), 23 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index fd5582e..7f50da1 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1569,6 +1569,7 @@ dependencies = [
  "anyhow",
  "chacha20poly1305",
  "dotenv",
+ "futures",
  "log",
  "migration",
  "pbkdf2",
diff --git a/Cargo.toml b/Cargo.toml
index cd06d38..5a236bb 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -15,6 +15,7 @@ members = [".", "migration"]
 anyhow = "1.0.70"
 chacha20poly1305 = { version = "0.10.1", features = ["std"] }
 dotenv = "0.15.0"
+futures = "0.3.28"
 log = "0.4.17"
 migration = { version = "0.1.0", path = "migration" }
 pbkdf2 = "0.12.1"
diff --git a/src/entity/account.rs b/src/entity/account.rs
index 2ba1ea7..9147915 100644
--- a/src/entity/account.rs
+++ b/src/entity/account.rs
@@ -1,6 +1,7 @@
 //! `SeaORM` Entity. Generated by sea-orm-codegen 0.11.2
 
 use chacha20poly1305::{aead::Aead, AeadCore, ChaCha20Poly1305, KeyInit};
+use futures::{Stream, StreamExt};
 use pbkdf2::pbkdf2_hmac_array;
 use rand::{rngs::OsRng, RngCore};
 use sea_orm::{prelude::*, ActiveValue::Set, QuerySelect};
@@ -88,14 +89,18 @@ impl Model {
 
 impl Entity {
     /// Gets a list of account names of a user
-    pub async fn get_names(user_id: u64, db: &DatabaseConnection) -> crate::Result<Vec<String>> {
-        Self::find()
+    pub async fn get_names(
+        user_id: u64,
+        db: &DatabaseConnection,
+    ) -> crate::Result<impl Stream<Item = crate::Result<String>> + '_> {
+        let result = Self::find()
             .select_only()
             .column(Column::Name)
             .filter(Column::UserId.eq(user_id))
             .into_tuple()
-            .all(db)
-            .await
-            .map_err(|err| err.into())
+            .stream(db)
+            .await?
+            .map(|item| item.map_err(Into::into));
+        Ok(result)
     }
 }
diff --git a/src/handlers/commands/add_account.rs b/src/handlers/commands/add_account.rs
index 26d9a7c..627bbe4 100644
--- a/src/handlers/commands/add_account.rs
+++ b/src/handlers/commands/add_account.rs
@@ -2,6 +2,7 @@ use crate::entity::account;
 use crate::handlers::{utils::package_handler, MainDialogue, State};
 use sea_orm::prelude::*;
 use teloxide::{adaptors::Throttle, prelude::*};
+use tokio::task::spawn_blocking;
 
 async fn get_master_pass(
     bot: Throttle<Bot>,
@@ -13,15 +14,13 @@ async fn get_master_pass(
     password: String,
     master_pass: String,
 ) -> crate::Result<()> {
+    let user_id = msg.from().unwrap().id.0;
     let _ = bot.delete_message(msg.chat.id, msg.id).await;
     dialogue.exit().await?;
-    let account = account::ActiveModel::from_unencrypted(
-        msg.from().unwrap().id.0,
-        name,
-        &login,
-        &password,
-        &master_pass,
-    )?;
+    let account = spawn_blocking(move || {
+        account::ActiveModel::from_unencrypted(user_id, name, &login, &password, &master_pass)
+    })
+    .await??;
     account.insert(&db).await?;
     bot.send_message(msg.chat.id, "Success").await?;
     Ok(())
diff --git a/src/handlers/commands/get_account.rs b/src/handlers/commands/get_account.rs
index c6f91ad..64cc1a2 100644
--- a/src/handlers/commands/get_account.rs
+++ b/src/handlers/commands/get_account.rs
@@ -1,8 +1,8 @@
 use crate::entity::{account, prelude::Account};
+use crate::handlers::{utils::package_handler, MainDialogue, State};
 use sea_orm::prelude::*;
 use teloxide::{adaptors::Throttle, prelude::*, types::ParseMode};
-
-use crate::handlers::{utils::package_handler, MainDialogue, State};
+use tokio::task::spawn_blocking;
 
 async fn get_master_pass(
     bot: Throttle<Bot>,
@@ -20,7 +20,7 @@ async fn get_master_pass(
         .one(&db)
         .await?
         .unwrap();
-    let (login, password) = account.decrypt(&master_pass)?;
+    let (login, password) = spawn_blocking(move || account.decrypt(&master_pass)).await??;
     let message = format!("Name:\n`{name}`\nLogin:\n`{login}`\nPassword:\n`{password}`");
     bot.send_message(msg.chat.id, message)
         .parse_mode(ParseMode::MarkdownV2)
diff --git a/src/handlers/commands/get_accounts.rs b/src/handlers/commands/get_accounts.rs
index 821c893..e275740 100644
--- a/src/handlers/commands/get_accounts.rs
+++ b/src/handlers/commands/get_accounts.rs
@@ -1,4 +1,5 @@
 use crate::entity::prelude::Account;
+use futures::TryStreamExt;
 use sea_orm::prelude::*;
 use teloxide::{adaptors::Throttle, prelude::*, types::ParseMode};
 
@@ -8,20 +9,20 @@ pub async fn get_accounts(
     db: DatabaseConnection,
 ) -> crate::Result<()> {
     let user_id = msg.from().unwrap().id.0;
-    let mut account_names = Account::get_names(user_id, &db).await?.into_iter();
-    let mut result = match account_names.next() {
+    let mut account_names = Account::get_names(user_id, &db).await?;
+    let mut result = match account_names.try_next().await? {
         Some(name) => format!("Accounts:\n`{name}`"),
         None => {
             bot.send_message(msg.chat.id, "No accounts found").await?;
             return Ok(());
         }
     };
-    for name in account_names {
-        result.reserve(name.len() + 3);
-        result.push_str("\n`");
-        result.push_str(&name);
-        result.push('`')
-    }
+    account_names
+        .try_for_each(|name| {
+            result.push_str(&format!("\n`{}`", name));
+            async { Ok(()) }
+        })
+        .await?;
     bot.send_message(msg.chat.id, result)
         .parse_mode(ParseMode::MarkdownV2)
         .await?;