Spaces:
Running on Zero
Running on Zero
File size: 3,318 Bytes
599eb60 | 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 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 | """
memory/db.py — Connection/session management for the memory Postgres instance.
This is COMPLETELY SEPARATE from the mock_infra mock_db. It connects to the
real Postgres+pgvector instance defined in POSTGRES_URL. The environment
FSM reset() never touches this connection or any tables it manages.
"""
from __future__ import annotations
from contextlib import asynccontextmanager
from typing import AsyncGenerator
import structlog
from sqlalchemy.ext.asyncio import (
AsyncEngine,
AsyncSession,
async_sessionmaker,
create_async_engine,
)
from sqlalchemy.orm import DeclarativeBase
from config import settings
log = structlog.get_logger(__name__)
# ── SQLAlchemy base ───────────────────────────────────────────────────────────
class Base(DeclarativeBase):
pass
# ── Engine + session factory (module-level singletons) ───────────────────────
_engine: AsyncEngine | None = None
_session_factory: async_sessionmaker[AsyncSession] | None = None
def get_engine() -> AsyncEngine:
global _engine
if _engine is None:
_engine = create_async_engine(
settings.postgres_url,
pool_size=10,
max_overflow=20,
pool_pre_ping=True,
echo=False,
)
return _engine
def get_session_factory() -> async_sessionmaker[AsyncSession]:
global _session_factory
if _session_factory is None:
_session_factory = async_sessionmaker(
bind=get_engine(),
expire_on_commit=False,
class_=AsyncSession,
)
return _session_factory
@asynccontextmanager
async def get_db_session() -> AsyncGenerator[AsyncSession, None]:
"""Async context manager for a DB session. Handles commit/rollback."""
factory = get_session_factory()
async with factory() as session:
try:
yield session
await session.commit()
except Exception:
await session.rollback()
raise
# ── Schema init ───────────────────────────────────────────────────────────────
async def init_db() -> None:
"""
Create all tables if they don't exist. In production, init.sql is run by
docker-entrypoint-initdb.d; this is a fallback for local dev without Docker.
"""
from sqlalchemy import text
engine = get_engine()
async with engine.begin() as conn:
# Ensure pgvector extension is available
await conn.execute(text("CREATE EXTENSION IF NOT EXISTS vector"))
# Create tables via SQLAlchemy metadata
await conn.run_sync(Base.metadata.create_all)
log.info("memory.db.init_complete", url=settings.postgres_url)
async def health_check() -> bool:
"""Returns True if the memory DB is reachable."""
from sqlalchemy import text
try:
async with get_db_session() as session:
await session.execute(text("SELECT 1"))
return True
except Exception as exc:
log.error("memory.db.health_check_failed", error=str(exc))
return False
|