/**
* postgres-mcp - Connection Pool Manager
*
* Wraps pg connection pooling with health monitoring,
* statistics tracking, and graceful shutdown support.
*/
import pg from "pg";
import type { PoolClient, QueryResult as PgQueryResult } from "pg";
import type { PoolConfig, PoolStats, HealthStatus } from "../types/index.js";
import { PoolError, ConnectionError } from "../types/index.js";
import { logger } from "../utils/logger.js";
/**
* Connection pool configuration with defaults
*/
export interface ConnectionPoolConfig {
host: string;
port: number;
user: string;
password: string;
database: string;
pool?: PoolConfig | undefined;
ssl?: pg.ConnectionConfig["ssl"] | undefined;
statementTimeout?: number | undefined;
applicationName?: string | undefined;
}
/**
* Connection pool wrapper with statistics and health monitoring
*/
export class ConnectionPool {
private pool: pg.Pool | null = null;
private config: ConnectionPoolConfig;
private stats: PoolStats = {
total: 0,
active: 0,
idle: 0,
waiting: 0,
totalQueries: 0,
};
private shuttingDown = false;
constructor(config: ConnectionPoolConfig) {
this.config = config;
}
/**
* Initialize the connection pool
*/
async initialize(): Promise<void> {
if (this.pool !== null) {
logger.warn("Connection pool already initialized");
return;
}
logger.info("Initializing PostgreSQL connection pool", {
host: this.config.host,
port: this.config.port,
database: this.config.database,
});
try {
const poolConfig: pg.PoolConfig = {
host: this.config.host,
port: this.config.port,
user: this.config.user,
password: this.config.password,
database: this.config.database,
max: this.config.pool?.max ?? 10,
min: this.config.pool?.min ?? 0,
idleTimeoutMillis: this.config.pool?.idleTimeoutMillis ?? 10000,
connectionTimeoutMillis:
this.config.pool?.connectionTimeoutMillis ?? 10000,
allowExitOnIdle: this.config.pool?.allowExitOnIdle ?? true,
application_name: this.config.applicationName ?? "postgres-mcp",
};
if (this.config.ssl === true) {
poolConfig.ssl = { rejectUnauthorized: false };
} else if (this.config.ssl !== undefined && this.config.ssl !== false) {
poolConfig.ssl = this.config.ssl;
}
if (
this.config.statementTimeout !== undefined &&
this.config.statementTimeout > 0
) {
poolConfig.statement_timeout = this.config.statementTimeout;
}
this.pool = new pg.Pool(poolConfig);
// Set up event handlers
this.pool.on("connect", () => {
this.stats.total++;
logger.debug("New connection established");
});
this.pool.on("acquire", () => {
this.stats.active++;
this.stats.idle = Math.max(0, this.stats.idle - 1);
});
this.pool.on("release", () => {
this.stats.active = Math.max(0, this.stats.active - 1);
this.stats.idle++;
});
this.pool.on("remove", () => {
this.stats.total = Math.max(0, this.stats.total - 1);
this.stats.idle = Math.max(0, this.stats.idle - 1);
});
this.pool.on("error", (err) => {
logger.error("Pool error", { error: err.message });
});
// Test connection
const client = await this.pool.connect();
const result = await client.query("SELECT version()");
client.release();
const version = result.rows[0] as { version?: string } | undefined;
logger.info("PostgreSQL connection pool initialized", {
version: version?.version ?? "unknown",
});
} catch (error) {
// Clean up pool on initialization failure
if (this.pool !== null) {
try {
await this.pool.end();
} catch {
// Ignore cleanup errors
}
this.pool = null;
}
const message = error instanceof Error ? error.message : "Unknown error";
logger.error("Failed to initialize connection pool", { error: message });
throw new ConnectionError(`Failed to connect to PostgreSQL: ${message}`);
}
}
/**
* Get a connection from the pool
*/
async getConnection(): Promise<PoolClient> {
if (this.pool === null) {
throw new PoolError("Connection pool not initialized");
}
if (this.shuttingDown) {
throw new PoolError("Connection pool is shutting down");
}
try {
this.stats.waiting++;
const client = await this.pool.connect();
this.stats.waiting = Math.max(0, this.stats.waiting - 1);
return client;
} catch (error) {
this.stats.waiting = Math.max(0, this.stats.waiting - 1);
const message = error instanceof Error ? error.message : "Unknown error";
throw new PoolError(`Failed to acquire connection: ${message}`);
}
}
/**
* Release a connection back to the pool
*/
releaseConnection(client: PoolClient): void {
try {
client.release();
} catch (error) {
logger.warn("Error releasing connection", {
error: error instanceof Error ? error.message : "Unknown error",
});
}
}
/**
* Execute a query using a pooled connection
*/
async query<T extends Record<string, unknown>[]>(
sql: string,
params?: unknown[],
): Promise<PgQueryResult<T[number]>> {
if (this.pool === null) {
throw new PoolError("Connection pool not initialized");
}
const startTime = Date.now();
this.stats.totalQueries++;
try {
const result = await this.pool.query<T[number]>(sql, params);
logger.debug("Query executed", {
sql: sql.substring(0, 100),
rowCount: result.rowCount,
durationMs: Date.now() - startTime,
});
return result;
} catch (error) {
const message = error instanceof Error ? error.message : "Unknown error";
logger.error("Query failed", {
sql: sql.substring(0, 100),
error: message,
});
throw error;
}
}
/**
* Get pool statistics
*/
getStats(): PoolStats {
if (this.pool !== null) {
// Update stats from pool if available
this.stats.total = this.pool.totalCount;
this.stats.idle = this.pool.idleCount;
this.stats.waiting = this.pool.waitingCount;
this.stats.active = this.stats.total - this.stats.idle;
}
return { ...this.stats };
}
/**
* Check pool health
*/
async checkHealth(): Promise<HealthStatus> {
if (this.pool === null || this.shuttingDown) {
return {
connected: false,
error: this.shuttingDown
? "Pool is shutting down"
: "Pool not initialized",
};
}
const startTime = Date.now();
try {
const result = await this.pool.query(
"SELECT version(), current_database()",
);
const latencyMs = Date.now() - startTime;
const row = result.rows[0] as
| { version?: string; current_database?: string }
| undefined;
return {
connected: true,
latencyMs,
version: row?.version ?? undefined,
poolStats: this.getStats(),
details: {
database: row?.current_database,
},
};
} catch (error) {
return {
connected: false,
error: error instanceof Error ? error.message : "Unknown error",
latencyMs: Date.now() - startTime,
};
}
}
/**
* Gracefully shutdown the pool
*/
async shutdown(): Promise<void> {
if (this.pool === null) {
return;
}
logger.info("Shutting down connection pool...");
this.shuttingDown = true;
try {
await this.pool.end();
this.pool = null;
logger.info("Connection pool shut down successfully");
} catch (error) {
logger.error("Error during pool shutdown", {
error: error instanceof Error ? error.message : "Unknown error",
});
throw error;
}
}
/**
* Check if pool is initialized
*/
isInitialized(): boolean {
return this.pool !== null && !this.shuttingDown;
}
/**
* Check if pool is shutting down
*/
isClosing(): boolean {
return this.shuttingDown;
}
}
