121 lines
3.7 KiB
Python
121 lines
3.7 KiB
Python
"""
|
|
Database Connection Manager for Arthur Agent.
|
|
|
|
Handles PostgreSQL connections for audit logging with support
|
|
for Docker Secrets and connection pooling.
|
|
"""
|
|
|
|
import os
|
|
import logging
|
|
from typing import Optional
|
|
from contextlib import asynccontextmanager
|
|
|
|
import asyncpg
|
|
from asyncpg import Pool
|
|
|
|
from src.security.secrets_manager import SecretsManager
|
|
|
|
logger = logging.getLogger("ArthurDB")
|
|
|
|
|
|
class DatabaseManager:
|
|
"""
|
|
Manages PostgreSQL database connections.
|
|
|
|
Supports both Docker Secrets (production) and environment
|
|
variables (development) for credential management.
|
|
"""
|
|
|
|
def __init__(self):
|
|
self._pool: Optional[Pool] = None
|
|
self._secrets = SecretsManager()
|
|
|
|
@property
|
|
def dsn(self) -> str:
|
|
"""Build PostgreSQL connection DSN from secrets."""
|
|
return (
|
|
f"postgresql://{self._secrets.get('POSTGRES_USER')}:"
|
|
f"{self._secrets.get('POSTGRES_PASSWORD')}@"
|
|
f"{self._secrets.get('POSTGRES_HOST')}:"
|
|
f"{self._secrets.get('POSTGRES_PORT')}/"
|
|
f"{self._secrets.get('POSTGRES_DB')}"
|
|
)
|
|
|
|
async def connect(self, min_size: int = 2, max_size: int = 10) -> None:
|
|
"""
|
|
Initialize the connection pool.
|
|
|
|
Args:
|
|
min_size: Minimum number of connections in pool
|
|
max_size: Maximum number of connections in pool
|
|
"""
|
|
if self._pool is not None:
|
|
logger.warning("Connection pool already initialized")
|
|
return
|
|
|
|
try:
|
|
self._pool = await asyncpg.create_pool(
|
|
dsn=self.dsn,
|
|
min_size=min_size,
|
|
max_size=max_size,
|
|
command_timeout=60
|
|
)
|
|
logger.info("Database connection pool initialized successfully")
|
|
except Exception as e:
|
|
logger.error(f"Failed to connect to database: {e}")
|
|
raise
|
|
|
|
async def disconnect(self) -> None:
|
|
"""Close the connection pool."""
|
|
if self._pool:
|
|
await self._pool.close()
|
|
self._pool = None
|
|
logger.info("Database connection pool closed")
|
|
|
|
@asynccontextmanager
|
|
async def acquire(self):
|
|
"""
|
|
Acquire a connection from the pool.
|
|
|
|
Usage:
|
|
async with db.acquire() as conn:
|
|
await conn.execute(...)
|
|
"""
|
|
if self._pool is None:
|
|
raise RuntimeError("Database not connected. Call connect() first.")
|
|
|
|
async with self._pool.acquire() as connection:
|
|
yield connection
|
|
|
|
async def execute(self, query: str, *args) -> str:
|
|
"""Execute a query and return status."""
|
|
async with self.acquire() as conn:
|
|
return await conn.execute(query, *args)
|
|
|
|
async def fetch(self, query: str, *args) -> list:
|
|
"""Execute a query and return all rows."""
|
|
async with self.acquire() as conn:
|
|
return await conn.fetch(query, *args)
|
|
|
|
async def fetchrow(self, query: str, *args) -> Optional[asyncpg.Record]:
|
|
"""Execute a query and return a single row."""
|
|
async with self.acquire() as conn:
|
|
return await conn.fetchrow(query, *args)
|
|
|
|
async def fetchval(self, query: str, *args):
|
|
"""Execute a query and return a single value."""
|
|
async with self.acquire() as conn:
|
|
return await conn.fetchval(query, *args)
|
|
|
|
|
|
# Singleton instance
|
|
_db_manager: Optional[DatabaseManager] = None
|
|
|
|
|
|
def get_db_manager() -> DatabaseManager:
|
|
"""Get the global database manager instance."""
|
|
global _db_manager
|
|
if _db_manager is None:
|
|
_db_manager = DatabaseManager()
|
|
return _db_manager
|