migrate_database.py

#!/usr/bin/env python3
"""
Database Migration Script for GSSOC 25 Dashboard Update

This script handles the migration from GSSOC 24 to GSSOC 25 data by:
1. Creating backups of existing data (optional)
2. Clearing existing GSSOC 24 data from MongoDB collections
3. Providing logging for migration progress and errors

Requirements: 3.1, 3.2
"""

import os
import sys
import json
import logging
import pymongo
import toml
from datetime import datetime, timezone
from dotenv import load_dotenv
from typing import Optional, Dict, Any

# Load environment variables
load_dotenv()

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('migration.log'),
        logging.StreamHandler(sys.stdout)
    ]
)

logger = logging.getLogger(__name__)

class DatabaseMigrator:
    """Handles database migration from GSSOC 24 to GSSOC 25"""
    
    def __init__(self, mongodb_uri: Optional[str] = None, database_name: str = "gssoc"):
        """
        Initialize the database migrator
        
        Args:
            mongodb_uri: MongoDB connection URI (defaults to secrets.toml or environment variable)
            database_name: Name of the database to migrate (defaults to "gssoc")
        """
        self.mongodb_uri = mongodb_uri or self._load_mongodb_uri()
        self.database_name = database_name
        self.client = None
        self.db = None
    
    def _load_mongodb_uri(self) -> str:
        """
        Load MongoDB URI from secrets.toml or environment variables
        
        Returns:
            str: MongoDB connection URI
        """
        # Try to load from .streamlit/secrets.toml first
        try:
            secrets_path = os.path.join('.streamlit', 'secrets.toml')
            if os.path.exists(secrets_path):
                with open(secrets_path, 'r') as f:
                    secrets = toml.load(f)
                    if 'MONGO_URI' in secrets:
                        logger.info("Loaded MongoDB URI from .streamlit/secrets.toml")
                        return secrets['MONGO_URI']
        except Exception as e:
            logger.warning(f"Failed to load from secrets.toml: {str(e)}")
        
        # Fallback to environment variables
        mongodb_uri = os.getenv('MONGODB_URI') or os.getenv('MONGO_URI')
        if mongodb_uri:
            logger.info("Loaded MongoDB URI from environment variables")
            return mongodb_uri
        
        # Final fallback to localhost
        logger.warning("No MongoDB URI found in secrets.toml or environment variables, using localhost")
        return 'mongodb://localhost:27017/'
        
    def connect(self) -> bool:
        """
        Connect to MongoDB database
        
        Returns:
            bool: True if connection successful, False otherwise
        """
        try:
            logger.info(f"Connecting to MongoDB at {self.mongodb_uri}")
            self.client = pymongo.MongoClient(self.mongodb_uri)
            self.db = self.client[self.database_name]
            
            # Test connection
            self.client.admin.command('ping')
            logger.info(f"Successfully connected to database: {self.database_name}")
            return True
            
        except Exception as e:
            logger.error(f"Failed to connect to MongoDB: {str(e)}")
            return False
    
    def disconnect(self):
        """Close MongoDB connection"""
        if self.client:
            self.client.close()
            logger.info("Disconnected from MongoDB")
    
    def backup_collection(self, collection_name: str, backup_dir: str = "backups") -> bool:
        """
        Create backup of a collection
        
        Args:
            collection_name: Name of the collection to backup
            backup_dir: Directory to store backup files
            
        Returns:
            bool: True if backup successful, False otherwise
        """
        try:
            # Create backup directory if it doesn't exist
            os.makedirs(backup_dir, exist_ok=True)
            
            collection = self.db[collection_name]
            document_count = collection.count_documents({})
            
            if document_count == 0:
                logger.info(f"Collection '{collection_name}' is empty, skipping backup")
                return True
            
            logger.info(f"Backing up collection '{collection_name}' ({document_count} documents)")
            
            # Create backup filename with timestamp
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            backup_filename = f"{collection_name}_backup_{timestamp}.json"
            backup_path = os.path.join(backup_dir, backup_filename)
            
            # Export collection data
            documents = list(collection.find({}))
            
            # Convert ObjectId to string for JSON serialization
            for doc in documents:
                if '_id' in doc:
                    doc['_id'] = str(doc['_id'])
            
            with open(backup_path, 'w', encoding='utf-8') as f:
                json.dump(documents, f, indent=2, default=str)
            
            logger.info(f"Backup created successfully: {backup_path}")
            return True
            
        except Exception as e:
            logger.error(f"Failed to backup collection '{collection_name}': {str(e)}")
            return False
    
    def clear_collection(self, collection_name: str) -> bool:
        """
        Clear all documents from a collection
        
        Args:
            collection_name: Name of the collection to clear
            
        Returns:
            bool: True if clearing successful, False otherwise
        """
        try:
            collection = self.db[collection_name]
            document_count = collection.count_documents({})
            
            if document_count == 0:
                logger.info(f"Collection '{collection_name}' is already empty")
                return True
            
            logger.info(f"Clearing collection '{collection_name}' ({document_count} documents)")
            
            result = collection.delete_many({})
            deleted_count = result.deleted_count
            
            logger.info(f"Successfully deleted {deleted_count} documents from '{collection_name}'")
            return True
            
        except Exception as e:
            logger.error(f"Failed to clear collection '{collection_name}': {str(e)}")
            return False
    
    def get_collection_stats(self, collection_name: str) -> Dict[str, Any]:
        """
        Get statistics for a collection
        
        Args:
            collection_name: Name of the collection
            
        Returns:
            dict: Collection statistics
        """
        try:
            collection = self.db[collection_name]
            stats = {
                'name': collection_name,
                'document_count': collection.count_documents({}),
                'exists': collection_name in self.db.list_collection_names()
            }
            return stats
        except Exception as e:
            logger.error(f"Failed to get stats for collection '{collection_name}': {str(e)}")
            return {'name': collection_name, 'document_count': 0, 'exists': False, 'error': str(e)}
    
    def load_gssoc25_projects(self, projects_file: str = "gssoc25_projects_valid.json") -> bool:
        """
        Load GSSOC 2025 projects from JSON file and insert into database
        
        Args:
            projects_file: Path to the JSON file containing GSSOC 2025 projects
            
        Returns:
            bool: True if loading successful, False otherwise
        """
        try:
            if not os.path.exists(projects_file):
                logger.error(f"Projects file not found: {projects_file}")
                return False
            
            logger.info(f"Loading GSSOC 2025 projects from {projects_file}")
            
            with open(projects_file, 'r', encoding='utf-8') as f:
                projects_data = json.load(f)
            
            if not isinstance(projects_data, list):
                logger.error("Projects file should contain a JSON array")
                return False
            
            projects_collection = self.db['projects']
            
            # Add metadata to each project
            for project in projects_data:
                project['created_at'] = datetime.now(timezone.utc)
                project['gssoc_year'] = 2025
            
            # Insert projects in batches for better performance
            batch_size = 50
            total_projects = len(projects_data)
            inserted_count = 0
            
            logger.info(f"Inserting {total_projects} GSSOC 2025 projects...")
            
            for i in range(0, total_projects, batch_size):
                batch = projects_data[i:i + batch_size]
                try:
                    result = projects_collection.insert_many(batch)
                    inserted_count += len(result.inserted_ids)
                    logger.info(f"Inserted batch {i//batch_size + 1}: {len(result.inserted_ids)} projects")
                except Exception as e:
                    logger.error(f"Failed to insert batch {i//batch_size + 1}: {str(e)}")
                    return False
            
            logger.info(f"Successfully inserted {inserted_count} GSSOC 2025 projects")
            
            # Verify insertion
            final_count = projects_collection.count_documents({'gssoc_year': 2025})
            if final_count == total_projects:
                logger.info(f"Verification successful: {final_count} projects in database")
                return True
            else:
                logger.error(f"Verification failed: Expected {total_projects}, found {final_count}")
                return False
                
        except Exception as e:
            logger.error(f"Failed to load GSSOC 2025 projects: {str(e)}")
            return False

    def migrate_database(self, create_backup: bool = True, load_projects: bool = True, projects_file: str = "gssoc25_projects_valid.json") -> bool:
        """
        Perform complete database migration
        
        Args:
            create_backup: Whether to create backups before clearing data
            load_projects: Whether to load GSSOC 2025 projects after clearing
            projects_file: Path to the JSON file containing GSSOC 2025 projects
            
        Returns:
            bool: True if migration successful, False otherwise
        """
        logger.info("Starting GSSOC 24 to GSSOC 25 database migration")
        
        # Collections to migrate
        collections_to_clear = ['projects', 'repo_stats']
        
        try:
            # Get initial statistics
            logger.info("Getting initial database statistics...")
            for collection_name in collections_to_clear:
                stats = self.get_collection_stats(collection_name)
                logger.info(f"Collection '{collection_name}': {stats['document_count']} documents")
            
            # Create backups if requested
            if create_backup:
                logger.info("Creating backups of existing data...")
                for collection_name in collections_to_clear:
                    if not self.backup_collection(collection_name):
                        logger.warning(f"Backup failed for '{collection_name}', continuing with migration")
            else:
                logger.info("Skipping backup creation as requested")
            
            # Clear collections
            logger.info("Clearing existing GSSOC 24 data...")
            migration_success = True
            
            for collection_name in collections_to_clear:
                if not self.clear_collection(collection_name):
                    logger.error(f"Failed to clear collection '{collection_name}'")
                    migration_success = False
            
            # Verify collections are empty
            logger.info("Verifying migration results...")
            for collection_name in collections_to_clear:
                stats = self.get_collection_stats(collection_name)
                if stats['document_count'] > 0:
                    logger.error(f"Collection '{collection_name}' still contains {stats['document_count']} documents")
                    migration_success = False
                else:
                    logger.info(f"Collection '{collection_name}' successfully cleared")
            
            # Load GSSOC 2025 projects if requested and clearing was successful
            if migration_success and load_projects:
                logger.info("Loading GSSOC 2025 projects...")
                if not self.load_gssoc25_projects(projects_file):
                    logger.error("Failed to load GSSOC 2025 projects")
                    migration_success = False
                else:
                    logger.info("GSSOC 2025 projects loaded successfully")
            elif not load_projects:
                logger.info("Skipping GSSOC 2025 projects loading as requested")
            
            if migration_success:
                logger.info("Database migration completed successfully!")
                if load_projects:
                    logger.info("Database is now populated with GSSOC 2025 projects")
                else:
                    logger.info("Ready for GSSOC 25 data population")
            else:
                logger.error("Database migration completed with errors")
            
            return migration_success
            
        except Exception as e:
            logger.error(f"Migration failed with exception: {str(e)}")
            return False

