Alembic experiments

src/alpine_bits_python/db_setup.py

@@ -0,0 +1,86 @@
+"""Database setup and initialization.
+
+This module handles all database setup tasks that should run once at startup,
+before the application starts accepting requests. It includes:
+- Schema migrations via Alembic
+- One-time data cleanup/backfill tasks (e.g., hashing existing customers)
+"""
+
+import asyncio
+from typing import Any
+
+from sqlalchemy.ext.asyncio import AsyncEngine, async_sessionmaker
+
+from .customer_service import CustomerService
+from .db import create_database_engine
+from .logging_config import get_logger
+
+_LOGGER = get_logger(__name__)
+
+
+async def setup_database(config: dict[str, Any] | None = None) -> tuple[AsyncEngine, async_sessionmaker]:
+    """Set up the database and prepare for application use.
+
+    This function should be called once at application startup, after
+    migrations have been run but before the app starts accepting requests. It:
+    1. Creates the async engine
+    2. Creates the sessionmaker
+    3. Performs one-time startup tasks (e.g., hashing existing customers)
+
+    NOTE: Database migrations should be run BEFORE calling this function,
+    typically using `uv run alembic upgrade head` or via run_migrations.py.
+
+    Args:
+        config: Application configuration dictionary
+
+    Returns:
+        Tuple of (engine, async_sessionmaker) for use in the application
+
+    Raises:
+        Any database-related exceptions that occur during setup
+    """
+    _LOGGER.info("Starting database setup...")
+
+    # Create database engine
+    engine = create_database_engine(config=config, echo=False)
+
+    try:
+        # Create sessionmaker for the application to use
+        AsyncSessionLocal = async_sessionmaker(engine, expire_on_commit=False)
+
+        # Perform startup tasks (NOT migrations)
+        _LOGGER.info("Running startup tasks...")
+        await run_startup_tasks(AsyncSessionLocal, config)
+        _LOGGER.info("Startup tasks completed successfully")
+
+        _LOGGER.info("Database setup completed successfully")
+        return engine, AsyncSessionLocal
+
+    except Exception as e:
+        _LOGGER.exception("Database setup failed: %s", e)
+        await engine.dispose()
+        raise
+
+
+async def run_startup_tasks(
+    sessionmaker: async_sessionmaker, config: dict[str, Any] | None = None
+) -> None:
+    """Run one-time startup tasks.
+
+    These are tasks that need to run at startup but are NOT schema migrations.
+    Examples: data backfills, hashing existing records, etc.
+
+    Args:
+        sessionmaker: SQLAlchemy async sessionmaker
+        config: Application configuration dictionary
+    """
+    # Hash any existing customers that don't have hashed data
+    async with sessionmaker() as session:
+        customer_service = CustomerService(session)
+        hashed_count = await customer_service.hash_existing_customers()
+        if hashed_count > 0:
+            _LOGGER.info(
+                "Backfilled hashed data for %d existing customers", hashed_count
+            )
+        else:
+            _LOGGER.info("All existing customers already have hashed data")