Enhance database interaction in coupon and server management. Updated check_unique_server_name to accept session and cluster name parameters for improved uniqueness checks. Introduced new functions: get_coupon_by_code, delete_server, create_coupon_usage, check_coupon_usage, and update_coupon_usage_count for better coupon management. Refactored handlers to utilize these new functions, improving code clarity and maintainability.

2025-01-23 11:28:21 +03:00
parent 906601f5cd
commit 7dafb47c38
4 changed files with 167 additions and 63 deletions
@@ -75,18 +75,26 @@ async def init_db(file_path: str = "assets/schema.sql"):
        await conn.close()


-async def check_unique_server_name(server_name: str) -> bool:
+async def check_unique_server_name(server_name: str, session: Any, cluster_name: str | None = None) -> bool:
    """
    Проверяет уникальность имени сервера.

    :param server_name: Имя сервера.
+    :param session: Сессия базы данных.
+    :param cluster_name: Имя кластера (опционально).
    :return: True, если имя сервера уникально, False, если уже существует.
    """
-    conn = await asyncpg.connect(DATABASE_URL)
-
-    result = await conn.fetchrow("SELECT 1 FROM servers WHERE server_name = $1 LIMIT 1", server_name)
-
-    await conn.close()
+    if cluster_name:
+        result = await session.fetchrow(
+            "SELECT 1 FROM servers WHERE server_name = $1 AND cluster_name = $2 LIMIT 1",
+            server_name,
+            cluster_name
+        )
+    else:
+        result = await session.fetchrow(
+            "SELECT 1 FROM servers WHERE server_name = $1 LIMIT 1",
+            server_name
+        )

    return result is None

