Add hook integration: cluster balancing, subscription routing, tariff menu

This commit is contained in:
Capybara-z
2025-09-16 22:09:32 +03:00
parent c1e26ca1ce
commit 4b8bea65be
4 changed files with 46 additions and 24 deletions
+19 -14
View File
@@ -92,21 +92,26 @@ async def key_cluster_mode(
if tariff.get("traffic_limit") is not None: if tariff.get("traffic_limit") is not None:
traffic_limit_gb = int(tariff["traffic_limit"]) traffic_limit_gb = int(tariff["traffic_limit"])
try: forced_cluster_results = await run_hooks("cluster_override", tg_id=tg_id, state_data=data, session=session, plan=plan)
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: if forced_cluster_results and forced_cluster_results[0]:
await edit_or_send_message( least_loaded_cluster = forced_cluster_results[0]
target_message=target_message, else:
text=error_message, try:
reply_markup=None, least_loaded_cluster = await get_least_loaded_cluster(session)
) except ValueError as e:
else: logger.error(f"Нет доступных кластеров: {e}")
await bot.send_message(chat_id=tg_id, text=error_message) error_message = str(e)
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( await create_key_on_cluster(
cluster_id=least_loaded_cluster, cluster_id=least_loaded_cluster,
+17 -10
View File
@@ -78,16 +78,23 @@ async def key_country_mode(
target_message = message_or_query target_message = message_or_query
safe_to_edit = True safe_to_edit = True
try: data = await state.get_data() if state else {}
least_loaded_cluster = await get_least_loaded_cluster(session)
except ValueError as e: forced_cluster_results = await run_hooks("cluster_override", tg_id=tg_id, state_data=data, session=session, plan=plan)
logger.error(f"Нет доступных кластеров: {e}")
text = str(e) if forced_cluster_results and forced_cluster_results[0]:
if safe_to_edit: least_loaded_cluster = forced_cluster_results[0]
await edit_or_send_message(target_message=target_message, text=text, reply_markup=None) else:
else: try:
await bot.send_message(chat_id=tg_id, text=text) least_loaded_cluster = await get_least_loaded_cluster(session)
return 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( result = await session.execute(
select( select(
+5
View File
@@ -39,6 +39,8 @@ from handlers.texts import (
SELECT_TARIFF_PLAN_MSG, SELECT_TARIFF_PLAN_MSG,
) )
from handlers.utils import edit_or_send_message, format_discount_time_left, get_least_loaded_cluster 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 logger import logger
from .key_cluster_mode import key_cluster_mode 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")) 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 target_message = message_or_query.message if isinstance(message_or_query, CallbackQuery) else message_or_query
+5
View File
@@ -21,6 +21,7 @@ from bot import bot
from config import ADMIN_ID from config import ADMIN_ID
from database import get_servers from database import get_servers
from database.models import Key, Notification, Server from database.models import Key, Notification, Server
from hooks.hooks import run_hooks
from logger import logger from logger import logger
@@ -67,6 +68,10 @@ async def get_least_loaded_cluster(session: AsyncSession) -> str:
else: else:
continue 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: if not available_clusters:
logger.warning("❌ Нет доступных кластеров с лимитом ключей!") logger.warning("❌ Нет доступных кластеров с лимитом ключей!")
raise ValueError("⚠️ Сервисы временно недоступны. Попробуйте позже.") raise ValueError("⚠️ Сервисы временно недоступны. Попробуйте позже.")