forked from sagnik/Project_Velocity
59 lines
1.8 KiB
Python
59 lines
1.8 KiB
Python
"""
|
|
backend/db/pool.py — asyncpg Connection Pool
|
|
|
|
Initialises a PostgreSQL connection pool from environment variables.
|
|
All credentials are sourced from the environment only — never hardcoded.
|
|
|
|
Environment variables required:
|
|
VELOCITY_DB_HOST PostgreSQL host (default: localhost)
|
|
VELOCITY_DB_PORT PostgreSQL port (default: 5432)
|
|
VELOCITY_DB_NAME Database name
|
|
VELOCITY_DB_USER Database user
|
|
VELOCITY_DB_PASSWORD Database password (injected from AWS SSM at service start)
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
from contextlib import asynccontextmanager
|
|
from typing import AsyncGenerator
|
|
|
|
import asyncpg
|
|
from fastapi import Request
|
|
|
|
_pool: asyncpg.Pool | None = None
|
|
|
|
|
|
async def create_pool() -> asyncpg.Pool:
|
|
"""Creates and returns the application-wide asyncpg connection pool."""
|
|
global _pool
|
|
_pool = await asyncpg.create_pool(
|
|
host=os.environ.get("VELOCITY_DB_HOST", "localhost"),
|
|
port=int(os.environ.get("VELOCITY_DB_PORT", "5432")),
|
|
database=os.environ["VELOCITY_DB_NAME"],
|
|
user=os.environ["VELOCITY_DB_USER"],
|
|
password=os.environ["VELOCITY_DB_PASSWORD"],
|
|
min_size=2,
|
|
max_size=10,
|
|
command_timeout=30,
|
|
# Set app_name for easier identification in pg_stat_activity
|
|
server_settings={"application_name": "velocity-backend"},
|
|
)
|
|
return _pool
|
|
|
|
|
|
async def close_pool() -> None:
|
|
"""Closes the connection pool on application shutdown."""
|
|
global _pool
|
|
if _pool:
|
|
await _pool.close()
|
|
_pool = None
|
|
|
|
|
|
def get_pool(request: Request) -> asyncpg.Pool:
|
|
"""FastAPI dependency: returns the pool stored in app.state."""
|
|
pool: asyncpg.Pool = request.app.state.db_pool
|
|
if pool is None:
|
|
raise RuntimeError("Database pool is not initialised.")
|
|
return pool
|