Spaces:
Sleeping
Sleeping
File size: 1,673 Bytes
db284b9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 | """SQLModel engine and session dependency for the Postgres history store."""
from __future__ import annotations
import logging
import os
from collections.abc import Iterator
from pathlib import Path
from sqlmodel import Session, SQLModel, create_engine
logger = logging.getLogger("tracerag.database")
# load .env so APP_DATABASE_URL is available; no-op if absent
try:
from dotenv import load_dotenv
load_dotenv(Path(__file__).resolve().parent / ".env")
except ImportError: # pragma: no cover
pass
_engine = None # lazily created
def get_engine():
"""Return the process-wide SQLModel engine, creating it on first use."""
global _engine
if _engine is None:
url = os.getenv("APP_DATABASE_URL")
if not url:
raise RuntimeError(
"APP_DATABASE_URL is not set. Add your Neon connection string "
"(postgresql://USER:PASSWORD@HOST/DB?sslmode=require) to .env."
)
_engine = create_engine(url, echo=False, pool_pre_ping=True)
return _engine
def init_db() -> None:
"""Create the history tables if they don't exist (idempotent)."""
# import models so their tables register before create
from models import ChatSession, TraceLog, User # noqa: F401
from sqlalchemy.orm import configure_mappers
# validate relationships at startup, not on first query
configure_mappers()
SQLModel.metadata.create_all(get_engine())
logger.info("Postgres history schema ready.")
def get_session() -> Iterator[Session]:
"""FastAPI dependency yielding a session, closed afterwards."""
with Session(get_engine()) as session:
yield session
|