/**
* Database connection and query utilities
*/
import pg from 'pg';
import { DatabaseConfig, QueryResult } from './types.js';
const { Pool } = pg;
/**
* Create a PostgreSQL connection pool
*/
export function createPool(config: DatabaseConfig): pg.Pool {
// Parse SSL configuration
let sslConfig: boolean | { rejectUnauthorized: boolean };
if (config.connectionString.includes('sslmode=')) {
// SSL mode is in connection string
sslConfig = config.ssl;
} else {
// Add SSL mode to connection string if needed
if (typeof config.ssl === 'object') {
sslConfig = config.ssl;
} else {
sslConfig = config.ssl;
}
}
const pool = new Pool({
connectionString: config.connectionString,
ssl: sslConfig,
max: 10, // Maximum number of clients in the pool
idleTimeoutMillis: 30000, // Close idle clients after 30 seconds
connectionTimeoutMillis: 5000, // Return an error after 5 seconds if connection could not be established
statement_timeout: config.timeout * 1000, // Query timeout in milliseconds
});
// Handle pool errors
pool.on('error', (err) => {
console.error('Unexpected error on idle client', err);
});
return pool;
}
/**
* Validate SQL query for read-only mode
* @param query SQL query to validate
* @param readOnly Whether read-only mode is enabled
* @throws Error if query contains disallowed operations in read-only mode
*/
export function validateQuery(query: string, readOnly: boolean): void {
if (!readOnly) {
return; // No restrictions in non-read-only mode
}
// Normalize query: remove comments and convert to uppercase
const normalizedQuery = query
.replace(/--[^\n]*/g, '') // Remove single-line comments
.replace(/\/\*[\s\S]*?\*\//g, '') // Remove multi-line comments
.trim()
.toUpperCase();
// List of disallowed operations in read-only mode
const disallowedOperations = [
'INSERT',
'UPDATE',
'DELETE',
'DROP',
'CREATE',
'ALTER',
'TRUNCATE',
'GRANT',
'REVOKE',
'EXECUTE',
'CALL',
];
// Check for disallowed operations at the start of statements
const statements = normalizedQuery.split(';').filter(s => s.trim());
for (const statement of statements) {
const trimmedStatement = statement.trim();
for (const operation of disallowedOperations) {
// Check if statement starts with the operation
if (trimmedStatement.startsWith(operation)) {
throw new Error(
`Read-only mode: ${operation} operations are not allowed. Only SELECT and EXPLAIN queries are permitted.`
);
}
// Also check for operations after WITH clauses (CTEs)
const withRegex = new RegExp(`\\bWITH\\b[\\s\\S]*?\\b${operation}\\b`, 'i');
if (withRegex.test(trimmedStatement)) {
throw new Error(
`Read-only mode: ${operation} operations are not allowed, even in CTEs. Only SELECT and EXPLAIN queries are permitted.`
);
}
}
}
}
/**
* Execute a SQL query with timeout and row limit
*/
export async function executeQuery(
pool: pg.Pool,
query: string,
config: DatabaseConfig,
params?: any[]
): Promise<QueryResult> {
// Validate query if in read-only mode
validateQuery(query, config.readOnly);
// Add LIMIT clause if not present and maxQueryResults is set
let finalQuery = query.trim();
const upperQuery = finalQuery.toUpperCase();
// Only add LIMIT for SELECT queries without existing LIMIT
if (upperQuery.startsWith('SELECT') && !upperQuery.includes('LIMIT')) {
finalQuery = `${finalQuery} LIMIT ${config.maxQueryResults}`;
}
try {
const result = params
? await pool.query(finalQuery, params)
: await pool.query(finalQuery);
return {
rows: result.rows,
rowCount: result.rowCount || 0,
fields: result.fields?.map(f => ({
name: f.name,
dataTypeID: f.dataTypeID,
})),
};
} catch (error) {
if (error instanceof Error) {
// Check if it's a timeout error
if (error.message.includes('statement timeout') || error.message.includes('query_canceled')) {
throw new Error(`Query timeout: execution exceeded ${config.timeout} seconds`);
}
throw error;
}
throw new Error('Unknown database error occurred');
}
}
/**
* Test database connection
*/
export async function testConnection(pool: pg.Pool): Promise<void> {
try {
const result = await pool.query('SELECT 1 as test');
if (result.rows[0].test !== 1) {
throw new Error('Connection test failed: unexpected result');
}
} catch (error) {
if (error instanceof Error) {
throw new Error(`Database connection failed: ${error.message}`);
}
throw new Error('Database connection failed: unknown error');
}
}
/**
* Validate that a schema is in the allowed list
*/
export function validateSchema(schema: string, allowedSchemas: string[]): void {
if (!allowedSchemas.includes(schema)) {
throw new Error(
`Schema "${schema}" is not in the allowed schemas list: ${allowedSchemas.join(', ')}`
);
}
}
/**
* Safely escape identifiers (table names, column names) for use in SQL
* This helps prevent SQL injection in dynamic queries
*/
export function escapeIdentifier(identifier: string): string {
return `"${identifier.replace(/"/g, '""')}"`;
}
/**
* Build a WHERE clause for schema filtering
*/
export function buildSchemaFilter(allowedSchemas: string[]): string {
const schemaList = allowedSchemas.map(s => `'${s.replace(/'/g, "''")}'`).join(', ');
return `table_schema IN (${schemaList})`;
}
/**
* Close the database connection pool gracefully
*/
export async function closePool(pool: pg.Pool): Promise<void> {
try {
await pool.end();
} catch (error) {
console.error('Error closing database pool:', error);
throw error;
}
}
