major-domo-database/migrate_to_postgres.py

#!/usr/bin/env python3

import os
import logging
from datetime import datetime
from playhouse.shortcuts import model_to_dict
from peewee import SqliteDatabase, PostgresqlDatabase

logger = logging.getLogger(f'{__name__}.migrate_to_postgres')

def setup_databases():
    """Setup both SQLite source and PostgreSQL target databases"""
    
    # SQLite source database
    sqlite_db = SqliteDatabase(
        'storage/sba_master.db',
        pragmas={
            'journal_mode': 'wal',
            'cache_size': -1 * 64000,
            'synchronous': 0
        }
    )
    
    # PostgreSQL target database
    postgres_db = PostgresqlDatabase(
        os.environ.get('POSTGRES_DB', 'sba_master'),
        user=os.environ.get('POSTGRES_USER', 'sba_admin'),
        password=os.environ.get('POSTGRES_PASSWORD', 'sba_dev_password_2024'),
        host=os.environ.get('POSTGRES_HOST', 'localhost'),
        port=int(os.environ.get('POSTGRES_PORT', '5432'))
    )
    
    return sqlite_db, postgres_db

def get_all_models():
    """Get all models in dependency order for migration"""
    # Set temporary environment to load models
    os.environ['DATABASE_TYPE'] = 'sqlite'
    
    from app.db_engine import (
        Current, Manager, Division, SbaPlayer,  # No dependencies
        Team,  # Depends on Manager, Division
        Player,  # Depends on Team, SbaPlayer
        Result, Schedule, Transaction,  # Depend on Team, Player
        BattingStat, PitchingStat,  # Depend on Player, Team
        Standings,  # Depends on Team
        BattingCareer, PitchingCareer, FieldingCareer,  # No dependencies
        BattingSeason, PitchingSeason, FieldingSeason,  # Depend on Player, Career tables
        DraftPick, DraftData, DraftList,  # Depend on Team, Player
        Award,  # Depends on Manager, Player, Team
        DiceRoll,  # Depends on Team
        Keeper, Injury,  # Depend on Team, Player
        StratGame,  # Depends on Team, Manager
        StratPlay, Decision,  # Depend on StratGame, Player, Team
        CustomCommandCreator, CustomCommand  # CustomCommand depends on Creator
    )
    
    # Return in dependency order
    return [
        # Base tables (no dependencies)
        Current, Manager, Division, SbaPlayer,
        BattingCareer, PitchingCareer, FieldingCareer,
        CustomCommandCreator,
        
        # First level dependencies
        Team, DraftData,
        
        # Second level dependencies
        Player, CustomCommand,
        
        # Third level dependencies
        Result, Schedule, Transaction, BattingStat, PitchingStat,
        Standings, DraftPick, DraftList, Award, DiceRoll,
        Keeper, Injury, StratGame,
        
        # Fourth level dependencies
        BattingSeason, PitchingSeason, FieldingSeason,
        StratPlay, Decision
    ]

def migrate_table_data(model_class, sqlite_db, postgres_db, batch_size=1000):
    """Migrate data from SQLite to PostgreSQL for a specific model"""
    table_name = model_class._meta.table_name
    logger.info(f"Migrating table: {table_name}")
    
    try:
        # Connect to SQLite and count records
        model_class._meta.database = sqlite_db
        sqlite_db.connect()
        
        total_records = model_class.select().count()
        if total_records == 0:
            logger.info(f"  No records in {table_name}, skipping")
            sqlite_db.close()
            return True
            
        logger.info(f"  Found {total_records} records")
        sqlite_db.close()
        
        # Connect to PostgreSQL and prepare
        model_class._meta.database = postgres_db
        postgres_db.connect()
        
        # Create table if it doesn't exist
        model_class.create_table(safe=True)
        
        # Migrate data in batches
        migrated = 0
        sqlite_db.connect()
        
        for batch_start in range(0, total_records, batch_size):
            # Get batch from SQLite
            model_class._meta.database = sqlite_db
            batch = list(model_class.select().offset(batch_start).limit(batch_size))
            
            if not batch:
                break
                
            # Convert to dicts and prepare for PostgreSQL
            batch_data = []
            for record in batch:
                data = model_to_dict(record, recurse=False)
                # Remove auto-increment ID if present to let PostgreSQL handle it
                if 'id' in data and hasattr(model_class, 'id'):
                    data.pop('id', None)
                batch_data.append(data)
            
            # Insert into PostgreSQL
            model_class._meta.database = postgres_db
            if batch_data:
                model_class.insert_many(batch_data).execute()
                migrated += len(batch_data)
                
            logger.info(f"  Migrated {migrated}/{total_records} records")
        
        sqlite_db.close()
        postgres_db.close()
        
        logger.info(f"✓ Successfully migrated {table_name}: {migrated} records")
        return True
        
    except Exception as e:
        logger.error(f"✗ Failed to migrate {table_name}: {e}")
        try:
            sqlite_db.close()
        except:
            pass
        try:
            postgres_db.close()
        except:
            pass
        return False

def migrate_all_data():
    """Migrate all data from SQLite to PostgreSQL"""
    logger.info("Starting full data migration from SQLite to PostgreSQL...")
    
    # Setup databases
    sqlite_db, postgres_db = setup_databases()
    
    # Test connections
    try:
        sqlite_db.connect()
        sqlite_db.execute_sql("SELECT 1").fetchone()
        sqlite_db.close()
        logger.info("✓ SQLite source database connection OK")
    except Exception as e:
        logger.error(f"✗ SQLite connection failed: {e}")
        return False
    
    try:
        postgres_db.connect()
        postgres_db.execute_sql("SELECT 1").fetchone()
        postgres_db.close()
        logger.info("✓ PostgreSQL target database connection OK")
    except Exception as e:
        logger.error(f"✗ PostgreSQL connection failed: {e}")
        return False
    
    # Get models in dependency order
    all_models = get_all_models()
    logger.info(f"Found {len(all_models)} models to migrate")
    
    # Migrate each table
    successful_migrations = 0
    failed_migrations = []
    
    for model in all_models:
        success = migrate_table_data(model, sqlite_db, postgres_db)
        if success:
            successful_migrations += 1
        else:
            failed_migrations.append(model._meta.table_name)
    
    # Report results
    logger.info(f"\nMigration completed:")
    logger.info(f"✓ Successful: {successful_migrations}/{len(all_models)} tables")
    
    if failed_migrations:
        logger.error(f"✗ Failed: {len(failed_migrations)} tables")
        for table in failed_migrations:
            logger.error(f"  - {table}")
        return False
    else:
        logger.info("🎉 All tables migrated successfully!")
        return True

def main():
    # Configure logging
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    )
    
    # Set PostgreSQL environment variables
    os.environ['POSTGRES_DB'] = 'sba_master'
    os.environ['POSTGRES_USER'] = 'sba_admin' 
    os.environ['POSTGRES_PASSWORD'] = 'sba_dev_password_2024'
    os.environ['POSTGRES_HOST'] = 'localhost'
    os.environ['POSTGRES_PORT'] = '5432'
    
    success = migrate_all_data()
    return 0 if success else 1

if __name__ == "__main__":
    exit(main())