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.

2025-01-23 11:58:35 +03:00
parent 56af0a7885
commit bb803ab145
4 changed files with 137 additions and 29 deletions
@@ -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()
@@ -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.")
@@ -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} уже существует.")
@@ -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: