From bb803ab1450de3c9ff1f2b73663d3f27caef9f30 Mon Sep 17 00:00:00 2001 From: Zakhar Izmaylov Date: Thu, 23 Jan 2025 11:58:35 +0300 Subject: [PATCH] Implement new database functions for server and coupon management. Added `check_server_name_by_cluster`, `get_coupon_details`, `get_referral_by_referred_id`, and `get_all_keys` to enhance data retrieval capabilities. Updated handlers to utilize these new functions, improving code clarity and maintainability. Refactored existing database calls in `start.py`, `utils.py`, and `keys.py` for consistency and better session management. --- database.py | 130 +++++++++++++++++++++++++++++++++++++++++- handlers/keys/keys.py | 20 ++----- handlers/start.py | 12 ++-- handlers/utils.py | 4 +- 4 files changed, 137 insertions(+), 29 deletions(-) diff --git a/database.py b/database.py index e1304589..ed8cb718 100644 --- a/database.py +++ b/database.py @@ -94,6 +94,37 @@ async def check_unique_server_name(server_name: str, session: Any, cluster_name: return result is None +async def check_server_name_by_cluster(server_name: str, session: Any) -> dict | None: + """ + Проверяет принадлежность сервера к кластеру. + + Args: + server_name (str): Имя сервера для проверки + session (Any): Сессия базы данных + + Returns: + dict | None: Словарь с информацией о кластере или None, если сервер не найден + - cluster_name (str): Название кластера + """ + try: + cluster_info = await session.fetchrow( + """ + SELECT cluster_name + FROM servers + WHERE server_name = $1 + """, + server_name, + ) + if cluster_info: + logger.info(f"Найден кластер для сервера {server_name}") + return dict(cluster_info) + logger.info(f"Кластер для сервера {server_name} не найден") + return None + except Exception as e: + logger.error(f"Ошибка при поиске кластера для сервера {server_name}: {e}") + raise + + async def create_coupon(coupon_code: str, amount: float, usage_limit: int, session: Any): """ Создает новый купон в базе данных. @@ -1409,8 +1440,6 @@ async def update_coupon_usage_count(coupon_id: int, session: Any): raise - - async def get_last_payments(tg_id: int, session: Any): """ Получает последние 3 платежа пользователя. @@ -1434,7 +1463,7 @@ async def get_last_payments(tg_id: int, session: Any): ORDER BY created_at DESC LIMIT 3 """, - tg_id + tg_id, ) logger.info(f"Успешно получены последние платежи для пользователя {tg_id}") return records @@ -1443,4 +1472,99 @@ async def get_last_payments(tg_id: int, session: Any): raise +async def get_coupon_details(coupon_id: str, session: Any): + """ + Получает детали купона по его ID. + Args: + coupon_id (str): ID купона + session (Any): Сессия базы данных + + Returns: + dict: Словарь с деталями купона или None если купон не найден + + Raises: + Exception: В случае ошибки при выполнении запроса + """ + try: + record = await session.fetchrow( + """ + SELECT id, code, discount, usage_count, usage_limit, is_used + FROM coupons + WHERE id = $1 + """, + coupon_id, + ) + + if record: + logger.info(f"Успешно получены детали купона {coupon_id}") + return dict(record) + + logger.warning(f"Купон {coupon_id} не найден") + return None + + except Exception as e: + logger.error(f"Ошибка при получении деталей купона {coupon_id}: {e}") + raise + + +async def get_referral_by_referred_id(referred_tg_id: int, session: Any): + """ + Получает информацию о реферале по ID приглашенного пользователя. + + Args: + referred_tg_id (int): ID приглашенного пользователя + session (Any): Сессия базы данных + + Returns: + dict: Словарь с информацией о реферале или None если не найден + + Raises: + Exception: В случае ошибки при выполнении запроса + """ + try: + record = await session.fetchrow( + """ + SELECT * FROM referrals + WHERE referred_tg_id = $1 + """, + referred_tg_id, + ) + + if record: + logger.info(f"Успешно получена информация о реферале для пользователя {referred_tg_id}") + return dict(record) + + logger.info(f"Реферал для пользователя {referred_tg_id} не найден") + return None + + except Exception as e: + logger.error(f"Ошибка при получении информации о реферале для пользователя {referred_tg_id}: {e}") + raise + + +async def get_all_keys(session: Any = None): + """ + Получает все записи из таблицы keys. + + Args: + session (Any, optional): Сессия базы данных. По умолчанию None. + + Returns: + list: Список всех записей из таблицы keys + + Raises: + Exception: В случае ошибки при выполнении запроса + """ + conn = None + try: + conn = session if session is not None else await asyncpg.connect(DATABASE_URL) + keys = await conn.fetch("SELECT * FROM keys") + logger.info(f"Успешно получены все записи из таблицы keys. Количество: {len(keys)}") + return keys + except Exception as e: + logger.error(f"Ошибка при получении записей из таблицы keys: {e}") + raise + finally: + if conn is not None and session is None: + await conn.close() diff --git a/handlers/keys/keys.py b/handlers/keys/keys.py index 8213f7eb..bece4e4f 100644 --- a/handlers/keys/keys.py +++ b/handlers/keys/keys.py @@ -27,9 +27,11 @@ from config import ( USE_NEW_PAYMENT_FLOW, ) from database import ( + check_server_name_by_cluster, delete_key, get_balance, get_key_details, + get_keys, get_keys_by_server, get_servers, create_temporary_data, @@ -83,14 +85,7 @@ async def process_callback_or_message_view_keys( send_photo = callback_query_or_message.answer_photo try: - records = await session.fetch( - """ - SELECT email, client_id, expiry_time - FROM keys - WHERE tg_id = $1 - """, - chat_id, - ) + records = await get_keys(chat_id, session) inline_keyboard, response_message = build_keys_response(records) @@ -512,14 +507,7 @@ async def complete_key_renewal(tg_id, client_id, email, new_expiry_time, total_g server_id = key_info["server_id"] if USE_COUNTRY_SELECTION: - cluster_info = await conn.fetchrow( - """ - SELECT cluster_name - FROM servers - WHERE server_name = $1 - """, - server_id, - ) + cluster_info = await check_server_name_by_cluster(server_id, conn) if not cluster_info: logger.error(f"[RENEW] Сервер {server_id} не найден в таблице servers.") diff --git a/handlers/start.py b/handlers/start.py index 25b83bfd..de10c227 100644 --- a/handlers/start.py +++ b/handlers/start.py @@ -28,6 +28,8 @@ from database import ( add_connection, add_referral, check_connection_exists, + get_coupon_details, + get_referral_by_referred_id, get_trial, set_trial, ) @@ -84,10 +86,7 @@ async def start_command(message: Message, state: FSMContext, session: Any, admin recipient_tg_id = message.chat.id - gift_info = await session.fetchrow( - "SELECT * FROM gifts WHERE gift_id = $1 AND is_used = FALSE", - gift_id, - ) + gift_info = await get_coupon_details(gift_id, session) if gift_info is None: logger.warning(f"Подарок с ID {gift_id} уже был использован или не существует.") @@ -137,10 +136,7 @@ async def start_command(message: Message, state: FSMContext, session: Any, admin await message.answer("❌ Вы не можете быть рефералом самого себя.") return await show_start_menu(message, admin, session) - existing_referral = await session.fetchrow( - "SELECT * FROM referrals WHERE referred_tg_id = $1", - message.chat.id, - ) + existing_referral = await get_referral_by_referred_id(message.chat.id, session) if existing_referral: logger.info(f"Реферал с ID {message.chat.id} уже существует.") diff --git a/handlers/utils.py b/handlers/utils.py index dba6fb7b..29d525f7 100644 --- a/handlers/utils.py +++ b/handlers/utils.py @@ -7,7 +7,7 @@ import asyncpg from bot import bot from config import DATABASE_URL -from database import get_servers +from database import get_all_keys, get_servers from logger import logger @@ -65,7 +65,7 @@ async def get_least_loaded_cluster() -> str: async with asyncpg.create_pool(DATABASE_URL) as pool: async with pool.acquire() as conn: - keys = await conn.fetch("SELECT server_id FROM keys") + keys = await get_all_keys(conn) for key in keys: cluster_id = key["server_id"] if cluster_id in cluster_loads: