instrument-data-mcp/instrument_config_mcp/db.py

from __future__ import annotations

import os
from typing import Any

from sqlalchemy import Integer, JSON, String, create_engine, select
from sqlalchemy.engine import Engine
from sqlalchemy.orm import DeclarativeBase, Mapped, Session, mapped_column


def normalize_database_url(raw: str) -> str:
    return raw.strip()


class Base(DeclarativeBase):
    pass


class SysConfig(Base):
    __tablename__ = "sys_config"

    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
    key: Mapped[str] = mapped_column(String(128), unique=True, nullable=False)
    value: Mapped[Any] = mapped_column(JSON, nullable=False)


_ENGINE: Engine | None = None
_ENGINE_URL: str | None = None


def database_url() -> str:
    raw = str(os.getenv("DATABASE_URL") or "sqlite:///llm_proxy.db").strip()
    return normalize_database_url(raw)


def sql_engine() -> Engine:
    global _ENGINE, _ENGINE_URL

    effective_url = database_url()
    if _ENGINE is not None and _ENGINE_URL == effective_url:
        return _ENGINE

    engine_options: dict[str, Any] = {}
    connect_args: dict[str, Any] = {}
    if effective_url.startswith("sqlite:"):
        connect_args["check_same_thread"] = False
    if connect_args:
        engine_options["connect_args"] = connect_args

    _ENGINE = create_engine(effective_url, future=True, **engine_options)
    _ENGINE_URL = effective_url
    return _ENGINE


def read_sys_config_value(config_key: str) -> Any | None:
    try:
        with Session(sql_engine()) as session:
            return session.scalar(select(SysConfig.value).where(SysConfig.key == config_key))
    except Exception as exc:
        raise ValueError(f"failed to read sys_config key={config_key}: {exc}") from exc


def write_sys_config_value(config_key: str, value: Any) -> None:
    try:
        with Session(sql_engine()) as session:
            row = session.scalar(select(SysConfig).where(SysConfig.key == config_key))
            if row is None:
                session.add(SysConfig(key=config_key, value=value))
            else:
                row.value = value
            session.commit()
    except Exception as exc:
        raise ValueError(f"failed to write sys_config key={config_key}: {exc}") from exc