"""Database configuration and session management for Homelab Automation. Uses SQLAlchemy 2.x async engine with SQLite + aiosqlite driver. """ from __future__ import annotations import asyncio import os from pathlib import Path from typing import AsyncGenerator from urllib.parse import urlparse from alembic import command from alembic.config import Config from sqlalchemy import event, MetaData from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker, create_async_engine from sqlalchemy.orm import declarative_base # Naming convention to keep Alembic happy with constraints NAMING_CONVENTION = { "ix": "ix_%(column_0_label)s", "uq": "uq_%(table_name)s_%(column_0_name)s", "ck": "ck_%(table_name)s_%(constraint_name)s", "fk": "fk_%(table_name)s_%(column_0_name)s_%(referred_table_name)s", "pk": "pk_%(table_name)s", } metadata_obj = MetaData(naming_convention=NAMING_CONVENTION) Base = declarative_base(metadata=metadata_obj) # Resolve base path (project root) ROOT_DIR = Path(__file__).resolve().parents[2] DEFAULT_DB_PATH = Path(os.environ.get("DB_PATH") or (ROOT_DIR / "data" / "homelab.db")) DATABASE_URL = os.environ.get("DATABASE_URL", f"sqlite+aiosqlite:///{DEFAULT_DB_PATH}") # Ensure SQLite directory exists even if DATABASE_URL overrides DB_PATH def _ensure_sqlite_dir(db_url: str) -> None: if not db_url.startswith("sqlite"): return # Extraire le chemin après sqlite+aiosqlite:/// # Sur Windows, le chemin peut être C:\... donc on ne peut pas utiliser urlparse prefix = "sqlite+aiosqlite:///" if db_url.startswith(prefix): path_str = db_url[len(prefix):] # Sur Windows, le chemin peut commencer par une lettre de lecteur (C:) db_path = Path(path_str) if db_path.parent and str(db_path.parent) != ".": db_path.parent.mkdir(parents=True, exist_ok=True) DEFAULT_DB_PATH.parent.mkdir(parents=True, exist_ok=True) _ensure_sqlite_dir(DATABASE_URL) def _debug_db_paths() -> None: try: print( "[DB] DATABASE_URL=%s, DEFAULT_DB_PATH=%s, parent_exists=%s, parent=%s" % ( DATABASE_URL, DEFAULT_DB_PATH, DEFAULT_DB_PATH.parent.exists(), DEFAULT_DB_PATH.parent, ) ) except Exception: # Debug logging should never break startup pass _debug_db_paths() engine: AsyncEngine = create_async_engine( DATABASE_URL, echo=False, pool_pre_ping=True, future=True, ) # Ensure SQLite pragmas (WAL + FK) when using SQLite if DATABASE_URL.startswith("sqlite"): @event.listens_for(engine.sync_engine, "connect") def _set_sqlite_pragmas(dbapi_connection, connection_record): # type: ignore[override] cursor = dbapi_connection.cursor() cursor.execute("PRAGMA foreign_keys=ON") # WAL mode can fail on some Docker volume mounts (e.g., NFS, CIFS, overlay issues) # Fall back to DELETE mode if WAL fails try: cursor.execute("PRAGMA journal_mode=WAL") except Exception: # WAL not supported, use DELETE mode instead try: cursor.execute("PRAGMA journal_mode=DELETE") except Exception: pass # Ignore if this also fails cursor.close() async_session_maker = async_sessionmaker( bind=engine, autoflush=False, expire_on_commit=False, class_=AsyncSession, ) async def get_db() -> AsyncGenerator[AsyncSession, None]: """FastAPI dependency that yields an AsyncSession with automatic rollback on error.""" async with async_session_maker() as session: # type: AsyncSession try: yield session except Exception: await session.rollback() raise finally: await session.close() async def init_db() -> None: """Create all tables (mostly for dev/tests; migrations should be handled by Alembic).""" from . import ( host, task, schedule, schedule_run, log, alert, app_setting, docker_container, docker_image, docker_volume, docker_alert, playbook_lint, favorite_group, favorite_container, ) # noqa: F401 def _to_sync_database_url(db_url: str) -> str: return db_url.replace("sqlite+aiosqlite:", "sqlite:") def _run_alembic_upgrade() -> None: # Try multiple locations for alembic.ini (dev vs Docker) alembic_ini_paths = [ ROOT_DIR / "alembic.ini", # Dev: /path/to/project/alembic.ini Path("/alembic.ini"), # Docker: /alembic.ini Path("/app/alembic.ini"), # Docker alternative ] alembic_ini = None for path in alembic_ini_paths: if path.exists(): alembic_ini = path print(f"[DB] Found alembic.ini at: {alembic_ini}") break if not alembic_ini: print(f"[DB] alembic.ini not found in any of: {[str(p) for p in alembic_ini_paths]}") return try: cfg = Config(str(alembic_ini)) cfg.set_main_option("sqlalchemy.url", _to_sync_database_url(DATABASE_URL)) print(f"[DB] Running Alembic upgrade to head...") command.upgrade(cfg, "head") print(f"[DB] Alembic upgrade completed successfully") except Exception as e: print(f"[DB] Alembic upgrade failed: {e}") raise try: await asyncio.to_thread(_run_alembic_upgrade) except Exception as e: print(f"[DB] Exception during Alembic migration: {e}") async with engine.begin() as conn: await conn.run_sync(Base.metadata.create_all)