"""Database configuration and session management for Homelab Automation. Uses SQLAlchemy 2.x async engine with SQLite + aiosqlite driver. """ 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, text 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 ( # noqa: F401 — register ALL models with Base.metadata host, host_metrics, task, schedule, schedule_run, log, alert, app_setting, docker_container, docker_image, docker_volume, docker_alert, playbook_lint, favorite_group, favorite_container, user, terminal_session, terminal_command_log, bootstrap_status, container_customization, ) def _to_sync_database_url(db_url: str) -> str: if db_url.startswith("sqlite+aiosqlite:"): return db_url.replace("sqlite+aiosqlite:", "sqlite:") elif db_url.startswith("mysql+aiomysql:"): return db_url.replace("mysql+aiomysql:", "mysql+pymysql:") return db_url 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: # For MySQL, pre-create alembic_version with VARCHAR(255) to prevent # "Data too long for column 'version_num'" error on long revision IDs. # SQLite doesn't enforce VARCHAR length, so this is only needed for MySQL. sync_url = _to_sync_database_url(DATABASE_URL) if "mysql" in sync_url: import sqlalchemy as sa sync_engine = sa.create_engine(sync_url) with sync_engine.connect() as conn: # Case 1: Table doesn't exist conn.execute(sa.text( "CREATE TABLE IF NOT EXISTS alembic_version " "(version_num VARCHAR(255) NOT NULL, " "PRIMARY KEY (version_num))" )) # Case 2: Table exists but column might be too short (VARCHAR(32)) # Descriptive revisions like '0011_add_docker_management_tables' are > 30 chars try: conn.execute(sa.text( "ALTER TABLE alembic_version MODIFY version_num VARCHAR(255) NOT NULL" )) except Exception: pass # Might fail if not MySQL or column already long (SQLite has no MODIFY) conn.commit() sync_engine.dispose() print("[DB] MySQL: alembic_version table ensured with VARCHAR(255)") cfg = Config(str(alembic_ini)) cfg.set_main_option("sqlalchemy.url", sync_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 print(f"[DB] Initializing database with URL: {DATABASE_URL}") await asyncio.to_thread(_run_alembic_upgrade) print("[DB] Ensuring all tables exist (metadata create_all)...") async with engine.begin() as conn: await conn.run_sync(Base.metadata.create_all) print("[DB] Database initialization complete.")