"""MongoDB database client and utilities with connection pooling and Redis caching.""" import os from pymongo import MongoClient, ASCENDING, TEXT from pymongo.errors import ServerSelectionTimeoutError, BulkWriteError import logging from typing import Optional, Dict, Any, List from datetime import datetime, timedelta logger = logging.getLogger(__name__) # MongoDB connection with connection pooling _client: Optional[MongoClient] = None _db = None # Optional Redis client for session-based caching _redis_client = None try: import redis HAS_REDIS = True except ImportError: HAS_REDIS = False logger.debug("Redis not installed - session caching disabled") def get_mongo_client() -> MongoClient: """Get or create MongoDB client with connection pooling.""" global _client if _client is None: mongo_uri = os.getenv("MONGO_URI", "mongodb://localhost:27017") max_pool_size = int(os.getenv("MONGO_MAX_POOL_SIZE", "10")) try: _client = MongoClient( mongo_uri, serverSelectionTimeoutMS=5000, maxPoolSize=max_pool_size, # Connection pooling minPoolSize=2, # Keep 2 connections ready maxIdleTimeMS=45000, # Close idle connections after 45s socketTimeoutMS=20000, # Socket timeout connectTimeoutMS=10000, # Connection timeout ) # Test connection _client.admin.command('ping') logger.info(f"Successfully connected to MongoDB (pool size: {max_pool_size})") except ServerSelectionTimeoutError: logger.error(f"Failed to connect to MongoDB at {mongo_uri}") raise return _client def get_redis_client(): """Get or create Redis client for session caching (optional).""" global _redis_client if not HAS_REDIS: return None if _redis_client is None: redis_host = os.getenv("REDIS_HOST", "localhost") redis_port = int(os.getenv("REDIS_PORT", "6379")) redis_db = int(os.getenv("REDIS_DB", "0")) redis_password = os.getenv("REDIS_PASSWORD") try: _redis_client = redis.Redis( host=redis_host, port=redis_port, db=redis_db, password=redis_password, decode_responses=True, socket_timeout=5, socket_connect_timeout=5, ) # Test connection _redis_client.ping() logger.info(f"Successfully connected to Redis at {redis_host}:{redis_port}") except Exception as e: logger.warning(f"Failed to connect to Redis: {e} - session caching disabled") _redis_client = None return _redis_client def get_database(): """Get database instance.""" global _db if _db is None: client = get_mongo_client() db_name = os.getenv("MONGO_DB_NAME", "grant_analyst") _db = client[db_name] return _db def close_mongo_connection(): """Close MongoDB connection.""" global _client, _redis_client if _client is not None: _client.close() _client = None logger.info("MongoDB connection closed") if _redis_client is not None: _redis_client.close() _redis_client = None logger.info("Redis connection closed") class SummaryStore: """Handle pre-computed grant summaries with optimized indexing and bulk operations.""" def __init__(self): """Initialize summary store with compound indexes and text search.""" self.db = get_database() self.collection = self.db["summaries"] self.redis = get_redis_client() # Create compound index for cache lookups (most common query pattern) self.collection.create_index( [("grant_id", ASCENDING), ("summary_type", ASCENDING)], unique=True, name="grant_summary_lookup" ) # Create text index for full-text search on summaries try: self.collection.create_index( [("summary_text", TEXT)], name="summary_text_search", default_language="english" ) except Exception as e: logger.debug(f"Text index may already exist: {e}") # Create index on created_at for time-based queries self.collection.create_index("created_at", name="created_at_idx") # Create index on metadata fields for analytics self.collection.create_index( [("metadata.model", ASCENDING)], name="model_idx", sparse=True ) def save_summary( self, grant_id: str, summary_type: str, summary_text: str, metadata: Optional[Dict[str, Any]] = None ) -> bool: """ Save a pre-computed summary to database. Args: grant_id: Grant ID (e.g., "competition-2315") summary_type: Type of summary ("layman", "technical", "exec") summary_text: The summary content metadata: Optional metadata (model used, tokens, etc.) Returns: True if saved successfully """ try: from datetime import datetime doc_key = f"{grant_id}_{summary_type}" document = { "grant_id": grant_id, "summary_type": summary_type, "summary_text": summary_text, "created_at": datetime.utcnow(), "metadata": metadata or {} } # Upsert: update if exists, insert if not self.collection.update_one( {"grant_id": grant_id, "summary_type": summary_type}, {"$set": document}, upsert=True ) logger.info(f"Summary saved: {doc_key}") return True except Exception as e: logger.error(f"Error saving summary for {grant_id}: {e}") return False def get_summary(self, grant_id: str, summary_type: str = "layman") -> Optional[str]: """ Retrieve a pre-computed summary from cache (Redis → MongoDB). Args: grant_id: Grant ID summary_type: Type of summary to retrieve Returns: Summary text if found, None otherwise """ cache_key = f"summary:{grant_id}:{summary_type}" try: # Try Redis first (if available) if self.redis is not None: try: cached = self.redis.get(cache_key) if cached: logger.info(f"⚡ Redis HIT: {grant_id}_{summary_type}") # Record cache hit try: from src.monitoring import record_cache_hit record_cache_hit() except Exception: pass return cached except Exception as redis_err: logger.debug(f"Redis read error: {redis_err}") # Fall back to MongoDB doc = self.collection.find_one( {"grant_id": grant_id, "summary_type": summary_type}, {"summary_text": 1, "_id": 0} # Project only needed field ) if doc: summary_text = doc.get("summary_text") logger.info(f"📦 MongoDB HIT: {grant_id}_{summary_type}") # Record cache hit try: from src.monitoring import record_cache_hit record_cache_hit() except Exception: pass # Cache in Redis for next time (TTL: 1 hour) if self.redis is not None and summary_text: try: self.redis.setex(cache_key, 3600, summary_text) except Exception as redis_err: logger.debug(f"Redis write error: {redis_err}") return summary_text # Record cache miss try: from src.monitoring import record_cache_miss record_cache_miss() except Exception: pass return None except Exception as e: logger.error(f"Error retrieving summary for {grant_id}: {e}") return None def get_all_summaries(self, grant_id: str) -> Dict[str, str]: """ Get all summary types for a grant. Returns: Dict with keys: layman, technical, exec (if available) """ try: docs = self.collection.find({"grant_id": grant_id}) return {doc["summary_type"]: doc["summary_text"] for doc in docs} except Exception as e: logger.error(f"Error retrieving summaries for {grant_id}: {e}") return {} def bulk_save_summaries(self, summaries: List[Dict[str, Any]]) -> int: """ Bulk save multiple summaries using bulk write operations. Args: summaries: List of dicts with keys: grant_id, summary_type, summary_text, metadata Returns: Number of summaries saved Example: summaries = [ {"grant_id": "comp-123", "summary_type": "layman", "summary_text": "...", "metadata": {}}, {"grant_id": "comp-124", "summary_type": "layman", "summary_text": "...", "metadata": {}}, ] store.bulk_save_summaries(summaries) """ if not summaries: return 0 try: from pymongo import UpdateOne operations = [] for summary in summaries: grant_id = summary.get("grant_id") summary_type = summary.get("summary_type", "layman") summary_text = summary.get("summary_text", "") metadata = summary.get("metadata", {}) if not grant_id or not summary_text: logger.warning(f"Skipping invalid summary: {summary}") continue document = { "grant_id": grant_id, "summary_type": summary_type, "summary_text": summary_text, "created_at": datetime.utcnow(), "metadata": metadata } # Upsert operation operations.append( UpdateOne( {"grant_id": grant_id, "summary_type": summary_type}, {"$set": document}, upsert=True ) ) if not operations: return 0 # Execute bulk write result = self.collection.bulk_write(operations, ordered=False) saved_count = result.upserted_count + result.modified_count logger.info(f"💾 Bulk saved {saved_count} summaries ({result.upserted_count} new, {result.modified_count} updated)") return saved_count except BulkWriteError as bwe: # Log errors but don't fail completely logger.error(f"Bulk write errors: {bwe.details}") # Return count of successful writes return bwe.details.get("nInserted", 0) + bwe.details.get("nModified", 0) except Exception as e: logger.error(f"Error in bulk save: {e}") return 0 def search_summaries(self, query: str, summary_type: Optional[str] = None, limit: int = 10) -> List[Dict[str, Any]]: """ Full-text search across summaries using text index. Args: query: Search query string summary_type: Optional filter by summary type limit: Maximum results to return Returns: List of matching summaries with grant_id and summary_text """ try: filter_dict = {"$text": {"$search": query}} if summary_type: filter_dict["summary_type"] = summary_type # Text search with relevance score results = self.collection.find( filter_dict, {"grant_id": 1, "summary_type": 1, "summary_text": 1, "score": {"$meta": "textScore"}} ).sort([("score", {"$meta": "textScore"})]).limit(limit) return list(results) except Exception as e: logger.error(f"Error searching summaries: {e}") return [] class GrantStore: """Handle grant data operations with bulk write support.""" def __init__(self): """Initialize grant store with text indexes.""" self.db = get_database() self.collection = self.db["grants"] # Compound index for common queries self.collection.create_index( [("grant_id", ASCENDING), ("status", ASCENDING)], name="grant_status_lookup" ) # Text index for full-text search on grant titles and descriptions try: self.collection.create_index( [("title", TEXT), ("summary", TEXT)], name="grant_text_search", default_language="english" ) except Exception as e: logger.debug(f"Text index may already exist: {e}") # Index on deadline for sorting self.collection.create_index("deadline", name="deadline_idx") def bulk_update_grants(self, grants: List[Dict[str, Any]]) -> int: """ Bulk update grants from crawler. Args: grants: List of grant documents to upsert Returns: Number of grants updated """ if not grants: return 0 try: from pymongo import UpdateOne operations = [] for grant in grants: grant_id = grant.get("id") or grant.get("grant_id") if not grant_id: logger.warning(f"Skipping grant without ID: {grant.get('title', 'unknown')}") continue # Upsert operation operations.append( UpdateOne( {"grant_id": grant_id}, {"$set": grant}, upsert=True ) ) if not operations: return 0 # Execute bulk write result = self.collection.bulk_write(operations, ordered=False) updated_count = result.upserted_count + result.modified_count logger.info(f"💾 Bulk updated {updated_count} grants ({result.upserted_count} new, {result.modified_count} updated)") return updated_count except BulkWriteError as bwe: logger.error(f"Bulk write errors: {bwe.details}") return bwe.details.get("nInserted", 0) + bwe.details.get("nModified", 0) except Exception as e: logger.error(f"Error in bulk grant update: {e}") return 0 class FeedbackStore: """Handle feedback data operations with optimized indexing.""" def __init__(self): """Initialize feedback store with compound indexes.""" self.db = get_database() self.collection = self.db["feedback"] # Compound index for user feedback history self.collection.create_index( [("user_id", ASCENDING), ("created_at", ASCENDING)], name="user_feedback_history" ) # Index on rating for statistics self.collection.create_index("rating", name="rating_idx") # Index on created_at for time-based queries self.collection.create_index("created_at", name="feedback_created_at_idx") def save_feedback(self, feedback_data: Dict[str, Any]) -> str: """ Save feedback to database. Args: feedback_data: Dictionary with feedback information Returns: String ID of the inserted feedback """ try: result = self.collection.insert_one(feedback_data) logger.info(f"Feedback saved with ID: {result.inserted_id}") return str(result.inserted_id) except Exception as e: logger.error(f"Error saving feedback: {e}") raise def get_feedback_stats(self) -> Dict[str, Any]: """ Get feedback statistics. Returns: Dictionary with feedback statistics """ try: total_feedback = self.collection.count_documents({}) # Calculate average rating (only count feedback with ratings) pipeline = [ {"$match": {"rating": {"$exists": True, "$ne": None}}}, {"$group": {"_id": None, "avg_rating": {"$avg": "$rating"}}} ] result = list(self.collection.aggregate(pipeline)) avg_rating = result[0]["avg_rating"] if result else 0.0 return { "total_feedback": total_feedback, "average_rating": round(avg_rating, 2) } except Exception as e: logger.error(f"Error getting feedback stats: {e}") return { "total_feedback": 0, "average_rating": 0.0 }