scripts/migrate_to_postgres.py

#!/usr/bin/env python3
"""
SQLite to PostgreSQL Migration Script
Migrates data from SQLite database to PostgreSQL
"""

import sys
import os
import sqlite3
import psycopg2
from psycopg2 import sql
from datetime import datetime
from pathlib import Path

# Add parent directory to path
sys.path.insert(0, str(Path(__file__).parent.parent))

def connect_sqlite(db_path):
    """Connect to SQLite database"""
    try:
        conn = sqlite3.connect(db_path)
        conn.row_factory = sqlite3.Row
        print(f"✓ Connected to SQLite: {db_path}")
        return conn
    except Exception as e:
        print(f"✗ SQLite connection failed: {e}")
        sys.exit(1)

def connect_postgres(config):
    """Connect to PostgreSQL database"""
    try:
        conn = psycopg2.connect(
            host=config.get('host', 'localhost'),
            port=config.get('port', 5432),
            database=config.get('database'),
            user=config.get('user'),
            password=config.get('password')
        )
        print(f"✓ Connected to PostgreSQL: {config['database']}")
        return conn
    except Exception as e:
        print(f"✗ PostgreSQL connection failed: {e}")
        sys.exit(1)

def get_tables(sqlite_conn):
    """Get all table names from SQLite"""
    cursor = sqlite_conn.cursor()
    cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%'")
    tables = [row[0] for row in cursor.fetchall()]
    cursor.close()
    return tables

def migrate_table(table_name, sqlite_conn, pg_conn):
    """Migrate a single table from SQLite to PostgreSQL"""
    print(f"\n→ Migrating table: {table_name}")
    
    # Get data from SQLite
    sqlite_cursor = sqlite_conn.cursor()
    sqlite_cursor.execute(f"SELECT * FROM {table_name}")
    rows = sqlite_cursor.fetchall()
    
    if not rows:
        print(f"  ⓘ Table {table_name} is empty, skipping...")
        return 0
    
    # Get column names
    columns = [description[0] for description in sqlite_cursor.description]
    
    # Prepare PostgreSQL insert
    pg_cursor = pg_conn.cursor()
    
    # Build insert query
    placeholders = ', '.join(['%s'] * len(columns))
    columns_str = ', '.join([f'"{col}"' for col in columns])
    insert_query = f'INSERT INTO {table_name} ({columns_str}) VALUES ({placeholders}) ON CONFLICT DO NOTHING'
    
    # Insert data
    migrated_count = 0
    failed_count = 0
    
    for row in rows:
        try:
            pg_cursor.execute(insert_query, tuple(row))
            migrated_count += 1
        except Exception as e:
            failed_count += 1
            print(f"  ✗ Failed to insert row: {e}")
            pg_conn.rollback()
            continue
    
    pg_conn.commit()
    pg_cursor.close()
    sqlite_cursor.close()
    
    print(f"  ✓ Migrated {migrated_count} rows")
    if failed_count > 0:
        print(f"  ⚠ Failed to migrate {failed_count} rows")
    
    return migrated_count

def update_sequences(pg_conn, tables):
    """Update PostgreSQL sequences after migration"""
    print("\n→ Updating sequences...")
    pg_cursor = pg_conn.cursor()
    
    for table in tables:
        try:
            # Find primary key column (assuming it's named 'id')
            pg_cursor.execute(f"""
                SELECT MAX(id) FROM {table}
            """)
            max_id = pg_cursor.fetchone()[0]
            
            if max_id:
                # Update sequence
                pg_cursor.execute(f"""
                    SELECT setval('{table}_id_seq', {max_id}, true)
                """)
                print(f"  ✓ Updated sequence for {table} to {max_id}")
        except Exception as e:
            print(f"  ⓘ No sequence for {table} or error: {e}")
            pg_conn.rollback()
    
    pg_conn.commit()
    pg_cursor.close()

def verify_migration(sqlite_conn, pg_conn, tables):
    """Verify migration by comparing row counts"""
    print("\n→ Verifying migration...")
    all_match = True
    
    for table in tables:
        sqlite_cursor = sqlite_conn.cursor()
        pg_cursor = pg_conn.cursor()
        
        sqlite_cursor.execute(f"SELECT COUNT(*) FROM {table}")
        sqlite_count = sqlite_cursor.fetchone()[0]
        
        pg_cursor.execute(f"SELECT COUNT(*) FROM {table}")
        pg_count = pg_cursor.fetchone()[0]
        
        if sqlite_count == pg_count:
            print(f"  ✓ {table}: {sqlite_count} rows (match)")
        else:
            print(f"  ✗ {table}: SQLite={sqlite_count}, PostgreSQL={pg_count} (mismatch)")
            all_match = False
        
        sqlite_cursor.close()
        pg_cursor.close()
    
    return all_match

def main():
    """Main migration function"""
    print("=" * 60)
    print("SQLite to PostgreSQL Migration Tool")
    print("AI Fitness Coach Platform")
    print("=" * 60)
    
    # Configuration
    SQLITE_DB = os.getenv('SQLITE_DB', 'backend/database/fitness_coach.db')
    
    PG_CONFIG = {
        'host': os.getenv('DB_HOST', 'localhost'),
        'port': int(os.getenv('DB_PORT', 5432)),
        'database': os.getenv('DB_NAME', 'fitness_coach'),
        'user': os.getenv('DB_USER', 'fitness_user'),
        'password': os.getenv('DB_PASSWORD', 'fitness_password')
    }
    
    # Check if SQLite database exists
    if not os.path.exists(SQLITE_DB):
        print(f"✗ SQLite database not found: {SQLITE_DB}")
        sys.exit(1)
    
    # Connect to databases
    sqlite_conn = connect_sqlite(SQLITE_DB)
    pg_conn = connect_postgres(PG_CONFIG)
    
    try:
        # Get tables to migrate
        tables = get_tables(sqlite_conn)
        print(f"\nFound {len(tables)} tables to migrate:")
        for table in tables:
            print(f"  • {table}")
        
        # Confirm migration
        response = input("\nProceed with migration? (yes/no): ")
        if response.lower() != 'yes':
            print("Migration cancelled.")
            return
        
        # Migrate each table
        print("\nStarting migration...")
        total_rows = 0
        for table in tables:
            rows = migrate_table(table, sqlite_conn, pg_conn)
            total_rows += rows
        
        # Update sequences
        update_sequences(pg_conn, tables)
        
        # Verify migration
        if verify_migration(sqlite_conn, pg_conn, tables):
            print("\n" + "=" * 60)
            print(f"✓ Migration completed successfully!")
            print(f"✓ Total rows migrated: {total_rows}")
            print("=" * 60)
        else:
            print("\n" + "=" * 60)
            print("⚠ Migration completed with mismatches!")
            print("  Please review the verification results above.")
            print("=" * 60)
        
    except Exception as e:
        print(f"\n✗ Migration failed: {e}")
        pg_conn.rollback()
        sys.exit(1)
    
    finally:
        sqlite_conn.close()
        pg_conn.close()
        print("\nDatabase connections closed.")

if __name__ == '__main__':
    main()