houseassassin/Solo_bot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

servers in clusters/total_gb

Browse Source
This commit is contained in:
Vladless
2024-11-14 03:03:32 +03:00
parent bd051f81a6
commit e23676b74d
9 changed files with 303 additions and 183 deletions
Download Patch File Download Diff File Expand all files Collapse all files
client.py
+13 -13
View File
@@ -47,7 +47,9 @@ async def add_client(
return {"status": "failed", "error": str(e)}
async def extend_client_key(xui, email: str, new_expiry_time: int, client_id: str):
async def extend_client_key(
xui, email: str, new_expiry_time: int, client_id: str, total_gb: int
):
"""
Функция для обновления срока действия ключа клиента по email.
"""
@@ -71,12 +73,7 @@ async def extend_client_key(xui, email: str, new_expiry_time: int, client_id: st
client.expiry_time = new_expiry_time
client.flow = "xtls-rprx-vision"
client.sub_id = email
if TOTAL_GB > 0:
client.total_gb = TOTAL_GB
logger.info(
f"Установлен объем трафика для клиента {client.email}: {TOTAL_GB} ГБ"
)
client.total_gb = total_gb
await xui.client.update(client.id, client)
logger.info(
@@ -91,23 +88,26 @@ async def delete_client(
xui,
email: str,
client_id: str,
):
) -> bool:
"""
Функция для удаления клиента с сервера 3x-ui
Функция для удаления клиента с сервера 3x-ui.
Возвращает True при успешном удалении, иначе False.
"""
await xui.login()
try:
client = await xui.client.get_by_email(email)
client.id = client_id
if not client:
logger.warning(f"Клиент с email {client_id} не найден.")
return
logger.warning(f"Клиент с email {email} и ID {client_id} не найден.")
return False
client.id = client_id
inbound_id = 1
await xui.client.delete(inbound_id, client.id)
await xui.client.delete(inbound_id, client.id)
logger.info(f"Клиент с ID {client_id} был удален успешно.")
return True
except Exception as e:
logger.error(f"Ошибка при удалении клиента с ID {client_id}: {e}")
return False
handlers/admin/admin_user_editor.py
+24 -12
View File
@@ -10,11 +10,11 @@ from aiogram.utils.keyboard import InlineKeyboardBuilder
from loguru import logger
from bot import bot
from config import DATABASE_URL, SERVERS
from config import CLUSTERS, DATABASE_URL
from database import get_client_id_by_email, restore_trial, update_key_expiry
from filters.admin import IsAdminFilter
from handlers.admin.admin_panel import back_to_admin_menu
from handlers.keys.key_utils import delete_key_from_server, renew_server_key
from handlers.keys.key_utils import delete_key_from_cluster, renew_key_in_cluster
from handlers.utils import sanitize_key_name
router = Router()
@@ -178,9 +178,14 @@ async def process_key_edit(callback_query: CallbackQuery):
key = record["key"]
expiry_time = record["expiry_time"]
server_id = record["server_id"]
server_name = SERVERS.get(server_id, {}).get(
"name", "Неизвестный сервер"
)
server_name = "Неизвестный сервер"
for cluster in CLUSTERS.values():
if server_id in cluster:
server_name = cluster[server_id].get(
"name", "Неизвестный сервер"
)
break
expiry_date = datetime.utcfromtimestamp(expiry_time / 1000)
current_date = datetime.utcnow()
@@ -284,7 +289,12 @@ async def handle_key_name_input(message: types.Message, state: FSMContext):
key = record["key"]
expiry_time = record["expiry_time"]
server_id = record["server_id"]
server_name = SERVERS.get(server_id, {}).get("name", "Неизвестный сервер")
server_name = "Неизвестный сервер"
for cluster in CLUSTERS.values():
if server_id in cluster:
server_name = cluster[server_id].get("name", "Неизвестный сервер")
break
expiry_date = datetime.utcfromtimestamp(expiry_time / 1000).strftime(
"%d %B %Y"
@@ -370,10 +380,12 @@ async def handle_expiry_time_input(message: types.Message, state: FSMContext):
async def update_key_on_all_servers():
tasks = []
for server_id in SERVERS:
for cluster_id in CLUSTERS: # Разворачиваем на cluster_id и cluster
tasks.append(
asyncio.create_task(
renew_server_key(server_id, email, client_id, expiry_time)
renew_key_in_cluster(
cluster_id, email, client_id, expiry_time
)
)
)
await asyncio.gather(*tasks)
@@ -468,15 +480,15 @@ async def process_callback_confirm_delete(callback_query: types.CallbackQuery):
InlineKeyboardButton(text="⬅️ Назад", callback_data="view_keys")
)
async def delete_key_from_servers():
async def delete_key_from_servers(email, client_id):
tasks = []
for server_id in SERVERS:
for cluster_id in CLUSTERS:
tasks.append(
delete_key_from_server(server_id, email, client_id)
delete_key_from_cluster(cluster_id, email, client_id)
)
await asyncio.gather(*tasks)
await delete_key_from_servers()
await delete_key_from_servers(email, client_id)
await delete_key_from_db(client_id)
await bot.edit_message_text(
handlers/keys/key_management.py
+15 -10
View File
@@ -10,13 +10,13 @@ from aiogram.types import CallbackQuery, InlineKeyboardButton, InlineKeyboardMar
from loguru import logger
from bot import bot, dp
from config import CONNECT_ANDROID, CONNECT_IOS, DATABASE_URL, DOWNLOAD_ANDROID, DOWNLOAD_IOS, PUBLIC_LINK, SERVERS
from config import CONNECT_ANDROID, CONNECT_IOS, DATABASE_URL, DOWNLOAD_ANDROID, DOWNLOAD_IOS, PUBLIC_LINK
from database import add_connection, get_balance, store_key, update_balance
from handlers.instructions.instructions import send_instructions
from handlers.keys.key_utils import create_key_on_server
from handlers.keys.key_utils import create_key_on_cluster
from handlers.profile import process_callback_view_profile
from handlers.texts import KEY, KEY_TRIAL, NULL_BALANCE, RENEWAL_PLANS, key_message_success
from handlers.utils import sanitize_key_name
from handlers.utils import get_least_loaded_cluster, sanitize_key_name
router = Router()
@@ -238,15 +238,20 @@ async def handle_key_name_input(message: Message, state: FSMContext):
)
try:
least_loaded_cluster = await get_least_loaded_cluster()
tasks = []
for server_id in SERVERS:
tasks.append(
asyncio.create_task(
create_key_on_server(
server_id, tg_id, client_id, email, expiry_timestamp
)
tasks.append(
asyncio.create_task(
create_key_on_cluster(
least_loaded_cluster,
tg_id,
client_id,
email,
expiry_timestamp,
)
)
)
await asyncio.gather(*tasks)
@@ -267,7 +272,7 @@ async def handle_key_name_input(message: Message, state: FSMContext):
logger.info(f"Storing key for user {tg_id} in the database.")
await store_key(
tg_id, client_id, email, expiry_timestamp, public_link, "all_servers"
tg_id, client_id, email, expiry_timestamp, public_link, least_loaded_cluster
)
except Exception as e:
handlers/keys/key_utils.py
+118 -64
View File
@@ -1,61 +1,88 @@
import asyncio
import asyncpg
from loguru import logger
from py3xui import AsyncApi
from client import add_client, delete_client, extend_client_key
from config import ADMIN_PASSWORD, ADMIN_USERNAME, DATABASE_URL, SERVERS
from config import ADMIN_PASSWORD, ADMIN_USERNAME, CLUSTERS, DATABASE_URL
async def create_key_on_server(server_id, tg_id, client_id, email, expiry_timestamp):
async def create_key_on_cluster(cluster_id, tg_id, client_id, email, expiry_timestamp):
try:
xui = AsyncApi(
SERVERS[server_id]["API_URL"],
username=ADMIN_USERNAME,
password=ADMIN_PASSWORD,
)
tasks = []
cluster = CLUSTERS.get(cluster_id)
conn = await asyncpg.connect(DATABASE_URL)
existing_key = await conn.fetchrow("SELECT 1 FROM keys WHERE email = $1", email)
if not cluster:
raise ValueError(f"Кластер с ID {cluster_id} не найден.")
if existing_key:
raise ValueError(f"Email {email} уже существует в базе данных.")
for server_info in cluster.values():
xui = AsyncApi(
server_info["API_URL"],
username=ADMIN_USERNAME,
password=ADMIN_PASSWORD,
)
await add_client(
xui,
client_id,
email,
tg_id,
limit_ip=1,
total_gb=0,
expiry_time=expiry_timestamp,
enable=True,
flow="xtls-rprx-vision",
)
conn = await asyncpg.connect(DATABASE_URL)
existing_key = await conn.fetchrow(
"SELECT 1 FROM keys WHERE email = $1", email
)
await conn.close()
if existing_key:
raise ValueError(f"Email {email} уже существует в базе данных.")
tasks.append(
add_client(
xui,
client_id,
email,
tg_id,
limit_ip=1,
total_gb=0,
expiry_time=expiry_timestamp,
enable=True,
flow="xtls-rprx-vision",
)
)
await conn.close()
await asyncio.gather(*tasks)
except Exception as e:
logger.error(f"Ошибка на сервере {server_id}: {e}")
logger.error(f"Ошибка при создании ключа: {e}")
raise e
async def renew_server_key(server_id, email, client_id, new_expiry_time):
async def renew_key_in_cluster(cluster_id, email, client_id, new_expiry_time, total_gb):
"""
Функция для продления срока действия ключа на сервере и сброса трафика, если необходимо.
Функция для продления срока действия ключа на всех серверах в кластере
и обновления объема трафика.
"""
try:
xui = AsyncApi(
SERVERS[server_id]["API_URL"],
username=ADMIN_USERNAME,
password=ADMIN_PASSWORD,
)
cluster = CLUSTERS.get(cluster_id)
await extend_client_key(xui, email, new_expiry_time, client_id)
if not cluster:
raise ValueError(f"Кластер с ID {cluster_id} не найден.")
tasks = []
for server_info in cluster.values():
xui = AsyncApi(
server_info["API_URL"],
username=ADMIN_USERNAME,
password=ADMIN_PASSWORD,
)
tasks.append(
extend_client_key(xui, email, new_expiry_time, client_id, total_gb)
)
await asyncio.gather(*tasks)
except Exception as e:
logger.error(
f"Не удалось продлить ключ {client_id} и сбросить трафик на сервере {server_id}: {e}"
f"Не удалось продлить ключ {client_id} в кластере {cluster_id}: {e}"
)
raise e
async def delete_key_from_db(client_id):
@@ -69,43 +96,70 @@ async def delete_key_from_db(client_id):
await conn.close()
async def delete_key_from_server(server_id, email, client_id):
"""Удаление ключа с сервера"""
async def delete_key_from_cluster(cluster_id, email, client_id):
"""Удаление ключа с серверов в кластере"""
try:
xui = AsyncApi(
SERVERS[server_id]["API_URL"],
username=ADMIN_USERNAME,
password=ADMIN_PASSWORD,
)
cluster = CLUSTERS.get(cluster_id)
await delete_client(xui, email, client_id)
if not cluster:
raise ValueError(f"Кластер с ID {cluster_id} не найден.")
except Exception as e:
logger.error(f"Не удалось удалить ключ {client_id} на сервере {server_id}: {e}")
tasks = []
for server_info in cluster.values():
xui = AsyncApi(
server_info["API_URL"],
username=ADMIN_USERNAME,
password=ADMIN_PASSWORD,
)
tasks.append(delete_client(xui, email, client_id))
async def update_key_on_server(tg_id, client_id, email, expiry_time, server_id):
try:
xui = AsyncApi(
SERVERS[server_id]["API_URL"],
username=ADMIN_USERNAME,
password=ADMIN_PASSWORD,
)
await add_client(
xui,
client_id,
email,
tg_id,
limit_ip=1,
total_gb=0,
expiry_time=expiry_time,
enable=True,
flow="xtls-rprx-vision",
)
logger.info(f"Ключ успешно обновлен на сервере {server_id} для {client_id}")
await asyncio.gather(*tasks)
except Exception as e:
logger.error(
f"Ошибка при обновлении ключа на сервере {server_id} для {client_id}: {e}"
f"Не удалось удалить ключ {client_id} в кластере {cluster_id}: {e}"
)
raise e
async def update_key_on_cluster(tg_id, client_id, email, expiry_time, cluster_id):
try:
cluster = CLUSTERS.get(cluster_id)
if not cluster:
raise ValueError(f"Кластер с ID {cluster_id} не найден.")
tasks = []
for server_info in cluster.values():
xui = AsyncApi(
server_info["API_URL"],
username=ADMIN_USERNAME,
password=ADMIN_PASSWORD,
)
tasks.append(
add_client(
xui,
client_id,
email,
tg_id,
limit_ip=1,
total_gb=0,
expiry_time=expiry_time,
enable=True,
flow="xtls-rprx-vision",
)
)
await asyncio.gather(*tasks)
logger.info(
f"Ключ успешно обновлен для {client_id} на всех серверах в кластере {cluster_id}"
)
except Exception as e:
logger.error(
f"Ошибка при обновлении ключа на серверах кластера {cluster_id} для {client_id}: {e}"
)
raise e
handlers/keys/keys.py
+27 -23
View File
@@ -9,11 +9,11 @@ from aiogram.types import BufferedInputFile
from loguru import logger
from bot import bot
from config import CONNECT_ANDROID, CONNECT_IOS, DATABASE_URL, DOWNLOAD_ANDROID, DOWNLOAD_IOS, PUBLIC_LINK, SERVERS
from config import CLUSTERS, CONNECT_ANDROID, CONNECT_IOS, DATABASE_URL, DOWNLOAD_ANDROID, DOWNLOAD_IOS, PUBLIC_LINK, TOTAL_GB
from database import delete_key, get_balance, store_key, update_balance, update_key_expiry
from handlers.keys.key_utils import delete_key_from_db, delete_key_from_server, renew_server_key, update_key_on_server
from handlers.keys.key_utils import delete_key_from_cluster, delete_key_from_db, renew_key_in_cluster, update_key_on_cluster
from handlers.texts import INSUFFICIENT_FUNDS_MSG, KEY_NOT_FOUND_MSG, NO_KEYS, PLAN_SELECTION_MSG, RENEWAL_PLANS, SUCCESS_RENEWAL_MSG, key_message
from handlers.utils import handle_error
from handlers.utils import get_least_loaded_cluster, handle_error
locale.setlocale(locale.LC_TIME, "ru_RU.UTF-8")
@@ -53,7 +53,7 @@ async def process_callback_view_keys(callback_query: types.CallbackQuery):
inline_keyboard = types.InlineKeyboardMarkup(inline_keyboard=buttons)
response_message = (
"<b>🔑 Список ваших устройств</b>\n\n"
"<i>👆 Выберите устройство для управления подпиской:</i>"
"<i>👇 Выберите устройство для управления подпиской:</i>"
)
image_path = os.path.join(os.path.dirname(__file__), "pic_keys.jpg")
@@ -167,7 +167,14 @@ async def process_callback_view_key(callback_query: types.CallbackQuery):
expiry_time = record["expiry_time"]
server_id = record["server_id"]
server_name = SERVERS.get(server_id, {}).get("name", "мультисервер")
if server_id == "все кластеры":
server_name = "мультикластер"
else:
cluster_name = CLUSTERS.get(server_id, {}).get(
"name", "кластер неизвестен"
)
server_name = cluster_name
expiry_date = datetime.utcfromtimestamp(expiry_time / 1000)
current_date = datetime.utcnow()
time_left = expiry_date - current_date
@@ -303,13 +310,14 @@ async def process_callback_update_subscription(callback_query: types.CallbackQue
)
return
least_loaded_cluster_id = await get_least_loaded_cluster()
tasks = []
for server_id in SERVERS:
tasks.append(
update_key_on_server(
tg_id, client_id, email, expiry_time, server_id
)
tasks.append(
update_key_on_cluster(
tg_id, client_id, email, expiry_time, least_loaded_cluster_id
)
)
await asyncio.gather(*tasks)
@@ -319,7 +327,7 @@ async def process_callback_update_subscription(callback_query: types.CallbackQue
email,
expiry_time,
public_link,
server_id="все сервера",
server_id=least_loaded_cluster_id,
)
try:
@@ -527,9 +535,9 @@ async def process_callback_confirm_delete(callback_query: types.CallbackQuery):
async def delete_key_from_servers():
try:
tasks = []
for server_id in SERVERS:
for cluster_id, cluster in CLUSTERS.items():
tasks.append(
delete_key_from_server(server_id, email, client_id)
delete_key_from_cluster(cluster_id, email, client_id)
)
await asyncio.gather(*tasks)
@@ -577,14 +585,10 @@ async def process_callback_renew_plan(callback_query: types.CallbackQuery):
)
days_to_extend = 30 * int(plan)
try:
try:
await bot.delete_message(
chat_id=tg_id, message_id=callback_query.message.message_id
)
except Exception:
pass
gb_multiplier = {"1": 1, "3": 3, "6": 6, "12": 12}
total_gb = TOTAL_GB * gb_multiplier.get(plan, 1) if TOTAL_GB > 0 else 0
try:
conn = await asyncpg.connect(DATABASE_URL)
try:
record = await conn.fetchrow(
@@ -643,10 +647,10 @@ async def process_callback_renew_plan(callback_query: types.CallbackQuery):
async def renew_key_on_servers():
tasks = []
for server_id in SERVERS:
for cluster_id in CLUSTERS:
task = asyncio.create_task(
renew_server_key(
server_id, email, client_id, new_expiry_time
renew_key_in_cluster(
cluster_id, email, client_id, new_expiry_time, total_gb
)
)
tasks.append(task)
handlers/keys/subscriptions.py
+5 -4
View File
@@ -4,7 +4,7 @@ import aiohttp
from aiohttp import web
from loguru import logger
from config import SERVERS
from config import CLUSTERS
async def fetch_url_content(url):
@@ -50,9 +50,10 @@ async def handle_subscription(request):
logger.info(f"Получен запрос на подписку для email: {email}")
urls = []
for server in SERVERS.values():
server_subscription_url = f"{server['SUBSCRIPTION']}/{email}"
urls.append(server_subscription_url)
for cluster in CLUSTERS.values():
for server in cluster.values():
server_subscription_url = f"{server['SUBSCRIPTION']}/{email}"
urls.append(server_subscription_url)
query_string = request.query_string
logger.debug(f"Извлечен query string: {query_string}")
handlers/keys/trial_key.py
+27 -13
View File
@@ -6,10 +6,10 @@ import asyncpg
from py3xui import AsyncApi
from client import add_client
from config import ADMIN_PASSWORD, ADMIN_USERNAME, DATABASE_URL, PUBLIC_LINK, SERVERS, TRIAL_TIME
from config import ADMIN_PASSWORD, ADMIN_USERNAME, CLUSTERS, DATABASE_URL, PUBLIC_LINK, TRIAL_TIME
from database import store_key
from handlers.texts import INSTRUCTIONS
from handlers.utils import generate_random_email
from handlers.utils import generate_random_email, get_least_loaded_cluster
async def create_trial_key(tg_id: int):
@@ -39,13 +39,20 @@ async def generate_and_store_keys(
conn = await asyncpg.connect(DATABASE_URL)
try:
current_time = datetime.utcnow()
expiry_time = current_time + timedelta(days={TRIAL_TIME}, hours=3)
expiry_time = current_time + timedelta(days=TRIAL_TIME, hours=3)
expiry_timestamp = int(expiry_time.timestamp() * 1000)
least_loaded_cluster = await get_least_loaded_cluster()
tasks = []
for server_id in SERVERS:
for server_id, server in CLUSTERS[least_loaded_cluster].items():
task = create_key_on_server(
server_id, client_id, email, tg_id, expiry_timestamp
least_loaded_cluster,
server_id,
client_id,
email,
tg_id,
expiry_timestamp,
)
tasks.append(task)
@@ -57,15 +64,15 @@ async def generate_and_store_keys(
email,
expiry_timestamp,
public_link,
server_id="all_servers",
server_id=least_loaded_cluster,
)
await conn.execute(
"""
INSERT INTO connections (tg_id, trial)
VALUES ($1, 1)
ON CONFLICT (tg_id)
DO UPDATE SET trial = 1
INSERT INTO connections (tg_id, trial)
VALUES ($1, 1)
ON CONFLICT (tg_id)
DO UPDATE SET trial = 1
""",
tg_id,
)
@@ -74,12 +81,19 @@ async def generate_and_store_keys(
async def create_key_on_server(
server_id: str, client_id: str, email: str, tg_id: int, expiry_timestamp: int
cluster_id: str,
server_id: str,
client_id: str,
email: str,
tg_id: int,
expiry_timestamp: int,
):
"""Создает ключ на сервере и возвращает результат."""
"""Создает ключ на сервере в указанном кластере и возвращает результат."""
xui = AsyncApi(
SERVERS[server_id]["API_URL"], username=ADMIN_USERNAME, password=ADMIN_PASSWORD
CLUSTERS[cluster_id][server_id]["API_URL"],
username=ADMIN_USERNAME,
password=ADMIN_PASSWORD,
)
response = await add_client(
handlers/notifications.py
+28 -30
View File
@@ -7,7 +7,7 @@ from loguru import logger
from py3xui import AsyncApi
from client import delete_client, extend_client_key
from config import ADMIN_PASSWORD, ADMIN_USERNAME, DATABASE_URL, SERVERS
from config import ADMIN_PASSWORD, ADMIN_USERNAME, CLUSTERS, DATABASE_URL
from database import delete_key, get_balance, update_balance, update_key_expiry
from handlers.texts import KEY_EXPIRY_10H, KEY_EXPIRY_24H, KEY_RENEWED, RENEWAL_PLANS
@@ -259,23 +259,28 @@ async def handle_expired_keys(bot: Bot, conn: asyncpg.Connection, current_time:
)
await update_key_expiry(client_id, new_expiry_time)
all_success = True
for server_id in SERVERS:
xui = AsyncApi(
SERVERS[server_id]["API_URL"],
username=ADMIN_USERNAME,
password=ADMIN_PASSWORD,
)
success = await extend_client_key(
xui, email, new_expiry_time, client_id
)
if not success:
all_success = False
logger.error(
f"Не удалось продлить ключ для пользователя {tg_id} на сервере {server_id}."
renewal_success = False
for cluster_id, cluster in CLUSTERS.items():
for server_id, server in cluster.items():
xui = AsyncApi(
server["API_URL"],
username=ADMIN_USERNAME,
password=ADMIN_PASSWORD,
)
success = await extend_client_key(
xui, email, new_expiry_time, client_id
)
if success:
renewal_success = True
logger.info(
f"Ключ для пользователя {tg_id} успешно продлен на сервере {server_id} в кластере {cluster_id}."
)
else:
logger.error(
f"Не удалось продлить ключ для пользователя {tg_id} на сервере {server_id} в кластере {cluster_id}."
)
if all_success:
if renewal_success:
try:
await bot.send_message(
tg_id, KEY_RENEWED, reply_markup=keyboard
@@ -299,21 +304,14 @@ async def handle_expired_keys(bot: Bot, conn: asyncpg.Connection, current_time:
await delete_key(client_id)
for server_id in SERVERS:
xui = AsyncApi(
SERVERS[server_id]["API_URL"],
username=ADMIN_USERNAME,
password=ADMIN_PASSWORD,
)
success = await delete_client(xui, email, client_id)
if success:
logger.info(
f"Ключ для клиента {tg_id} успешно удален с сервера {server_id}."
)
else:
logger.error(
f"Не удалось удалить ключ для клиента {tg_id} на сервере {server_id}."
for cluster_id, cluster in CLUSTERS.items():
for server_id, server in cluster.items():
xui = AsyncApi(
server["API_URL"],
username=ADMIN_USERNAME,
password=ADMIN_PASSWORD,
)
await delete_client(xui, email, client_id)
except Exception as e:
logger.error(f"Ошибка при обработке ключа для клиента {tg_id}: {e}")
handlers/utils.py
+46 -14
View File
@@ -1,10 +1,11 @@
import random
import re
import asyncpg
from loguru import logger
from bot import bot
from config import SERVERS
from config import CLUSTERS, DATABASE_URL
def sanitize_key_name(key_name: str) -> str:
@@ -17,21 +18,52 @@ def generate_random_email():
return random_string
async def get_least_loaded_server(conn):
"""Находит сервер с наименьшей загрузкой."""
least_loaded_server_id = None
min_load_percentage = float("inf")
async def get_least_loaded_cluster():
"""
Функция для получения кластера с наименьшей загрузкой (по количеству ключей).
Возвращает идентификатор кластера с наименьшей загрузкой или первый кластер из конфигурации,
если загруженность не определяется. В случае отсутствия кластеров с номером, возвращает 'cluster1'.
"""
cluster_loads = {}
for server_id, server in SERVERS.items():
count = await conn.fetchval(
"SELECT COUNT(*) FROM keys WHERE server_id = $1", server_id
)
percent_full = (count / 60) * 100 if count <= 60 else 100
if percent_full < min_load_percentage:
min_load_percentage = percent_full
least_loaded_server_id = server_id
conn = await asyncpg.connect(DATABASE_URL)
try:
keys = await conn.fetch("SELECT * FROM keys")
for key in keys:
cluster_id = key["server_id"]
return least_loaded_server_id
if re.match(r"^cluster\d+$", cluster_id):
if cluster_id not in cluster_loads:
cluster_loads[cluster_id] = 0
cluster_loads[cluster_id] += 1
finally:
await conn.close()
logger.info(f"Cluster loads: {cluster_loads}")
if not cluster_loads:
available_clusters = [
cluster_id
for cluster_id in CLUSTERS.keys()
if re.match(r"^cluster\d+$", cluster_id)
]
logger.info(f"Available clusters from config: {available_clusters}")
if available_clusters:
logger.info(
f"Returning the first available cluster: {available_clusters[0]}"
)
return available_clusters[0]
else:
logger.warning("No valid clusters found in config, returning 'cluster1'.")
return "cluster1"
least_loaded_cluster = min(cluster_loads, key=cluster_loads.get)
logger.info(f"Least loaded cluster selected: {least_loaded_cluster}")
return least_loaded_cluster
async def handle_error(tg_id, callback_query, message):