HabitForge-zepp/scripts/migrate_to_postgres.py

"""
Migration script from SQLite to PostgreSQL.

This script migrates data from the existing SQLite database to PostgreSQL.

Usage:
    python scripts/migrate_to_postgres.py

Prerequisites:
    - PostgreSQL database must be running
    - DATABASE_URL environment variable must point to PostgreSQL
    - SQLite database must exist at data/habitforge.db
"""
import os
import sys
from datetime import datetime

# Add parent directory to path
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from sqlalchemy import create_engine, text
from sqlalchemy.orm import sessionmaker
from backend import models
from backend.database import Base, SessionLocal as SQLiteSessionLocal

# PostgreSQL connection
POSTGRES_URL = os.getenv("DATABASE_URL", "postgresql://habitforge:habitforge@localhost:5432/habitforge")

def migrate_data():
    """Migrate data from SQLite to PostgreSQL."""
    print("Starting migration from SQLite to PostgreSQL...")
    
    # Create PostgreSQL engine
    postgres_engine = create_engine(POSTGRES_URL)
    PostgresSessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=postgres_engine)
    
    # Create tables in PostgreSQL
    print("Creating tables in PostgreSQL...")
    Base.metadata.create_all(bind=postgres_engine)
    
    # Get SQLite session
    sqlite_session = SQLiteSessionLocal()
    
    # Get PostgreSQL session
    postgres_session = PostgresSessionLocal()
    
    try:
        # Migrate Users
        print("Migrating Users...")
        users = sqlite_session.query(models.User).all()
        for user in users:
            new_user = models.User(
                id=user.id,
                username=user.username,
                hashed_password=user.hashed_password
            )
            postgres_session.add(new_user)
        postgres_session.commit()
        print(f"Migrated {len(users)} users")
        
        # Migrate Challenges
        print("Migrating Challenges...")
        challenges = sqlite_session.query(models.Challenge).all()
        for challenge in challenges:
            new_challenge = models.Challenge(
                id=challenge.id,
                name=challenge.name,
                period_type=challenge.period_type,
                daily_target=challenge.daily_target,
                start_date=challenge.start_date,
                end_date=challenge.end_date,
                display_order=challenge.display_order,
                user_id=challenge.user_id,
                icon=challenge.icon,
                image_url=challenge.image_url,
                description=challenge.description,
                unit_type=challenge.unit_type,
                rest_days=challenge.rest_days,
                frequency=challenge.frequency
            )
            postgres_session.add(new_challenge)
        postgres_session.commit()
        print(f"Migrated {len(challenges)} challenges")
        
        # Migrate Tracking
        print("Migrating Tracking...")
        trackings = sqlite_session.query(models.Tracking).all()
        for tracking in trackings:
            new_tracking = models.Tracking(
                id=tracking.id,
                challenge_id=tracking.challenge_id,
                date=tracking.date,
                reps=tracking.reps,
                completed=tracking.completed,
                notes=tracking.notes
            )
            postgres_session.add(new_tracking)
        postgres_session.commit()
        print(f"Migrated {len(trackings)} trackings")
        
        # Migrate DailyHealthMetrics
        print("Migrating DailyHealthMetrics...")
        metrics = sqlite_session.query(models.DailyHealthMetrics).all()
        for metric in metrics:
            new_metric = models.DailyHealthMetrics(
                id=metric.id,
                user_id=metric.user_id,
                date=metric.date,
                step_count=metric.step_count,
                calories_burned=metric.calories_burned,
                distance_meters=metric.distance_meters,
                sleep_duration_minutes=metric.sleep_duration_minutes,
                deep_sleep_minutes=metric.deep_sleep_minutes,
                light_sleep_minutes=metric.light_sleep_minutes,
                rem_sleep_minutes=metric.rem_sleep_minutes,
                awake_duration_minutes=metric.awake_duration_minutes,
                avg_heart_rate=metric.avg_heart_rate,
                min_heart_rate=metric.min_heart_rate,
                max_heart_rate=metric.max_heart_rate,
                avg_spo2=metric.avg_spo2,
                pai_score=metric.pai_score,
                weight=metric.weight,
                resting_heart_rate=metric.resting_heart_rate,
                hrv=metric.hrv
            )
            postgres_session.add(new_metric)
        postgres_session.commit()
        print(f"Migrated {len(metrics)} daily health metrics")
        
        print("\nMigration completed successfully!")
        print(f"Total records migrated:")
        print(f"  - Users: {len(users)}")
        print(f"  - Challenges: {len(challenges)}")
        print(f"  - Trackings: {len(trackings)}")
        print(f"  - DailyHealthMetrics: {len(metrics)}")
        
    except Exception as e:
        print(f"\nError during migration: {e}")
        postgres_session.rollback()
        raise
    finally:
        sqlite_session.close()
        postgres_session.close()
        postgres_engine.dispose()


if __name__ == "__main__":
    # Check if SQLite database exists
    if not os.path.exists("data/habitforge.db"):
        print("Error: SQLite database not found at data/habitforge.db")
        sys.exit(1)
    
    # Confirm migration
    print("This will migrate data from SQLite to PostgreSQL.")
    print(f"PostgreSQL URL: {POSTGRES_URL}")
    response = input("Do you want to continue? (yes/no): ")
    
    if response.lower() != "yes":
        print("Migration cancelled.")
        sys.exit(0)
    
    migrate_data()