import os from motor.motor_asyncio import AsyncIOMotorClient from pymongo.errors import DuplicateKeyError, NetworkTimeout 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): """Добавляет новый тарифный план в коллекцию.""" try: result = await self.plans_collection.insert_one(plan_data) self.logger.debug(f"Тарифный план добавлен с ID: {result.inserted_id}") return result.inserted_id except DuplicateKeyError: self.logger.error("Дублирующий ключ.") except NetworkTimeout: self.logger.error("Сетевой таймаут.") async def get_subscription_plan(self, plan_name): """Получает тарифный план по его имени.""" try: 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 except DuplicateKeyError: self.logger.error("Дублирующий ключ.") except NetworkTimeout: self.logger.error("Сетевой таймаут.") async def add_server(self, server_data): """Добавляет новый VPN сервер в коллекцию.""" try: result = await self.collection.insert_one(server_data) self.logger.debug(f"VPN сервер добавлен с ID: {result.inserted_id}") return result.inserted_id except DuplicateKeyError: self.logger.error("Дублирующий ключ.") except NetworkTimeout: self.logger.error("Сетевой таймаут.") async def get_server(self, server_name: str): """Получает сервер VPN по его ID.""" try: 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 except DuplicateKeyError: self.logger.error("Дублирующий ключ.") except NetworkTimeout: self.logger.error("Сетевой таймаут.") async def get_server_with_least_clients(self): """Возвращает сервер с наименьшим количеством подключенных клиентов.""" try: 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 except DuplicateKeyError: self.logger.error("Дублирующий ключ.") except NetworkTimeout: self.logger.error("Сетевой таймаут.") async def update_server(self, server_name, update_data): """Обновляет данные VPN сервера.""" try: result = await self.collection.update_one({"server_name": server_name}, {"$set": update_data}) if result.matched_count > 0: self.logger.debug(f"VPN сервер с ID {server_name} обновлен.") else: self.logger.debug(f"VPN сервер с ID {server_name} не найден.") return result.matched_count > 0 except DuplicateKeyError: self.logger.error("Дублирующий ключ.") except NetworkTimeout: self.logger.error("Сетевой таймаут.") async def delete_server(self, server_name): """Удаляет VPN сервер по его ID.""" try: result = await self.collection.delete_one({"name": server_name}) if result.deleted_count > 0: self.logger.debug(f"VPN сервер с ID {server_name} удален.") else: self.logger.debug(f"VPN сервер с ID {server_name} не найден.") return result.deleted_count > 0 except DuplicateKeyError: self.logger.error("Дублирующий ключ.") except NetworkTimeout: self.logger.error("Сетевой таймаут.") async def list_servers(self): """Возвращает список всех VPN серверов.""" try: servers = await self.collection.find().to_list(length=1000) # Получить до 1000 серверов (можно настроить) self.logger.debug(f"Найдено {len(servers)} VPN серверов.") return servers except DuplicateKeyError: self.logger.error("Дублирующий ключ.") except NetworkTimeout: self.logger.error("Сетевой таймаут.") async def __aenter__(self): """ Метод вызывается при входе в блок with. """ self.logger.debug("Контекстный менеджер: подключение открыто.") return self async def __aexit__(self, exc_type, exc_value, traceback): """ Метод вызывается при выходе из блока with. """ await self.close_connection() if exc_type: self.logger.error(f"Контекстный менеджер завершён с ошибкой: {exc_value}") else: self.logger.debug("Контекстный менеджер: подключение закрыто.")