src/database/connection.py

import os
from dotenv import load_dotenv
from pymongo import MongoClient
from src.utils.constants import Constants
from src.utils.logging import get_logger

class DatabaseConnection:
    """
    Class to handle the connection to MongoDB Atlas and provide methods for database operations.
    """

    def __init__(self):
        """
        Initialize the database connection using the MongoDB URI from environment variables.
        """
        self.logger = get_logger()
        self.logger.info("Initializing database connection")

        load_dotenv()
        self.mongodb_uri = os.getenv("MONGODB_URI")
        if not self.mongodb_uri:
            self.logger.error("MONGODB_URI environment variable is not set")
            raise ValueError("MONGODB_URI environment variable is not set")

        try:
            self.client = MongoClient(self.mongodb_uri)
            self.db = self.client[Constants.DB_NAME]
            self.collection = self.db[Constants.COLLECTION_NAME]
            self.logger.info(f"Connected to MongoDB database: {Constants.DB_NAME}, collection: {Constants.COLLECTION_NAME}")
        except Exception as e:
            self.logger.error(f"Failed to connect to MongoDB: {str(e)}")
            raise

    def get_collection(self):
        """
        Get the MongoDB collection.

        Returns:
            pymongo.collection.Collection: The MongoDB collection.
        """
        self.logger.debug(f"Getting collection: {Constants.COLLECTION_NAME}")
        return self.collection

    def close_connection(self):
        """
        Close the MongoDB connection.
        """
        self.logger.info("Closing MongoDB connection")
        if self.client:
            try:
                self.client.close()
                self.logger.info("MongoDB connection closed successfully")
            except Exception as e:
                self.logger.error(f"Error closing MongoDB connection: {str(e)}")