Spaces:
Running
Running
File size: 5,040 Bytes
4276a62 | 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 | """
MongoDB Connection Module
Handles all database connections and configurations
FastAPI-compatible version (no Streamlit dependencies)
"""
import os
from pymongo import MongoClient
from pymongo.errors import ConnectionFailure, ServerSelectionTimeoutError
from datetime import datetime
import logging
# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
class MongoDBConnection:
"""MongoDB Connection Manager"""
def __init__(self):
self.client = None
self.database = None
self.connection_string = self._get_connection_string()
def _get_connection_string(self):
"""Get MongoDB connection string from environment variables"""
try:
# Get from environment variables (for FastAPI/HF Spaces)
# Fallback to environment variables
mongodb_uri = os.getenv('MONGODB_URI')
if mongodb_uri:
return mongodb_uri
# Default local MongoDB connection
host = os.getenv('MONGODB_HOST', 'localhost')
port = os.getenv('MONGODB_PORT', '27017')
username = os.getenv('MONGODB_USERNAME')
password = os.getenv('MONGODB_PASSWORD')
database_name = os.getenv('MONGODB_DATABASE', 'virus_prediction')
if username and password:
return f"mongodb://{username}:{password}@{host}:{port}/{database_name}"
else:
return f"mongodb://{host}:{port}/{database_name}"
except Exception as e:
logger.warning(f"Could not load MongoDB configuration: {e}")
return "mongodb://localhost:27017/virus_prediction" # Default fallback
def connect(self):
"""Establish connection to MongoDB"""
try:
if not self.client:
self.client = MongoClient(
self.connection_string,
serverSelectionTimeoutMS=5000, # 5 seconds timeout
connectTimeoutMS=10000, # 10 seconds timeout
maxPoolSize=10, # Maximum connection pool size
retryWrites=True
)
# Test the connection
self.client.admin.command('ping')
# Get database name from connection string or use default
db_name = os.getenv('MONGODB_DATABASE', 'virus_prediction')
self.database = self.client[db_name]
logger.info("Successfully connected to MongoDB")
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 disconnect(self):
"""Close MongoDB connection"""
try:
if self.client:
self.client.close()
self.client = None
self.database = None
logger.info("MongoDB connection closed")
except Exception as e:
logger.error(f"Error closing MongoDB connection: {e}")
def get_database(self):
"""Get database instance"""
if not self.database:
if self.connect():
return self.database
else:
return None
return self.database
def test_connection(self):
"""Test MongoDB connection and return status"""
try:
if self.connect():
# Test with a simple operation
db = self.get_database()
if db:
db.list_collection_names()
return {
'status': 'success',
'message': 'Successfully connected to MongoDB',
'timestamp': datetime.now().isoformat()
}
return {
'status': 'error',
'message': 'Failed to connect to MongoDB',
'timestamp': datetime.now().isoformat()
}
except Exception as e:
return {
'status': 'error',
'message': f'Connection test failed: {str(e)}',
'timestamp': datetime.now().isoformat()
}
# Global connection instance
mongo_connection = MongoDBConnection()
def get_db():
"""Get database instance - use this function in your app"""
return mongo_connection.get_database()
def test_db_connection():
"""Test database connection - use this function to check status"""
return mongo_connection.test_connection()
def close_db_connection():
"""Close database connection - call this when app shuts down"""
mongo_connection.disconnect() |