From dba8752436100bb0ed7163cfc37bb7ca163a8787 Mon Sep 17 00:00:00 2001 From: Boris Kovalskii <36034823+JustYay@users.noreply.github.com> Date: Sun, 16 Feb 2025 17:08:38 +1000 Subject: [PATCH 1/2] optimize database --- database.py | 42 ++++++++++++++++++++++++------------------ 1 file changed, 24 insertions(+), 18 deletions(-) diff --git a/database.py b/database.py index 3b29e7f8..9b84db9e 100644 --- a/database.py +++ b/database.py @@ -537,18 +537,18 @@ async def get_balance(tg_id: int) -> float: await conn.close() -async def update_balance(tg_id: int, amount: float, session: Any = None, is_admin: bool = False): +async def update_balance( + tg_id: int, + amount: float, + session: Any = None, + is_admin: bool = False, + skip_referral: bool = False, # <- флаг "пропустить реферальное начисление" + skip_cashback: bool = False # <- флаг "пропустить кэшбэк" +): """ Обновляет баланс пользователя в базе данных. - Кэшбек применяется только для положительных сумм, если пополнение НЕ через админку. - - Args: - tg_id (int): Telegram ID пользователя. - amount (float): Сумма для обновления баланса. - session (Any, optional): Сессия базы данных. Если не передана, создается новая. - is_admin (bool, optional): Флаг, указывающий, что пополнение идёт через админку. По умолчанию False. - Raises: - Exception: В случае ошибки при подключении к базе данных или обновлении баланса. + - Кэшбек применяется только для положительных сумм, если пополнение НЕ через админку и не пропущен явно. + - Реферальный бонус тоже не срабатывает, если явно попросили пропустить (например, при начислении за купон). """ conn = None try: @@ -556,15 +556,20 @@ async def update_balance(tg_id: int, amount: float, session: Any = None, is_admi conn = await asyncpg.connect(DATABASE_URL) session = conn - extra = amount * (CASHBACK / 100.0) if (CASHBACK > 0 and amount > 0 and not is_admin) else 0 + # Если пополнение не от админа и не сказали пропустить кэшбэк + if (CASHBACK > 0 and amount > 0 and not is_admin and not skip_cashback): + extra = amount * (CASHBACK / 100.0) + else: + extra = 0 + total_amount = int(amount + extra) - current_balance = await session.fetchval("SELECT balance FROM connections WHERE tg_id = $1", tg_id) + current_balance = await session.fetchval( + "SELECT balance FROM connections WHERE tg_id = $1", + tg_id + ) or 0 - if current_balance is None: - current_balance = 0 - - new_balance = int(current_balance) + total_amount + new_balance = current_balance + total_amount await session.execute( """ @@ -580,7 +585,8 @@ async def update_balance(tg_id: int, amount: float, session: Any = None, is_admi f"({'+ кешбэк' if extra > 0 else 'без кешбэка'}), стало: {new_balance}" ) - if not is_admin: + # Если не админ и не пропустили реферальное начисление — обрабатываем реферальную цепочку + if not is_admin and not skip_referral: await handle_referral_on_balance_update(tg_id, int(amount)) except Exception as e: @@ -729,7 +735,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) + await update_balance(referrer_tg_id, bonus, skip_referral=True, skip_cashback=True) if CHECK_REFERRAL_REWARD_ISSUED: await conn.execute( From 770f40f4ffdd7b71533b6ef30b58a8ca392b5c29 Mon Sep 17 00:00:00 2001 From: Boris Kovalskii <36034823+JustYay@users.noreply.github.com> Date: Sun, 16 Feb 2025 17:09:16 +1000 Subject: [PATCH 2/2] Opimize and fix coupons and refferals --- handlers/keys/key_utils.py | 27 ++++++++-- handlers/keys/subscriptions.py | 96 ++++++++++++++++++++++------------ 2 files changed, 88 insertions(+), 35 deletions(-) diff --git a/handlers/keys/key_utils.py b/handlers/keys/key_utils.py index 91b0bfac..9e286f72 100644 --- a/handlers/keys/key_utils.py +++ b/handlers/keys/key_utils.py @@ -14,14 +14,23 @@ async def create_key_on_cluster( cluster_id: str, tg_id: int, client_id: str, email: str, expiry_timestamp: int, plan: int = None ): """ - Создает ключ на всех серверах указанного кластера. + Создает ключ на всех серверах указанного кластера (или на конкретном сервере, если cluster_id — это имя сервера). """ try: servers = await get_servers() cluster = servers.get(cluster_id) + # Если не нашли кластер по ключу, ищем сервер по имени (аналогично renew_key_in_cluster, delete_key_from_cluster) if not cluster: - raise ValueError(f"Кластер с ID {cluster_id} не найден.") + found_servers = [] + for _key, server_list in servers.items(): + for server_info in server_list: + if server_info.get("server_name", "").lower() == cluster_id.lower(): + found_servers.append(server_info) + if found_servers: + cluster = found_servers + else: + raise ValueError(f"Кластер или сервер с ID/именем {cluster_id} не найден.") semaphore = asyncio.Semaphore(2) @@ -197,12 +206,24 @@ async def delete_key_from_cluster(cluster_id, email, client_id): async def update_key_on_cluster(tg_id, client_id, email, expiry_time, cluster_id): + """ + Обновляет ключ на всех серверах указанного кластера (или сервера, если передано имя). + """ try: servers = await get_servers() cluster = servers.get(cluster_id) + # Аналогичная логика поиска кластера или конкретного сервера if not cluster: - raise ValueError(f"Кластер с ID {cluster_id} не найден.") + found_servers = [] + for _key, server_list in servers.items(): + for server_info in server_list: + if server_info.get("server_name", "").lower() == cluster_id.lower(): + found_servers.append(server_info) + if found_servers: + cluster = found_servers + else: + raise ValueError(f"Кластер или сервер с ID/именем {cluster_id} не найден.") tasks = [] for server_info in cluster: diff --git a/handlers/keys/subscriptions.py b/handlers/keys/subscriptions.py index 2d66e99b..a4a6ce64 100644 --- a/handlers/keys/subscriptions.py +++ b/handlers/keys/subscriptions.py @@ -6,10 +6,14 @@ import aiohttp import asyncpg from aiohttp import web -from config import DATABASE_URL, PROJECT_NAME, SUB_MESSAGE, SUPERNODE, TRANSITION_DATE_STR, USE_COUNTRY_SELECTION +from config import ( + DATABASE_URL, PROJECT_NAME, SUB_MESSAGE, SUPERNODE, + TRANSITION_DATE_STR, USE_COUNTRY_SELECTION +) from database import get_key_details, get_servers from logger import logger + db_pool = None @@ -53,7 +57,10 @@ async def combine_unique_lines(urls, tg_id, query_string): logger.info(f"Начинаем объединение подписок для tg_id: {tg_id}, запрос: {query_string}") - urls_with_query = [f"{url}?{query_string}" if query_string else url for url in urls] + urls_with_query = [ + f"{url}?{query_string}" if query_string else url + for url in urls + ] logger.info(f"Составлены URL-адреса: {urls_with_query}") tasks = [fetch_url_content(url, tg_id) for url in urls_with_query] @@ -64,7 +71,6 @@ async def combine_unique_lines(urls, tg_id, query_string): all_lines.update(filter(None, lines)) logger.info(f"Объединено {len(all_lines)} строк после фильтрации и удаления дубликатов для tg_id: {tg_id}") - return list(all_lines) @@ -75,6 +81,40 @@ transition_timestamp_ms_adjusted = transition_timestamp_ms - (3 * 60 * 60 * 1000 logger.info(f"Время перехода (с поправкой на часовой пояс): {transition_timestamp_ms_adjusted}") +async def get_subscription_urls(server_id: str, email: str, conn) -> list: + """ + Универсальная функция, которая в зависимости от флага USE_COUNTRY_SELECTION + получает список URL-адресов для подписки. Возвращает пустой список, если + нужные данные не найдены. + """ + if USE_COUNTRY_SELECTION: + logger.info(f"Режим выбора страны активен. Ищем сервер {server_id} в БД.") + server_data = await conn.fetchrow( + "SELECT subscription_url FROM servers WHERE server_name = $1", + server_id + ) + if not server_data: + logger.warning(f"Не найден сервер {server_id} в БД!") + return [] + subscription_url = server_data["subscription_url"] + urls = [f"{subscription_url}/{email}"] + logger.info(f"Используем подписку {urls[0]}") + return urls + + # Если режим выбора страны выключен, получаем все сервера кластера + servers = await get_servers() + logger.info(f"Режим выбора страны отключен. Используем кластер {server_id}.") + cluster_servers = servers.get(server_id, []) + + if not cluster_servers: + logger.warning(f"Не найдены сервера для {server_id}") + return [] + + urls = [f"{server['subscription_url']}/{email}" for server in cluster_servers] + logger.info(f"Найдено {len(urls)} URL-адресов в кластере {server_id}") + return urls + + async def handle_subscription(request, old_subscription=False): """Обрабатывает запрос на подписку (старую или новую).""" email = request.match_info.get("email") @@ -85,7 +125,8 @@ async def handle_subscription(request, old_subscription=False): return web.Response(text="❌ Неверные параметры запроса.", status=400) logger.info( - f"Обработка запроса для {'старого' if old_subscription else 'нового'} клиента: email={email}, tg_id={tg_id}" + f"Обработка запроса для {'старого' if old_subscription else 'нового'} клиента: " + f"email={email}, tg_id={tg_id}" ) await init_db_pool() @@ -100,6 +141,7 @@ async def handle_subscription(request, old_subscription=False): stored_tg_id = client_data.get("tg_id") server_id = client_data["server_id"] + # Проверяем, что tg_id из запроса совпадает с сохранённым в БД (для новой подписки) if not old_subscription and str(tg_id) != str(stored_tg_id): logger.warning(f"Неверный tg_id для клиента с email {email}.") return web.Response(text="❌ Неверные данные. Получите свой ключ в боте.", status=403) @@ -112,39 +154,29 @@ async def handle_subscription(request, old_subscription=False): if created_at_ms >= transition_timestamp_ms_adjusted: logger.info(f"Клиент с email {email} является новым.") - return web.Response(text="❌ Эта ссылка устарела. Пожалуйста, обновите ссылку.", status=400) + return web.Response( + text="❌ Эта ссылка устарела. Пожалуйста, обновите ссылку.", + status=400 + ) - urls = [] - - if USE_COUNTRY_SELECTION: - logger.info(f"Режим выбора страны активен. Ищем сервер {server_id} в БД.") - server_data = await conn.fetchrow("SELECT subscription_url FROM servers WHERE server_name = $1", server_id) - - if not server_data: - logger.warning(f"Не найден сервер {server_id} в БД!") - return web.Response(text="❌ Сервер не найден.", status=404) - - subscription_url = server_data["subscription_url"] - urls = [f"{subscription_url}/{email}"] - logger.info(f"Используем подписку {urls[0]}") - - else: - servers = await get_servers() - logger.info(f"Режим выбора страны отключен. Используем кластер {server_id}.") - cluster_servers = servers.get(server_id, []) - - if not cluster_servers: - logger.warning(f"Не найдены сервера для {server_id}") - return web.Response(text="❌ Сервер не найден.", status=404) - - urls = [f"{server['subscription_url']}/{email}" for server in cluster_servers] + # Получаем список URL-адресов для подписки через универсальную функцию + urls = await get_subscription_urls(server_id, email, conn) + if not urls: + # Сообщаем, что сервер не найден или неверные данные + return web.Response(text="❌ Сервер не найден.", status=404) query_string = request.query_string if not old_subscription else "" - combined_subscriptions = await combine_unique_lines(urls, tg_id or email, query_string) + combined_subscriptions = await combine_unique_lines( + urls, + tg_id or email, # Если tg_id нет, для лога используем email + query_string + ) + + base64_encoded = base64.b64encode( + "\n".join(combined_subscriptions).encode("utf-8") + ).decode("utf-8") - base64_encoded = base64.b64encode("\n".join(combined_subscriptions).encode("utf-8")).decode("utf-8") encoded_project_name = f"{PROJECT_NAME} - {SUB_MESSAGE}" - headers = { "Content-Type": "text/plain; charset=utf-8", "Content-Disposition": "inline",