Refactor key management functions to improve database interactions by consolidating key deletion logic into a single function. Introduce timezone-aware expiry date calculations for keys and enhance key retrieval processes across various handlers. Update the start command to utilize a dedicated function for adding connections, ensuring cleaner code and better session management.

2025-01-23 07:49:21 +03:00
parent 311c48edcd
commit 2d638b0c82
10 changed files with 95 additions and 158 deletions
@@ -3,6 +3,7 @@ from datetime import datetime
 from typing import Any

 import asyncpg
+import pytz

 from config import DATABASE_URL, REFERRAL_BONUS_PERCENTAGES
 from logger import logger
@@ -57,7 +58,7 @@ async def remove_blocked_user(tg_id: int | list[int], conn: asyncpg.Connection):


 async def init_db(file_path: str = "assets/schema.sql"):
-    with open(file_path, mode="r") as file:
+    with open(file_path) as file:
        sql_content = file.read()

    statements = [stmt.strip() for stmt in sql_content.split(";") if stmt.strip()]
@@ -843,39 +844,6 @@ async def update_key_expiry(client_id: str, new_expiry_time: int):
            logger.info("Закрытие подключения к базе данных")


-async def delete_key(client_id: str):
-    """
-    Удаление ключа из базы данных для указанного клиента.
-
-    Args:
-        client_id (str): Уникальный идентификатор клиента, ключ которого будет удален
-
-    Raises:
-        Exception: В случае ошибки при подключении к базе данных или удалении ключа
-    """
-    conn = None
-    try:
-        conn = await asyncpg.connect(DATABASE_URL)
-        logger.info(f"Установлено подключение к базе данных для удаления ключа клиента {client_id}")
-
-        await conn.execute(
-            """
-            DELETE FROM keys
-            WHERE client_id = $1
-            """,
-            client_id,
-        )
-        logger.info(f"Успешно удален ключ для клиента {client_id}")
-
-    except Exception as e:
-        logger.error(f"Ошибка при удалении ключа для клиента {client_id}: {e}")
-        raise
-    finally:
-        if conn:
-            await conn.close()
-            logger.info("Закрытие подключения к базе данных")
-
-
 async def add_balance_to_client(client_id: str, amount: float):
    """
    Добавление баланса клиенту по его идентификатору Telegram.
@@ -1202,7 +1170,7 @@ async def delete_user_data(session: Any, tg_id: int):
    await session.execute("DELETE FROM payments WHERE tg_id = $1", tg_id)
    await session.execute("DELETE FROM users WHERE tg_id = $1", tg_id)
    await session.execute("DELETE FROM connections WHERE tg_id = $1", tg_id)
-    await session.execute("DELETE FROM keys WHERE tg_id = $1", tg_id)
+    await delete_key(tg_id, session)
    await session.execute("DELETE FROM referrals WHERE referrer_tg_id = $1", tg_id)


@@ -1261,3 +1229,69 @@ async def store_gift_link(
    finally:
        if conn is not None and session is None:
            await conn.close()
+
+
+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
+        FROM keys k
+        JOIN connections c ON k.tg_id = c.tg_id
+        WHERE k.email = $1
+        """,
+        email,
+    )
+
+    if not record:
+        return None
+
+    cluster_name = record["server_id"]
+
+    moscow_tz = pytz.timezone("Europe/Moscow")
+    expiry_date = datetime.fromtimestamp(record["expiry_time"] / 1000, tz=moscow_tz)
+    current_date = datetime.now(moscow_tz)
+    time_left = expiry_date - current_date
+
+    if time_left.total_seconds() <= 0:
+        days_left_message = "<b>Ключ истек.</b>"
+    elif time_left.days > 0:
+        days_left_message = f"Осталось дней: <b>{time_left.days}</b>"
+    else:
+        hours_left = time_left.seconds // 3600
+        days_left_message = f"Осталось часов: <b>{hours_left}</b>"
+
+    return {
+        "key": record["key"],
+        "created_at": record["created_at"],
+        "expiry_time": record["expiry_time"],
+        "client_id": record["client_id"],
+        "expiry_date": expiry_date.strftime("%d %B %Y года %H:%M"),
+        "days_left_message": days_left_message,
+        "server_name": cluster_name,
+        "balance": record["balance"],
+        "tg_id": record["tg_id"],
+    }
+
+
+async def delete_key(identifier, session):
+    """
+    Удаляет ключ из базы данных по client_id или tg_id
+
+    Args:
+        identifier (str): client_id или tg_id для удаления
+        session: Сессия базы данных
+
+    Raises:
+        Exception: В случае ошибки при удалении ключа
+    """
+    try:
+        # Проверяем, является ли идентификатор числом (tg_id) или строкой (client_id)
+        if identifier.isdigit():
+            query = "DELETE FROM keys WHERE tg_id = $1"
+        else:
+            query = "DELETE FROM keys WHERE client_id = $1"
+
+        await session.execute(query, identifier)
+        logger.info(f"Ключ с идентификатором {identifier} успешно удалён")
+    except Exception as e:
+        logger.error(f"Ошибка при удалении ключа с идентификатором {identifier} из базы данных: {e}")
@@ -13,6 +13,7 @@ from config import TOTAL_GB
 from database import (
    delete_user_data,
    get_client_id_by_email,
+    get_key_details,
    get_servers_from_db,
    restore_trial,
    update_key_expiry,
@@ -257,45 +258,6 @@ async def handle_new_balance_input(message: types.Message, state: FSMContext, se
    await state.clear()


-async def get_key_details(email, session):
-    record = await session.fetchrow(
-        """
-        SELECT k.key, k.expiry_time, k.server_id, c.tg_id, c.balance
-        FROM keys k
-        JOIN connections c ON k.tg_id = c.tg_id
-        WHERE k.email = $1
-        """,
-        email,
-    )
-
-    if not record:
-        return None
-
-    cluster_name = record["server_id"]
-
-    moscow_tz = pytz.timezone("Europe/Moscow")
-    expiry_date = datetime.fromtimestamp(record["expiry_time"] / 1000, tz=moscow_tz)
-    current_date = datetime.now(moscow_tz)
-    time_left = expiry_date - current_date
-
-    if time_left.total_seconds() <= 0:
-        days_left_message = "<b>Ключ истек.</b>"
-    elif time_left.days > 0:
-        days_left_message = f"Осталось дней: <b>{time_left.days}</b>"
-    else:
-        hours_left = time_left.seconds // 3600
-        days_left_message = f"Осталось часов: <b>{hours_left}</b>"
-
-    return {
-        "key": record["key"],
-        "expiry_date": expiry_date.strftime("%d %B %Y года %H:%M"),
-        "days_left_message": days_left_message,
-        "server_name": cluster_name,
-        "balance": record["balance"],
-        "tg_id": record["tg_id"],
-    }
-
-
@router.callback_query(F.data.startswith("edit_key_"), IsAdminFilter())
 async def process_key_edit(callback_query: CallbackQuery, session: Any):
    email = callback_query.data.split("_", 2)[2]
@@ -492,9 +454,9 @@ async def handle_expiry_time_input(message: types.Message, state: FSMContext, se
@router.callback_query(F.data.startswith("delete_key_admin|"), IsAdminFilter())
 async def process_callback_delete_key(callback_query: types.CallbackQuery, session: Any):
    email = callback_query.data.split("|")[1]
-    client_id = await session.fetchval("SELECT client_id FROM keys WHERE email = $1", email)
+    key_details = await get_key_details(email, session)

-    if client_id is None:
+    if key_details is None:
        builder = InlineKeyboardBuilder()
        builder.row(InlineKeyboardButton(text="🔙 Назад", callback_data="user_editor"))
        await callback_query.message.answer("🔍 Ключ не найден. 🚫", reply_markup=builder.as_markup())
@@ -504,7 +466,7 @@ async def process_callback_delete_key(callback_query: types.CallbackQuery, sessi
    builder.row(
        types.InlineKeyboardButton(
            text="✅ Да, удалить",
-            callback_data=f"confirm_delete_admin|{client_id}",
+            callback_data=f"confirm_delete_admin|{key_details['client_id']}",
        )
    )
    builder.row(types.InlineKeyboardButton(text="❌ Нет, отменить", callback_data="user_editor"))
@@ -2,12 +2,12 @@ import os
 from typing import Any

 import aiofiles
-import asyncpg
 from aiogram import F, Router, types
 from aiogram.types import BufferedInputFile, InlineKeyboardButton
 from aiogram.utils.keyboard import InlineKeyboardBuilder

-from config import CONNECT_MACOS, CONNECT_WINDOWS, DATABASE_URL, SUPPORT_CHAT_URL
+from config import CONNECT_MACOS, CONNECT_WINDOWS, SUPPORT_CHAT_URL
+from database import get_key_details
 from handlers.texts import (
    CONNECT_TV_TEXT,
    INSTRUCTION_PC,
@@ -60,15 +60,7 @@ async def process_connect_pc(callback_query: types.CallbackQuery, session: Any):
    tg_id = callback_query.message.chat.id
    key_name = callback_query.data.split("|")[1]

-    record = await session.fetchrow(
-        """
-        SELECT k.key
-        FROM keys k
-        WHERE k.tg_id = $1 AND k.email = $2
-        """,
-        tg_id,
-        key_name,
-    )
+    record = await get_key_details(key_name, session)

    if not record:
        await callback_query.message.answer("❌ <b>Ключ не найден. Проверьте имя ключа.</b> 🔍")
@@ -107,27 +99,15 @@ async def process_connect_tv(callback_query: types.CallbackQuery):


@router.callback_query(F.data.startswith("continue_tv|"))
-async def process_continue_tv(callback_query: types.CallbackQuery):
+async def process_continue_tv(callback_query: types.CallbackQuery, session: Any):
    key_name = callback_query.data.split("|")[1]
    tg_id = callback_query.from_user.id

    logger.info(f"tg_id: {tg_id}, key_name: {key_name}")

-    conn = await asyncpg.connect(DATABASE_URL)
-    try:
-        record = await conn.fetchrow(
-            """
-            SELECT k.key
-            FROM keys k
-            WHERE k.tg_id = $1 AND k.email = $2
-            """,
-            tg_id,
-            key_name,
-        )
+    record = await get_key_details(key_name, session)

-        logger.info(f"Query result: {record}")
-    finally:
-        await conn.close()
+    logger.info(f"Query result: {record}")

    subscription_link = record["key"]

@@ -25,6 +25,7 @@ from config import (
 )
 from database import (
    get_balance,
+    get_key_details,
    get_trial,
    save_temporary_data,
    store_key,
@@ -352,11 +353,7 @@ async def finalize_key_creation(
        key_name = generate_random_email()
        logger.info(f"Generated random key name for user {tg_id}: {key_name}")

-        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 name '{key_name}' already exists for user {tg_id}. Generating a new one.")
@@ -142,13 +142,6 @@ async def renew_key_in_cluster(cluster_id, email, client_id, new_expiry_time, to
        raise e


-async def delete_key_from_db(client_id, session):
-    try:
-        await session.execute("DELETE FROM keys WHERE client_id = $1", client_id)
-    except Exception as e:
-        logger.error(f"Ошибка при удалении ключа {client_id} из базы данных: {e}")
-
-
 async def delete_key_from_cluster(cluster_id, email, client_id):
    """Удаление ключа с серверов в кластере"""
    try:
@@ -29,6 +29,7 @@ from config import (
 from database import (
    delete_key,
    get_balance,
+    get_key_details,
    get_servers_from_db,
    save_temporary_data,
    store_key,
@@ -163,15 +164,7 @@ async def process_callback_view_key(callback_query: types.CallbackQuery, session
    tg_id = callback_query.message.chat.id
    key_name = callback_query.data.split("|")[1]
    try:
-        record = await session.fetchrow(
-            """
-            SELECT k.expiry_time, k.server_id, k.key
-            FROM keys k
-            WHERE k.tg_id = $1 AND k.email = $2
-            """,
-            tg_id,
-            key_name,
-        )
+        record = await get_key_details(key_name, session)

        if record:
            key = record["key"]
@@ -264,15 +257,7 @@ async def process_callback_update_subscription(callback_query: types.CallbackQue
    tg_id = callback_query.message.chat.id
    email = callback_query.data.split("|")[1]
    try:
-        record = await session.fetchrow(
-            """
-            SELECT k.key, k.expiry_time, k.email, k.server_id, k.client_id
-            FROM keys k
-            WHERE k.tg_id = $1 AND k.email = $2
-            """,
-            tg_id,
-            email,
-        )
+        record = await get_key_details(email, session)

        if record:
            expiry_time = record["expiry_time"]
@@ -353,14 +338,7 @@ async def process_callback_renew_key(callback_query: types.CallbackQuery, sessio
    tg_id = callback_query.message.chat.id
    key_name = callback_query.data.split("|")[1]
    try:
-        record = await session.fetchrow(
-            """
-            SELECT client_id, expiry_time 
-            FROM keys 
-            WHERE email = $1
-            """,
-            key_name,
-        )
+        record = await get_key_details(key_name, session)

        if record:
            client_id = record["client_id"]
@@ -405,7 +383,7 @@ async def process_callback_renew_key(callback_query: types.CallbackQuery, sessio
 async def process_callback_confirm_delete(callback_query: types.CallbackQuery, session: Any):
    email = callback_query.data.split("|")[1]
    try:
-        record = await session.fetchrow("SELECT client_id FROM keys WHERE email = $1", email)
+        record = await get_key_details(email, session)

        if record:
            client_id = record["client_id"]
@@ -413,7 +391,7 @@ async def process_callback_confirm_delete(callback_query: types.CallbackQuery, s
            back_button = types.InlineKeyboardButton(text="Назад", callback_data="view_keys")
            keyboard = types.InlineKeyboardMarkup(inline_keyboard=[[back_button]])

-            await delete_key(client_id)
+            await delete_key(client_id, session)
            await callback_query.message.answer(
                response_message,
                reply_markup=keyboard,
@@ -7,7 +7,7 @@ import asyncpg
 from aiohttp import web

 from config import DATABASE_URL, PROJECT_NAME, SUB_MESSAGE, SUPERNODE, TRANSITION_DATE_STR
-from database import get_servers_from_db
+from database import get_key_details, get_servers_from_db
 from logger import logger

 # Глобальная переменная для пула соединений
@@ -92,7 +92,7 @@ async def handle_old_subscription(request):
    await init_db_pool()

    async with db_pool.acquire() as conn:
-        key_info = await conn.fetchrow("SELECT created_at, server_id FROM keys WHERE email = $1", email)
+        key_info = await get_key_details(email, conn)

        if not key_info:
            logger.warning(f"Клиент с email {email} не найден в базе.")
@@ -161,7 +161,7 @@ async def handle_new_subscription(request):
    await init_db_pool()

    async with db_pool.acquire() as conn:
-        client_data = await conn.fetchrow("SELECT tg_id, server_id FROM keys WHERE email = $1", email)
+        client_data = await get_key_details(email, conn)

        if not client_data:
            logger.warning(f"Клиент с email {email} не найден в базе.")
@@ -8,7 +8,7 @@ from py3xui import AsyncApi

 from client import add_client
 from config import ADMIN_PASSWORD, ADMIN_USERNAME, LIMIT_IP, PUBLIC_LINK, SUPERNODE, TOTAL_GB, TRIAL_TIME
-from database import get_servers_from_db, store_key, use_trial
+from database import get_servers_from_db, get_trial, store_key, use_trial
 from handlers.texts import INSTRUCTIONS
 from handlers.utils import generate_random_email, get_least_loaded_cluster
 from logger import logger
@@ -16,7 +16,7 @@ from logger import logger

 async def create_trial_key(tg_id: int, session: Any):
    try:
-        trial_status = await session.fetchval("SELECT trial FROM connections WHERE tg_id = $1", tg_id)
+        trial_status = await get_trial(tg_id, session)
        if trial_status == 1:
            return {"error": "Вы уже использовали пробную версию."}
    except Exception as e:
@@ -401,7 +401,7 @@ async def handle_expired_keys(bot: Bot, conn: asyncpg.Connection, current_time:
    for record in expired_keys:
        try:
            await delete_key_from_cluster(record["server_id"], record["email"], record["email"])
-            await conn.execute("DELETE FROM keys WHERE client_id = $1", record["client_id"])
+            await delete_key(record["client_id"], conn)
            logger.info(f"Удалён истёкший ключ {record['client_id']} пользователя {record['tg_id']}.")
        except Exception as e:
            logger.error(f"Ошибка при удалении истёкшего ключа {record['client_id']}: {e}")
@@ -497,7 +497,7 @@ async def process_key(record, bot, conn):
                        logger.error(f"Ошибка при удалении клиента {client_id} из кластера {cluster_id}: {e}")

                try:
-                    await delete_key(client_id)
+                    await delete_key(client_id, conn)
                    logger.info(f"Ключ {client_id} удалён из базы данных.")
                except Exception as e:
                    logger.error(f"Ошибка при удалении ключа {client_id} из базы данных: {e}")
@@ -101,14 +101,7 @@ async def start_command(message: Message, state: FSMContext, session: Any, admin
                    await message.answer("❌ Вы не можете получить подарок от самого себя.")
                    return await show_start_menu(message, admin, session)

-                await session.execute(
-                    """
-                    INSERT INTO connections (tg_id, balance, trial)
-                    VALUES ($1, 0, 1)
-                    ON CONFLICT (tg_id) DO UPDATE SET trial = 1
-                    """,
-                    recipient_tg_id,
-                )
+                await add_connection(tg_id=recipient_tg_id, session=session)

                selected_months = gift_info["selected_months"]
                expiry_time = gift_info["expiry_time"]