dss/dss-mvp1/scripts/run_migrations.py

#!/usr/bin/env python3
"""
DSS Database Migration Runner

This script runs SQL migrations in the correct order, with proper error handling
and transaction safety.

Usage:
    python run_migrations.py                 # Run all pending migrations
    python run_migrations.py --check         # Show pending migrations only
    python run_migrations.py --rollback 0001 # Rollback specific migration (CAREFUL!)
"""

import os
import sys
import sqlite3
import argparse
from pathlib import Path
from datetime import datetime
import json

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


class MigrationRunner:
    """Manages database migrations with version tracking and rollback support"""

    def __init__(self, db_path: Path = None):
        """
        Initialize migration runner

        Args:
            db_path: Path to database file. If None, uses default DSS location.
        """
        if db_path is None:
            # Default DSS database location
            db_path = Path.cwd() / ".dss" / "dss.db"

        self.db_path = Path(db_path)
        self.migrations_dir = Path(__file__).parent.parent / "dss" / "storage" / "migrations"
        self.migrations_table = "_dss_migrations"

    def _ensure_migrations_table(self, conn: sqlite3.Connection):
        """Create migrations tracking table if it doesn't exist"""
        conn.execute(f"""
            CREATE TABLE IF NOT EXISTS {self.migrations_table} (
                id TEXT PRIMARY KEY,
                applied_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                description TEXT,
                checksum TEXT,
                status TEXT DEFAULT 'applied'
            )
        """)
        conn.commit()

    def _get_migration_checksum(self, migration_file: Path) -> str:
        """Calculate checksum of migration file for integrity verification"""
        import hashlib
        content = migration_file.read_text()
        return hashlib.sha256(content.encode()).hexdigest()

    def _get_applied_migrations(self, conn: sqlite3.Connection) -> dict:
        """Get dictionary of applied migrations"""
        cursor = conn.execute(f"SELECT id, checksum FROM {self.migrations_table}")
        return {row[0]: row[1] for row in cursor.fetchall()}

    def _get_pending_migrations(self) -> list:
        """Get list of pending migrations in order"""
        applied = {}
        try:
            conn = sqlite3.connect(self.db_path)
            self._ensure_migrations_table(conn)
            applied = self._get_applied_migrations(conn)
            conn.close()
        except Exception as e:
            print(f"Warning: Could not read migration history: {e}")

        pending = []
        if self.migrations_dir.exists():
            for migration_file in sorted(self.migrations_dir.glob("*.sql")):
                migration_id = migration_file.stem  # e.g., "0002_add_uuid_columns"
                if migration_id not in applied:
                    pending.append({
                        'id': migration_id,
                        'file': migration_file,
                        'checksum': self._get_migration_checksum(migration_file),
                        'status': 'pending'
                    })

        return pending

    def check(self) -> bool:
        """Check for pending migrations without applying them"""
        pending = self._get_pending_migrations()

        if not pending:
            print("✓ No pending migrations - database is up to date")
            return True

        print(f"Found {len(pending)} pending migration(s):\n")
        for migration in pending:
            print(f"  - {migration['id']}")
            print(f"    File: {migration['file'].name}")
            print(f"    Checksum: {migration['checksum'][:16]}...")

        return False

    def run(self, dry_run: bool = False) -> bool:
        """
        Run all pending migrations

        Args:
            dry_run: If True, show migrations but don't apply them

        Returns:
            True if successful, False if any migration failed
        """
        pending = self._get_pending_migrations()

        if not pending:
            print("✓ No pending migrations - database is up to date")
            return True

        print(f"Found {len(pending)} pending migration(s)")
        if dry_run:
            print("\nDRY RUN - No changes will be applied\n")
        else:
            print("Running migrations...\n")

        # Backup database before running migrations
        if not dry_run:
            backup_path = self.db_path.with_suffix(f".backup-{datetime.now().strftime('%Y%m%d-%H%M%S')}")
            import shutil
            try:
                shutil.copy2(self.db_path, backup_path)
                print(f"✓ Database backed up to: {backup_path}\n")
            except Exception as e:
                print(f"✗ Failed to create backup: {e}")
                print("  Aborting migration")
                return False

        conn = sqlite3.connect(self.db_path)
        conn.row_factory = sqlite3.Row
        self._ensure_migrations_table(conn)

        try:
            for migration in pending:
                migration_id = migration['id']
                migration_file = migration['file']

                print(f"Running: {migration_id}")

                # Read migration SQL
                sql_content = migration_file.read_text()

                if not dry_run:
                    try:
                        # Execute migration with transaction
                        conn.executescript(sql_content)

                        # Record migration as applied
                        conn.execute(f"""
                            INSERT INTO {self.migrations_table}
                            (id, description, checksum, status)
                            VALUES (?, ?, ?, 'applied')
                        """, (migration_id, migration_file.name, migration['checksum']))

                        conn.commit()
                        print(f"  ✓ Migration applied successfully\n")
                    except sqlite3.Error as e:
                        conn.rollback()
                        print(f"  ✗ Migration failed: {e}")
                        print(f"  ✗ Changes rolled back")
                        return False
                else:
                    # Dry run: just show what would happen
                    print("  (DRY RUN - Would execute)")
                    lines = sql_content.split('\n')[:5]  # Show first 5 lines
                    for line in lines:
                        if line.strip() and not line.strip().startswith('--'):
                            print(f"    {line[:70]}")
                    if len(sql_content.split('\n')) > 5:
                        print(f"    ... ({len(sql_content.split(chr(10)))} lines total)")
                    print()

            if not dry_run:
                print("\n✓ All migrations applied successfully")
                return True
            else:
                print("✓ Dry run complete - no changes made")
                return True

        finally:
            conn.close()

    def status(self) -> None:
        """Show migration status"""
        try:
            conn = sqlite3.connect(self.db_path)
            self._ensure_migrations_table(conn)

            cursor = conn.execute(f"""
                SELECT id, applied_at, status FROM {self.migrations_table}
                ORDER BY applied_at
            """)

            applied = cursor.fetchall()
            print(f"Applied migrations ({len(applied)}):\n")

            if applied:
                for row in applied:
                    print(f"  ✓ {row[0]}")
                    print(f"    Applied at: {row[1]}")
            else:
                print("  (none)")

            pending = self._get_pending_migrations()
            print(f"\nPending migrations ({len(pending)}):\n")

            if pending:
                for migration in pending:
                    print(f"  ⏳ {migration['id']}")
            else:
                print("  (none)")

            conn.close()
        except Exception as e:
            print(f"Error reading migration status: {e}")


