Spaces:
Sleeping
Sleeping
File size: 4,742 Bytes
f6278c5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 | """
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") |