From 13eaac4c9238859cc4b701480f17ba13d416f538 Mon Sep 17 00:00:00 2001 From: Eugene Date: Sun, 14 Jun 2026 08:40:28 -0400 Subject: [PATCH] feat: per-query timeout that cancels over-running queries (#22) Long-running queries previously couldn't be cancelled. This adds a user-configurable per-query timeout, surfaced in the editor toolbar and persisted in localStorage (default 30s, "No timeout" available). The client sends `timeoutMs` with each /api/query request. When the deadline passes, each driver cancels the statement on the server using its idiomatic mechanism and rejects with a `transient` DriverError whose message ("Query cancelled: exceeded the Ns timeout.") shows up in the results error banner: - MySQL: KILL QUERY on a separate pooled connection - Postgres: SELECT pg_cancel_backend() on a separate pooled client - MongoDB: maxTimeMS on the read operations (find/findOne/aggregate/count) On timeout the SQL drivers discard the connection (mysql2 destroy / pg release(true)) instead of returning it to the pool, since the KILLed statement is still draining on that socket and a reused connection could hand the next caller a half-read result set. A shared withQueryTimeout helper races the work promise against the deadline; Promise.race keeps the late statement rejection consumed so it never surfaces as an unhandled rejection. Co-Authored-By: Claude Opus 4.8 (1M context) --- server/src/drivers/interface.ts | 14 +++-- server/src/drivers/mongodb.test.ts | 6 +- server/src/drivers/mongodb.ts | 23 ++++++-- server/src/drivers/mysql.ts | 56 +++++++++++++++--- server/src/drivers/postgres.ts | 68 +++++++++++++++++----- server/src/drivers/timeout.test.ts | 91 ++++++++++++++++++++++++++++++ server/src/drivers/timeout.ts | 46 +++++++++++++++ server/src/routes/query.test.ts | 78 ++++++++++++++++++++++++- server/src/routes/query.ts | 13 ++++- src/App.tsx | 26 ++++++++- src/api.ts | 10 ++-- src/components/QueryEditor.tsx | 35 ++++++++++++ 12 files changed, 422 insertions(+), 44 deletions(-) create mode 100644 server/src/drivers/timeout.test.ts create mode 100644 server/src/drivers/timeout.ts diff --git a/server/src/drivers/interface.ts b/server/src/drivers/interface.ts index 2e1518d..7fded8b 100644 --- a/server/src/drivers/interface.ts +++ b/server/src/drivers/interface.ts @@ -84,13 +84,19 @@ export interface DbDriver { * accept a MQL request object. Routes branch on this to validate the body shape. */ readonly queryMode: 'sql' | 'mql'; - /** Run a query, optionally switching schema first (atomically on one connection). */ - query(sql: string, params?: unknown[], schema?: string): Promise; + /** + * Run a query, optionally switching schema first (atomically on one connection). + * When `timeoutMs` is set the driver cancels the statement on the server after + * that many ms (MySQL `KILL QUERY`, Postgres `pg_cancel_backend`, MongoDB + * `maxTimeMS`) and rejects with a `transient` DriverError. + */ + query(sql: string, params?: unknown[], schema?: string, timeoutMs?: number): Promise; /** * Run one or more semicolon-separated statements and return a result per statement. - * Implemented by sql-mode drivers; mql-mode drivers can omit it. + * Implemented by sql-mode drivers; mql-mode drivers can omit it. `timeoutMs` has + * the same cancel-on-deadline semantics as {@link query}. */ - queryAll?(sql: string, schema?: string): Promise; + queryAll?(sql: string, schema?: string, timeoutMs?: number): Promise; getSchemas(): Promise; getSchema(schema: string): Promise; /** Targeted lookup for a single table — returns null if it doesn't exist. */ diff --git a/server/src/drivers/mongodb.test.ts b/server/src/drivers/mongodb.test.ts index 8a63d1f..64b71ea 100644 --- a/server/src/drivers/mongodb.test.ts +++ b/server/src/drivers/mongodb.test.ts @@ -235,7 +235,7 @@ describe('MongoDBDriver.query – find', () => { ); expect(mockClient.db).toHaveBeenCalledWith('app'); expect(mockDb.collection).toHaveBeenCalledWith('users'); - expect(mockCollection.find).toHaveBeenCalledWith({ active: true }); + expect(mockCollection.find).toHaveBeenCalledWith({ active: true }, {}); expect(r.rows).toEqual([ { _id: '507f1f77bcf86cd799439011', name: 'Alice', when: '2024-01-02 03:04:05' }, ]); @@ -402,7 +402,7 @@ describe('MongoDBDriver.query – findOne / aggregate / count', () => { ); expect(mockCollection.aggregate).toHaveBeenCalledWith([ { $group: { _id: '$status', n: { $sum: 1 } } }, - ]); + ], {}); expect(r.rows).toEqual([{ _id: 'A', n: 2 }]); }); @@ -411,7 +411,7 @@ describe('MongoDBDriver.query – findOne / aggregate / count', () => { const r = await makeDriver().query( JSON.stringify({ collection: 'users', operation: 'count', filter: { active: true } }), ); - expect(mockCollection.countDocuments).toHaveBeenCalledWith({ active: true }); + expect(mockCollection.countDocuments).toHaveBeenCalledWith({ active: true }, {}); expect(r.rows).toEqual([{ count: 42 }]); }); }); diff --git a/server/src/drivers/mongodb.ts b/server/src/drivers/mongodb.ts index f4f24b2..498a41f 100644 --- a/server/src/drivers/mongodb.ts +++ b/server/src/drivers/mongodb.ts @@ -1,6 +1,7 @@ import { MongoClient, ObjectId, MongoNetworkError, MongoServerSelectionError, MongoParseError } from 'mongodb'; import type { Db, Document } from 'mongodb'; import { DriverError } from './interface.js'; +import { timeoutError } from './timeout.js'; import type { DbDriver, ConnectionConfig, @@ -253,7 +254,15 @@ export class MongoDBDriver implements DbDriver { } } - async query(sql: string, _params?: unknown[], schema?: string): Promise { + async query(sql: string, _params?: unknown[], schema?: string, timeoutMs?: number): Promise { + // MongoDB has no out-of-band "kill query" like SQL's KILL/pg_cancel_backend; + // `maxTimeMS` is the idiomatic equivalent — the server aborts the operation + // itself once the deadline passes. Only the read operations honour it; the + // single-document writes return promptly. + const maxTimeMS = timeoutMs && timeoutMs > 0 ? timeoutMs : undefined; + // Spread into the read operations' options; empty (no `maxTimeMS` key) when + // no timeout is armed, so the deadline is opt-in per query. + const readOpts = maxTimeMS ? { maxTimeMS } : {}; try { await this.ensureConnected(); const req = parseRequest(sql); @@ -262,7 +271,7 @@ export class MongoDBDriver implements DbDriver { switch (req.operation) { case 'find': { - let cursor = coll.find(req.filter ?? {}); + let cursor = coll.find(req.filter ?? {}, readOpts); if (req.projection) cursor = cursor.project(req.projection) as typeof cursor; if (req.sort) cursor = cursor.sort(req.sort); if (typeof req.skip === 'number') cursor = cursor.skip(req.skip); @@ -274,15 +283,16 @@ export class MongoDBDriver implements DbDriver { const doc = await coll.findOne(req.filter ?? {}, { projection: req.projection, sort: req.sort as Document | undefined, + ...readOpts, }); return this.toQueryResult(doc ? [doc] : []); } case 'aggregate': { - const docs = await coll.aggregate(req.pipeline ?? []).toArray(); + const docs = await coll.aggregate(req.pipeline ?? [], readOpts).toArray(); return this.toQueryResult(docs); } case 'count': { - const n = await coll.countDocuments(req.filter ?? {}); + const n = await coll.countDocuments(req.filter ?? {}, readOpts); return { rows: [{ count: n }], columnMeta: [ @@ -343,6 +353,11 @@ export class MongoDBDriver implements DbDriver { } } catch (err) { if (err instanceof DriverError) throw err; + // MaxTimeMSExpired (code 50) means our deadline fired server-side; surface + // it with the same wording the SQL drivers use for a timed-out query. + if (maxTimeMS && (err as { code?: number }).code === 50) { + throw timeoutError(maxTimeMS, err); + } throw classifyMongoError(err); } } diff --git a/server/src/drivers/mysql.ts b/server/src/drivers/mysql.ts index 86bf60e..4ba0c29 100644 --- a/server/src/drivers/mysql.ts +++ b/server/src/drivers/mysql.ts @@ -1,6 +1,7 @@ import mysql from 'mysql2/promise'; import type { RowDataPacket, FieldPacket, ResultSetHeader } from 'mysql2/promise'; import type { DbDriver, ConnectionConfig, QueryResult, ColumnMeta, ColumnInfo, SchemaInfo, TableInfo } from './interface.js'; +import { withQueryTimeout } from './timeout.js'; function buildMysqlPoolOptions(config: ConnectionConfig): mysql.PoolOptions { return { @@ -162,14 +163,49 @@ export class MysqlDriver implements DbDriver { await this.pool.end(); } - async query(sql: string, params?: unknown[], schema?: string): Promise { + /** + * `KILL QUERY` aborts the statement running on `threadId` without dropping the + * connection. Issued on a separate pooled connection because the one running + * the timed-out query is busy. Best-effort: the statement may have already + * finished, or no pool slot may be free — the timeout error surfaces either way. + */ + private async killQuery(threadId: number): Promise { + let conn: mysql.PoolConnection | undefined; + try { + conn = await this.pool.getConnection(); + await conn.query(`KILL QUERY ${threadId}`); + } catch { + /* best-effort cancel */ + } finally { + conn?.release(); + } + } + + /** + * Return `conn` to the pool, but `destroy` it instead when the query timed out: + * the `KILL`ed statement is still draining on that socket, and a released + * connection could be handed to the next caller mid-interrupt with a stale + * result set. mysql2 opens a fresh connection on the next request. + */ + private finishConnection(conn: mysql.PoolConnection, timedOut: boolean): void { + if (timedOut) conn.destroy(); + else conn.release(); + } + + async query(sql: string, params?: unknown[], schema?: string, timeoutMs?: number): Promise { const conn = await this.pool.getConnection(); + const threadId = conn.threadId; + let timedOut = false; try { if (schema) { await conn.query(`USE \`${schema.replace(/`/g, '')}\``); } - const [rows, fields] = await conn.query(sql, params) as [RowDataPacket[] | ResultSetHeader, FieldPacket[]]; + const work = conn.query(sql, params) as Promise<[RowDataPacket[] | ResultSetHeader, FieldPacket[]]>; + const [rows, fields] = await withQueryTimeout(work, timeoutMs, () => { + timedOut = true; + if (threadId != null) void this.killQuery(threadId); + }); // With multipleStatements enabled, mysql2 returns arrays-of-arrays when // the SQL contains more than one statement. Internal callers (insertRow, // updateCell, deleteRow, …) only ever pass single statements, so flatten @@ -182,20 +218,26 @@ export class MysqlDriver implements DbDriver { } return buildMysqlResult(rows, fields); } finally { - conn.release(); + this.finishConnection(conn, timedOut); } } - async queryAll(sql: string, schema?: string): Promise { + async queryAll(sql: string, schema?: string, timeoutMs?: number): Promise { const conn = await this.pool.getConnection(); + const threadId = conn.threadId; + let timedOut = false; try { if (schema) { await conn.query(`USE \`${schema.replace(/`/g, '')}\``); } - const [rawRows, rawFields] = await conn.query(sql) as [ + const work = conn.query(sql) as Promise<[ RowDataPacket[] | ResultSetHeader | (RowDataPacket[] | ResultSetHeader)[], FieldPacket[] | FieldPacket[][], - ]; + ]>; + const [rawRows, rawFields] = await withQueryTimeout(work, timeoutMs, () => { + timedOut = true; + if (threadId != null) void this.killQuery(threadId); + }); // Single statement: mysql2 returns the result directly. Multi-statement: it // returns parallel arrays — one rowset and one fields array per statement. const isMulti = Array.isArray(rawRows) && rawRows.length > 0 && Array.isArray((rawRows as unknown[])[0]); @@ -206,7 +248,7 @@ export class MysqlDriver implements DbDriver { const fieldSets = rawFields as FieldPacket[][]; return rowSets.map((r, i) => buildMysqlResult(r, fieldSets[i] ?? [])); } finally { - conn.release(); + this.finishConnection(conn, timedOut); } } diff --git a/server/src/drivers/postgres.ts b/server/src/drivers/postgres.ts index 0da3605..d33d1c1 100644 --- a/server/src/drivers/postgres.ts +++ b/server/src/drivers/postgres.ts @@ -1,6 +1,7 @@ import pg from 'pg'; import { DriverError } from './interface.js'; import type { DbDriver, ConnectionConfig, QueryResult, ColumnMeta, ColumnInfo, SchemaInfo, TableInfo } from './interface.js'; +import { withQueryTimeout } from './timeout.js'; // Return DATE / TIME / TIMESTAMP / TIMESTAMPTZ / TIMETZ as raw strings rather // than JS Dates. pg's default parsers route through `new Date(...)`, which @@ -142,10 +143,29 @@ export class PostgresDriver implements DbDriver { await this.pool.end(); } - async query(sql: string, params?: unknown[], schema?: string): Promise { + /** + * `pg_cancel_backend` cancels the statement running on `pid` without dropping + * the connection. Run on a separate pooled client because the one running the + * timed-out query is busy. Best-effort — the timeout error surfaces regardless. + */ + private async cancelBackend(pid: number): Promise { + let client: pg.PoolClient | undefined; + try { + client = await this.pool.connect(); + await client.query('SELECT pg_cancel_backend($1)', [pid]); + } catch { + /* best-effort cancel */ + } finally { + client?.release(); + } + } + + async query(sql: string, params?: unknown[], schema?: string, timeoutMs?: number): Promise { try { const client = await this.pool.connect(); + const pid = (client as { processID?: number }).processID; let searchPathSet = false; + let timedOut = false; try { if (schema) { await client.query(`SET search_path TO ${this.escapeIdent(schema)}`); @@ -159,7 +179,11 @@ export class PostgresDriver implements DbDriver { let n = 0; pgSql = sql.replace(/\?/g, () => `$${++n}`); } - const result = await client.query({ text: pgSql, values: params?.length ? params : undefined }); + const work = client.query({ text: pgSql, values: params?.length ? params : undefined }); + const result = await withQueryTimeout(work, timeoutMs, () => { + timedOut = true; + if (typeof pid === 'number') void this.cancelBackend(pid); + }); // node-pg's simple query protocol returns an array of results when the SQL // contains multiple statements. `query()` keeps the legacy single-result // contract — `queryAll` is the multi-statement entry point — so collapse @@ -167,12 +191,7 @@ export class PostgresDriver implements DbDriver { const single = Array.isArray(result) ? (result as pg.QueryResult[])[result.length - 1] : result; return buildPgResult(single); } finally { - // pg.Pool reuses clients without resetting session state, so search_path - // would leak to the next caller on this same client. - if (searchPathSet) { - try { await client.query('SET search_path TO DEFAULT'); } catch { /* fall through to release */ } - } - client.release(); + await this.finishClient(client, searchPathSet, timedOut); } } catch (err) { if (err instanceof DriverError) throw err; @@ -180,10 +199,12 @@ export class PostgresDriver implements DbDriver { } } - async queryAll(sql: string, schema?: string): Promise { + async queryAll(sql: string, schema?: string, timeoutMs?: number): Promise { try { const client = await this.pool.connect(); + const pid = (client as { processID?: number }).processID; let searchPathSet = false; + let timedOut = false; try { if (schema) { await client.query(`SET search_path TO ${this.escapeIdent(schema)}`); @@ -191,14 +212,15 @@ export class PostgresDriver implements DbDriver { } // Always go through the simple query protocol (no values) so multi-statement // SQL fans out to one result per statement. - const raw = await client.query(sql); + const work = client.query(sql); + const raw = await withQueryTimeout(work, timeoutMs, () => { + timedOut = true; + if (typeof pid === 'number') void this.cancelBackend(pid); + }); const results = Array.isArray(raw) ? (raw as pg.QueryResult[]) : [raw]; return results.map(buildPgResult); } finally { - if (searchPathSet) { - try { await client.query('SET search_path TO DEFAULT'); } catch { /* fall through to release */ } - } - client.release(); + await this.finishClient(client, searchPathSet, timedOut); } } catch (err) { if (err instanceof DriverError) throw err; @@ -206,6 +228,24 @@ export class PostgresDriver implements DbDriver { } } + /** + * Reset session state and return the client to the pool — but on timeout, skip + * the reset (it would queue behind the statement still being cancelled) and + * `release(true)` to discard the client so the pool never reuses it mid-cancel. + * pg.Pool reuses clients without resetting state, so search_path would + * otherwise leak to the next caller on this same client. + */ + private async finishClient(client: pg.PoolClient, searchPathSet: boolean, timedOut: boolean): Promise { + if (timedOut) { + client.release(true); + return; + } + if (searchPathSet) { + try { await client.query('SET search_path TO DEFAULT'); } catch { /* fall through to release */ } + } + client.release(); + } + async getSchemas(): Promise { const result = await this.pool.query<{ name: string }>( `SELECT schema_name AS name diff --git a/server/src/drivers/timeout.test.ts b/server/src/drivers/timeout.test.ts new file mode 100644 index 0000000..aef3d01 --- /dev/null +++ b/server/src/drivers/timeout.test.ts @@ -0,0 +1,91 @@ +import { describe, it, expect, vi } from 'vitest'; +import { withQueryTimeout, timeoutError, timeoutLabel } from './timeout.js'; +import { DriverError } from './interface.js'; + +describe('timeoutLabel', () => { + it.each([ + [500, '500ms'], + [1000, '1s'], + [1500, '2s'], + [30_000, '30s'], + [60_000, '1m'], + [300_000, '5m'], + [90_000, '90s'], // not a whole number of minutes → seconds + ])('renders %dms as %s', (ms, label) => { + expect(timeoutLabel(ms)).toBe(label); + }); +}); + +describe('timeoutError', () => { + it('is a transient DriverError carrying the timeout in its message', () => { + const err = timeoutError(5000); + expect(err).toBeInstanceOf(DriverError); + expect(err.errorClass).toBe('transient'); + expect(err.message).toBe('Query cancelled: exceeded the 5s timeout.'); + }); +}); + +describe('withQueryTimeout', () => { + it('returns the work result when it settles before the deadline', async () => { + const onTimeout = vi.fn(); + const result = await withQueryTimeout(Promise.resolve('ok'), 1000, onTimeout); + expect(result).toBe('ok'); + expect(onTimeout).not.toHaveBeenCalled(); + }); + + it('returns the work promise untouched when timeout is falsy or non-positive', async () => { + const onTimeout = vi.fn(); + await expect(withQueryTimeout(Promise.resolve(1), 0, onTimeout)).resolves.toBe(1); + await expect(withQueryTimeout(Promise.resolve(2), undefined, onTimeout)).resolves.toBe(2); + await expect(withQueryTimeout(Promise.resolve(3), -5, onTimeout)).resolves.toBe(3); + expect(onTimeout).not.toHaveBeenCalled(); + }); + + it('propagates the work rejection (not a timeout) when work fails first', async () => { + const onTimeout = vi.fn(); + const boom = new Error('boom'); + await expect(withQueryTimeout(Promise.reject(boom), 1000, onTimeout)).rejects.toBe(boom); + expect(onTimeout).not.toHaveBeenCalled(); + }); + + it('fires onTimeout and rejects with a transient DriverError when the deadline wins', async () => { + vi.useFakeTimers(); + try { + const onTimeout = vi.fn(); + // A promise that never settles on its own — only the deadline can resolve the race. + const pending = new Promise(() => {}); + const raced = withQueryTimeout(pending, 5000, onTimeout); + const assertion = expect(raced).rejects.toMatchObject({ + errorClass: 'transient', + message: 'Query cancelled: exceeded the 5s timeout.', + }); + await vi.advanceTimersByTimeAsync(5000); + await assertion; + expect(onTimeout).toHaveBeenCalledTimes(1); + } finally { + vi.useRealTimers(); + } + }); + + it('does not leave the late work rejection unhandled after a timeout', async () => { + vi.useFakeTimers(); + const unhandled = vi.fn(); + process.on('unhandledRejection', unhandled); + try { + let rejectWork!: (e: unknown) => void; + const work = new Promise((_, reject) => { rejectWork = reject; }); + const raced = withQueryTimeout(work, 1000, () => { /* would KILL here */ }); + const assertion = expect(raced).rejects.toBeInstanceOf(DriverError); + await vi.advanceTimersByTimeAsync(1000); + await assertion; + // The statement's error arrives only after the timeout already won the race. + rejectWork(new Error('query interrupted')); + await vi.advanceTimersByTimeAsync(0); + await Promise.resolve(); + expect(unhandled).not.toHaveBeenCalled(); + } finally { + process.off('unhandledRejection', unhandled); + vi.useRealTimers(); + } + }); +}); diff --git a/server/src/drivers/timeout.ts b/server/src/drivers/timeout.ts new file mode 100644 index 0000000..bbd0109 --- /dev/null +++ b/server/src/drivers/timeout.ts @@ -0,0 +1,46 @@ +import { DriverError } from './interface.js'; + +/** Human-readable rendering of a timeout used in the surfaced error message. */ +export function timeoutLabel(timeoutMs: number): string { + if (timeoutMs >= 60_000 && timeoutMs % 60_000 === 0) return `${timeoutMs / 60_000}m`; + if (timeoutMs >= 1000) return `${Math.round(timeoutMs / 1000)}s`; + return `${timeoutMs}ms`; +} + +/** The error surfaced to the client when a query is cancelled for exceeding its deadline. */ +export function timeoutError(timeoutMs: number, cause?: unknown): DriverError { + return new DriverError( + `Query cancelled: exceeded the ${timeoutLabel(timeoutMs)} timeout.`, + 'transient', + cause === undefined ? undefined : { cause }, + ); +} + +/** + * Race a query promise against a deadline. When the timer wins, `onTimeout` is + * invoked — drivers use it to cancel the in-flight statement on the server + * (MySQL `KILL QUERY`, Postgres `pg_cancel_backend`) — and the returned promise + * rejects with a `transient` DriverError carrying the reason. A falsy or + * non-positive `timeoutMs` disables the timeout and `work` is returned as-is. + * + * `Promise.race` keeps a rejection handler attached to `work`, so the statement's + * eventual error (it settles shortly after `onTimeout` cancels it) is consumed + * here and never surfaces as an unhandled rejection. + */ +export function withQueryTimeout( + work: Promise, + timeoutMs: number | undefined, + onTimeout: () => void, +): Promise { + if (!timeoutMs || timeoutMs <= 0) return work; + + let timer: ReturnType; + const deadline = new Promise((_, reject) => { + timer = setTimeout(() => { + try { onTimeout(); } catch { /* best-effort cancel — the timeout still surfaces */ } + reject(timeoutError(timeoutMs)); + }, timeoutMs); + }); + + return Promise.race([work, deadline]).finally(() => clearTimeout(timer)); +} diff --git a/server/src/routes/query.test.ts b/server/src/routes/query.test.ts index 2fe3492..9df772f 100644 --- a/server/src/routes/query.test.ts +++ b/server/src/routes/query.test.ts @@ -39,7 +39,7 @@ describe('postQuery – driver delegation (sql mode)', () => { .send({ sql: 'SELECT * FROM users', schema: 'mydb' }); expect(res.status).toBe(200); - expect(mockDriver.query).toHaveBeenCalledWith('SELECT * FROM users', [], 'mydb'); + expect(mockDriver.query).toHaveBeenCalledWith('SELECT * FROM users', [], 'mydb', undefined); expect(res.body.columns).toEqual(['id', 'name']); expect(res.body.rows).toEqual([{ id: 1, name: 'Alice' }]); expect(res.body.columnMeta[0]).toMatchObject({ name: 'id', pk: true }); @@ -61,7 +61,7 @@ describe('postQuery – driver delegation (sql mode)', () => { .send({ sql: 'SELECT 1 AS one' }); expect(res.status).toBe(200); - expect(mockDriver.query).toHaveBeenCalledWith('SELECT 1 AS one', [], undefined); + expect(mockDriver.query).toHaveBeenCalledWith('SELECT 1 AS one', [], undefined, undefined); }); it('returns affectedRows and insertId for DML (empty columnMeta)', async () => { @@ -102,7 +102,7 @@ describe('postQuery – driver delegation (sql mode)', () => { .send({ sql: 'SELECT 1 AS a; SELECT 2 AS b', schema: 'mydb' }); expect(res.status).toBe(200); - expect(mockDriver.queryAll).toHaveBeenCalledWith('SELECT 1 AS a; SELECT 2 AS b', 'mydb'); + expect(mockDriver.queryAll).toHaveBeenCalledWith('SELECT 1 AS a; SELECT 2 AS b', 'mydb', undefined); expect(mockDriver.query).not.toHaveBeenCalled(); // Legacy fields mirror the first statement so old clients keep working. expect(res.body.columns).toEqual(['a']); @@ -174,6 +174,78 @@ describe('postQuery – driver delegation (sql mode)', () => { }); }); +describe('postQuery – timeout forwarding', () => { + beforeEach(() => vi.clearAllMocks()); + + it('forwards a positive timeoutMs to driver.queryAll', async () => { + const mockDriver = { + queryMode: 'sql' as const, + query: vi.fn(), + queryAll: vi.fn().mockResolvedValue([{ rows: [], columnMeta: [] }]), + }; + vi.mocked(getDriver).mockReturnValue(mockDriver as any); + + await request(makeApp()) + .post('/api/query') + .send({ sql: 'SELECT 1', schema: 'mydb', timeoutMs: 5000 }); + + expect(mockDriver.queryAll).toHaveBeenCalledWith('SELECT 1', 'mydb', 5000); + }); + + it('forwards a positive timeoutMs to driver.query in mql mode', async () => { + const mockDriver = { + queryMode: 'mql' as const, + query: vi.fn().mockResolvedValue({ rows: [], columnMeta: [] }), + }; + vi.mocked(getDriver).mockReturnValue(mockDriver as any); + + await request(makeApp()) + .post('/api/query') + .send({ mql: { collection: 'c', operation: 'find' }, timeoutMs: 1500 }); + + const [, , , timeoutArg] = mockDriver.query.mock.calls[0]; + expect(timeoutArg).toBe(1500); + }); + + it.each([ + ['absent', undefined], + ['zero', 0], + ['negative', -1], + ['non-numeric', 'soon'], + ])('passes undefined when timeoutMs is %s', async (_label, value) => { + const mockDriver = { + queryMode: 'sql' as const, + query: vi.fn(), + queryAll: vi.fn().mockResolvedValue([{ rows: [], columnMeta: [] }]), + }; + vi.mocked(getDriver).mockReturnValue(mockDriver as any); + + await request(makeApp()) + .post('/api/query') + .send({ sql: 'SELECT 1', timeoutMs: value }); + + expect(mockDriver.queryAll).toHaveBeenCalledWith('SELECT 1', undefined, undefined); + }); + + it('surfaces a transient timeout DriverError as HTTP 503', async () => { + const mockDriver = { + queryMode: 'sql' as const, + query: vi.fn(), + queryAll: vi.fn().mockRejectedValue( + new DriverError('Query cancelled: exceeded the 5s timeout.', 'transient'), + ), + }; + vi.mocked(getDriver).mockReturnValue(mockDriver as any); + + const res = await request(makeApp()) + .post('/api/query') + .send({ sql: 'SELECT SLEEP(99)', timeoutMs: 5000 }); + + expect(res.status).toBe(503); + expect(res.body.error).toMatch(/exceeded the 5s timeout/i); + }); +}); + describe('postQuery – input validation (sql mode)', () => { beforeEach(() => vi.clearAllMocks()); diff --git a/server/src/routes/query.ts b/server/src/routes/query.ts index 5412292..602dc2e 100644 --- a/server/src/routes/query.ts +++ b/server/src/routes/query.ts @@ -4,12 +4,19 @@ import { DriverError } from '../drivers/interface.js'; export const postQuery: RequestHandler = async (req, res) => { const driver = getDriver(); - const { sql, mql, schema } = req.body as { + const { sql, mql, schema, timeoutMs: rawTimeout } = req.body as { sql?: unknown; mql?: unknown; schema?: string; + timeoutMs?: unknown; }; + // A positive, finite number arms the per-query deadline; anything else + // (absent, 0, NaN, the client's "no timeout" choice) leaves it disarmed. + const timeoutMs = typeof rawTimeout === 'number' && Number.isFinite(rawTimeout) && rawTimeout > 0 + ? rawTimeout + : undefined; + let queryText: string; // Both `sql` and `mql` present is treated as a mode mismatch — the client is @@ -47,8 +54,8 @@ export const postQuery: RequestHandler = async (req, res) => { // sql-mode drivers expose `queryAll` for multi-statement support; mql-mode // payloads are always a single request and use `query`. const resultList = driver.queryMode === 'sql' && driver.queryAll - ? await driver.queryAll(queryText, schema) - : [await driver.query(queryText, [], schema)]; + ? await driver.queryAll(queryText, schema, timeoutMs) + : [await driver.query(queryText, [], schema, timeoutMs)]; const executionTime = Date.now() - start; const results = resultList.map(r => ({ diff --git a/src/App.tsx b/src/App.tsx index e8ae59a..91af6cd 100644 --- a/src/App.tsx +++ b/src/App.tsx @@ -65,6 +65,20 @@ function readStoredTheme(): ThemeName { } } +// Per-query timeout in seconds; 0 means "no timeout". Sent to the server, which +// cancels an over-running query (KILL QUERY / pg_cancel_backend / maxTimeMS). +const QUERY_TIMEOUT_KEY = 'helix.queryTimeoutSec'; +const DEFAULT_QUERY_TIMEOUT_SEC = 30; + +function readStoredQueryTimeout(): number { + try { + const v = Number(localStorage.getItem(QUERY_TIMEOUT_KEY)); + return Number.isFinite(v) && v >= 0 ? v : DEFAULT_QUERY_TIMEOUT_SEC; + } catch { + return DEFAULT_QUERY_TIMEOUT_SEC; + } +} + export default function App() { const [themeName, setThemeName] = useState(readStoredTheme); const t = themeName === 'dark' ? DARK : LIGHT; @@ -74,6 +88,11 @@ export default function App() { try { localStorage.setItem(THEME_KEY, themeName); } catch { /* quota or disabled storage */ } }, [themeName]); + const [queryTimeoutSec, setQueryTimeoutSec] = useState(readStoredQueryTimeout); + useEffect(() => { + try { localStorage.setItem(QUERY_TIMEOUT_KEY, String(queryTimeoutSec)); } catch { /* quota or disabled storage */ } + }, [queryTimeoutSec]); + const [connected, setConnected] = useState(false); const [queryMode, setQueryMode] = useState('sql'); const [dbType, setDbType] = useState(null); @@ -327,10 +346,11 @@ export default function App() { updateTab(targetTab, { isRunning: true, results: null, resultSets: undefined, activeResultIndex: 0, queryError: null, execTime: null, explainPlan: null, explainSql: null, explainError: null }); const started = Date.now(); + const timeoutMs = queryTimeoutSec > 0 ? queryTimeoutSec * 1000 : undefined; try { const res = queryMode === 'mql' - ? await api.queryMql(mqlPayload, activeSchema) - : await api.query(text, activeSchema); + ? await api.queryMql(mqlPayload, activeSchema, timeoutMs) + : await api.query(text, activeSchema, timeoutMs); // The route always emits the legacy fields from the first statement; `results` // is the new array (one entry per statement). Old responses without it are // synthesised into a single-element array so the rest of the app stays uniform. @@ -593,6 +613,8 @@ export default function App() { onDeleteSavedQuery={handleDeleteSavedQuery} onRenameSavedQuery={handleRenameSavedQuery} onReopenSavedQuery={handleReopenSavedQuery} + queryTimeoutSec={queryTimeoutSec} + onChangeQueryTimeout={setQueryTimeoutSec} t={t} /> diff --git a/src/api.ts b/src/api.ts index 0d50c02..076cc41 100644 --- a/src/api.ts +++ b/src/api.ts @@ -145,17 +145,19 @@ export const api = { return request<{ ddl: string }>(`/api/table-ddl?schema=${encodeURIComponent(schema)}&table=${encodeURIComponent(name)}${t}`); }, - query(sql: string, schema: string) { + // `timeoutMs`, when > 0, asks the server to cancel the query (KILL QUERY / + // pg_cancel_backend / maxTimeMS) after that many ms. Omit or pass 0 for none. + query(sql: string, schema: string, timeoutMs?: number) { return request( '/api/query', - { method: 'POST', body: JSON.stringify({ sql, schema }) } + { method: 'POST', body: JSON.stringify({ sql, schema, timeoutMs }) } ); }, - queryMql(mql: unknown, schema: string) { + queryMql(mql: unknown, schema: string, timeoutMs?: number) { return request( '/api/query', - { method: 'POST', body: JSON.stringify({ mql, schema }) } + { method: 'POST', body: JSON.stringify({ mql, schema, timeoutMs }) } ); }, diff --git a/src/components/QueryEditor.tsx b/src/components/QueryEditor.tsx index 2aa660d..842b712 100644 --- a/src/components/QueryEditor.tsx +++ b/src/components/QueryEditor.tsx @@ -48,9 +48,21 @@ interface QueryEditorProps { onDeleteSavedQuery?: (id: string) => void; onRenameSavedQuery?: (id: string, name: string) => void; onReopenSavedQuery?: (query: SavedQuery) => void; + /** Per-query timeout in seconds; 0 means no timeout. */ + queryTimeoutSec?: number; + onChangeQueryTimeout?: (sec: number) => void; t: Theme; } +// Options for the toolbar timeout picker. 0 = no timeout. +const TIMEOUT_OPTIONS: { sec: number; label: string }[] = [ + { sec: 0, label: 'No timeout' }, + { sec: 10, label: '10s' }, + { sec: 30, label: '30s' }, + { sec: 60, label: '1m' }, + { sec: 300, label: '5m' }, +]; + function formatRelative(ts: number, now: number): string { const s = Math.max(0, Math.round((now - ts) / 1000)); if (s < 60) return `${s}s ago`; @@ -82,6 +94,7 @@ export function QueryEditor({ value, onChange, onRun, isRunning, onExplain, isExplaining = false, queryMode = 'sql', activeSchema, schemaData, runtimeError, history = [], onReopenHistory, onDeleteHistoryEntry, onClearHistory, savedQueries = [], onSaveQuery, onDeleteSavedQuery, onRenameSavedQuery, onReopenSavedQuery, + queryTimeoutSec = 0, onChangeQueryTimeout, t, }: QueryEditorProps) { const isMql = queryMode === 'mql'; @@ -498,6 +511,8 @@ export function QueryEditor({ explainBtn: { display: 'flex', alignItems: 'center', gap: 6, height: 28, padding: '0 10px', background: t.bgElevated, color: t.textPrimary, border: `1px solid ${t.border}`, borderRadius: 5, fontSize: 12, fontWeight: 500, fontFamily: '"IBM Plex Sans", sans-serif', cursor: 'pointer' } as CSSProperties, sep: { width: 1, height: 18, background: t.border, margin: '0 4px' } as CSSProperties, schemaPill: { display: 'flex', alignItems: 'center', gap: 5, fontSize: 11, color: t.textSecondary, fontFamily: 'monospace', background: t.bgElevated, border: `1px solid ${t.border}`, padding: '3px 9px', borderRadius: 4 } as CSSProperties, + timeoutPill: { display: 'flex', alignItems: 'center', gap: 4, color: t.textSecondary, background: t.bgElevated, border: `1px solid ${t.border}`, padding: '2px 6px', borderRadius: 4 } as CSSProperties, + timeoutSelect: { background: 'transparent', color: t.textSecondary, border: 'none', outline: 'none', fontSize: 11, fontFamily: '"IBM Plex Sans", sans-serif', cursor: 'pointer' } as CSSProperties, editorWrap: { flex: 1, display: 'flex', overflow: 'hidden', minHeight: 0 } as CSSProperties, lineNums: { background: t.bgToolbar, borderRight: `1px solid ${t.borderSubtle}`, padding: '12px 0', minWidth: 40, textAlign: 'right', userSelect: 'none', flexShrink: 0, overflowY: 'hidden' } as CSSProperties, lineNum: { padding: '0 10px', fontSize: 11, lineHeight: '21px', color: t.textMuted, fontFamily: '"JetBrains Mono", monospace' } as CSSProperties, @@ -784,6 +799,26 @@ export function QueryEditor({
+ {onChangeQueryTimeout && ( + + )}