app/database.py

"""
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")