Spaces:
Sleeping
Sleeping
| """ | |
| Database connection and utilities for MongoDB integration | |
| """ | |
| import os | |
| from pymongo import MongoClient | |
| from pymongo.errors import ConnectionFailure, ServerSelectionTimeoutError | |
| import logging | |
| from typing import Optional | |
| from dotenv import load_dotenv | |
| # Load environment variables | |
| load_dotenv() | |
| # Configure logging | |
| logger = logging.getLogger(__name__) | |
| class DatabaseConnection: | |
| """MongoDB connection manager""" | |
| def __init__(self, mongodb_uri=None, database_name=None): | |
| self.client: Optional[MongoClient] = None | |
| self.db = None | |
| self.mongodb_url = mongodb_uri or os.getenv('MONGODB_URL') or os.getenv('MONGODB_URI') | |
| self.database_name = database_name or os.getenv('MONGODB_DATABASE', 'Atlas') | |
| if not self.mongodb_url: | |
| raise ValueError("MONGODB_URL environment variable is required") | |
| def connect(self): | |
| """Establish connection to MongoDB""" | |
| try: | |
| self.client = MongoClient( | |
| self.mongodb_url, | |
| serverSelectionTimeoutMS=5000, # 5 second timeout | |
| connectTimeoutMS=10000, # 10 second connection timeout | |
| socketTimeoutMS=20000, # 20 second socket timeout | |
| maxPoolSize=50, # Maximum connection pool size | |
| retryWrites=True | |
| ) | |
| # Test the connection | |
| self.client.admin.command('ping') | |
| self.db = self.client[self.database_name] | |
| logger.info(f"Successfully connected to MongoDB database: {self.database_name}") | |
| return True | |
| except (ConnectionFailure, ServerSelectionTimeoutError) as e: | |
| logger.error(f"Failed to connect to MongoDB: {e}") | |
| return False | |
| except Exception as e: | |
| logger.error(f"Unexpected error connecting to MongoDB: {e}") | |
| return False | |
| def get_database(self): | |
| """Get database instance""" | |
| if self.db is None: | |
| if not self.connect(): | |
| raise ConnectionError("Could not establish database connection") | |
| return self.db | |
| def get_collection(self, collection_name: str): | |
| """Get a specific collection""" | |
| db = self.get_database() | |
| return db[collection_name] | |
| def close(self): | |
| """Close database connection""" | |
| if self.client: | |
| self.client.close() | |
| logger.info("Database connection closed") | |
| # Global database connection instance | |
| db_connection = DatabaseConnection() | |
| def get_db(): | |
| """Get database instance - convenience function""" | |
| return db_connection.get_database() | |
| def get_users_collection(): | |
| """Get users collection""" | |
| return db_connection.get_collection('users') | |
| def get_chat_sessions_collection(): | |
| """Get chat sessions collection""" | |
| return db_connection.get_collection('chat_sessions') | |
| def init_database(): | |
| """Initialize database connection and create indexes""" | |
| try: | |
| db = get_db() | |
| if db is None: | |
| logger.error("Database connection is None") | |
| return False | |
| # Create indexes for better performance | |
| users_collection = get_users_collection() | |
| chat_sessions_collection = get_chat_sessions_collection() | |
| # Create unique index on email for users collection | |
| users_collection.create_index("email", unique=True) | |
| logger.info("Created unique index on users.email") | |
| # Create indexes on chat_sessions collection | |
| chat_sessions_collection.create_index("user_id") | |
| chat_sessions_collection.create_index("timestamp") | |
| chat_sessions_collection.create_index([("user_id", 1), ("timestamp", -1)]) | |
| logger.info("Created indexes on chat_sessions collection") | |
| return True | |
| except Exception as e: | |
| logger.error(f"Failed to initialize database: {e}") | |
| return False | |
| def test_connection(): | |
| """Test database connection""" | |
| try: | |
| db = get_db() | |
| # Simple test query | |
| result = db.command('ping') | |
| logger.info("Database connection test successful") | |
| return True | |
| except Exception as e: | |
| logger.error(f"Database connection test failed: {e}") | |
| return False | |
| if __name__ == "__main__": | |
| # Test the connection when run directly | |
| print("Testing MongoDB connection...") | |
| if test_connection(): | |
| print("✓ Database connection successful") | |
| if init_database(): | |
| print("✓ Database initialization successful") | |
| else: | |
| print("✗ Database initialization failed") | |
| else: | |
| print("✗ Database connection failed") |