@@ -122,6 +130,39 @@ async def create_coupon(coupon_code: str, amount: float, usage_limit: int, sessi
        logger.error(f"Ошибка при создании купона {coupon_code}: {e}")
        raise

+async def get_coupon_by_code(coupon_code: str, session: Any) -> dict | None:
+    """
+    Получает информацию о купоне по его коду.
+
+    Args:
+        coupon_code (str): Код купона для поиска
+        session (Any): Сессия базы данных
+
+    Returns:
+        dict | None: Словарь с информацией о купоне или None, если купон не найден
+            - id (int): ID купона
+            - usage_limit (int): Лимит использований
+            - usage_count (int): Текущее количество использований
+            - is_used (bool): Флаг использования
+            - amount (float): Сумма купона
+
+    Raises:
+        Exception: В случае ошибки при выполнении запроса
+    """
+    try:
+        result = await session.fetchrow(
+            """
+            SELECT id, usage_limit, usage_count, is_used, amount
+            FROM coupons
+            WHERE code = $1 AND (usage_count < usage_limit OR usage_limit = 0) AND is_used = FALSE
+            """,
+            coupon_code,
+        )
+        return dict(result) if result else None
+    except Exception as e:
+        logger.error(f"Ошибка при получении купона {coupon_code}: {e}")
+        raise
+

 async def get_all_coupons(session: Any, page: int = 1, per_page: int = 10):
    """
@@ -1259,3 +1300,110 @@ async def add_server_to_db(
    except Exception as e:
        logger.error(f"Ошибка при добавлении сервера {server_name} в кластер {cluster_name}: {e}")
        raise
+
+async def delete_server(server_name: str, session: Any):
+    """
+    Удаляет сервер из базы данных по его названию.
+
+    Args:
+        server_name (str): Название сервера для удаления
+        session (Any): Сессия базы данных
+
+    Raises:
+        Exception: В случае ошибки при удалении сервера
+    """
+    try:
+        await session.execute(
+            """
+            DELETE FROM servers WHERE server_name = $1
+            """,
+            server_name,
+        )
+        logger.info(f"Сервер {server_name} успешно удалён из базы данных")
+    except Exception as e:
+        logger.error(f"Ошибка при удалении сервера {server_name} из базы данных: {e}")
+        raise
+
+
+async def create_coupon_usage(coupon_id: int, user_id: int, session: Any):
+    """
+    Создаёт запись об использовании купона в базе данных.
+
+    Args:
+        coupon_id (int): ID купона
+        user_id (int): ID пользователя
+        session (Any): Сессия базы данных
+
+    Raises:
+        Exception: В случае ошибки при создании записи
+    """
+    try:
+        await session.execute(
+            """
+            INSERT INTO coupon_usages (coupon_id, user_id, used_at)
+            VALUES ($1, $2, $3)
+            """,
+            coupon_id,
+            user_id,
+            datetime.utcnow(),
+        )
+        logger.info(f"Создана запись об использовании купона {coupon_id} пользователем {user_id}")
+    except Exception as e:
+        logger.error(f"Ошибка при создании записи об использовании купона {coupon_id} пользователем {user_id}: {e}")
+        raise
+
+async def check_coupon_usage(coupon_id: int, user_id: int, session: Any) -> bool:
+    """
+    Проверяет, использовал ли пользователь данный купон.
+
+    Args:
+        coupon_id (int): ID купона для проверки
+        user_id (int): ID пользователя для проверки
+        session (Any): Сессия базы данных
+
+    Returns:
+        bool: True если купон уже использован, False если нет
+
+    Raises:
+        Exception: В случае ошибки при выполнении запроса
+    """
+    try:
+        result = await session.fetchrow(
+            """
+            SELECT 1 FROM coupon_usages WHERE coupon_id = $1 AND user_id = $2
+            """,
+            coupon_id,
+            user_id,
+        )
+        return result is not None
+    except Exception as e:
+        logger.error(f"Ошибка при проверке использования купона {coupon_id} пользователем {user_id}: {e}")
+        raise
+
+
+async def update_coupon_usage_count(coupon_id: int, session: Any):
+    """
+    Обновляет счетчик использования купона и его статус.
+
+    Args:
+        coupon_id (int): ID купона для обновления
+        session (Any): Сессия базы данных
+
+    Raises:
+        Exception: В случае ошибки при обновлении данных купона
+    """
+    try:
+        await session.execute(
+            """
+            UPDATE coupons
+            SET usage_count = usage_count + 1,
+                is_used = CASE WHEN usage_count + 1 >= usage_limit AND usage_limit > 0 THEN TRUE ELSE FALSE END
+            WHERE id = $1
+            """,
+            coupon_id,
+        )
+        logger.info(f"Успешно обновлен счетчик использования купона {coupon_id}")
+    except Exception as e:
+        logger.error(f"Ошибка при обновлении счетчика использования купона {coupon_id}: {e}")
+        raise
+
@@ -11,7 +11,7 @@ from py3xui import AsyncApi

 from backup import create_backup_and_send_to_admins
 from config import ADMIN_PASSWORD, ADMIN_USERNAME, DATABASE_URL
-from database import add_server_to_db, check_unique_server_name, get_keys_by_server, get_servers_from_db
+from database import add_server_to_db, check_unique_server_name, delete_server, get_keys_by_server, get_servers_from_db
 from filters.admin import IsAdminFilter
 from handlers.keys.key_utils import create_key_on_cluster
 from logger import logger
@@ -93,7 +93,7 @@ async def handle_cluster_name_input(message: types.Message, state: FSMContext):


@router.message(UserEditorState.waiting_for_server_name, IsAdminFilter())
-async def handle_server_name_input(message: types.Message, state: FSMContext):
+async def handle_server_name_input(message: types.Message, state: FSMContext, session: Any):
    server_name = message.text.strip()

    if server_name == "❌ Отменить":
@@ -110,7 +110,7 @@ async def handle_server_name_input(message: types.Message, state: FSMContext):
        await message.answer("❌ Имя сервера не может быть пустым. Попробуйте снова.")
        return

-    server_unique = await check_unique_server_name(server_name)
+    server_unique = await check_unique_server_name(server_name, session)
    if not server_unique:
        await message.answer("❌ Сервер с таким именем уже существует. Пожалуйста, выберите другое имя.")
        return
@@ -432,17 +432,10 @@ async def handle_delete_server(callback_query: types.CallbackQuery, state: FSMCo


@router.callback_query(F.data.startswith("confirm_delete_server|"), IsAdminFilter())
-async def handle_confirm_delete_server(callback_query: types.CallbackQuery, state: FSMContext):
+async def handle_confirm_delete_server(callback_query: types.CallbackQuery, state: FSMContext, session: Any):
    server_name = callback_query.data.split("|")[1]

-    conn = await asyncpg.connect(DATABASE_URL)
-    await conn.execute(
-        """
-        DELETE FROM servers WHERE server_name = $1
-        """,
-        server_name,
-    )
-    await conn.close()
+    await delete_server(server_name, session)

    builder = InlineKeyboardBuilder()
    builder.row(InlineKeyboardButton(text="🔙 Назад в управление кластерами", callback_data="servers_editor"))
@@ -7,7 +7,7 @@ from aiogram.fsm.state import State, StatesGroup
 from aiogram.types import InlineKeyboardButton
 from aiogram.utils.keyboard import InlineKeyboardBuilder

-from database import update_balance
+from database import check_coupon_usage, create_coupon_usage, get_coupon_by_code, update_balance, update_coupon_usage_count


 class CouponActivationState(StatesGroup):
@@ -43,51 +43,21 @@ async def process_coupon_code(message: types.Message, state: FSMContext, session


 async def activate_coupon(user_id: int, coupon_code: str, session: Any):
-    coupon_record = await session.fetchrow(
-        """
-        SELECT id, usage_limit, usage_count, is_used, amount
-        FROM coupons
-        WHERE code = $1 AND (usage_count < usage_limit OR usage_limit = 0) AND is_used = FALSE
-    """,
-        coupon_code,
-    )
+    coupon_record = await get_coupon_by_code(coupon_code, session)

    if not coupon_record:
        return "<b>❌ Купон не найден</b> 🚫 или его использование ограничено. 🔒 Пожалуйста, проверьте код и попробуйте снова. 🔍"

-    usage_exists = await session.fetchrow(
-        """
-        SELECT 1 FROM coupon_usages WHERE coupon_id = $1 AND user_id = $2
-    """,
-        coupon_record["id"],
-        user_id,
-    )
+    usage_exists = await check_coupon_usage(coupon_record["id"], user_id, session)
+

    if usage_exists:
        return "<b>❌ Вы уже активировали этот купон.</b> 🚫 Купоны могут быть активированы только один раз. 🔒"

    coupon_amount = coupon_record["amount"]

-    async with session.transaction():
-        await session.execute(
-            """
-            UPDATE coupons
-            SET usage_count = usage_count + 1,
-                is_used = CASE WHEN usage_count + 1 >= usage_limit AND usage_limit > 0 THEN TRUE ELSE FALSE END
-            WHERE id = $1
-        """,
-            coupon_record["id"],
-        )
-
-        await session.execute(
-            """
-            INSERT INTO coupon_usages (coupon_id, user_id, used_at)
-            VALUES ($1, $2, $3)
-        """,
-            coupon_record["id"],
-            user_id,
-            datetime.utcnow(),
-        )
+    await update_coupon_usage_count(coupon_record["id"], session)
+    await create_coupon_usage(coupon_record["id"], user_id, session)

    await update_balance(user_id, coupon_amount, session)
    return f"<b>✅ Купон успешно активирован! 🎉</b>\n\nНа ваш баланс добавлено <b>{coupon_amount} рублей</b> 💰."
@@ -9,7 +9,7 @@ from ping3 import ping

 from bot import bot
 from config import ADMIN_ID, DATABASE_URL, PING_TIME
-from database import add_server_to_db, get_servers_from_db
+from database import add_server_to_db, check_unique_server_name, get_servers_from_db
 from logger import logger

 try:
@@ -34,14 +34,7 @@ async def sync_servers_with_db():

        for cluster_name, servers in CLUSTERS.items():
            for server_key, server_info in servers.items():
-                exists = await conn.fetchval(
-                    """
-                    SELECT 1 FROM servers
-                    WHERE cluster_name = $1 AND server_name = $2
-                    """,
-                    cluster_name,
-                    server_info["name"],
-                )
+                exists = await check_unique_server_name(server_info["name"], conn, cluster_name)

                if not exists:
                    await add_server_to_db(