rail-game/scripts/init_demo_db.py

#!/usr/bin/env python3
"""
Initialize the database with demo data for the Rail Game.

This script automates the database setup process:
1. Validates environment setup
2. Runs database migrations
3. Loads OSM fixtures for demo data

Usage:
    python scripts/init_demo_db.py [--dry-run] [--region REGION]

Requirements:
- Virtual environment activated
- .env file configured with DATABASE_URL
- PostgreSQL with PostGIS running
"""

import argparse
import os
import subprocess
import sys
from pathlib import Path

try:
    from dotenv import load_dotenv
    load_dotenv()
except ImportError:
    print("WARNING: python-dotenv not installed. .env file will not be loaded automatically.")
    print("Install with: pip install python-dotenv")


def check_virtualenv():
    """Check if we're running in a virtual environment."""
    # Skip virtualenv check in Docker containers
    if os.getenv('INIT_DEMO_DB') == 'true':
        return

    if not hasattr(sys, 'real_prefix') and not (hasattr(sys, 'base_prefix') and sys.base_prefix != sys.prefix):
        print("ERROR: Virtual environment not activated. Run:")
        print("  .venv\\Scripts\\Activate.ps1  (PowerShell)")
        print("  source .venv/bin/activate     (Bash/macOS/Linux)")
        sys.exit(1)


def check_env_file():
    """Check if .env file exists."""
    env_file = Path('.env')
    if not env_file.exists():
        print("ERROR: .env file not found. Copy .env.example to .env and configure:")
        print("  Copy-Item .env.example .env  (PowerShell)")
        print("  cp .env.example .env          (Bash)")
        sys.exit(1)


def check_database_url():
    """Check if DATABASE_URL is set in environment."""
    database_url = os.getenv('DATABASE_URL')
    if not database_url:
        print("ERROR: DATABASE_URL not set. Check your .env file.")
        sys.exit(1)
    print(f"Using database: {database_url}")


def run_command(cmd, cwd=None, description="", env=None):
    """Run a shell command and return the result."""
    print(f"\n>>> {description}")
    print(f"Running: {' '.join(cmd)}")
    try:
        env_vars = os.environ.copy()
        if env:
            env_vars.update(env)
        env_vars.setdefault("PYTHONPATH", "/app")
        result = subprocess.run(
            cmd,
            cwd=cwd,
            check=True,
            capture_output=True,
            text=True,
            env=env_vars,
        )
        if result.stdout:
            print(result.stdout)
        return result
    except subprocess.CalledProcessError as e:
        print(f"ERROR: Command failed with exit code {e.returncode}")
        if e.stdout:
            print(e.stdout)
        if e.stderr:
            print(e.stderr)
        sys.exit(1)


def run_migrations():
    """Run database migrations using alembic."""
    run_command(
        ['alembic', 'upgrade', 'head'],
        cwd='backend',
        description="Running database migrations",
    )


def load_osm_fixtures(region, dry_run=False):
    """Load OSM fixtures for demo data."""
    cmd = ['python', '-m', 'backend.scripts.osm_refresh', '--region', region]
    if dry_run:
        cmd.append('--no-commit')
        description = f"Loading OSM fixtures (dry run) for region: {region}"
    else:
        description = f"Loading OSM fixtures for region: {region}"

    run_command(cmd, description=description)


def main():
    parser = argparse.ArgumentParser(
        description="Initialize database with demo data")
    parser.add_argument(
        '--region',
        default='all',
        help='OSM region to load (default: all)'
    )
    parser.add_argument(
        '--dry-run',
        action='store_true',
        help='Dry run: run migrations and load fixtures without committing'
    )
    parser.add_argument(
        '--skip-migrations',
        action='store_true',
        help='Skip running migrations'
    )
    parser.add_argument(
        '--skip-fixtures',
        action='store_true',
        help='Skip loading OSM fixtures'
    )

    args = parser.parse_args()

    print("Rail Game Database Initialization")
    print("=" * 40)

    # Pre-flight checks
    check_virtualenv()
    check_env_file()
    check_database_url()

    # Run migrations
    if not args.skip_migrations:
        run_migrations()
    else:
        print("Skipping migrations (--skip-migrations)")

    # Load fixtures
    if not args.skip_fixtures:
        load_osm_fixtures(args.region, args.dry_run)
    else:
        print("Skipping fixtures (--skip-fixtures)")

    print("\n✅ Database initialization completed successfully!")
    if args.dry_run:
        print("Note: This was a dry run. No data was committed to the database.")
    else:
        print("Demo data loaded. You can now start the backend server.")


if __name__ == '__main__':
    main()