CA/backend/database.py

"""
Database connection and session management for PostGIS
"""
import logging
import asyncpg
from typing import Optional
from contextlib import asynccontextmanager
from pydantic_settings import BaseSettings

logger = logging.getLogger("cbpoa.database")


class Settings(BaseSettings):
    """Database settings from environment variables"""
    POSTGRES_HOST: str = "localhost"
    POSTGRES_PORT: int = 5432
    POSTGRES_USER: str = ""
    POSTGRES_PASSWORD: str = ""
    POSTGRES_DB: str = ""

    class Config:
        env_file = ".env"


settings = Settings()

if not settings.POSTGRES_USER or not settings.POSTGRES_PASSWORD or not settings.POSTGRES_DB:
    raise RuntimeError(
        "Missing required database environment variables: POSTGRES_USER, POSTGRES_PASSWORD, POSTGRES_DB. "
        "Create a .env file or set them in your environment."
    )


class Database:
    """Async database connection pool manager"""

    def __init__(self):
        self.pool: Optional[asyncpg.Pool] = None

    async def connect(self):
        """Initialize database connection pool"""
        if self.pool is None:
            dsn = f"postgresql://{settings.POSTGRES_USER}:{settings.POSTGRES_PASSWORD}@{settings.POSTGRES_HOST}:{settings.POSTGRES_PORT}/{settings.POSTGRES_DB}"
            self.pool = await asyncpg.create_pool(
                dsn=dsn,
                min_size=5,
                max_size=20,
                command_timeout=60
            )
            logger.info("Database connection pool created successfully")

    async def disconnect(self):
        """Close database connection pool"""
        if self.pool:
            await self.pool.close()
            self.pool = None
            logger.info("Database connection pool closed")

    @asynccontextmanager
    async def get_connection(self):
        """Get a connection from the pool"""
        if self.pool is None:
            await self.connect()
        assert self.pool is not None
        async with self.pool.acquire() as connection:
            yield connection

    @asynccontextmanager
    async def get_transaction(self):
        """Get a transaction context"""
        if self.pool is None:
            await self.connect()
        assert self.pool is not None
        async with self.pool.acquire() as connection:
            async with connection.transaction():
                yield connection


# Global database instance
db = Database()


async def init_db():
    """Initialize database on startup - graceful degradation if unavailable"""
    try:
        await db.connect()
    except Exception as e:
        logger.warning("Database not available (%s). Running in demo mode.", e)
        logger.warning("Set POSTGRES_HOST/POSTGRES_USER/POSTGRES_PASSWORD environment variables for database access.")


async def close_db():
    """Close database on shutdown"""
    await db.disconnect()