Spaces:
Running
Running
| """ | |
| 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() |