yakpanel-core/YakPanel-server/backend/app/core/database.py

"""YakPanel - Database configuration"""
from pathlib import Path

from sqlalchemy.engine.url import make_url
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
from sqlalchemy.orm import DeclarativeBase

from app.core.config import get_settings


def _resolve_database_url(raw_url: str) -> str:
    """Anchor relative SQLite paths to the backend directory (not process CWD)."""
    url = make_url(raw_url)
    if not url.drivername.startswith("sqlite"):
        return raw_url
    db_path = url.database
    if db_path is None or db_path.startswith(":"):
        return raw_url
    path = Path(db_path)
    backend_dir = Path(__file__).resolve().parent.parent
    if not path.is_absolute():
        path = (backend_dir / path).resolve()
    else:
        path = path.resolve()
    path.parent.mkdir(parents=True, exist_ok=True)
    # POSIX path for SQLite URL (backslashes break aiosqlite on unexpected setups)
    posix = path.as_posix()
    return str(url.set(database=posix))


settings = get_settings()
DATABASE_URL = _resolve_database_url(settings.database_url)

_engine_kw = {"echo": settings.debug}
if make_url(DATABASE_URL).drivername.startswith("sqlite"):
    _engine_kw["connect_args"] = {"timeout": 30}

engine = create_async_engine(
    DATABASE_URL,
    **_engine_kw,
)

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


class Base(DeclarativeBase):
    """SQLAlchemy declarative base"""

    pass


async def get_db():
    """Dependency for async database sessions"""
    async with AsyncSessionLocal() as session:
        try:
            yield session
            await session.commit()
        except Exception:
            await session.rollback()
            raise
        finally:
            await session.close()


def _run_migrations(conn):
    """Add new columns to existing tables (SQLite)."""
    import sqlalchemy

    try:
        r = conn.execute(sqlalchemy.text("PRAGMA table_info(sites)"))
        cols = [row[1] for row in r.fetchall()]
        if "php_version" not in cols:
            conn.execute(
                sqlalchemy.text(
                    "ALTER TABLE sites ADD COLUMN php_version VARCHAR(16) DEFAULT '74'"
                )
            )
        if "force_https" not in cols:
            conn.execute(
                sqlalchemy.text(
                    "ALTER TABLE sites ADD COLUMN force_https INTEGER DEFAULT 0"
                )
            )
    except Exception:
        pass
    # Create backup_plans if not exists (create_all handles new installs)
    try:
        conn.execute(sqlalchemy.text("""
            CREATE TABLE IF NOT EXISTS backup_plans (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                name VARCHAR(128) NOT NULL,
                plan_type VARCHAR(32) NOT NULL,
                target_id INTEGER NOT NULL,
                schedule VARCHAR(64) NOT NULL,
                enabled BOOLEAN DEFAULT 1
            )
        """))
    except Exception:
        pass
    # Create custom_plugins if not exists
    try:
        conn.execute(sqlalchemy.text("""
            CREATE TABLE IF NOT EXISTS custom_plugins (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                plugin_id VARCHAR(64) UNIQUE NOT NULL,
                name VARCHAR(128) NOT NULL,
                version VARCHAR(32) DEFAULT '1.0',
                desc VARCHAR(512) DEFAULT '',
                source_url VARCHAR(512) DEFAULT '',
                enabled BOOLEAN DEFAULT 1
            )
        """))
    except Exception:
        pass


async def init_db():
    """Initialize database tables"""
    import app.models  # noqa: F401 - register all models with Base.metadata

    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)
        if "sqlite" in str(engine.url):
            await conn.run_sync(_run_migrations)