""" MongoDB database connection and logging utilities """ import os import logging from datetime import datetime from typing import Optional, Dict, Any from pymongo import MongoClient from pymongo.errors import ConnectionFailure, ServerSelectionTimeoutError logger = logging.getLogger(__name__) # MongoDB connection _client: Optional[MongoClient] = None _db = None def get_mongodb_client() -> Optional[MongoClient]: """Get or create MongoDB client""" global _client if _client is None: mongodb_uri = os.getenv("MONGODB_URI") if not mongodb_uri: logger.warning("MONGODB_URI environment variable not set. MongoDB features will be disabled.") return None try: _client = MongoClient( mongodb_uri, serverSelectionTimeoutMS=5000, connectTimeoutMS=5000 ) # Test connection _client.admin.command('ping') logger.info("MongoDB connection established successfully") except (ConnectionFailure, ServerSelectionTimeoutError) as e: logger.error("Failed to connect to MongoDB: %s", str(e)) _client = None return _client def get_database(): """Get database instance""" global _db if _db is None: client = get_mongodb_client() if client: db_name = os.getenv("MONGODB_DB_NAME", "colorization_db") _db = client[db_name] else: logger.warning("MongoDB client not available") return _db def log_api_call( endpoint: str, method: str, status_code: int = 200, request_data: Optional[Dict[str, Any]] = None, response_data: Optional[Dict[str, Any]] = None, error: Optional[str] = None, user_id: Optional[str] = None, ip_address: Optional[str] = None ) -> bool: """ Log API call to MongoDB Args: endpoint: API endpoint path method: HTTP method (GET, POST, etc.) status_code: HTTP status code request_data: Request data/parameters response_data: Response data error: Error message if any user_id: User ID if authenticated ip_address: Client IP address Returns: True if logged successfully, False otherwise """ try: db = get_database() if not db: logger.warning("MongoDB not available, skipping API log") return False collection = db["api_calls"] log_entry = { "endpoint": endpoint, "method": method, "status_code": status_code, "timestamp": datetime.utcnow(), "request_data": request_data or {}, "response_data": response_data or {}, "error": error, "user_id": user_id, "ip_address": ip_address } result = collection.insert_one(log_entry) logger.info("API call logged to MongoDB: %s", result.inserted_id) return True except Exception as e: logger.error("Failed to log API call to MongoDB: %s", str(e)) return False def log_image_upload( image_id: str, filename: str, file_size: int, content_type: str, user_id: Optional[str] = None, ip_address: Optional[str] = None ) -> bool: """ Log image upload to MongoDB Args: image_id: Unique image identifier filename: Original filename file_size: File size in bytes content_type: MIME type user_id: User ID if authenticated ip_address: Client IP address Returns: True if logged successfully, False otherwise """ try: db = get_database() if not db: logger.warning("MongoDB not available, skipping upload log") return False collection = db["image_uploads"] log_entry = { "image_id": image_id, "filename": filename, "file_size": file_size, "content_type": content_type, "uploaded_at": datetime.utcnow(), "user_id": user_id, "ip_address": ip_address } result = collection.insert_one(log_entry) logger.info("Image upload logged to MongoDB: %s", result.inserted_id) return True except Exception as e: logger.error("Failed to log image upload to MongoDB: %s", str(e)) return False def log_colorization( result_id: str, image_id: Optional[str] = None, prompt: Optional[str] = None, model_type: Optional[str] = None, processing_time: Optional[float] = None, user_id: Optional[str] = None, ip_address: Optional[str] = None ) -> bool: """ Log colorization request to MongoDB Args: result_id: Unique result identifier image_id: Original image identifier prompt: Text prompt used (if any) model_type: Model type used (fastai, pytorch, sdxl, etc.) processing_time: Time taken to process in seconds user_id: User ID if authenticated ip_address: Client IP address Returns: True if logged successfully, False otherwise """ try: db = get_database() if not db: logger.warning("MongoDB not available, skipping colorization log") return False collection = db["colorizations"] log_entry = { "result_id": result_id, "image_id": image_id, "prompt": prompt, "model_type": model_type, "processing_time": processing_time, "created_at": datetime.utcnow(), "user_id": user_id, "ip_address": ip_address } result = collection.insert_one(log_entry) logger.info("Colorization logged to MongoDB: %s", result.inserted_id) return True except Exception as e: logger.error("Failed to log colorization to MongoDB: %s", str(e)) return False def close_connection(): """Close MongoDB connection""" global _client, _db if _client: _client.close() _client = None _db = None logger.info("MongoDB connection closed")