database.py

"""
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()