1  """Migration runner for application startup."""
 2  import logging
 3  from alembic import command
 4  from alembic.config import Config
 5  import os
 6  import time
 7  from sqlalchemy import create_engine, text
 8  from sqlalchemy.exc import OperationalError
 9  
10  logger = logging.getLogger(__name__)
11  
12  
13  def wait_for_db(max_retries=30, retry_delay=2):
14      """Wait for database to be ready before running migrations."""
15      database_url = os.getenv(
16          "DATABASE_URL",
17          "postgresql://smeops:change-me@postgres:5432/smeops"
18      )
19      
20      for attempt in range(max_retries):
21          try:
22              engine = create_engine(database_url, pool_pre_ping=True)
23              with engine.connect() as conn:
24                  conn.execute(text("SELECT 1"))
25              logger.info("Database is ready")
26              return True
27          except OperationalError as e:
28              if attempt < max_retries - 1:
29                  logger.warning(f"Database not ready (attempt {attempt + 1}/{max_retries}): {e}")
30                  time.sleep(retry_delay)
31              else:
32                  logger.error(f"Database not ready after {max_retries} attempts")
33                  raise
34      return False
35  
36  
37  def run_migrations():
38      """Run Alembic migrations on application startup."""
39      try:
40          # Wait for database to be ready
41          wait_for_db()
42          
43          # Get path to alembic.ini
44          alembic_cfg = Config(os.path.join(os.path.dirname(os.path.dirname(__file__)), "alembic.ini"))
45          
46          # Override sqlalchemy.url with environment variable
47          alembic_cfg.set_main_option("sqlalchemy.url", os.getenv(
48              "DATABASE_URL",
49              "postgresql://smeops:change-me@postgres:5432/smeops"
50          ))
51          
52          logger.info("Running database migrations...")
53          command.upgrade(alembic_cfg, "head")
54          logger.info("Database migrations completed successfully")
55      except Exception as e:
56          logger.error(f"Failed to run database migrations: {e}")
57          raise