From ec9016f9589d2bf8854edea53fcbddc4d39d06cf Mon Sep 17 00:00:00 2001 From: Capybara-z Date: Sat, 5 Apr 2025 00:46:24 +0300 Subject: [PATCH] =?UTF-8?q?=D0=9F=D0=B5=D1=80=D0=B5=D0=BD=D0=BE=D1=81=20?= =?UTF-8?q?=D0=BE=D0=B1=D1=80=D0=B0=D1=82=D0=BE=D0=B1=D1=82=D0=BA=D0=B8=20?= =?UTF-8?q?=D1=81=D0=BE=D1=81=D1=82=D0=BE=D1=8F=D0=BD=D0=B8=D0=B9=20=D0=BF?= =?UTF-8?q?=D0=BE=D1=81=D0=BB=D0=B5=20=D1=81=D0=BE=D0=B7=D0=B4=D0=B0=D0=BD?= =?UTF-8?q?=D0=B8=D1=8F=20=D0=BA=D0=BB=D1=8E=D1=87=D0=B0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Перенос оплат и списания триала после создания ключа --- handlers/keys/key_management.py | 59 +++++++++++++++++---------------- 1 file changed, 31 insertions(+), 28 deletions(-) diff --git a/handlers/keys/key_management.py b/handlers/keys/key_management.py index 61817bce..5d08cf57 100644 --- a/handlers/keys/key_management.py +++ b/handlers/keys/key_management.py @@ -96,20 +96,14 @@ async def handle_key_creation( extra_days = NOTIFY_EXTRA_DAYS if trial_status == -1 else 0 expiry_time = current_time + timedelta(days=TRIAL_TIME + extra_days) logger.info(f"Доступен {TRIAL_TIME + extra_days}-дневный пробный период пользователю {tg_id}.") - updated = await update_trial(tg_id, 1, session) - if updated: - await edit_or_send_message( - target_message=message_or_query - if isinstance(message_or_query, Message) - else message_or_query.message, - text=CREATING_CONNECTION_MSG, - reply_markup=None, - ) - - await create_key(tg_id, expiry_time, state, session, message_or_query) - return - else: - logger.error(f"Не удалось обновить статус триального периода для пользователя {tg_id}.") + await edit_or_send_message( + target_message=message_or_query if isinstance(message_or_query, Message) else message_or_query.message, + text=CREATING_CONNECTION_MSG, + reply_markup=None, + ) + await state.update_data(is_trial=True) + await create_key(tg_id, expiry_time, state, session, message_or_query) + return builder = InlineKeyboardBuilder() for index, (plan_id, price) in enumerate(RENEWAL_PRICES.items()): @@ -144,7 +138,7 @@ async def handle_key_creation( @router.callback_query(F.data.startswith("select_plan_")) -async def select_tariff_plan(callback_query: CallbackQuery, session: Any): +async def select_tariff_plan(callback_query: CallbackQuery, session: Any, state: FSMContext): tg_id = callback_query.message.chat.id plan_id = callback_query.data.split("_")[-1] plan_price = RENEWAL_PRICES.get(plan_id) @@ -191,9 +185,8 @@ async def select_tariff_plan(callback_query: CallbackQuery, session: Any): ) expiry_time = datetime.now(moscow_tz) + timedelta(days=duration_days) - await create_key(tg_id, expiry_time, None, session, callback_query, None, plan_id) - await update_balance(tg_id, -plan_price, session) - + await state.update_data(plan_id=plan_id) + await create_key(tg_id, expiry_time, state, session, callback_query) async def create_key( tg_id: int, @@ -214,13 +207,8 @@ async def create_key( if USE_COUNTRY_SELECTION: logger.info("[Country Selection] USE_COUNTRY_SELECTION включен. Получение наименее загруженного кластера") least_loaded_cluster = await get_least_loaded_cluster() - logger.info( - f"[Country Selection] Наименее загруженный кластер: {least_loaded_cluster}. Получаем список серверов" - ) - servers = await session.fetch( - "SELECT server_name FROM servers WHERE cluster_name = $1", - least_loaded_cluster, - ) + logger.info(f"[Country Selection] Наименее загруженный кластер: {least_loaded_cluster}. Получаем список серверов") + servers = await session.fetch("SELECT server_name FROM servers WHERE cluster_name = $1", least_loaded_cluster) countries = [server["server_name"] for server in servers] logger.info(f"[Country Selection] Список серверов: {countries}") @@ -280,7 +268,14 @@ async def create_key( least_loaded_cluster, session, ) - await update_trial(tg_id, 1, session) + data = await state.get_data() + if data.get("is_trial"): + trial_status = await get_trial(tg_id, session) + if trial_status in [0, -1]: + await update_trial(tg_id, 1, session) + if data.get("plan_id"): + plan_price = RENEWAL_PRICES.get(data["plan_id"]) + await update_balance(tg_id, -plan_price, session) logger.info(f"[Database] Ключ сохранён в базе данных для пользователя {tg_id}") except Exception as e: logger.error(f"[Error] Ошибка при создании ключа для пользователя {tg_id}: {e}") @@ -375,7 +370,7 @@ async def change_location_callback(callback_query: CallbackQuery, session: Any): @router.callback_query(F.data.startswith("select_country|")) -async def handle_country_selection(callback_query: CallbackQuery, session: Any): +async def handle_country_selection(callback_query: CallbackQuery, session: Any, state: FSMContext): """ Обрабатывает выбор страны. Формат callback data: @@ -402,7 +397,7 @@ async def handle_country_selection(callback_query: CallbackQuery, session: Any): logger.info(f"Пользователь {tg_id} выбрал страну: {selected_country}") logger.info(f"Получено время истечения (timestamp): {ts}") - await finalize_key_creation(tg_id, expiry_time, selected_country, None, session, callback_query, old_key_name) + await finalize_key_creation(tg_id, expiry_time, selected_country, state, session, callback_query, old_key_name) async def finalize_key_creation( @@ -516,6 +511,14 @@ async def finalize_key_creation( public_link, selected_country, ) + data = await state.get_data() + if data.get("is_trial"): + trial_status = await get_trial(tg_id, session) + if trial_status in [0, -1]: + await update_trial(tg_id, 1, session) + if data.get("plan_id"): + plan_price = RENEWAL_PRICES.get(data["plan_id"]) + await update_balance(tg_id, -plan_price, session) except Exception as e: logger.error(f"Error while creating the key for user {tg_id}: {e}")