persistence/db.py

"""Database primitives for persistence models.

Tables are created only by scripts/init_db.py. Importing this module must not
modify the database.
"""

from __future__ import annotations

import base64
import os
from pathlib import Path
from typing import Iterator

from sqlalchemy import create_engine
from sqlalchemy.engine import make_url
from sqlalchemy.orm import DeclarativeBase, Session, sessionmaker
from dotenv import load_dotenv


load_dotenv()

DATABASE_URL = os.getenv("DATABASE_URL", "sqlite:///./rag_persistence.db")
DB_CA_CERT_B64 = os.getenv("DB_CA_CERT_B64")


def _bootstrap_db_ca_cert_from_b64() -> str | None:
    """Materialize the HF-provided CA certificate secret without logging it."""
    if not DB_CA_CERT_B64:
        return os.getenv("DB_CA_CERT")

    ca_path = Path("/tmp/aiven-ca.pem")
    try:
        ca_path.parent.mkdir(parents=True, exist_ok=True)
        ca_path.write_bytes(base64.b64decode(DB_CA_CERT_B64, validate=True))
        ca_path.chmod(0o600)
    except Exception as exc:
        raise RuntimeError("DB_CA_CERT_B64 is configured, but the CA certificate could not be prepared.") from exc

    os.environ["DB_CA_CERT"] = str(ca_path)
    return str(ca_path)


DB_CA_CERT = _bootstrap_db_ca_cert_from_b64()


def _is_mysql_url(database_url: str) -> bool:
    return make_url(database_url).drivername.startswith("mysql")


def _build_connect_args(database_url: str) -> dict:
    if database_url.startswith("sqlite:"):
        return {"check_same_thread": False}

    if not _is_mysql_url(database_url):
        return {}

    if not DB_CA_CERT:
        raise RuntimeError("Managed MySQL requires DB_CA_CERT or DB_CA_CERT_B64 for TLS certificate verification.")

    ca_path = Path(DB_CA_CERT)
    if not ca_path.exists():
        raise RuntimeError("DB_CA_CERT is configured, but the certificate file does not exist.")

    return {"ssl": {"ca": str(ca_path)}}


_connect_args = _build_connect_args(DATABASE_URL)
engine = create_engine(DATABASE_URL, connect_args=_connect_args, future=True)
SessionLocal = sessionmaker(bind=engine, autoflush=False, autocommit=False, future=True)


class Base(DeclarativeBase):
    pass


def get_session() -> Iterator[Session]:
    """Yield a SQLAlchemy session for future API dependencies."""
    session = SessionLocal()
    try:
        yield session
    finally:
        session.close()