From 97c03d0396cf45c9269594f7a5bfe87ca2456a84 Mon Sep 17 00:00:00 2001 From: Zakhar Izmaylov Date: Sun, 24 Nov 2024 09:15:09 +0300 Subject: [PATCH] Add logging --- database.py | 843 ++++++++++++++++++++++++-------- handlers/keys/key_management.py | 23 +- 2 files changed, 643 insertions(+), 223 deletions(-) diff --git a/database.py b/database.py index 4edc4d7d..c660a762 100644 --- a/database.py +++ b/database.py @@ -140,15 +140,20 @@ async def create_coupon(coupon_code: str, amount: float, usage_limit: int, sessi Example: await create_coupon('SALE50', 50.0, 5, session) """ - await session.execute( - """ - INSERT INTO coupons (code, amount, usage_limit, usage_count, is_used) - VALUES ($1, $2, $3, 0, FALSE) - """, - coupon_code, - amount, - usage_limit, - ) + try: + await session.execute( + """ + INSERT INTO coupons (code, amount, usage_limit, usage_count, is_used) + VALUES ($1, $2, $3, 0, FALSE) + """, + coupon_code, + amount, + usage_limit, + ) + logger.info(f"Успешно создан купон с кодом {coupon_code} на сумму {amount}") + except Exception as e: + logger.error(f"Ошибка при создании купона {coupon_code}: {e}") + raise async def get_all_coupons(session: Any): @@ -166,15 +171,22 @@ async def get_all_coupons(session: Any): Exception: В случае ошибки при получении данных из базы """ try: + # Выполняем запрос на получение всех купонов из базы данных coupons = await session.fetch( """ SELECT code, amount, usage_limit, usage_count FROM coupons """ ) + + # Логируем успешное получение списка купонов + logger.info(f"Успешно получено {len(coupons)} купонов из базы данных") + return coupons except Exception as e: - logger.error(f"Ошибка при получении купонов: {e}") + # Подробное логирование ошибки при получении купонов + logger.error(f"Критическая ошибка при получении списка купонов: {e}") + logger.exception("Трассировка стека ошибки получения купонов") return [] @@ -196,7 +208,7 @@ async def delete_coupon_from_db(coupon_code: str, session: Any): result = await delete_coupon_from_db('SALE50', session) """ try: - + # Проверяем существование купона в базе данных coupon_record = await session.fetchrow( """ SELECT id FROM coupons WHERE code = $1 @@ -204,9 +216,12 @@ async def delete_coupon_from_db(coupon_code: str, session: Any): coupon_code, ) + # Если купон не найден, возвращаем False if not coupon_record: + logger.info(f"Купон {coupon_code} не найден в базе данных") return False + # Удаляем купон из базы данных await session.execute( """ DELETE FROM coupons WHERE code = $1 @@ -214,10 +229,13 @@ async def delete_coupon_from_db(coupon_code: str, session: Any): coupon_code, ) + # Логируем успешное удаление купона + logger.info(f"Купон {coupon_code} успешно удален из базы данных") return True except Exception as e: - logger.error(f"Ошибка при удалении купона: {e}") + # Логируем ошибку на русском с подробным описанием + logger.error(f"Произошла ошибка при удалении купона {coupon_code}: {e}") return False @@ -233,7 +251,7 @@ async def restore_trial(tg_id: int, session: Any): bool: True, если триал успешно восстановлен, False в случае ошибки """ try: - response = await session.execute( + await session.execute( """ INSERT INTO connections (tg_id, trial) VALUES ($1, 0) @@ -242,10 +260,10 @@ async def restore_trial(tg_id: int, session: Any): """, tg_id, ) - logger.info(response) + logger.info(f"Триальный период успешно восстановлен для пользователя {tg_id}") return True except Exception as e: - logger.error(f"Ошибка при установке значения триала: {e}") + logger.error(f"Ошибка при восстановлении триального периода для пользователя {tg_id}: {e}") return False @@ -261,7 +279,7 @@ async def use_trial(tg_id: int, session: Any): bool: True, если триал успешно использован, False в случае ошибки """ try: - response = await session.execute( + await session.execute( """ INSERT INTO connections (tg_id, trial) VALUES ($1, 1) @@ -270,273 +288,643 @@ async def use_trial(tg_id: int, session: Any): """, tg_id, ) - logger.info(response) + logger.info(f"Триальный период успешно использован для пользователя {tg_id}") return True except Exception as e: - logger.error(f"Ошибка при использовании триала: {e}") + logger.error(f"Ошибка при использовании триала для пользователя {tg_id}: {e}") return False async def add_connection(tg_id: int, balance: float = 0.0, trial: int = 0, session: Any = None): - await session.execute( - """ - INSERT INTO connections (tg_id, balance, trial) - VALUES ($1, $2, $3) - """, - tg_id, - balance, - trial, - ) + """ + Добавляет новое подключение для пользователя в базу данных. + + Args: + tg_id (int): Telegram ID пользователя + balance (float, optional): Начальный баланс пользователя. По умолчанию 0.0. + trial (int, optional): Статус триального периода. По умолчанию 0. + session (Any, optional): Сессия базы данных. + + Raises: + Exception: Если возникает ошибка при добавлении подключения в базу данных. + """ + try: + await session.execute( + """ + INSERT INTO connections (tg_id, balance, trial) + VALUES ($1, $2, $3) + """, + tg_id, + balance, + trial, + ) + logger.info( + f"Успешно добавлено новое подключение для пользователя {tg_id} с балансом {balance} и статусом триала {trial}" + ) + except Exception as e: + logger.error(f"Не удалось добавить подключение для пользователя {tg_id}. Причина: {e}") + raise async def check_connection_exists(tg_id: int): - conn = await asyncpg.connect(DATABASE_URL) - exists = await conn.fetchval( - """ - SELECT EXISTS(SELECT 1 FROM connections WHERE tg_id = $1) - """, - tg_id, - ) - await conn.close() - return exists + """ + Проверяет существование подключения для указанного пользователя в базе данных. + + Args: + tg_id (int): Telegram ID пользователя для проверки. + + Returns: + bool: True, если подключение существует, иначе False. + + Raises: + Exception: В случае ошибки при подключении к базе данных. + """ + try: + conn = await asyncpg.connect(DATABASE_URL) + exists = await conn.fetchval( + """ + SELECT EXISTS(SELECT 1 FROM connections WHERE tg_id = $1) + """, + tg_id, + ) + logger.info( + f"Проверка существования подключения для пользователя {tg_id}: {'найдено' if exists else 'не найдено'}" + ) + return exists + except Exception as e: + logger.error(f"Ошибка при проверке подключения для пользователя {tg_id}: {e}") + raise + finally: + if conn: + await conn.close() -async def store_key( - tg_id: int, - client_id: str, - email: str, - expiry_time: int, - key: str, - server_id: str, -): - conn = await asyncpg.connect(DATABASE_URL) - await conn.execute( - """ - INSERT INTO keys (tg_id, client_id, email, created_at, expiry_time, key, server_id) - VALUES ($1, $2, $3, $4, $5, $6, $7) - """, - tg_id, - client_id, - email, - int(datetime.utcnow().timestamp() * 1000), - expiry_time, - key, - server_id, - ) - await conn.close() +async def store_key(tg_id: int, client_id: str, email: str, expiry_time: int, key: str, server_id: str, session: Any): + """ + Сохраняет информацию о ключе в базу данных. + + Args: + tg_id (int): Telegram ID пользователя + client_id (str): Уникальный идентификатор клиента + email (str): Электронная почта или имя устройства + expiry_time (int): Время истечения ключа в миллисекундах + key (str): Ключ доступа + server_id (str): Идентификатор сервера + + Raises: + Exception: Если возникает ошибка при сохранении ключа в базу данных + """ + try: + await session.execute( + """ + INSERT INTO keys (tg_id, client_id, email, created_at, expiry_time, key, server_id) + VALUES ($1, $2, $3, $4, $5, $6, $7) + """, + tg_id, + client_id, + email, + int(datetime.utcnow().timestamp() * 1000), + expiry_time, + key, + server_id, + ) + logger.info(f"Ключ успешно сохранен для пользователя {tg_id} на сервере {server_id}") + except Exception as e: + logger.error(f"Ошибка при сохранении ключа для пользователя {tg_id}: {e}") + raise async def get_keys(tg_id: int): - conn = await asyncpg.connect(DATABASE_URL) - records = await conn.fetch( - """ - SELECT client_id, email, created_at, key - FROM keys - WHERE tg_id = $1 - """, - tg_id, - ) - await conn.close() - return records + """ + Получает список ключей для указанного пользователя. + + Args: + tg_id (int): Telegram ID пользователя + + Returns: + list: Список записей ключей с информацией о клиенте, электронной почте, времени создания и ключе + + Raises: + Exception: В случае ошибки при подключении к базе данных или выполнении запроса + """ + conn = None + try: + conn = await asyncpg.connect(DATABASE_URL) + records = await conn.fetch( + """ + SELECT client_id, email, created_at, key + FROM keys + WHERE tg_id = $1 + """, + tg_id, + ) + logger.info(f"Успешно получено {len(records)} ключей для пользователя {tg_id}") + return records + except Exception as e: + logger.error(f"Ошибка при получении ключей для пользователя {tg_id}: {e}") + raise + finally: + if conn: + await conn.close() async def get_keys_by_server(tg_id: int, server_id: str): - conn = await asyncpg.connect(DATABASE_URL) - records = await conn.fetch( - """ - SELECT client_id, email, created_at, key - FROM keys - WHERE tg_id = $1 AND server_id = $2 - """, - tg_id, - server_id, - ) - await conn.close() - return records + """ + Получает список ключей для указанного пользователя на определенном сервере. + + Args: + tg_id (int): Telegram ID пользователя + server_id (str): Идентификатор сервера + + Returns: + list: Список записей ключей с информацией о клиенте, электронной почте, времени создания и ключе + + Raises: + Exception: В случае ошибки при подключении к базе данных или выполнении запроса + """ + conn = None + try: + conn = await asyncpg.connect(DATABASE_URL) + records = await conn.fetch( + """ + SELECT client_id, email, created_at, key + FROM keys + WHERE tg_id = $1 AND server_id = $2 + """, + tg_id, + server_id, + ) + logger.info(f"Успешно получено {len(records)} ключей для пользователя {tg_id} на сервере {server_id}") + return records + except Exception as e: + logger.error(f"Ошибка при получении ключей для пользователя {tg_id} на сервере {server_id}: {e}") + raise + finally: + if conn: + await conn.close() async def has_active_key(tg_id: int) -> bool: - conn = await asyncpg.connect(DATABASE_URL) - count = await conn.fetchval("SELECT COUNT(*) FROM keys WHERE tg_id = $1", tg_id) - await conn.close() - return count > 0 + """ + Проверяет наличие активных ключей для указанного пользователя. + + Args: + tg_id (int): Telegram ID пользователя + + Returns: + bool: True, если у пользователя есть активные ключи, иначе False + + Raises: + Exception: В случае ошибки при подключении к базе данных или выполнении запроса + """ + conn = None + try: + conn = await asyncpg.connect(DATABASE_URL) + count = await conn.fetchval("SELECT COUNT(*) FROM keys WHERE tg_id = $1", tg_id) + logger.info(f"Проверка наличия ключей для пользователя {tg_id}. Найдено ключей: {count}") + return count > 0 + except Exception as e: + logger.error(f"Ошибка при проверке наличия ключей для пользователя {tg_id}: {e}") + raise + finally: + if conn: + await conn.close() async def get_balance(tg_id: int) -> float: - conn = await asyncpg.connect(DATABASE_URL) - balance = await conn.fetchval("SELECT balance FROM connections WHERE tg_id = $1", tg_id) - await conn.close() - return balance if balance is not None else 0.0 + """ + Получает баланс пользователя из базы данных. + + Args: + tg_id (int): Telegram ID пользователя + + Returns: + float: Баланс пользователя, 0.0 если баланс не найден + + Raises: + Exception: В случае ошибки при подключении к базе данных или выполнении запроса + """ + conn = None + try: + conn = await asyncpg.connect(DATABASE_URL) + balance = await conn.fetchval("SELECT balance FROM connections WHERE tg_id = $1", tg_id) + logger.info(f"Получен баланс для пользователя {tg_id}: {balance}") + return balance if balance is not None else 0.0 + except Exception as e: + logger.error(f"Ошибка при получении баланса для пользователя {tg_id}: {e}") + return 0.0 + finally: + if conn: + await conn.close() async def update_balance(tg_id: int, amount: float): - conn = await asyncpg.connect(DATABASE_URL) - await conn.execute( - """ - UPDATE connections - SET balance = balance + $1 - WHERE tg_id = $2 - """, - amount, - tg_id, - ) + """ + Обновляет баланс пользователя в базе данных. - await handle_referral_on_balance_update(tg_id, amount) + Args: + tg_id (int): Telegram ID пользователя + amount (float): Сумма для обновления баланса - await conn.close() + Raises: + Exception: В случае ошибки при подключении к базе данных или обновлении баланса + """ + conn = None + try: + conn = await asyncpg.connect(DATABASE_URL) + await conn.execute( + """ + UPDATE connections + SET balance = balance + $1 + WHERE tg_id = $2 + """, + amount, + tg_id, + ) + logger.info(f"Баланс пользователя {tg_id} обновлен на сумму {amount}") + + await handle_referral_on_balance_update(tg_id, amount) + + except Exception as e: + logger.error(f"Ошибка при обновлении баланса для пользователя {tg_id}: {e}") + raise + finally: + if conn: + await conn.close() async def get_trial(tg_id: int, session: Any) -> int: - trial = await session.fetchval("SELECT trial FROM connections WHERE tg_id = $1", tg_id) - return trial if trial is not None else 0 + """ + Получает статус триала для пользователя из базы данных. + + Args: + tg_id (int): Telegram ID пользователя + session (Any): Сессия базы данных + + Returns: + int: Статус триала (0 - не использован, 1 - использован) + """ + try: + trial = await session.fetchval("SELECT trial FROM connections WHERE tg_id = $1", tg_id) + logger.info(f"Получен статус триала для пользователя {tg_id}: {trial}") + return trial if trial is not None else 0 + except Exception as e: + logger.error(f"Ошибка при получении статуса триала для пользователя {tg_id}: {e}") + return 0 async def get_key_count(tg_id: int) -> int: - conn = await asyncpg.connect(DATABASE_URL) - count = await conn.fetchval("SELECT COUNT(*) FROM keys WHERE tg_id = $1", tg_id) - await conn.close() - return count if count is not None else 0 + """ + Получает количество ключей для указанного пользователя. + + Args: + tg_id (int): Telegram ID пользователя + + Returns: + int: Количество ключей пользователя, 0 если ключей нет + + Raises: + Exception: В случае ошибки при подключении к базе данных + """ + conn = None + try: + conn = await asyncpg.connect(DATABASE_URL) + count = await conn.fetchval("SELECT COUNT(*) FROM keys WHERE tg_id = $1", tg_id) + logger.info(f"Получено количество ключей для пользователя {tg_id}: {count}") + return count if count is not None else 0 + except Exception as e: + logger.error(f"Ошибка при получении количества ключей для пользователя {tg_id}: {e}") + return 0 + finally: + if conn: + await conn.close() async def get_all_users(conn): - return await conn.fetch("SELECT tg_id FROM connections") + """ + Получает список всех пользователей из базы данных. + + Args: + conn: Подключение к базе данных + + Returns: + list: Список Telegram ID всех пользователей + + Raises: + Exception: В случае ошибки при получении данных + """ + try: + пользователи = await conn.fetch("SELECT tg_id FROM connections") + logger.info(f"Получен список всех пользователей. Количество: {len(пользователи)}") + return пользователи + except Exception as e: + logger.error(f"Ошибка при получении списка пользователей: {e}") + raise async def add_referral(referred_tg_id: int, referrer_tg_id: int, session: Any): - await session.execute( - """ - INSERT INTO referrals (referred_tg_id, referrer_tg_id) - VALUES ($1, $2) - """, - referred_tg_id, - referrer_tg_id, - ) + try: + # Вставка записи о реферале в базу данных + await session.execute( + """ + INSERT INTO referrals (referred_tg_id, referrer_tg_id) + VALUES ($1, $2) + """, + referred_tg_id, + referrer_tg_id, + ) + logger.info(f"Добавлена реферальная связь: приглашенный {referred_tg_id}, пригласивший {referrer_tg_id}") + except Exception as e: + logger.error(f"Ошибка при добавлении реферала: {e}") + raise async def handle_referral_on_balance_update(tg_id: int, amount: float): - conn = await asyncpg.connect(DATABASE_URL) + """ + Обработка реферальной системы при обновлении баланса. - referral = await conn.fetchrow( - """ - SELECT referrer_tg_id FROM referrals WHERE referred_tg_id = $1 - """, - tg_id, - ) + Эта функция проверяет, есть ли у пользователя реферер, + и если есть, начисляет бонус за пополнение баланса. - if referral: - referrer_tg_id = referral["referrer_tg_id"] + Args: + tg_id (int): Telegram ID пользователя + amount (float): Сумма пополнения баланса - bonus = amount * BONUS_PERCENT + Returns: + None + """ + conn = None + try: + conn = await asyncpg.connect(DATABASE_URL) + logger.info(f"Установлено подключение к базе данных для обработки реферала пользователя {tg_id}") - if bonus < 0: - bonus = 0 - - await update_balance(referrer_tg_id, bonus) - - await conn.execute( + referral = await conn.fetchrow( """ - UPDATE referrals SET reward_issued = TRUE - WHERE referrer_tg_id = $1 AND referred_tg_id = $2 - """, - referrer_tg_id, + SELECT referrer_tg_id FROM referrals WHERE referred_tg_id = $1 + """, tg_id, ) - await conn.close() + if referral: + referrer_tg_id = referral["referrer_tg_id"] + logger.info(f"Найден реферер {referrer_tg_id} для пользователя {tg_id}") + + bonus = amount * BONUS_PERCENT + bonus = max(bonus, 0) # Гарантируем неотрицательный бонус + + logger.info(f"Начисление бонуса {bonus} рублей рефереру {referrer_tg_id}") + await update_balance(referrer_tg_id, bonus) + + await conn.execute( + """ + UPDATE referrals SET reward_issued = TRUE + WHERE referrer_tg_id = $1 AND referred_tg_id = $2 + """, + referrer_tg_id, + tg_id, + ) + logger.info(f"Обновлен статус реферала для {referrer_tg_id}") + + except Exception as e: + logger.error(f"Ошибка при обработке реферала для пользователя {tg_id}: {e}") + finally: + if conn: + await conn.close() + logger.info("Закрытие подключения к базе данных") async def get_referral_stats(referrer_tg_id: int): - conn = await asyncpg.connect(DATABASE_URL) - total_referrals = await conn.fetchval( - """ - SELECT COUNT(*) FROM referrals WHERE referrer_tg_id = $1 - """, - referrer_tg_id, - ) + """ + Получение статистики рефералов для указанного пользователя. - active_referrals = await conn.fetchval( - """ - SELECT COUNT(*) FROM referrals WHERE referrer_tg_id = $1 AND reward_issued = TRUE - """, - referrer_tg_id, - ) + Args: + referrer_tg_id (int): Telegram ID пользователя, для которого запрашивается статистика рефералов. - await conn.close() + Returns: + dict: Словарь со статистикой рефералов, содержащий: + - total_referrals (int): Общее количество рефералов + - active_referrals (int): Количество активных рефералов (с начисленным бонусом) - return { - "total_referrals": total_referrals, - "active_referrals": active_referrals, - } + Raises: + Exception: В случае ошибки при подключении к базе данных или выполнении запроса + """ + conn = None + try: + conn = await asyncpg.connect(DATABASE_URL) + logger.info( + f"Установлено подключение к базе данных для получения статистики рефералов пользователя {referrer_tg_id}" + ) + + total_referrals = await conn.fetchval( + """ + SELECT COUNT(*) FROM referrals WHERE referrer_tg_id = $1 + """, + referrer_tg_id, + ) + logger.debug(f"Получено общее количество рефералов: {total_referrals}") + + active_referrals = await conn.fetchval( + """ + SELECT COUNT(*) FROM referrals WHERE referrer_tg_id = $1 AND reward_issued = TRUE + """, + referrer_tg_id, + ) + logger.debug(f"Получено количество активных рефералов: {active_referrals}") + + return { + "total_referrals": total_referrals, + "active_referrals": active_referrals, + } + + except Exception as e: + logger.error(f"Ошибка при получении статистики рефералов для пользователя {referrer_tg_id}: {e}") + raise + finally: + if conn: + await conn.close() + logger.info("Закрытие подключения к базе данных") async def update_key_expiry(client_id: str, new_expiry_time: int): """ - Обновление времени истечения ключа на новое значение. + Обновление времени истечения ключа для указанного клиента. + + Args: + client_id (str): Уникальный идентификатор клиента + new_expiry_time (int): Новое время истечения ключа + + Raises: + Exception: В случае ошибки при подключении к базе данных или обновлении ключа """ - conn = await asyncpg.connect(DATABASE_URL) - await conn.execute( - """ - UPDATE keys - SET expiry_time = $1, notified = FALSE, notified_24h = FALSE - WHERE client_id = $2 - """, - new_expiry_time, - client_id, - ) - await conn.close() + conn = None + try: + conn = await asyncpg.connect(DATABASE_URL) + logger.info(f"Установлено подключение к базе данных для обновления времени истечения ключа клиента {client_id}") + + await conn.execute( + """ + UPDATE keys + SET expiry_time = $1, notified = FALSE, notified_24h = FALSE + WHERE client_id = $2 + """, + new_expiry_time, + 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 delete_key(client_id: str): """ - Удаление ключа из базы данных. + Удаление ключа из базы данных для указанного клиента. + + Args: + client_id (str): Уникальный идентификатор клиента, ключ которого будет удален + + Raises: + Exception: В случае ошибки при подключении к базе данных или удалении ключа """ - conn = await asyncpg.connect(DATABASE_URL) - await conn.execute( - """ - DELETE FROM keys - WHERE client_id = $1 - """, - client_id, - ) - await conn.close() + 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): - conn = await asyncpg.connect(DATABASE_URL) - await conn.execute( - """ - UPDATE connections - SET balance = balance + $1 - WHERE tg_id = $2 - """, - amount, - client_id, - ) - await conn.close() + """ + Добавление баланса клиенту по его идентификатору Telegram. + + Args: + client_id (str): Идентификатор клиента в Telegram + amount (float): Сумма для пополнения баланса + + Raises: + Exception: В случае ошибки при подключении к базе данных или обновлении баланса + """ + conn = None + try: + conn = await asyncpg.connect(DATABASE_URL) + logger.info(f"Установлено подключение к базе данных для пополнения баланса клиента {client_id}") + + await conn.execute( + """ + UPDATE connections + SET balance = balance + $1 + WHERE tg_id = $2 + """, + amount, + client_id, + ) + logger.info(f"Успешно пополнен баланс клиента {client_id} на сумму {amount}") + + except Exception as e: + logger.error(f"Ошибка при пополнении баланса для клиента {client_id}: {e}") + raise + finally: + if conn: + await conn.close() + logger.info("Закрытие подключения к базе данных") async def get_client_id_by_email(email: str): """ - Получение client_id по email. + Получение идентификатора клиента по электронной почте. + + Args: + email (str): Электронная почта клиента + + Returns: + str: Идентификатор клиента или None, если клиент не найден + + Raises: + Exception: В случае ошибки при подключении к базе данных или выполнении запроса """ - conn = await asyncpg.connect(DATABASE_URL) - client_id = await conn.fetchval( - """ - SELECT client_id FROM keys WHERE email = $1 - """, - email, - ) - await conn.close() - return client_id + conn = None + try: + conn = await asyncpg.connect(DATABASE_URL) + logger.info(f"Установлено подключение к базе данных для поиска client_id по email: {email}") + + client_id = await conn.fetchval( + """ + SELECT client_id FROM keys WHERE email = $1 + """, + email, + ) + + if client_id: + logger.info(f"Найден client_id для email: {email}") + else: + logger.warning(f"Не найден client_id для email: {email}") + + return client_id + + except Exception as e: + logger.error(f"Ошибка при получении client_id для email {email}: {e}") + raise + finally: + if conn: + await conn.close() + logger.info("Закрытие подключения к базе данных") async def get_tg_id_by_client_id(client_id: str): - conn = await asyncpg.connect(DATABASE_URL) + """ + Получение Telegram ID по идентификатору клиента. + + Args: + client_id (str): Идентификатор клиента + + Returns: + int или None: Telegram ID клиента, если найден, иначе None + + Raises: + Exception: В случае ошибки при подключении к базе данных или выполнении запроса + """ + conn = None try: + conn = await asyncpg.connect(DATABASE_URL) + logger.info(f"Установлено подключение к базе данных для поиска Telegram ID по client_id: {client_id}") + result = await conn.fetchrow("SELECT tg_id FROM keys WHERE client_id = $1", client_id) - return result["tg_id"] if result else None + + if result: + logger.info(f"Найден Telegram ID для client_id: {client_id}") + return result["tg_id"] + else: + logger.warning(f"Не найден Telegram ID для client_id: {client_id}") + return None + + except Exception as e: + logger.error(f"Ошибка при получении Telegram ID для client_id {client_id}: {e}") + raise finally: - await conn.close() + if conn: + await conn.close() + logger.info("Закрытие подключения к базе данных") async def upsert_user( @@ -547,8 +935,25 @@ async def upsert_user( language_code: str = None, is_bot: bool = False, ): - conn = await asyncpg.connect(DATABASE_URL) + """ + Обновляет или вставляет информацию о пользователе в базу данных. + + Args: + tg_id (int): Идентификатор пользователя в Telegram + username (str, optional): Имя пользователя в Telegram + first_name (str, optional): Имя пользователя + last_name (str, optional): Фамилия пользователя + language_code (str, optional): Код языка пользователя + is_bot (bool, optional): Флаг, указывающий является ли пользователь ботом + + Raises: + Exception: В случае ошибки при работе с базой данных + """ + conn = None try: + conn = await asyncpg.connect(DATABASE_URL) + logger.info(f"Установлено подключение к базе данных для обновления пользователя {tg_id}") + await conn.execute( """ INSERT INTO users (tg_id, username, first_name, last_name, language_code, is_bot, created_at, updated_at) @@ -569,13 +974,33 @@ async def upsert_user( language_code, is_bot, ) + logger.info(f"Успешно обновлена информация о пользователе {tg_id}") + except Exception as e: + logger.error(f"Ошибка при обновлении информации о пользователе {tg_id}: {e}") + raise finally: - await conn.close() + if conn: + await conn.close() + logger.info("Закрытие подключения к базе данных") async def add_payment(tg_id: int, amount: float, payment_system: str): - conn = await asyncpg.connect(DATABASE_URL) + """ + Добавляет информацию о платеже в базу данных. + + Args: + tg_id (int): Идентификатор пользователя в Telegram + amount (float): Сумма платежа + payment_system (str): Система оплаты + + Raises: + Exception: В случае ошибки при добавлении платежа + """ + conn = None try: + conn = await asyncpg.connect(DATABASE_URL) + logger.info(f"Установлено подключение к базе данных для добавления платежа пользователя {tg_id}") + await conn.execute( """ INSERT INTO payments (tg_id, amount, payment_system, status) @@ -585,7 +1010,11 @@ async def add_payment(tg_id: int, amount: float, payment_system: str): amount, payment_system, ) + logger.info(f"Успешно добавлен платеж для пользователя {tg_id} на сумму {amount}") except Exception as e: - logger.error(f"Ошибка при добавлении платежа: {e}") + logger.error(f"Ошибка при добавлении платежа для пользователя {tg_id}: {e}") + raise finally: - await conn.close() + if conn: + await conn.close() + logger.info("Закрытие подключения к базе данных после добавления платежа") diff --git a/handlers/keys/key_management.py b/handlers/keys/key_management.py index a1afa07a..6cede92f 100644 --- a/handlers/keys/key_management.py +++ b/handlers/keys/key_management.py @@ -10,7 +10,7 @@ from aiogram.types import CallbackQuery, InlineKeyboardButton, Message from aiogram.utils.keyboard import InlineKeyboardBuilder from config import CONNECT_ANDROID, CONNECT_IOS, DOWNLOAD_ANDROID, DOWNLOAD_IOS, PUBLIC_LINK, SUPPORT_CHAT_URL -from database import add_connection, get_balance, get_trial, store_key, update_balance +from database import add_connection, get_balance, get_trial, store_key, update_balance, use_trial from handlers.keys.key_utils import create_key_on_cluster from handlers.texts import KEY, KEY_TRIAL, NULL_BALANCE, RENEWAL_PLANS, key_message_success from handlers.utils import get_least_loaded_cluster, sanitize_key_name @@ -34,7 +34,7 @@ async def process_callback_create_key(callback_query: CallbackQuery, state: FSMC async def select_server(callback_query: CallbackQuery, state: FSMContext, session: Any): - trial_status = get_trial(callback_query.message.from_user.id,session) + trial_status = get_trial(callback_query.message.from_user.id, session) logger.info(f'trial_status {trial_status}') if trial_status == 1: builder = InlineKeyboardBuilder() @@ -107,9 +107,7 @@ async def handle_key_name_input(message: Message, state: FSMContext, session: An expiry_time = None logger.info(f"Checking trial status for user {tg_id}.") - existing_connection = await session.fetchrow("SELECT trial FROM connections WHERE tg_id = $1", tg_id) - - trial_status = existing_connection["trial"] if existing_connection else 0 + trial_status = get_trial(message.from_user.id, session) if trial_status == 0: expiry_time = current_time + timedelta(days=1, hours=3) @@ -175,21 +173,14 @@ async def handle_key_name_input(message: Message, state: FSMContext, session: An await asyncio.gather(*tasks) logger.info(f"Updating trial status for user {tg_id} in the database.") - existing_connection = await session.fetchrow("SELECT * FROM connections WHERE tg_id = $1", tg_id) - if existing_connection: - await session.execute("UPDATE connections SET trial = 1 WHERE tg_id = $1", tg_id) + trial_status = await get_trial(message.from_user.id, session) + if trial_status == 0: + await use_trial(message.from_user.id, session) else: await add_connection(tg_id, 0, 1) logger.info(f"Storing key for user {tg_id} in the database.") - await store_key( - tg_id, - client_id, - email, - expiry_timestamp, - public_link, - least_loaded_cluster, - ) + await store_key(tg_id, client_id, email, expiry_timestamp, public_link, least_loaded_cluster, session) except Exception as e: logger.error(f"Error while creating the key for user {tg_id}: {e}")