_export2 no longer blocks event loop during decryption
Removed sorting in get_all_accounts Removed decrypt_multiple function because it is no longer used Now running decrytion of accounts in ProcessPoolExecutor
This commit is contained in:
		| @@ -1,5 +1,7 @@ | |||||||
|  | import asyncio | ||||||
| import functools | import functools | ||||||
| import gc | import gc | ||||||
|  | from concurrent.futures import ProcessPoolExecutor | ||||||
|  |  | ||||||
| import telebot | import telebot | ||||||
| from sqlalchemy.future import Engine | from sqlalchemy.future import Engine | ||||||
| @@ -534,7 +536,16 @@ async def _export2( | |||||||
|         ) |         ) | ||||||
|  |  | ||||||
|     accounts = database.get.get_all_accounts(engine, mes.from_user.id) |     accounts = database.get.get_all_accounts(engine, mes.from_user.id) | ||||||
|     accounts = encryption.other_accounts.decrypt_multiple(accounts, text) |     with ProcessPoolExecutor() as pool: | ||||||
|  |         loop = asyncio.get_running_loop() | ||||||
|  |         tasks = [] | ||||||
|  |         for account in accounts: | ||||||
|  |             function = functools.partial( | ||||||
|  |                 encryption.other_accounts.decrypt, account, text | ||||||
|  |             ) | ||||||
|  |             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) |     json_io = accounts_to_json(accounts) | ||||||
|     await bot.send_document( |     await bot.send_document( | ||||||
|         mes.chat.id, |         mes.chat.id, | ||||||
|   | |||||||
| @@ -36,13 +36,9 @@ def get_accounts( | |||||||
|  |  | ||||||
| def get_all_accounts(engine: Engine, user_id: int) -> list[models.Account]: | def get_all_accounts(engine: Engine, user_id: int) -> list[models.Account]: | ||||||
|     """Returns a list of accounts of a user""" |     """Returns a list of accounts of a user""" | ||||||
|     statement = ( |     statement = sqlmodel.select(models.Account).where( | ||||||
|         sqlmodel.select(models.Account) |  | ||||||
|         .where( |  | ||||||
|         models.Account.user_id == user_id, |         models.Account.user_id == user_id, | ||||||
|     ) |     ) | ||||||
|         .order_by(models.Account.name) |  | ||||||
|     ) |  | ||||||
|     with sqlmodel.Session(engine) as session: |     with sqlmodel.Session(engine) as session: | ||||||
|         result = session.exec(statement).fetchall() |         result = session.exec(statement).fetchall() | ||||||
|     return result |     return result | ||||||
|   | |||||||
| @@ -1,6 +1,5 @@ | |||||||
| import base64 | import base64 | ||||||
| import os | import os | ||||||
| from typing import Iterable, Iterator |  | ||||||
|  |  | ||||||
| from cryptography.fernet import Fernet | from cryptography.fernet import Fernet | ||||||
| from cryptography.hazmat.backends import default_backend | from cryptography.hazmat.backends import default_backend | ||||||
| @@ -71,12 +70,3 @@ def decrypt( | |||||||
|         login=login, |         login=login, | ||||||
|         password=password, |         password=password, | ||||||
|     ) |     ) | ||||||
|  |  | ||||||
|  |  | ||||||
| def decrypt_multiple( |  | ||||||
|     accounts: Iterable[Account], master_pass: str |  | ||||||
| ) -> Iterator[DecryptedAccount]: |  | ||||||
|     """Decrypts an iterable of accounts using master_pass and |  | ||||||
|     returns an Iterator of decrypted accounts""" |  | ||||||
|     for account in accounts: |  | ||||||
|         yield decrypt(account, master_pass) |  | ||||||
|   | |||||||
		Reference in New Issue
	
	Block a user