db/mongoDB.py

# 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.")