# db/mongoDB.py import logging from typing import Optional from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorDatabase, AsyncIOMotorCollection from pymongo.errors import ConnectionFailure, ConfigurationError import config # Import từ file config tập trung logger = logging.getLogger(__name__) class MongoDatabase: """ Một lớp singleton để quản lý kết nối và các collection của MongoDB. Điều này đảm bảo chúng ta chỉ có một kết nối duy nhất trong toàn bộ ứng dụng. """ client: Optional[AsyncIOMotorClient] = None db: Optional[AsyncIOMotorDatabase] = None # Khai báo các collection bạn sẽ sử dụng users: Optional[AsyncIOMotorCollection] = None token_blacklist: Optional[AsyncIOMotorCollection] = None conversations: Optional[AsyncIOMotorCollection] = None processed_documents: Optional[AsyncIOMotorCollection] = None # Tạo một instance duy nhất để import và sử dụng trong toàn bộ ứng dụng mongo_db = MongoDatabase() async def connect_to_mongo(): """ Hàm khởi tạo kết nối đến MongoDB Atlas và gán vào object mongo_db. Hàm này sẽ được gọi từ lifespan của FastAPI. """ if mongo_db.client: logger.info("✅ MongoDB connection already established.") return logger.info(f"🔸 Connecting to MongoDB Atlas...") if not config.MONGODB_CLOUD_URI or not config.DB_NAME: logger.error("❌ MONGODB_CLOUD_URI hoặc DB_NAME chưa được thiết lập trong biến môi trường.") raise ConfigurationError("MONGODB_CLOUD_URI and DB_NAME must be set.") try: # 1. Khởi tạo client bất đồng bộ mongo_db.client = AsyncIOMotorClient( config.MONGODB_CLOUD_URI, serverSelectionTimeoutMS=30000 # Thời gian chờ kết nối là 5 giây ) # 2. Kiểm tra kết nối một cách rõ ràng await mongo_db.client.admin.command('ping') # 3. Gán các đối tượng database và collection mongo_db.db = mongo_db.client[config.DB_NAME] mongo_db.users = mongo_db.db["users"] mongo_db.token_blacklist = mongo_db.db["token_blacklist"] mongo_db.conversations = mongo_db.db["conversations"] mongo_db.processed_documents = mongo_db.db["processed_documents"] # 4. Tạo TTL index một cách an toàn # Lấy danh sách index hiện có index_info = await mongo_db.token_blacklist.index_information() if "expires_at_1" not in index_info: # Chỉ tạo index nếu nó chưa tồn tại await mongo_db.token_blacklist.create_index("expires_at", expireAfterSeconds=0) logger.info("🔸 Successfully created TTL index for 'expires_at' in 'token_blacklist'.") logger.info("✅ MongoDB connection successful and collections are ready.") except ConnectionFailure as e: logger.error(f"❌ Failed to connect to MongoDB: Connection Failure. Check your URI and network access rules in Atlas. Error: {e}", exc_info=True) raise e except ConfigurationError as e: logger.error(f"❌ Failed to connect to MongoDB: Configuration Error. Check your connection string format. Error: {e}", exc_info=True) raise e except Exception as e: logger.error(f"❌ An unexpected error occurred while connecting to MongoDB: {e}", exc_info=True) raise e async def close_mongo_connection(): """Hàm đóng kết nối MongoDB khi ứng dụng tắt.""" if mongo_db.client: mongo_db.client.close() logger.info("✅ MongoDB connection closed.")