High-performance asynchronous applications (such as FastAPI web services, event-driven microservices, or concurrent worker queues) frequently interact with
High-performance asynchronous applications (such as FastAPI web services, event-driven microservices, or concurrent worker queues) frequently interact with relational databases. Opening a new PostgreSQL connection for every incoming request or concurrent task introduces massive performance bottlenecks. PostgreSQL spawns a dedicated backend process for each TCP connection, which consumes significant CPU and memory resources while adding 20–50ms of network handshake and authentication overhead per request.
Under high load, failing to reuse connections results in database connection exhaustion (too many clients already), severe latency spikes, and eventual application failure. This script solves the problem by implementing a production-grade connection pool using asyncpg. It manages connection creation, reuse, health checks, error handling, and graceful teardown across multiple concurrent asynchronous tasks.
asyncpg is built specifically for asyncio and communicates directly with PostgreSQL using its binary wire protocol, making it significantly faster than traditional drivers like psycopg2 or aiopg. The core design centers around asyncpg.Pool, an asynchronous thread-safe connection pool manager.
The script encapsulates database lifecycle logic inside a DatabaseManager class:
asyncpg.create_pool() pre-allocates a minimum number of idle connections (min_size) and dynamically scales up to a ceiling limit (max_size) as demand increases.async with db.pool.acquire() as conn: context manager. When the block exits, the connection is automatically returned to the pool rather than closed.async with conn.transaction():) ensuring atomic operations with automatic rollback on error.finally block ensures that database schemas are cleaned up and the pool is gracefully drained on shutdown. [ Async Worker 1 ] ----\
[ Async Worker 2 ] -----+--> [ asyncpg Pool ] === (Pre-warmed Connections) ===> PostgreSQL
[ Async Worker N ] ----/ (Min: 5, Max: 20)
#!/usr/bin/env python3
"""
PostgreSQL Connection Pool with asyncpg
Demonstrates robust PostgreSQL connection pooling, health checks,
concurrent query execution, and lifecycle management using asyncpg.
"""
import argparse
import asyncio
import logging
import os
import sys
import time
from typing import Any, Optional
import asyncpg
# Configure structured logging
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
handlers=[logging.StreamHandler(sys.stdout)],
)
logger = logging.getLogger("async_pg_pool")
class DatabaseManager:
"""Manages asyncpg connection pool lifecycle and query execution."""
def __init__(
self,
dsn: str,
min_size: int = 5,
max_size: int = 20,
command_timeout: float = 10.0,
) -> None:
self.dsn = dsn
self.min_size = min_size
self.max_size = max_size
self.command_timeout = command_timeout
self.pool: Optional[asyncpg.Pool] = None
async def initialize(self) -> None:
"""Create and initialize the connection pool."""
logger.info(
"Initializing connection pool (min_size=%d, max_size=%d)...",
self.min_size,
self.max_size,
)
try:
self.pool = await asyncpg.create_pool(
dsn=self.dsn,
min_size=self.min_size,
max_size=self.max_size,
command_timeout=self.command_timeout,
max_inactive_connection_lifetime=300.0,
)
logger.info("Connection pool successfully created.")
except (asyncpg.PostgresError, OSError) as exc:
logger.error("Failed to create connection pool: %s", exc)
raise
async def close(self) -> None:
"""Gracefully close all connections in the pool."""
if self.pool:
logger.info("Closing connection pool...")
await self.pool.close()
logger.info("Connection pool closed.")
async def health_check(self) -> bool:
"""Validate database connectivity by executing a simple ping query."""
if not self.pool:
logger.warning("Health check failed: Pool is not initialized.")
return False
try:
async with self.pool.acquire(timeout=5.0) as conn:
result = await conn.fetchval("SELECT 1;")
return result == 1
except Exception as exc:
logger.error("Health check ping failed: %s", exc)
return False
async def setup_schema(self) -> None:
"""Create demonstration table."""
if not self.pool:
raise RuntimeError("Pool is not initialized.")
query = """
CREATE TABLE IF NOT EXISTS asyncpg_demo (
id SERIAL PRIMARY KEY,
worker_id INT NOT NULL,
payload TEXT NOT NULL,
created_at TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP
);
"""
async with self.pool.acquire() as conn:
await conn.execute(query)
logger.info("Database schema initialized (asyncpg_demo table ready).")
async def teardown_schema(self) -> None:
"""Clean up demonstration table."""
if not self.pool:
return
async with self.pool.acquire() as conn:
await conn.execute("DROP TABLE IF EXISTS asyncpg_demo;")
logger.info("Database schema cleaned up.")
async def execute_worker_task(self, worker_id: int, items_count: int) -> int:
"""Simulate worker task acquiring pool connections to execute queries."""
if not self.pool:
raise RuntimeError("Pool is not initialized.")
inserted_records = 0
for i in range(items_count):
payload = f"Worker-{worker_id}-Item-{i}"
# Acquire connection from pool
async with self.pool.acquire() as conn:
# Execute inside transaction block
async with conn.transaction():
await conn.execute(
"INSERT INTO asyncpg_demo (worker_id, payload) VALUES ($1, $2);",
worker_id,
payload,
)
inserted_records += 1
# Yield control back to asyncio event loop
await asyncio.sleep(0.01)
logger.info("Worker %d finished inserting %d records.", worker_id, inserted_records)
return inserted_records
async def fetch_summary(self) -> list[dict[str, Any]]:
"""Fetch aggregate summary from the database using pool."""
if not self.pool:
raise RuntimeError("Pool is not initialized.")
query = """
SELECT worker_id, COUNT(*) as total_items, MAX(created_at) as last_created
FROM asyncpg_demo
GROUP BY worker_id
ORDER BY worker_id;
"""
async with self.pool.acquire() as conn:
rows = await conn.fetch(query)
return [dict(row) for row in rows]
def parse_args() -> argparse.Namespace:
"""Parse command-line arguments and environment variables."""
parser = argparse.ArgumentParser(
description="PostgreSQL Connection Pool Demo with asyncpg"
)
parser.add_argument(
"--dsn",
type=str,
default=os.environ.get(
"POSTGRES_DSN", "postgresql://postgres:postgres@localhost:5432/postgres"
),
help="PostgreSQL DSN string (default: env POSTGRES_DSN or localhost default)",
)
parser.add_argument(
"--min-size",
type=int,
default=int(os.environ.get("POSTGRES_MIN_POOL", "5")),
help="Minimum connections in pool (default: 5)",
)
parser.add_argument(
"--max-size",
type=int,
default=int(os.environ.get("POSTGRES_MAX_POOL", "20")),
help="Maximum connections in pool (default: 20)",
)
parser.add_argument(
"--workers",
type=int,
default=10,
help="Number of concurrent worker tasks to run (default: 10)",
)
parser.add_argument(
"--items-per-worker",
type=int,
default=5,
help="Number of DB operations per worker (default: 5)",
)
return parser.parse_args()
async def main() -> None:
"""Main application entry point."""
args = parse_args()
db = DatabaseManager(
dsn=args.dsn,
min_size=args.min_size,
max_size=args.max_size,
)
try:
await db.initialize()
# Perform health check
if not await db.health_check():
logger.error("Database health check failed. Aborting operations.")
sys.exit(1)
logger.info("Health check passed.")
# Prepare demo table
await db.setup_schema()
# Run concurrent workers borrowing from the connection pool
start_time = time.perf_counter()
logger.info("Launching %d concurrent workers...", args.workers)
tasks = [
asyncio.create_task(
db.execute_worker_task(worker_id=i, items_count=args.items_per_worker)
)
for i in range(1, args.workers + 1)
]
results = await asyncio.gather(*tasks, return_exceptions=True)
total_inserted = sum(r for r in results if isinstance(r, int))
elapsed = time.perf_counter() - start_time
logger.info(
"Workload completed in %.2f seconds. Total inserted records: %d",
elapsed,
total_inserted,
)
# Print aggregate summary
summary = await db.fetch_summary()
logger.info("Database Workload Summary:")
for row in summary:
logger.info(
" Worker %d: %d items recorded (Last active: %s)",
row["worker_id"],
row["total_items"],
row