95 lines
2.8 KiB
Python
95 lines
2.8 KiB
Python
|
"""
|
||
|
|
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()
|
||
|
|
|
||
|
|
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()
|
||
|
|
|
||
|
|
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()
|