homelable/backend/app/db/database.py

import logging
import shutil
from collections.abc import AsyncGenerator
from contextlib import suppress
from pathlib import Path

from sqlalchemy.exc import OperationalError
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
from sqlalchemy.orm import DeclarativeBase

from app.core.config import APP_VERSION, settings

logger = logging.getLogger(__name__)

# Ensure the data directory exists before SQLite tries to open the file
Path(settings.sqlite_path).parent.mkdir(parents=True, exist_ok=True)

engine = create_async_engine(
    f"sqlite+aiosqlite:///{settings.sqlite_path}",
    echo=False,
)

AsyncSessionLocal = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)


class Base(DeclarativeBase):
    pass


def _backup_db() -> None:
    db_path = Path(settings.sqlite_path)
    if not db_path.exists():
        return
    backup_path = db_path.with_suffix(f".db.back-{APP_VERSION}")
    if backup_path.exists():
        return
    try:
        shutil.copy2(db_path, backup_path)
        logger.info("DB backup created: %s", backup_path.name)
    except OSError:
        logger.warning("Could not create DB backup at %s", backup_path)


async def init_db() -> None:
    _backup_db()
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)
        # Add columns introduced after initial schema (idempotent)
        with suppress(OperationalError):
            await conn.exec_driver_sql("ALTER TABLE nodes ADD COLUMN container_mode BOOLEAN NOT NULL DEFAULT 0")
        with suppress(OperationalError):
            await conn.exec_driver_sql("ALTER TABLE nodes ADD COLUMN custom_colors JSON")
        with suppress(OperationalError):
            await conn.exec_driver_sql("ALTER TABLE edges ADD COLUMN custom_color TEXT")
        with suppress(OperationalError):
            await conn.exec_driver_sql("ALTER TABLE edges ADD COLUMN path_style TEXT")
        with suppress(OperationalError):
            await conn.exec_driver_sql("ALTER TABLE nodes ADD COLUMN custom_icon TEXT")
        with suppress(OperationalError):
            await conn.exec_driver_sql("ALTER TABLE edges ADD COLUMN source_handle TEXT")
        with suppress(OperationalError):
            await conn.exec_driver_sql("ALTER TABLE edges ADD COLUMN target_handle TEXT")
        with suppress(OperationalError):
            await conn.exec_driver_sql("ALTER TABLE edges ADD COLUMN animated BOOLEAN NOT NULL DEFAULT 0")
        with suppress(OperationalError):
            await conn.exec_driver_sql("ALTER TABLE nodes ADD COLUMN cpu_count INTEGER")
        with suppress(OperationalError):
            await conn.exec_driver_sql("ALTER TABLE nodes ADD COLUMN cpu_model TEXT")
        with suppress(OperationalError):
            await conn.exec_driver_sql("ALTER TABLE nodes ADD COLUMN ram_gb REAL")
        with suppress(OperationalError):
            await conn.exec_driver_sql("ALTER TABLE nodes ADD COLUMN disk_gb REAL")
        with suppress(OperationalError):
            await conn.exec_driver_sql("ALTER TABLE nodes ADD COLUMN show_hardware BOOLEAN NOT NULL DEFAULT 0")
        with suppress(OperationalError):
            await conn.exec_driver_sql("ALTER TABLE nodes ADD COLUMN width REAL")
        with suppress(OperationalError):
            await conn.exec_driver_sql("ALTER TABLE nodes ADD COLUMN height REAL")
        with suppress(OperationalError):
            await conn.exec_driver_sql("ALTER TABLE nodes ADD COLUMN bottom_handles INTEGER NOT NULL DEFAULT 1")
        with suppress(OperationalError):
            await conn.exec_driver_sql("ALTER TABLE pending_devices ADD COLUMN discovery_source TEXT")
        with suppress(OperationalError):
            await conn.exec_driver_sql("ALTER TABLE edges ADD COLUMN waypoints JSON")
        with suppress(OperationalError):
            await conn.exec_driver_sql("ALTER TABLE nodes ADD COLUMN properties JSON")
        # Migrate hardware columns → properties JSON (idempotent: only runs on nodes where properties IS NULL)
        with suppress(OperationalError):
            rows = await conn.exec_driver_sql(
                "SELECT id, cpu_model, cpu_count, ram_gb, disk_gb, show_hardware "
                "FROM nodes WHERE properties IS NULL"
            )
            for row in rows.fetchall():
                node_id, cpu_model, cpu_count, ram_gb, disk_gb, show_hardware = row
                props = []
                visible = bool(show_hardware)
                if cpu_model:
                    props.append({"key": "CPU Model", "value": str(cpu_model), "icon": "Cpu", "visible": visible})
                if cpu_count is not None:
                    props.append({"key": "CPU Cores", "value": str(cpu_count), "icon": "Cpu", "visible": visible})
                if ram_gb is not None:
                    props.append({"key": "RAM", "value": f"{ram_gb} GB", "icon": "MemoryStick", "visible": visible})
                if disk_gb is not None:
                    props.append({"key": "Disk", "value": f"{disk_gb} GB", "icon": "HardDrive", "visible": visible})
                import json as _json
                await conn.exec_driver_sql(
                    "UPDATE nodes SET properties = ? WHERE id = ?",
                    (_json.dumps(props), node_id),
                )
        # Migrate animated column from boolean (0/1) to string ('none'/'snake')
        with suppress(OperationalError):
            await conn.exec_driver_sql("UPDATE edges SET animated = 'snake' WHERE animated = '1' OR animated = 1")
        with suppress(OperationalError):
            sql = "UPDATE edges SET animated = 'none' WHERE animated = '0' OR animated = 0 OR animated IS NULL"
            await conn.exec_driver_sql(sql)


async def get_db() -> AsyncGenerator[AsyncSession, None]:
    async with AsyncSessionLocal() as session:
        yield session