Spaces:
Running
Running
| from contextlib import contextmanager | |
| from typing import Any | |
| from sqlalchemy import create_engine, text | |
| from sqlalchemy.orm import Session, sessionmaker | |
| from .config import get_settings | |
| settings = get_settings() | |
| engine = None | |
| SessionLocal = None | |
| database_url = settings.database_url | |
| if database_url: | |
| if database_url.startswith("postgresql://"): | |
| database_url = database_url.replace("postgresql://", "postgresql+psycopg://", 1) | |
| engine = create_engine( | |
| database_url, | |
| future=True, | |
| pool_pre_ping=True, | |
| ) | |
| SessionLocal = sessionmaker(bind=engine, autoflush=False, autocommit=False, future=True) | |
| else: | |
| print("WARNING: DATABASE_URL is empty. Database features will be disabled.") | |
| def get_session() -> Session: | |
| if not SessionLocal: | |
| raise ValueError("Database connection not established. Please check your DATABASE_URL.") | |
| session = SessionLocal() | |
| try: | |
| yield session | |
| finally: | |
| session.close() | |
| def run_sql(session: Session, sql: str, params: dict[str, Any] | None = None) -> list[dict[str, Any]]: | |
| result = session.execute(text(sql), params or {}) | |
| return [dict(row) for row in result.mappings().all()] | |
| def init_app_tables() -> None: | |
| if not engine: | |
| return | |
| create_chat_history_sql = """ | |
| CREATE TABLE IF NOT EXISTS chat_messages ( | |
| id BIGSERIAL PRIMARY KEY, | |
| session_id TEXT NOT NULL, | |
| role TEXT NOT NULL, | |
| content TEXT NOT NULL, | |
| created_at TIMESTAMPTZ NOT NULL DEFAULT NOW() | |
| ); | |
| """ | |
| with engine.begin() as conn: | |
| conn.execute(text(create_chat_history_sql)) | |