bakery-ia/scripts/run_migrations.py

#!/usr/bin/env python3
"""
Enhanced Migration Runner

Handles automatic table creation and Alembic migrations for Kubernetes deployments.
Supports both first-time deployments and incremental migrations.
"""

import os
import sys
import asyncio
import argparse
import structlog
from pathlib import Path

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

from shared.database.base import DatabaseManager
from shared.database.init_manager import initialize_service_database

# Configure logging
structlog.configure(
    processors=[
        structlog.stdlib.filter_by_level,
        structlog.stdlib.add_logger_name,
        structlog.stdlib.add_log_level,
        structlog.stdlib.PositionalArgumentsFormatter(),
        structlog.processors.TimeStamper(fmt="iso"),
        structlog.processors.StackInfoRenderer(),
        structlog.processors.format_exc_info,
        structlog.processors.JSONRenderer()
    ],
    context_class=dict,
    logger_factory=structlog.stdlib.LoggerFactory(),
    wrapper_class=structlog.stdlib.BoundLogger,
    cache_logger_on_first_use=True,
)

logger = structlog.get_logger()


async def run_service_migration(service_name: str, force_recreate: bool = False) -> bool:
    """
    Run migration for a specific service

    Args:
        service_name: Name of the service (e.g., 'auth', 'inventory')
        force_recreate: Whether to force recreate tables (development mode)

    Returns:
        True if successful, False otherwise
    """
    logger.info("Starting migration for service", service=service_name, force_recreate=force_recreate)

    try:
        # Get database URL from environment (try both constructed and direct approaches)
        db_url_key = f"{service_name.upper().replace('-', '_')}_DATABASE_URL"
        database_url = os.getenv(db_url_key) or os.getenv("DATABASE_URL")

        # If no direct URL, construct from components
        if not database_url:
            host = os.getenv("POSTGRES_HOST")
            port = os.getenv("POSTGRES_PORT")
            db_name = os.getenv("POSTGRES_DB")
            user = os.getenv("POSTGRES_USER")
            password = os.getenv("POSTGRES_PASSWORD")

            if all([host, port, db_name, user, password]):
                database_url = f"postgresql+asyncpg://{user}:{password}@{host}:{port}/{db_name}"
                logger.info("Constructed database URL from components", host=host, port=port, db=db_name)
            else:
                logger.error("Database connection details not found",
                           db_url_key=db_url_key,
                           host=bool(host),
                           port=bool(port),
                           db=bool(db_name),
                           user=bool(user),
                           password=bool(password))
                return False

        # Create database manager
        db_manager = DatabaseManager(database_url=database_url)

        # Initialize the database
        result = await initialize_service_database(
            database_manager=db_manager,
            service_name=service_name,
            force_recreate=force_recreate
        )

        logger.info("Migration completed successfully", service=service_name, result=result)
        return True

    except Exception as e:
        logger.error("Migration failed", service=service_name, error=str(e))
        return False

    finally:
        # Cleanup database connections
        try:
            await db_manager.close_connections()
        except:
            pass


async def main():
    """Main migration runner"""
    parser = argparse.ArgumentParser(description="Enhanced Migration Runner")
    parser.add_argument("service", help="Service name (e.g., auth, inventory)")
    parser.add_argument("--force-recreate", action="store_true",
                       help="Force recreate tables (development mode)")
    parser.add_argument("--verbose", "-v", action="store_true", help="Verbose logging")

    args = parser.parse_args()

    if args.verbose:
        logger.info("Starting migration runner", service=args.service,
                   force_recreate=args.force_recreate)

    # Run the migration
    success = await run_service_migration(args.service, args.force_recreate)

    if success:
        logger.info("Migration runner completed successfully")
        sys.exit(0)
    else:
        logger.error("Migration runner failed")
        sys.exit(1)


if __name__ == "__main__":
    asyncio.run(main())