104 lines
4.7 KiB
Python
104 lines
4.7 KiB
Python
import os
|
||
from motor.motor_asyncio import AsyncIOMotorClient
|
||
import logging
|
||
|
||
|
||
class MongoDBRepository:
|
||
def __init__(self):
|
||
# Настройки MongoDB из переменных окружения
|
||
mongo_uri = os.getenv("MONGO_URL")
|
||
database_name = os.getenv("DB_NAME")
|
||
server_collection = os.getenv("SERVER_COLLECTION", "servers")
|
||
plan_collection = os.getenv("PLAN_COLLECTION", "plans")
|
||
|
||
# Подключение к базе данных и коллекциям
|
||
self.client = AsyncIOMotorClient(mongo_uri)
|
||
self.db = self.client[database_name]
|
||
self.collection = self.db[server_collection] # Коллекция серверов
|
||
self.plans_collection = self.db[plan_collection] # Коллекция тарифных планов
|
||
self.logger = logging.getLogger(__name__)
|
||
|
||
async def add_subscription_plan(self, plan_data):
|
||
"""Добавляет новый тарифный план в коллекцию."""
|
||
result = await self.plans_collection.insert_one(plan_data)
|
||
self.logger.debug(f"Тарифный план добавлен с ID: {result.inserted_id}")
|
||
return result.inserted_id
|
||
|
||
async def get_subscription_plan(self, plan_name):
|
||
"""Получает тарифный план по его имени."""
|
||
plan = await self.plans_collection.find_one({"name": plan_name})
|
||
if plan:
|
||
self.logger.debug(f"Найден тарифный план: {plan}")
|
||
else:
|
||
self.logger.error(f"Тарифный план {plan_name} не найден.")
|
||
return plan
|
||
|
||
async def add_server(self, server_data):
|
||
"""Добавляет новый VPN сервер в коллекцию."""
|
||
result = await self.collection.insert_one(server_data)
|
||
self.logger.debug(f"VPN сервер добавлен с ID: {result.inserted_id}")
|
||
return result.inserted_id
|
||
|
||
async def get_server(self, server_name: str):
|
||
"""Получает сервер VPN по его ID."""
|
||
server = await self.collection.find_one({"server.name": server_name})
|
||
if server:
|
||
self.logger.debug(f"Найден VPN сервер: {server}")
|
||
else:
|
||
self.logger.debug(f"VPN сервер с ID {server_name} не найден.")
|
||
return server
|
||
|
||
async def get_server_with_least_clients(self):
|
||
"""Возвращает сервер с наименьшим количеством подключенных клиентов."""
|
||
pipeline = [
|
||
{
|
||
"$addFields": {
|
||
"current_clients": {"$size": {"$ifNull": ["$clients", []]}}
|
||
}
|
||
},
|
||
{
|
||
"$sort": {"current_clients": 1}
|
||
},
|
||
{
|
||
"$limit": 1
|
||
}
|
||
]
|
||
|
||
result = await self.collection.aggregate(pipeline).to_list(length=1)
|
||
if result:
|
||
server = result[0]
|
||
self.logger.debug(f"Найден сервер с наименьшим количеством клиентов: {server}")
|
||
return server
|
||
else:
|
||
self.logger.debug("Не найдено серверов.")
|
||
return None
|
||
|
||
async def update_server(self, server_id, update_data):
|
||
"""Обновляет данные VPN сервера."""
|
||
result = await self.collection.update_one({"_id": server_id}, {"$set": update_data})
|
||
if result.matched_count > 0:
|
||
self.logger.debug(f"VPN сервер с ID {server_id} обновлен.")
|
||
else:
|
||
self.logger.debug(f"VPN сервер с ID {server_id} не найден.")
|
||
return result.matched_count > 0
|
||
|
||
async def delete_server(self, server_id):
|
||
"""Удаляет VPN сервер по его ID."""
|
||
result = await self.collection.delete_one({"_id": server_id})
|
||
if result.deleted_count > 0:
|
||
self.logger.debug(f"VPN сервер с ID {server_id} удален.")
|
||
else:
|
||
self.logger.debug(f"VPN сервер с ID {server_id} не найден.")
|
||
return result.deleted_count > 0
|
||
|
||
async def list_servers(self):
|
||
"""Возвращает список всех VPN серверов."""
|
||
servers = await self.collection.find().to_list(length=1000) # Получить до 1000 серверов (можно настроить)
|
||
self.logger.debug(f"Найдено {len(servers)} VPN серверов.")
|
||
return servers
|
||
|
||
async def close_connection(self):
|
||
"""Закрывает подключение к базе данных MongoDB."""
|
||
self.client.close()
|
||
self.logger.debug("Подключение к MongoDB закрыто.")
|