From 052124d483bdefb4f3a4d9bc203036256c51f290 Mon Sep 17 00:00:00 2001 From: Zakhar Izmaylov Date: Sun, 24 Nov 2024 20:29:22 +0300 Subject: [PATCH] Refactor database and notification handlers for improved readability and consistency. Clean up import statements, enhance logging messages, and streamline function calls across multiple files. Ensure proper session handling in payment processing functions. --- database.py | 58 ++++++++++++++++-------------- handlers/admin/admin_coupons.py | 4 +-- handlers/keys/key_management.py | 10 +++++- handlers/notifications.py | 20 +++++------ handlers/payments/cryprobot_pay.py | 5 +-- handlers/payments/robokassa_pay.py | 4 +-- handlers/payments/stars_pay.py | 5 +-- handlers/payments/yookassa_pay.py | 4 +-- middlewares/delete.py | 3 +- 9 files changed, 65 insertions(+), 48 deletions(-) diff --git a/database.py b/database.py index 99f35b46..08b934ce 100644 --- a/database.py +++ b/database.py @@ -3,7 +3,7 @@ from typing import Any import asyncpg -from config import REFERRAL_BONUS_PERCENTAGES, DATABASE_URL +from config import DATABASE_URL, REFERRAL_BONUS_PERCENTAGES from logger import logger @@ -611,9 +611,9 @@ async def get_all_users(conn): Exception: В случае ошибки при получении данных """ try: - пользователи = await conn.fetch("SELECT tg_id FROM connections") - logger.info(f"Получен список всех пользователей. Количество: {len(пользователи)}") - return пользователи + users = await conn.fetch("SELECT tg_id FROM connections") + logger.info(f"Получен список всех пользователей. Количество: {len(users)}") + return users except Exception as e: logger.error(f"Ошибка при получении списка пользователей: {e}") raise @@ -679,10 +679,7 @@ async def handle_referral_on_balance_update(tg_id: int, amount: float): break referrer_tg_id = referral['referrer_tg_id'] - referral_chain.append({ - 'tg_id': referrer_tg_id, - 'level': level - }) + referral_chain.append({'tg_id': referrer_tg_id, 'level': level}) # Переходим к следующему уровню current_tg_id = referrer_tg_id @@ -691,7 +688,7 @@ async def handle_referral_on_balance_update(tg_id: int, amount: float): for referral in referral_chain: referrer_tg_id = referral['tg_id'] level = referral['level'] - + # Расчет бонуса для текущего уровня bonus_percent = REFERRAL_BONUS_PERCENTAGES.get(level, 0) bonus = amount * bonus_percent @@ -699,7 +696,7 @@ async def handle_referral_on_balance_update(tg_id: int, amount: float): if bonus > 0: logger.info(f"Начисление бонуса {bonus} рублей рефереру {referrer_tg_id} на уровне {level}") - + # Обновляем баланс реферера await update_balance(referrer_tg_id, bonus) @@ -777,13 +774,11 @@ async def get_referral_stats(referrer_tg_id: int): """, referrer_tg_id, ) - + # Преобразование результатов в словарь referrals_by_level = { - record['level']: { - 'total': record['level_count'], - 'active': record['active_level_count'] - } for record in referrals_by_level_records + record['level']: {'total': record['level_count'], 'active': record['active_level_count']} + for record in referrals_by_level_records } logger.debug(f"Получена статистика рефералов по уровням: {referrals_by_level}") @@ -1139,7 +1134,13 @@ async def check_notification_time(tg_id: int, notification_type: str, hours: int Exception: В случае ошибки при проверке времени уведомления """ try: - result = await session.fetchrow( + # Если сессия не передана, создаем новое подключение + if session is None: + conn = await asyncpg.connect(DATABASE_URL) + else: + conn = session + + result = await conn.fetchrow( """ SELECT CASE @@ -1149,23 +1150,28 @@ async def check_notification_time(tg_id: int, notification_type: str, hours: int END as can_notify FROM notifications WHERE tg_id = $2 AND notification_type = $3 - """, - hours, - tg_id, - notification_type + """, + str(hours), # Преобразуем hours в строку + tg_id, + notification_type, ) + # Если сессия не была передана, закрываем подключение + if session is None and conn: + await conn.close() + # Если записи нет, значит уведомление можно отправить if result is None: return True can_notify = result['can_notify'] - - logger.info(f"Проверка уведомления типа {notification_type} для пользователя {tg_id}: {'можно отправить' if can_notify else 'слишком рано'}") - + + logger.info( + f"Проверка уведомления типа {notification_type} для пользователя {tg_id}: {'можно отправить' if can_notify else 'слишком рано'}" + ) + return can_notify - + except Exception as e: logger.error(f"Ошибка при проверке времени уведомления для пользователя {tg_id}: {e}") - raise - + return True # По умолчанию разрешаем отправку уведомления diff --git a/handlers/admin/admin_coupons.py b/handlers/admin/admin_coupons.py index d632a9eb..5ea3ca98 100644 --- a/handlers/admin/admin_coupons.py +++ b/handlers/admin/admin_coupons.py @@ -76,12 +76,12 @@ async def handle_delete_coupon(callback_query: types.CallbackQuery, session: Any result = await delete_coupon_from_db(coupon_code, session) if result: - await show_coupon_list(callback_query,session) + await show_coupon_list(callback_query, session) else: await callback_query.message.answer( f"❌ Купон с кодом {coupon_code} не найден.", ) - await show_coupon_list(callback_query,session) + await show_coupon_list(callback_query, session) except Exception as e: logger.error(f"Ошибка при удалении купона: {e}") diff --git a/handlers/keys/key_management.py b/handlers/keys/key_management.py index 2870aadf..780a7e71 100644 --- a/handlers/keys/key_management.py +++ b/handlers/keys/key_management.py @@ -10,7 +10,15 @@ 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, check_connection_exists, get_balance, get_trial, store_key, update_balance, use_trial +from database import ( + add_connection, + check_connection_exists, + 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 diff --git a/handlers/notifications.py b/handlers/notifications.py index 5a8049e7..613f940a 100644 --- a/handlers/notifications.py +++ b/handlers/notifications.py @@ -8,7 +8,14 @@ from py3xui import AsyncApi from client import delete_client from config import ADMIN_PASSWORD, ADMIN_USERNAME, CLUSTERS, DATABASE_URL, TOTAL_GB, TRIAL_TIME -from database import delete_key, get_balance, update_balance, update_key_expiry,add_notification,check_notification_time +from database import ( + add_notification, + check_notification_time, + delete_key, + get_balance, + update_balance, + update_key_expiry, +) from handlers.keys.key_utils import renew_key_in_cluster from handlers.texts import KEY_EXPIRY_10H, KEY_EXPIRY_24H, KEY_RENEWED, RENEWAL_PLANS from logger import logger @@ -220,10 +227,7 @@ async def notify_inactive_trial_users(bot: Bot, conn: asyncpg.Connection): try: # Проверяем, можно ли отправить уведомление can_notify = await check_notification_time( - tg_id, - 'inactive_trial', - hours=24, # Уведомление не чаще, чем раз в 24 часа - session=conn + tg_id, 'inactive_trial', hours=24, session=conn # Уведомление не чаще, чем раз в 24 часа ) if can_notify and not await is_bot_blocked(bot, tg_id): @@ -245,11 +249,7 @@ async def notify_inactive_trial_users(bot: Bot, conn: asyncpg.Connection): logger.info(f"Отправлено уведомление неактивному пользователю {tg_id}.") # Добавляем запись о notification - await add_notification( - tg_id, - 'inactive_trial', - session=conn - ) + await add_notification(tg_id, 'inactive_trial', session=conn) except Exception as e: logger.error(f"Ошибка при отправке уведомления неактивному пользователю {tg_id}: {e}") diff --git a/handlers/payments/cryprobot_pay.py b/handlers/payments/cryprobot_pay.py index acda2a1a..d25c263d 100644 --- a/handlers/payments/cryprobot_pay.py +++ b/handlers/payments/cryprobot_pay.py @@ -1,4 +1,5 @@ from typing import Any + from aiocryptopay import AioCryptoPay, Networks from aiogram import F, Router, types from aiogram.fsm.context import FSMContext @@ -26,7 +27,7 @@ class ReplenishBalanceState(StatesGroup): @router.callback_query(F.data == "pay_cryptobot") -async def process_callback_pay_cryptobot(callback_query: types.CallbackQuery, state: FSMContext,session:Any): +async def process_callback_pay_cryptobot(callback_query: types.CallbackQuery, state: FSMContext, session: Any): builder = InlineKeyboardBuilder() for i in range(0, len(PAYMENT_OPTIONS), 2): if i + 1 < len(PAYMENT_OPTIONS): @@ -58,7 +59,7 @@ async def process_callback_pay_cryptobot(callback_query: types.CallbackQuery, st if key_count == 0: exists = await check_connection_exists(callback_query.message.chat.id) if not exists: - await add_connection(tg_id=callback_query.message.chat.id, balance=0.0, trial=0,session=session) + await add_connection(tg_id=callback_query.message.chat.id, balance=0.0, trial=0, session=session) await callback_query.message.answer( "Выберите сумму пополнения:", reply_markup=builder.as_markup(), diff --git a/handlers/payments/robokassa_pay.py b/handlers/payments/robokassa_pay.py index 4a21eac8..f9248aec 100644 --- a/handlers/payments/robokassa_pay.py +++ b/handlers/payments/robokassa_pay.py @@ -49,7 +49,7 @@ def generate_payment_link(amount, inv_id, description, tg_id): @router.callback_query(F.data == "pay_robokassa") -async def process_callback_pay_robokassa(callback_query: types.CallbackQuery, state: FSMContext,session:Any): +async def process_callback_pay_robokassa(callback_query: types.CallbackQuery, state: FSMContext, session: Any): tg_id = callback_query.message.chat.id logger.info(f"User {tg_id} initiated Robokassa payment.") @@ -86,7 +86,7 @@ async def process_callback_pay_robokassa(callback_query: types.CallbackQuery, st if key_count == 0: exists = await check_connection_exists(tg_id) if not exists: - await add_connection(tg_id, balance=0.0, trial=0,session=session) + await add_connection(tg_id, balance=0.0, trial=0, session=session) logger.info(f"Created new connection for user {tg_id} with balance 0.0.") await callback_query.message.answer( diff --git a/handlers/payments/stars_pay.py b/handlers/payments/stars_pay.py index b55bff13..7e5730a5 100644 --- a/handlers/payments/stars_pay.py +++ b/handlers/payments/stars_pay.py @@ -1,4 +1,5 @@ from typing import Any + from aiogram import F, Router, types from aiogram.fsm.context import FSMContext from aiogram.fsm.state import State, StatesGroup @@ -21,7 +22,7 @@ class ReplenishBalanceState(StatesGroup): @router.callback_query(F.data == "pay_stars") -async def process_callback_pay_stars(callback_query: types.CallbackQuery, state: FSMContext,session:Any): +async def process_callback_pay_stars(callback_query: types.CallbackQuery, state: FSMContext, session: Any): tg_id = callback_query.message.chat.id builder = InlineKeyboardBuilder() @@ -59,7 +60,7 @@ async def process_callback_pay_stars(callback_query: types.CallbackQuery, state: if key_count == 0: exists = await check_connection_exists(tg_id) if not exists: - await add_connection(tg_id, balance=0.0, trial=0,session=session) + await add_connection(tg_id, balance=0.0, trial=0, session=session) try: await callback_query.message.delete() diff --git a/handlers/payments/yookassa_pay.py b/handlers/payments/yookassa_pay.py index c8ec2ebf..edc9f071 100644 --- a/handlers/payments/yookassa_pay.py +++ b/handlers/payments/yookassa_pay.py @@ -31,7 +31,7 @@ class ReplenishBalanceState(StatesGroup): @router.callback_query(F.data == "pay_yookassa") -async def process_callback_pay_yookassa(callback_query: types.CallbackQuery, state: FSMContext,session:Any): +async def process_callback_pay_yookassa(callback_query: types.CallbackQuery, state: FSMContext, session: Any): tg_id = callback_query.message.chat.id builder = InlineKeyboardBuilder() @@ -68,7 +68,7 @@ async def process_callback_pay_yookassa(callback_query: types.CallbackQuery, sta if key_count == 0: exists = await check_connection_exists(tg_id) if not exists: - await add_connection(tg_id, balance=0.0, trial=0,session=session) + await add_connection(tg_id, balance=0.0, trial=0, session=session) await callback_query.message.answer( text="Выберите сумму пополнения:", diff --git a/middlewares/delete.py b/middlewares/delete.py index 4ac4d812..e3911cf8 100644 --- a/middlewares/delete.py +++ b/middlewares/delete.py @@ -1,8 +1,9 @@ from typing import Any, Awaitable, Callable, Dict -from logger import logger + from aiogram import BaseMiddleware from aiogram.types import CallbackQuery, Message, TelegramObject + class DeleteMessageMiddleware(BaseMiddleware): async def __call__( self,