Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 12 additions & 0 deletions src/connectors/__tests__/dsn-parser.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -99,6 +99,18 @@ describe('DSN Parser - PostgreSQL SSL Modes', () => {
});
});

describe('DSN Parser - PostgreSQL query timeout', () => {
it('configures a server-side statement timeout before the client fallback', async () => {
const parser = new PostgresConnector().dsnParser;
const config = await parser.parse('postgres://user:pass@localhost:5432/db', {
queryTimeoutSeconds: 30,
});

expect(config.statement_timeout).toBe(30_000);
expect(config.query_timeout).toBe(35_000);
});
});

describe('DSN Parser - AWS IAM Authentication', () => {
describe('MySQL', () => {
const connector = new MySQLConnector();
Expand Down
61 changes: 60 additions & 1 deletion src/connectors/__tests__/postgres.integration.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -214,6 +214,65 @@ describe('PostgreSQL Connector Integration Tests', () => {
postgresTest.createErrorHandlingTests();
postgresTest.createSSLTests();
describe('PostgreSQL-specific Features', () => {
it('should cancel a timed-out query on the PostgreSQL server', async () => {
const timedConnector = new PostgresConnector();
const observer = new PostgresConnector();
const probe = 'dbhub_query_timeout_probe';

const runningProbeCount = async (): Promise<number> => {
const result = await observer.executeSQL(
`SELECT count(*)::int AS count
FROM pg_stat_activity
WHERE state = 'active'
AND query LIKE '%${probe}%'
AND query NOT LIKE '%pg_stat_activity%'`,
{}
);
return result.resultSets[0].rows[0].count;
};

const waitFor = async (
predicate: () => Promise<boolean>,
timeoutMs: number
): Promise<boolean> => {
const deadline = Date.now() + timeoutMs;
while (Date.now() < deadline) {
if (await predicate()) return true;
await new Promise((resolve) => setTimeout(resolve, 50));
}
return false;
};

try {
await timedConnector.connect(postgresTest.connectionString, undefined, {
queryTimeoutSeconds: 1,
});
await observer.connect(postgresTest.connectionString);

const query = timedConnector.executeSQL(
`SELECT pg_sleep(10), '${probe}'`,
{ readonly: true }
);
const settled = query.then(
() => null,
(error) => error as NodeJS.ErrnoException
);

expect(await waitFor(async () => (await runningProbeCount()) === 1, 5_000)).toBe(true);

const error = await settled;
expect(error).toBeInstanceOf(Error);
expect(error?.code).toBe('57014');
expect(await waitFor(async () => (await runningProbeCount()) === 0, 2_000)).toBe(true);

const after = await timedConnector.executeSQL('SELECT 1 AS ok', { readonly: true });
expect(after.resultSets[0].rows[0].ok).toBe(1);
} finally {
await timedConnector.disconnect();
await observer.disconnect();
}
}, 20_000);

it('should execute multiple statements with transaction support', async () => {
const result = await postgresTest.connector.executeSQL(`
INSERT INTO users (name, email, age) VALUES ('Multi User 1', 'multi1@example.com', 30);
Expand Down Expand Up @@ -756,4 +815,4 @@ describe('PostgreSQL Connector Integration Tests', () => {
}
});
});
});
});
12 changes: 9 additions & 3 deletions src/connectors/postgres/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,8 @@ import { splitSQLStatements } from "../../utils/sql-parser.js";
import { FailedToReadCertificate } from "./failed-to-read-certificate.js";
import { closeQuietly } from "../../utils/resource-cleanup.js";

const POSTGRES_CLIENT_QUERY_TIMEOUT_GRACE_MS = 5_000;

/**
* PostgreSQL DSN Parser
* Handles DSN strings like: postgres://user:password@localhost:5432/dbname?sslmode=disable
Expand Down Expand Up @@ -113,10 +115,14 @@ class PostgresDSNParser implements DSNParser {
poolConfig.connectionTimeoutMillis = connectionTimeoutSeconds * 1000;
}

// Apply query timeout if specified (client-side timeout)
// Apply the configured limit on the server so a timed-out statement does
// not keep running after DBHub stops waiting for it. Retain the client-side
// timeout as a fallback, with enough grace for PostgreSQL's cancellation
// response to arrive first.
if (queryTimeoutSeconds !== undefined) {
// pg library expects query_timeout in milliseconds
poolConfig.query_timeout = queryTimeoutSeconds * 1000;
const queryTimeoutMs = queryTimeoutSeconds * 1000;
poolConfig.statement_timeout = queryTimeoutMs;
poolConfig.query_timeout = queryTimeoutMs + POSTGRES_CLIENT_QUERY_TIMEOUT_GRACE_MS;
}

return poolConfig;
Expand Down
Loading