Added return statement after sending a message about that there is no active message

Now running encrytion of accounts in ProcessPoolExecutor

src/account_parsing.py

@@ -35,7 +35,7 @@ class _Accounts(pydantic.BaseModel):
 def _accounts_list_to_json(accounts: Iterable[DecryptedAccount]) -> str:
     result = _Accounts(
         accounts=[_Account.from_usual_account(i) for i in accounts],
-    ).json(ensure_ascii=False)
+    ).json(ensure_ascii=False, indent=2)
     return result
 
 

src/bot/__init__.py

@@ -3,13 +3,13 @@ import functools
 from sqlalchemy.future import Engine
 from telebot.async_telebot import AsyncTeleBot
 
-from . import callback_handlers, message_handlers
+from . import callback_handlers, exception_handler, message_handlers
 
-__all__ = ["callback_handlers", "message_handlers"]
+__all__ = ["callback_handlers", "exception_handler", "message_handlers"]
 
 
 def create_bot(token: str, engine: Engine) -> AsyncTeleBot:
-    bot = AsyncTeleBot(token)
+    bot = AsyncTeleBot(token, exception_handler=exception_handler.Handler)
     bot.register_message_handler(
         functools.partial(message_handlers.set_master_password, bot, engine),
         commands=["set_master_pass"],

src/bot/exception_handler.py

@@ -0,0 +1,8 @@
+import traceback
+from typing import Type
+
+
+class Handler:
+    @staticmethod
+    def handle(exc: Type[BaseException]) -> None:
+        traceback.print_exception(exc)

src/bot/helper_functions.py

@@ -18,7 +18,7 @@ states: dict[tuple[int, int], Handler] = {}
 
 def register_state(
     message: Message,
-    handler: Callable[[Message], Any],
+    handler: Handler,
 ) -> None:
     states[(message.chat.id, message.from_user.id)] = handler
 
@@ -40,9 +40,8 @@ async def delete_message(
     *,
     sleep_time: int = 0,
 ) -> bool:
+    await asyncio.sleep(sleep_time)
     try:
-        if sleep_time != 0:
-            await asyncio.sleep(sleep_time)
         await bot.delete_message(mes.chat.id, mes.id)
     except telebot.apihelper.ApiException:
         return False

src/bot/message_handlers.py

@@ -1,6 +1,7 @@
 import asyncio
 import functools
 import gc
+import itertools
 from concurrent.futures import ProcessPoolExecutor
 
 import telebot
@@ -540,7 +541,6 @@ async def _export2(
             )
             tasks.append(loop.run_in_executor(pool, function))
         accounts = await asyncio.gather(*tasks)
-        accounts.sort(key=lambda account: account.name)
     json_io = accounts_to_json(accounts)
     await bot.send_document(
         mes.chat.id,
@@ -640,22 +640,35 @@ async def _import3(
     # List of names of accounts, which failed to be added to the database
     # or failed the tests
     failed: list[str] = []
-    for account in accounts:
-        if not check_account(account):
-            failed.append(account.name)
-            continue
-        account = encryption.accounts.encrypt(account, text)
-        result = db.add.add_account(engine, account)
-        if not result:
-            failed.append(account.name)
+    tasks: list[asyncio.Future[db.models.Account]] = []
+    loop = asyncio.get_running_loop()
+    with ProcessPoolExecutor() as pool:
+        for account in accounts:
+            if not check_account(account):
+                failed.append(account.name)
+                continue
+            function = functools.partial(
+                encryption.accounts.encrypt,
+                account,
+                text,
+            )
+            tasks.append(loop.run_in_executor(pool, function))
+        enc_accounts: list[db.models.Account] = await asyncio.gather(*tasks)
+    results = db.add.add_accounts(engine, enc_accounts)
+
+    failed_accounts = itertools.compress(
+        enc_accounts, (not result for result in results)
+    )
+    failed.extend((account.name for account in failed_accounts))
 
     if failed:
-        mes_text = "Не удалось добавить:\n" + "\n".join(failed)
+        await send_deleteable_message(
+            bot, mes.chat.id, "Не удалось добавить:\n" + "\n".join(failed)
+        )
     else:
-        mes_text = "Успех"
+        await send_tmp_message(bot, mes.chat.id, "Успех")
 
-    await send_tmp_message(bot, mes.chat.id, mes_text, 10)
-    del text, mes, accounts
+    del text, mes, accounts, function, tasks, failed_accounts
     gc.collect()
 
 
@@ -677,6 +690,7 @@ async def message_handler(bot: AsyncTeleBot, mes: Message) -> None:
         await delete_message(bot, mes)
         if mes.text.strip() == "/cancel":
             await send_tmp_message(bot, mes.chat.id, "Нет активного действия")
+            return
         await send_tmp_message(
             bot,
             mes.chat.id,

src/db/add.py

@@ -5,27 +5,42 @@ from sqlalchemy.future import Engine
 from . import models
 
 
-def add_account(engine: Engine, account: models.Account) -> bool:
-    """Adds account to the database. Returns true on success,
+def _add_model(
+    session: sqlmodel.Session, model: models.Account | models.MasterPass
+) -> bool:
+    """Adds model to the session. Returns true on success,
     false otherwise"""
     try:
-        with sqlmodel.Session(engine) as session:
-            session.add(account)
-            session.commit()
+        session.add(model)
     except IntegrityError:
         return False
     else:
         return True
 
 
+def add_account(engine: Engine, account: models.Account) -> bool:
+    """Adds account to the database. Returns true on success,
+    false otherwise"""
+    with sqlmodel.Session(engine) as session:
+        result = _add_model(session, account)
+        session.commit()
+    return result
+
+
 def add_master_pass(engine: Engine, master_pass: models.MasterPass) -> bool:
     """Adds master password the database. Returns true on success,
     false otherwise"""
-    try:
-        with sqlmodel.Session(engine) as session:
-            session.add(master_pass)
-            session.commit()
-    except IntegrityError:
-        return False
-    else:
-        return True
+    with sqlmodel.Session(engine) as session:
+        result = _add_model(session, master_pass)
+        session.commit()
+    return result
+
+
+def add_accounts(
+    engine: Engine,
+    accounts: list[models.Account],
+) -> list[bool]:
+    with sqlmodel.Session(engine) as session:
+        result = [_add_model(session, account) for account in accounts]
+        session.commit()
+    return result

src/db/get.py

@@ -36,8 +36,12 @@ def get_account_names(
 
 def get_accounts(engine: Engine, user_id: int) -> list[models.Account]:
     """Returns a list of accounts of a user"""
-    statement = sqlmodel.select(models.Account).where(
-        models.Account.user_id == user_id,
+    statement = (
+        sqlmodel.select(models.Account)
+        .where(
+            models.Account.user_id == user_id,
+        )
+        .order_by(models.Account.name)
     )
     with sqlmodel.Session(engine) as session:
         result = session.exec(statement).fetchall()

src/db/models.py

@@ -4,7 +4,11 @@ import sqlmodel
 class MasterPass(sqlmodel.SQLModel, table=True):
     __tablename__ = "master_passwords"
     user_id: int = sqlmodel.Field(
-        sa_column=sqlmodel.Column(sqlmodel.INT(), primary_key=True)
+        sa_column=sqlmodel.Column(
+            sqlmodel.INT(),
+            primary_key=True,
+            autoincrement=False,
+        )
     )
     salt: bytes = sqlmodel.Field(
         sa_column=sqlmodel.Column(sqlmodel.BINARY(64), nullable=False)