""" Database configuration and session management for Silver Table Assistant backend. Uses SQLModel with AsyncEngine for Supabase PostgreSQL with pgvector support. """ import os from typing import AsyncGenerator from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, create_async_engine from sqlalchemy.orm import sessionmaker from sqlmodel import SQLModel # Environment variables SUPABASE_URL = os.getenv("SUPABASE_URL") SUPABASE_SERVICE_ROLE_KEY = os.getenv("SUPABASE_SERVICE_ROLE_KEY") DATABASE_URL = os.getenv("DATABASE_URL") # Validate required environment variables if not SUPABASE_URL: raise ValueError("SUPABASE_URL environment variable is required") if not SUPABASE_SERVICE_ROLE_KEY: raise ValueError("SUPABASE_SERVICE_ROLE_KEY environment variable is required") # Construct database URL for async Supabase connection # Using service role key for server-side operations if DATABASE_URL: # Use provided DATABASE_URL if available ASYNC_DATABASE_URL = DATABASE_URL.replace("postgresql://", "postgresql+asyncpg://") else: # Construct from Supabase URL # Format: postgresql+asyncpg://postgres:[password]@db.[project-ref].supabase.co:5432/postgres base_url = SUPABASE_URL.replace("https://", "").replace("http://", "") project_ref = base_url.split(".")[0] ASYNC_DATABASE_URL = f"postgresql+asyncpg://postgres:{SUPABASE_SERVICE_ROLE_KEY}@db.{project_ref}.supabase.co:5432/postgres" from sqlalchemy.pool import NullPool # Create async engine with pgvector support # Using NullPool because we are connecting through PGBouncer in transaction mode (port 6543) engine: AsyncEngine = create_async_engine( ASYNC_DATABASE_URL, echo=False, # Set to True for SQL query logging poolclass=NullPool, # Additional settings for Supabase/PGBouncer connect_args={ "statement_cache_size": 0, # Disable prepared statement cache for PGBouncer "server_settings": { "jit": "off", # Disable JIT for pgvector compatibility } }, ) # Create session factory SessionLocal = sessionmaker( bind=engine, class_=AsyncSession, expire_on_commit=False, autocommit=False, autoflush=False, ) async def create_db_and_tables() -> None: """Create database tables if they don't exist.""" async with engine.begin() as conn: await conn.run_sync(SQLModel.metadata.create_all) async def get_session() -> AsyncGenerator[AsyncSession, None]: """ Dependency function to get database session. Should be used with FastAPI's Depends() decorator. """ async with SessionLocal() as session: try: yield session finally: await session.close() def get_db_session() -> AsyncSession: """ Get a database session for use in functions. Remember to close the session after use. """ return SessionLocal() async def close_db_connections() -> None: """Close all database connections.""" await engine.dispose()