def main():
    parser = argparse.ArgumentParser(
        description="DSS Database Migration Runner",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python run_migrations.py              # Run all pending migrations
  python run_migrations.py --check      # Check for pending migrations
  python run_migrations.py --dry-run    # Show what would be applied
  python run_migrations.py --status     # Show migration status
"""
    )

    parser.add_argument(
        '--check',
        action='store_true',
        help='Check for pending migrations without applying them'
    )
    parser.add_argument(
        '--dry-run',
        action='store_true',
        help='Show migrations that would be applied without applying them'
    )
    parser.add_argument(
        '--status',
        action='store_true',
        help='Show migration status (applied and pending)'
    )
    parser.add_argument(
        '--db',
        type=Path,
        help='Path to database file (default: .dss/dss.db)'
    )

    args = parser.parse_args()

    runner = MigrationRunner(db_path=args.db)

    try:
        if args.status:
            runner.status()
            return 0
        elif args.check:
            success = runner.check()
            return 0 if success else 1
        elif args.dry_run:
            success = runner.run(dry_run=True)
            return 0 if success else 1
        else:
            # Run migrations
            success = runner.run(dry_run=False)
            return 0 if success else 1
    except KeyboardInterrupt:
        print("\nMigration cancelled by user")
        return 1
    except Exception as e:
        print(f"Error: {e}")
        import traceback
        traceback.print_exc()
        return 1


if __name__ == '__main__':
    sys.exit(main())