scripts/init_db.py

"""CLI script to initialize Airflow database."""
import sys
from pathlib import Path
import subprocess
import os

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

from mlpipeline.logging.logger import get_logger

logger = get_logger(__name__)


def main():
    """Initialize Airflow database and create admin user."""
    try:
        airflow_home = project_root / "airflow"
        
        # Set AIRFLOW_HOME environment variable
        os.environ["AIRFLOW_HOME"] = str(airflow_home)
        logger.info(f"AIRFLOW_HOME set to: {airflow_home}")
        
        # Initialize database
        logger.info("Initializing Airflow database...")
        result = subprocess.run(
            ["airflow", "db", "init"],
            capture_output=True,
            text=True
        )
        
        if result.returncode != 0:
            logger.error(f"Failed to initialize database: {result.stderr}")
            return 1
        
        logger.info("✅ Database initialized successfully!")
        
        # Create admin user
        logger.info("Creating admin user...")
        result = subprocess.run(
            [
                "airflow", "users", "create",
                "--username", "admin",
                "--firstname", "Admin",
                "--lastname", "User",
                "--role", "Admin",
                "--email", "admin@example.com",
                "--password", "admin"
            ],
            capture_output=True,
            text=True
        )
        
        if result.returncode != 0 and "already exists" not in result.stderr:
            logger.warning(f"User creation note: {result.stderr}")
        else:
            logger.info("✅ Admin user created (username: admin, password: admin)")
        
        logger.info("\n🚀 Next steps:")
        logger.info("  1. Start scheduler: airflow scheduler")
        logger.info("  2. Start webserver: airflow webserver")
        logger.info("  3. Access UI: http://localhost:8080")
        
        return 0
    except FileNotFoundError:
        logger.error("❌ Airflow not found. Install with: pip install apache-airflow")
        return 1
    except Exception as e:
        logger.error(f"❌ Initialization failed: {e}")
        return 1


if __name__ == "__main__":
    exit(main())