SneakyScan/app/init_db.py

#!/usr/bin/env python3
"""
Database initialization script for SneakyScanner.

This script:
1. Creates the database schema using Alembic migrations
2. Initializes default settings
3. Optionally sets up an initial admin password

Usage:
    python3 init_db.py [--password PASSWORD]
"""

import argparse
import os
import sys
from pathlib import Path

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

from alembic import command
from alembic.config import Config
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker

from web.models import Base
from web.utils.settings import PasswordManager, SettingsManager


def init_database(db_url: str = "sqlite:///./sneakyscanner.db", run_migrations: bool = True):
    """
    Initialize the database schema and settings.

    Args:
        db_url: Database URL (defaults to SQLite in current directory)
        run_migrations: Whether to run Alembic migrations (True) or create all tables directly (False)
    """
    print(f"Initializing SneakyScanner database at: {db_url}")

    # Create database directory if it doesn't exist (for SQLite)
    if db_url.startswith('sqlite:///'):
        db_path = db_url.replace('sqlite:///', '')
        db_dir = Path(db_path).parent
        if not db_dir.exists():
            print(f"Creating database directory: {db_dir}")
            db_dir.mkdir(parents=True, exist_ok=True)

    if run_migrations:
        # Run Alembic migrations
        print("Running Alembic migrations...")
        alembic_cfg = Config("alembic.ini")
        alembic_cfg.set_main_option("sqlalchemy.url", db_url)

        try:
            # Upgrade to head (latest migration)
            command.upgrade(alembic_cfg, "head")
            print("✓ Database schema created successfully via migrations")
        except Exception as e:
            print(f"✗ Migration failed: {e}")
            print("Falling back to direct table creation...")
            run_migrations = False

    if not run_migrations:
        # Create tables directly using SQLAlchemy (fallback or if migrations disabled)
        print("Creating database schema directly...")
        engine = create_engine(db_url, echo=False)
        Base.metadata.create_all(engine)
        print("✓ Database schema created successfully")

    # Initialize settings
    print("\nInitializing default settings...")
    engine = create_engine(db_url, echo=False)
    Session = sessionmaker(bind=engine)
    session = Session()

    try:
        settings_manager = SettingsManager(session)
        settings_manager.init_defaults()
        print("✓ Default settings initialized")
    except Exception as e:
        print(f"✗ Failed to initialize settings: {e}")
        session.rollback()
        raise
    finally:
        session.close()

    print("\n✓ Database initialization complete!")
    return True


def set_password(db_url: str, password: str):
    """
    Set the application password.

    Args:
        db_url: Database URL
        password: Password to set
    """
    print("Setting application password...")

    engine = create_engine(db_url, echo=False)
    Session = sessionmaker(bind=engine)
    session = Session()

    try:
        settings_manager = SettingsManager(session)
        PasswordManager.set_app_password(settings_manager, password)
        print("✓ Password set successfully")
    except Exception as e:
        print(f"✗ Failed to set password: {e}")
        session.rollback()
        raise
    finally:
        session.close()


def verify_database(db_url: str):
    """
    Verify database schema and settings.

    Args:
        db_url: Database URL
    """
    print("\nVerifying database...")

    engine = create_engine(db_url, echo=False)
    Session = sessionmaker(bind=engine)
    session = Session()

    try:
        # Check if tables exist by querying settings
        from web.models import Setting
        count = session.query(Setting).count()
        print(f"✓ Settings table accessible ({count} settings found)")

        # Display current settings (sanitized)
        settings_manager = SettingsManager(session)
        settings = settings_manager.get_all(decrypt=False, sanitize=True)
        print("\nCurrent settings:")
        for key, value in sorted(settings.items()):
            print(f"  {key}: {value}")

    except Exception as e:
        print(f"✗ Database verification failed: {e}")
        raise
    finally:
        session.close()


def main():
    """Main entry point for database initialization."""
    parser = argparse.ArgumentParser(
        description="Initialize SneakyScanner database",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Initialize database with default settings
  python3 init_db.py

  # Initialize and set password
  python3 init_db.py --password mysecretpassword

  # Use custom database URL
  python3 init_db.py --db-url postgresql://user:pass@localhost/sneakyscanner

  # Verify existing database
  python3 init_db.py --verify-only
        """
    )

    parser.add_argument(
        '--db-url',
        default='sqlite:///./sneakyscanner.db',
        help='Database URL (default: sqlite:///./sneakyscanner.db)'
    )

    parser.add_argument(
        '--password',
        help='Set application password'
    )

    parser.add_argument(
        '--verify-only',
        action='store_true',
        help='Only verify database, do not initialize'
    )

    parser.add_argument(
        '--no-migrations',
        action='store_true',
        help='Create tables directly instead of using migrations'
    )

    args = parser.parse_args()

    # Check if database already exists
    db_exists = False
    if args.db_url.startswith('sqlite:///'):
        db_path = args.db_url.replace('sqlite:///', '')
        db_exists = Path(db_path).exists()

    if db_exists and not args.verify_only:
        response = input(f"\nDatabase already exists at {db_path}. Reinitialize? (y/N): ")
        if response.lower() != 'y':
            print("Aborting.")
            return

    try:
        if args.verify_only:
            verify_database(args.db_url)
        else:
            # Initialize database
            init_database(args.db_url, run_migrations=not args.no_migrations)

            # Set password if provided
            if args.password:
                set_password(args.db_url, args.password)

            # Verify
            verify_database(args.db_url)

        print("\n✓ All done! Database is ready to use.")

        if not args.password and not args.verify_only:
            print("\n⚠ WARNING: No password set. Run with --password to set one:")
            print(f"  python3 init_db.py --db-url {args.db_url} --password YOUR_PASSWORD")

    except Exception as e:
        print(f"\n✗ Initialization failed: {e}")
        sys.exit(1)


if __name__ == '__main__':
    main()