def main():
    """Main function to run the migration"""
    import argparse
    
    parser = argparse.ArgumentParser(description='Migrate GSSOC dashboard database from 2024 to 2025')
    parser.add_argument('--no-backup', action='store_true', 
                       help='Skip creating backups (not recommended)')
    parser.add_argument('--no-load-projects', action='store_true',
                       help='Skip loading GSSOC 2025 projects after clearing')
    parser.add_argument('--projects-file', type=str, default='gssoc25_projects_valid.json',
                       help='Path to GSSOC 2025 projects JSON file (default: gssoc25_projects_valid.json)')
    parser.add_argument('--mongodb-uri', type=str, 
                       help='MongoDB connection URI (overrides environment variable)')
    parser.add_argument('--database', type=str, default='gssoc',
                       help='Database name (default: gssoc)')
    
    args = parser.parse_args()
    
    # Initialize migrator
    migrator = DatabaseMigrator(
        mongodb_uri=args.mongodb_uri,
        database_name=args.database
    )
    
    try:
        # Connect to database
        if not migrator.connect():
            logger.error("Failed to connect to database. Exiting.")
            sys.exit(1)
        
        # Perform migration
        create_backup = not args.no_backup
        load_projects = not args.no_load_projects
        success = migrator.migrate_database(
            create_backup=create_backup,
            load_projects=load_projects,
            projects_file=args.projects_file
        )
        
        if success:
            logger.info("Migration completed successfully!")
            sys.exit(0)
        else:
            logger.error("Migration failed!")
            sys.exit(1)
            
    except KeyboardInterrupt:
        logger.info("Migration interrupted by user")
        sys.exit(1)
    except Exception as e:
        logger.error(f"Unexpected error: {str(e)}")
        sys.exit(1)
    finally:
        migrator.disconnect()

if __name__ == "__main__":
    main()