""" Database Initialization Manager Handles automatic table creation and Alembic integration for both: 1. Production deployments (first-time table creation) 2. Development workflows (reset to clean slate) """ import os import asyncio import structlog from typing import Optional, List, Dict, Any from pathlib import Path from sqlalchemy import text, inspect from sqlalchemy.ext.asyncio import AsyncSession from alembic.config import Config from alembic import command from alembic.runtime.migration import MigrationContext from alembic.script import ScriptDirectory from .base import DatabaseManager, Base logger = structlog.get_logger() class DatabaseInitManager: """ Manages database initialization with support for: - Automatic table creation from SQLAlchemy models - Alembic integration and version management - Development reset capabilities """ def __init__( self, database_manager: DatabaseManager, service_name: str, alembic_ini_path: Optional[str] = None, models_module: Optional[str] = None, force_recreate: bool = False ): self.database_manager = database_manager self.service_name = service_name self.alembic_ini_path = alembic_ini_path self.models_module = models_module self.force_recreate = force_recreate self.logger = logger.bind(service=service_name) async def initialize_database(self) -> Dict[str, Any]: """ Main initialization method that handles all scenarios: 1. Check if database exists and has tables 2. Create tables if needed (first-time deployment) 3. Handle Alembic version management 4. Support development reset scenarios """ self.logger.info("Starting database initialization") try: # Check current database state db_state = await self._check_database_state() self.logger.info("Database state checked", state=db_state) # Handle different initialization scenarios if self.force_recreate: result = await self._handle_force_recreate() elif db_state["is_empty"]: result = await self._handle_first_time_deployment() elif db_state["has_alembic_version"]: result = await self._handle_existing_database_with_alembic() else: result = await self._handle_existing_database_without_alembic() self.logger.info("Database initialization completed", result=result) return result except Exception as e: self.logger.error("Database initialization failed", error=str(e)) raise async def _check_database_state(self) -> Dict[str, Any]: """Check the current state of the database""" state = { "is_empty": False, "has_alembic_version": False, "existing_tables": [], "alembic_version": None } try: async with self.database_manager.get_session() as session: # Check if database has any tables inspector = await self._get_inspector(session) existing_tables = await self._get_existing_tables(inspector) state["existing_tables"] = existing_tables state["is_empty"] = len(existing_tables) == 0 # Check if alembic_version table exists and has version if "alembic_version" in existing_tables: state["has_alembic_version"] = True result = await session.execute(text("SELECT version_num FROM alembic_version")) version = result.scalar() state["alembic_version"] = version except Exception as e: self.logger.warning("Error checking database state", error=str(e)) state["is_empty"] = True return state async def _handle_first_time_deployment(self) -> Dict[str, Any]: """Handle first-time deployment: create tables and set up Alembic""" self.logger.info("Handling first-time deployment") # Import models to ensure they're registered with Base if self.models_module: await self._import_models() # Create all tables from SQLAlchemy models await self._create_tables_from_models() # Initialize Alembic and stamp with the latest version if self.alembic_ini_path and os.path.exists(self.alembic_ini_path): await self._initialize_alembic() return { "action": "first_time_deployment", "tables_created": True, "alembic_initialized": True, "message": "Database initialized for first-time deployment" } async def _handle_force_recreate(self) -> Dict[str, Any]: """Handle development scenario: drop everything and recreate""" self.logger.info("Handling force recreate (development mode)") # Drop all tables await self._drop_all_tables() # Create tables from models if self.models_module: await self._import_models() await self._create_tables_from_models() # Re-initialize Alembic if self.alembic_ini_path and os.path.exists(self.alembic_ini_path): await self._initialize_alembic() return { "action": "force_recreate", "tables_dropped": True, "tables_created": True, "alembic_reinitialized": True, "message": "Database recreated from scratch (development mode)" } async def _handle_existing_database_with_alembic(self) -> Dict[str, Any]: """Handle existing database with Alembic version management""" self.logger.info("Handling existing database with Alembic") # Run pending migrations if self.alembic_ini_path and os.path.exists(self.alembic_ini_path): await self._run_migrations() return { "action": "existing_with_alembic", "migrations_run": True, "message": "Existing database updated with pending migrations" } async def _handle_existing_database_without_alembic(self) -> Dict[str, Any]: """Handle existing database without Alembic (legacy scenario)""" self.logger.info("Handling existing database without Alembic") # Initialize Alembic on existing database if self.alembic_ini_path and os.path.exists(self.alembic_ini_path): await self._initialize_alembic_on_existing() return { "action": "existing_without_alembic", "alembic_initialized": True, "message": "Alembic initialized on existing database" } async def _import_models(self): """Import models module to ensure models are registered with Base""" try: import importlib import sys from pathlib import Path # Add project root to Python path if not already there project_root = Path(__file__).parent.parent.parent if str(project_root) not in sys.path: sys.path.insert(0, str(project_root)) # Try to import the models module try: importlib.import_module(self.models_module) self.logger.info("Models imported successfully", module=self.models_module) except ImportError as import_error: # Try alternative import path (from app directly) alt_module = f"app.models" self.logger.warning("Primary import failed, trying alternative", primary=self.models_module, alternative=alt_module, error=str(import_error)) # Change working directory to service directory import os original_cwd = os.getcwd() service_dir = project_root / "services" / self.service_name if service_dir.exists(): os.chdir(service_dir) try: importlib.import_module(alt_module) self.logger.info("Models imported with alternative path", module=alt_module) finally: os.chdir(original_cwd) else: raise import_error except Exception as e: self.logger.error("Failed to import models", module=self.models_module, error=str(e)) # Don't raise for now, continue without models import self.logger.warning("Continuing without models import - tables may not be created") async def _create_tables_from_models(self): """Create all tables from registered SQLAlchemy models""" try: async with self.database_manager.async_engine.begin() as conn: await conn.run_sync(Base.metadata.create_all) self.logger.info("Tables created from SQLAlchemy models") except Exception as e: self.logger.error("Failed to create tables from models", error=str(e)) raise async def _drop_all_tables(self): """Drop all tables (for development reset)""" try: async with self.database_manager.async_engine.begin() as conn: await conn.run_sync(Base.metadata.drop_all) self.logger.info("All tables dropped") except Exception as e: self.logger.error("Failed to drop tables", error=str(e)) raise async def _initialize_alembic(self): """Initialize Alembic and stamp with latest version""" try: def run_alembic_init(): # Create Alembic config alembic_cfg = Config(self.alembic_ini_path) # Get the latest revision script_dir = ScriptDirectory.from_config(alembic_cfg) latest_revision = script_dir.get_current_head() if latest_revision: # Stamp the database with the latest revision command.stamp(alembic_cfg, latest_revision) return latest_revision return None # Run Alembic operations in async executor latest_revision = await asyncio.get_event_loop().run_in_executor(None, run_alembic_init) if latest_revision: self.logger.info("Alembic initialized and stamped", revision=latest_revision) else: self.logger.warning("No Alembic revisions found") except Exception as e: self.logger.error("Failed to initialize Alembic", error=str(e)) raise async def _initialize_alembic_on_existing(self): """Initialize Alembic on an existing database""" try: def run_alembic_stamp(): alembic_cfg = Config(self.alembic_ini_path) script_dir = ScriptDirectory.from_config(alembic_cfg) latest_revision = script_dir.get_current_head() if latest_revision: command.stamp(alembic_cfg, latest_revision) return latest_revision return None # Run Alembic operations in async executor latest_revision = await asyncio.get_event_loop().run_in_executor(None, run_alembic_stamp) if latest_revision: self.logger.info("Alembic initialized on existing database", revision=latest_revision) except Exception as e: self.logger.error("Failed to initialize Alembic on existing database", error=str(e)) raise async def _run_migrations(self): """Run pending Alembic migrations""" try: def run_alembic_upgrade(): alembic_cfg = Config(self.alembic_ini_path) command.upgrade(alembic_cfg, "head") # Run Alembic operations in async executor await asyncio.get_event_loop().run_in_executor(None, run_alembic_upgrade) self.logger.info("Alembic migrations completed") except Exception as e: self.logger.error("Failed to run migrations", error=str(e)) raise async def _get_inspector(self, session: AsyncSession): """Get SQLAlchemy inspector for the current connection""" def get_inspector_sync(connection): return inspect(connection) connection = await session.connection() return await connection.run_sync(get_inspector_sync) async def _get_existing_tables(self, inspector) -> List[str]: """Get list of existing tables in the database""" def get_tables_sync(connection): insp = inspect(connection) return insp.get_table_names() async with self.database_manager.get_session() as session: connection = await session.connection() return await connection.run_sync(get_tables_sync) def create_init_manager( database_manager: DatabaseManager, service_name: str, service_path: Optional[str] = None, force_recreate: bool = False ) -> DatabaseInitManager: """ Factory function to create a DatabaseInitManager with auto-detected paths Args: database_manager: DatabaseManager instance service_name: Name of the service service_path: Path to service directory (auto-detected if None) force_recreate: Whether to force recreate tables (development mode) """ # Auto-detect paths if not provided if service_path is None: # Try Docker container path first (service files at root level) if os.path.exists("alembic.ini"): service_path = "." else: # Fallback to development path service_path = f"services/{service_name}" # Set up paths based on environment if service_path == ".": # Docker container environment alembic_ini_path = "alembic.ini" models_module = "app.models" else: # Development environment alembic_ini_path = f"{service_path}/alembic.ini" models_module = f"services.{service_name}.app.models" # Check if paths exist if not os.path.exists(alembic_ini_path): logger.warning("Alembic config not found", path=alembic_ini_path) alembic_ini_path = None return DatabaseInitManager( database_manager=database_manager, service_name=service_name, alembic_ini_path=alembic_ini_path, models_module=models_module, force_recreate=force_recreate ) async def initialize_service_database( database_manager: DatabaseManager, service_name: str, force_recreate: bool = False ) -> Dict[str, Any]: """ Convenience function for service database initialization Args: database_manager: DatabaseManager instance service_name: Name of the service force_recreate: Whether to force recreate (development mode) Returns: Dict with initialization results """ init_manager = create_init_manager( database_manager=database_manager, service_name=service_name, force_recreate=force_recreate ) return await init_manager.initialize_database()