From 17d519d44016690806df7012f45b4e9dce7c9f8e Mon Sep 17 00:00:00 2001 From: Zakhar Izmaylov Date: Thu, 23 Jan 2025 08:41:44 +0300 Subject: [PATCH] Refactor key retrieval and management logic across multiple handlers. Update `get_key_details` to include email in the fetched record and replace direct SQL queries with dedicated functions for improved code clarity and maintainability. Add a TODO comment in `admin_servers.py` for future enhancements. --- database.py | 3 ++- handlers/admin/admin_servers.py | 1 + handlers/admin/admin_user_editor.py | 9 +++++---- handlers/keys/key_management.py | 6 +----- handlers/keys/keys.py | 25 ++----------------------- 5 files changed, 11 insertions(+), 33 deletions(-) diff --git a/database.py b/database.py index 0b21d734..ec8261d6 100644 --- a/database.py +++ b/database.py @@ -1234,7 +1234,7 @@ async def store_gift_link( async def get_key_details(email, session): record = await session.fetchrow( """ - SELECT k.key, k.expiry_time, k.server_id,k, k.client_id, k.created_at, c.tg_id, c.balance + SELECT k.key, k.email, k.expiry_time, k.server_id,k, k.client_id, k.created_at, c.tg_id, c.balance FROM keys k JOIN connections c ON k.tg_id = c.tg_id WHERE k.email = $1 @@ -1270,6 +1270,7 @@ async def get_key_details(email, session): "server_name": cluster_name, "balance": record["balance"], "tg_id": record["tg_id"], + "email": record["email"], } diff --git a/handlers/admin/admin_servers.py b/handlers/admin/admin_servers.py index 2e8aefc4..1fd5dcd9 100644 --- a/handlers/admin/admin_servers.py +++ b/handlers/admin/admin_servers.py @@ -300,6 +300,7 @@ async def sync_cluster_handler(callback_query: types.CallbackQuery): conn = await asyncpg.connect(DATABASE_URL) try: + #TODO query_keys = """ SELECT tg_id, client_id, email, expiry_time FROM keys diff --git a/handlers/admin/admin_user_editor.py b/handlers/admin/admin_user_editor.py index 64c7b6e1..58af54ab 100644 --- a/handlers/admin/admin_user_editor.py +++ b/handlers/admin/admin_user_editor.py @@ -14,6 +14,7 @@ from database import ( delete_user_data, get_client_id_by_email, get_key_details, + get_keys, get_servers_from_db, restore_trial, update_key_expiry, @@ -74,7 +75,7 @@ async def handle_username_input(message: types.Message, state: FSMContext, sessi tg_id = user_record["tg_id"] username = await session.fetchval("SELECT username FROM users WHERE tg_id = $1", tg_id) balance = await session.fetchval("SELECT balance FROM connections WHERE tg_id = $1", tg_id) - key_records = await session.fetch("SELECT email FROM keys WHERE tg_id = $1", tg_id) + key_records = await get_keys(tg_id, session) referral_count = await session.fetchval("SELECT COUNT(*) FROM referrals WHERE referrer_tg_id = $1", tg_id) if balance is None: @@ -154,7 +155,7 @@ async def handle_tg_id_input(message: types.Message, state: FSMContext, session: tg_id = int(message.text) username = await session.fetchval("SELECT username FROM users WHERE tg_id = $1", tg_id) balance = await session.fetchval("SELECT balance FROM connections WHERE tg_id = $1", tg_id) - key_records = await session.fetch("SELECT email FROM keys WHERE tg_id = $1", tg_id) + key_records = await get_keys(tg_id, session) referral_count = await session.fetchval("SELECT COUNT(*) FROM referrals WHERE referrer_tg_id = $1", tg_id) if balance is None: @@ -398,7 +399,7 @@ async def handle_expiry_time_input(message: types.Message, state: FSMContext, se await state.clear() return - record = await session.fetchrow("SELECT server_id FROM keys WHERE client_id = $1", client_id) + record = await get_key_details(email, session) if not record: builder = InlineKeyboardBuilder() builder.row(InlineKeyboardButton(text="🔙 Назад", callback_data="user_editor")) @@ -512,7 +513,7 @@ async def handle_user_info(callback_query: types.CallbackQuery, state: FSMContex tg_id = int(callback_query.data.split("|")[1]) username = await session.fetchval("SELECT username FROM users WHERE tg_id = $1", tg_id) balance = await session.fetchval("SELECT balance FROM connections WHERE tg_id = $1", tg_id) - key_records = await session.fetch("SELECT email FROM keys WHERE tg_id = $1", tg_id) + key_records = await get_keys(tg_id, session) referral_count = await session.fetchval("SELECT COUNT(*) FROM referrals WHERE referrer_tg_id = $1", tg_id) builder = InlineKeyboardBuilder() diff --git a/handlers/keys/key_management.py b/handlers/keys/key_management.py index a06f004f..70c775ec 100644 --- a/handlers/keys/key_management.py +++ b/handlers/keys/key_management.py @@ -228,11 +228,7 @@ async def create_key( key_name = generate_random_email() logger.info(f"[Key Generation] Сгенерировано имя ключа: {key_name} для пользователя {tg_id}") - existing_key = await session.fetchrow( - "SELECT * FROM keys WHERE email = $1 AND tg_id = $2", - key_name, - tg_id, - ) + existing_key = await get_key_details(key_name, session) if not existing_key: break logger.warning(f"[Key Generation] Имя ключа {key_name} уже существует. Генерация нового.") diff --git a/handlers/keys/keys.py b/handlers/keys/keys.py index 20c036fb..af4e6b56 100644 --- a/handlers/keys/keys.py +++ b/handlers/keys/keys.py @@ -264,20 +264,7 @@ async def process_callback_update_subscription(callback_query: types.CallbackQue client_id = record["client_id"] public_link = f"{PUBLIC_LINK}{email}/{tg_id}" - try: - await session.execute( - """ - DELETE FROM keys - WHERE tg_id = $1 AND email = $2 - """, - tg_id, - email, - ) - except Exception as delete_error: - await callback_query.message.answer( - f"Ошибка при удалении старой подписки: {delete_error}", - ) - return + await delete_key(client_id, session) least_loaded_cluster_id = await get_least_loaded_cluster() @@ -517,15 +504,7 @@ async def complete_key_renewal(tg_id, client_id, email, new_expiry_time, total_g await bot.send_message(tg_id, response_message, reply_markup=builder.as_markup()) conn = await asyncpg.connect(DATABASE_URL) - key_info = await conn.fetchrow( - """ - SELECT server_id - FROM keys - WHERE tg_id = $1 AND client_id = $2 - """, - tg_id, - client_id, - ) + key_info = await get_key_details(email, conn) if not key_info: logger.error(f"[RENEW] Ключ с client_id {client_id} для пользователя {tg_id} не найден.")