diff --git a/handlers/keys/key_mode/key_cluster_mode.py b/handlers/keys/key_mode/key_cluster_mode.py index 3be45878..a77703e9 100644 --- a/handlers/keys/key_mode/key_cluster_mode.py +++ b/handlers/keys/key_mode/key_cluster_mode.py @@ -92,21 +92,26 @@ async def key_cluster_mode( if tariff.get("traffic_limit") is not None: traffic_limit_gb = int(tariff["traffic_limit"]) - try: - least_loaded_cluster = await get_least_loaded_cluster(session) - except ValueError as e: - logger.error(f"Нет доступных кластеров: {e}") - error_message = str(e) + forced_cluster_results = await run_hooks("cluster_override", tg_id=tg_id, state_data=data, session=session, plan=plan) + + if forced_cluster_results and forced_cluster_results[0]: + least_loaded_cluster = forced_cluster_results[0] + else: + try: + least_loaded_cluster = await get_least_loaded_cluster(session) + except ValueError as e: + logger.error(f"Нет доступных кластеров: {e}") + error_message = str(e) - if safe_to_edit: - await edit_or_send_message( - target_message=target_message, - text=error_message, - reply_markup=None, - ) - else: - await bot.send_message(chat_id=tg_id, text=error_message) - return + if safe_to_edit: + await edit_or_send_message( + target_message=target_message, + text=error_message, + reply_markup=None, + ) + else: + await bot.send_message(chat_id=tg_id, text=error_message) + return await create_key_on_cluster( cluster_id=least_loaded_cluster, diff --git a/handlers/keys/key_mode/key_country_mode.py b/handlers/keys/key_mode/key_country_mode.py index 83e46f83..c9f91801 100644 --- a/handlers/keys/key_mode/key_country_mode.py +++ b/handlers/keys/key_mode/key_country_mode.py @@ -78,16 +78,23 @@ async def key_country_mode( target_message = message_or_query safe_to_edit = True - try: - least_loaded_cluster = await get_least_loaded_cluster(session) - except ValueError as e: - logger.error(f"Нет доступных кластеров: {e}") - text = str(e) - if safe_to_edit: - await edit_or_send_message(target_message=target_message, text=text, reply_markup=None) - else: - await bot.send_message(chat_id=tg_id, text=text) - return + data = await state.get_data() if state else {} + + forced_cluster_results = await run_hooks("cluster_override", tg_id=tg_id, state_data=data, session=session, plan=plan) + + if forced_cluster_results and forced_cluster_results[0]: + least_loaded_cluster = forced_cluster_results[0] + else: + try: + least_loaded_cluster = await get_least_loaded_cluster(session) + except ValueError as e: + logger.error(f"Нет доступных кластеров: {e}") + text = str(e) + if safe_to_edit: + await edit_or_send_message(target_message=target_message, text=text, reply_markup=None) + else: + await bot.send_message(chat_id=tg_id, text=text) + return result = await session.execute( select( diff --git a/handlers/keys/key_mode/key_create.py b/handlers/keys/key_mode/key_create.py index 3df33897..ddd5664e 100644 --- a/handlers/keys/key_mode/key_create.py +++ b/handlers/keys/key_mode/key_create.py @@ -39,6 +39,8 @@ from handlers.texts import ( SELECT_TARIFF_PLAN_MSG, ) from handlers.utils import edit_or_send_message, format_discount_time_left, get_least_loaded_cluster +from hooks.hook_buttons import insert_hook_buttons +from hooks.hooks import run_hooks from logger import logger from .key_cluster_mode import key_cluster_mode @@ -240,6 +242,9 @@ async def handle_key_creation( ) ) + tariff_menu_buttons = await run_hooks("tariff_menu", group_code=group_code, cluster_name=cluster_name, tg_id=tg_id, session=session) + builder = insert_hook_buttons(builder, tariff_menu_buttons) + builder.row(InlineKeyboardButton(text=MAIN_MENU, callback_data="profile")) target_message = message_or_query.message if isinstance(message_or_query, CallbackQuery) else message_or_query diff --git a/handlers/utils.py b/handlers/utils.py index 685d65a6..26fd0433 100644 --- a/handlers/utils.py +++ b/handlers/utils.py @@ -21,6 +21,7 @@ from bot import bot from config import ADMIN_ID from database import get_servers from database.models import Key, Notification, Server +from hooks.hooks import run_hooks from logger import logger @@ -67,6 +68,10 @@ async def get_least_loaded_cluster(session: AsyncSession) -> str: else: continue + cluster_filter_results = await run_hooks("cluster_balancer", available_clusters=available_clusters, session=session) + if cluster_filter_results and cluster_filter_results[0]: + available_clusters = cluster_filter_results[0] + if not available_clusters: logger.warning("❌ Нет доступных кластеров с лимитом ключей!") raise ValueError("⚠️ Сервисы временно недоступны. Попробуйте позже.")