From 298e8df04ca5c4f92e87d0ba782337bead3d21ab Mon Sep 17 00:00:00 2001 From: Johannes Fleck Date: Thu, 8 Oct 2026 13:28:50 +0200 Subject: [PATCH 1/7] test: cover Trino API route handlers (#345) Unit tests for the query, catalog and completion metadata endpoints: request validation, connection checks, generated SQL including identifier and literal escaping, result mapping and error status codes. --- .../(app)/api/trino/catalog/server.test.ts | 104 +++++++++ .../trino/completion/metadata/server.test.ts | 163 +++++++++++++++ .../(app)/api/trino/query/server.test.ts | 197 ++++++++++++++++++ 3 files changed, 464 insertions(+) create mode 100644 src/routes/(app)/api/trino/catalog/server.test.ts create mode 100644 src/routes/(app)/api/trino/completion/metadata/server.test.ts create mode 100644 src/routes/(app)/api/trino/query/server.test.ts diff --git a/src/routes/(app)/api/trino/catalog/server.test.ts b/src/routes/(app)/api/trino/catalog/server.test.ts new file mode 100644 index 00000000..37956e3d --- /dev/null +++ b/src/routes/(app)/api/trino/catalog/server.test.ts @@ -0,0 +1,104 @@ +import { describe, it, expect, vi, beforeEach } from 'vitest'; + +vi.mock('$lib/server/trino/client.js', () => ({ + resolveTrinoClient: vi.fn(), + trinoMetadataQuery: vi.fn() +})); + +import { GET } from './+server.js'; +import { resolveTrinoClient, trinoMetadataQuery } from '$lib/server/trino/client.js'; +import type { TrinoClient } from '$lib/server/trino/client.js'; + +const fakeClient = { fake: true } as unknown as TrinoClient; + +function mockEvent( + params: Record, + user: object | null = { id: 'u1', username: 'alice' } +) { + const url = new URL('http://localhost/api/trino/catalog'); + for (const [k, v] of Object.entries(params)) url.searchParams.set(k, v); + return { + url, + locals: { logger: { debug: vi.fn(), info: vi.fn() }, user } + } as unknown as Parameters[0]; +} + +/** SQL passed to Trino, with whitespace collapsed so assertions ignore formatting. */ +function submittedSql(): string { + return vi.mocked(trinoMetadataQuery).mock.calls[0][1].replace(/\s+/g, ' ').trim(); +} + +describe('GET /api/trino/catalog', () => { + beforeEach(() => { + vi.clearAllMocks(); + vi.mocked(resolveTrinoClient).mockReturnValue(fakeClient); + vi.mocked(trinoMetadataQuery).mockResolvedValue({ columns: [], rows: [] }); + }); + + it('throws 400 when no Trino connection is configured', async () => { + vi.mocked(resolveTrinoClient).mockReturnValue(null); + + await expect(GET(mockEvent({ level: 'catalogs' }))).rejects.toThrow( + expect.objectContaining({ status: 400 }) + ); + expect(trinoMetadataQuery).not.toHaveBeenCalled(); + }); + + it.each([ + ['level is missing', {}], + ['level is unknown', { level: 'functions' }], + ['schemas lacks catalog', { level: 'schemas' }], + ['tables lacks schema', { level: 'tables', catalog: 'tpch' }], + ['columns lacks table', { level: 'columns', catalog: 'tpch', schema: 'tiny' }] + ])('throws 400 when %s', async (_label, params) => { + await expect(GET(mockEvent(params))).rejects.toThrow(expect.objectContaining({ status: 400 })); + expect(trinoMetadataQuery).not.toHaveBeenCalled(); + }); + + it('lists catalogs and returns the rows', async () => { + vi.mocked(trinoMetadataQuery).mockResolvedValue({ columns: [], rows: [['tpch'], ['system']] }); + + const res = await GET(mockEvent({ level: 'catalogs' })); + + expect(await res.json()).toEqual([['tpch'], ['system']]); + expect(trinoMetadataQuery).toHaveBeenCalledWith(fakeClient, 'SHOW CATALOGS', { user: 'alice' }); + }); + + it('falls back to the anonymous user without a session user', async () => { + await GET(mockEvent({ level: 'catalogs' }, null)); + + expect(resolveTrinoClient).toHaveBeenCalledWith('anonymous'); + expect(trinoMetadataQuery).toHaveBeenCalledWith(fakeClient, 'SHOW CATALOGS', { + user: 'anonymous' + }); + }); + + it('quotes the catalog identifier when listing schemas', async () => { + await GET(mockEvent({ level: 'schemas', catalog: 'my"cat' })); + + expect(submittedSql()).toBe('SHOW SCHEMAS FROM "my""cat"'); + }); + + it('quotes identifiers and escapes string literals when listing tables', async () => { + await GET(mockEvent({ level: 'tables', catalog: 'it\'s"cat', schema: "o'schema" })); + + const sql = submittedSql(); + expect(sql).toContain(`FROM "it's""cat".information_schema.tables t`); + expect(sql).toContain(`mv.catalog_name = 'it''s"cat'`); + expect(sql).toContain(`WHERE t.table_schema = 'o''schema'`); + }); + + it('quotes every name part when describing columns', async () => { + await GET(mockEvent({ level: 'columns', catalog: 'tpch', schema: 'my-schema', table: 'a"b' })); + + expect(submittedSql()).toBe('DESCRIBE "tpch"."my-schema"."a""b"'); + }); + + it('throws 502 with the Trino error message when the query fails', async () => { + vi.mocked(trinoMetadataQuery).mockRejectedValue(new Error('Catalog does not exist')); + + await expect(GET(mockEvent({ level: 'schemas', catalog: 'nope' }))).rejects.toThrow( + expect.objectContaining({ status: 502, body: { message: 'Catalog does not exist' } }) + ); + }); +}); diff --git a/src/routes/(app)/api/trino/completion/metadata/server.test.ts b/src/routes/(app)/api/trino/completion/metadata/server.test.ts new file mode 100644 index 00000000..cce3e18b --- /dev/null +++ b/src/routes/(app)/api/trino/completion/metadata/server.test.ts @@ -0,0 +1,163 @@ +import { describe, it, expect, vi, beforeEach } from 'vitest'; + +const { flags } = vi.hoisted(() => ({ flags: { completionEnabled: true } })); + +vi.mock('$lib/server/feature-flags.js', () => ({ + get completionEnabled() { + return flags.completionEnabled; + } +})); + +vi.mock('$lib/server/trino/client.js', () => ({ + resolveTrinoClient: vi.fn(), + trinoMetadataQuery: vi.fn() +})); + +import { GET } from './+server.js'; +import { resolveTrinoClient, trinoMetadataQuery } from '$lib/server/trino/client.js'; +import type { TrinoClient } from '$lib/server/trino/client.js'; + +const fakeClient = { fake: true } as unknown as TrinoClient; + +function mockEvent(params: Record) { + const url = new URL('http://localhost/api/trino/completion/metadata'); + for (const [k, v] of Object.entries(params)) url.searchParams.set(k, v); + return { + url, + locals: { logger: { debug: vi.fn(), info: vi.fn() }, user: { id: 'u1', username: 'alice' } } + } as unknown as Parameters[0]; +} + +function submitted(): { sql: string; opts: unknown } { + const [, sql, opts] = vi.mocked(trinoMetadataQuery).mock.calls[0]; + return { sql: sql.replace(/\s+/g, ' ').trim(), opts }; +} + +describe('GET /api/trino/completion/metadata', () => { + beforeEach(() => { + vi.clearAllMocks(); + flags.completionEnabled = true; + vi.mocked(resolveTrinoClient).mockReturnValue(fakeClient); + vi.mocked(trinoMetadataQuery).mockResolvedValue({ columns: [], rows: [] }); + }); + + it('throws 404 when code completion is disabled', async () => { + flags.completionEnabled = false; + + await expect(GET(mockEvent({ level: 'catalogs' }))).rejects.toThrow( + expect.objectContaining({ status: 404 }) + ); + expect(resolveTrinoClient).not.toHaveBeenCalled(); + }); + + it('throws 400 when no Trino connection is configured', async () => { + vi.mocked(resolveTrinoClient).mockReturnValue(null); + + await expect(GET(mockEvent({ level: 'catalogs' }))).rejects.toThrow( + expect.objectContaining({ status: 400 }) + ); + }); + + it.each([ + ['level is missing', {}], + ['level is unknown', { level: 'views' }], + ['catalog is empty', { level: 'schemas', catalog: '' }], + ['schema is too long', { level: 'tables', schema: 'x'.repeat(1025) }], + ['columns lacks table', { level: 'columns', catalog: 'tpch', schema: 'tiny' }] + ])('throws 400 when %s', async (_label, params) => { + await expect(GET(mockEvent(params))).rejects.toThrow(expect.objectContaining({ status: 400 })); + expect(trinoMetadataQuery).not.toHaveBeenCalled(); + }); + + it('lists catalogs as a flat name array', async () => { + vi.mocked(trinoMetadataQuery).mockResolvedValue({ columns: [], rows: [['tpch'], ['system']] }); + + const res = await GET(mockEvent({ level: 'catalogs' })); + + expect(await res.json()).toEqual(['tpch', 'system']); + expect(submitted()).toEqual({ sql: 'SHOW CATALOGS', opts: { user: 'alice' } }); + }); + + it('lists schemas using the catalog as session context', async () => { + await GET(mockEvent({ level: 'schemas', catalog: 'tpch' })); + + expect(submitted()).toEqual({ sql: 'SHOW SCHEMAS', opts: { user: 'alice', catalog: 'tpch' } }); + }); + + it('lists tables without a schema filter when no schema is given', async () => { + await GET(mockEvent({ level: 'tables', catalog: 'tpch' })); + + const { sql, opts } = submitted(); + expect(sql).not.toContain('WHERE'); + expect(opts).toEqual({ user: 'alice', catalog: 'tpch' }); + }); + + it('filters tables by an escaped schema literal', async () => { + await GET(mockEvent({ level: 'tables', catalog: 'tpch', schema: "o'schema" })); + + expect(submitted().sql).toContain(`WHERE t.table_schema = 'o''schema'`); + }); + + it('maps table types to completion kinds', async () => { + vi.mocked(trinoMetadataQuery).mockResolvedValue({ + columns: [], + rows: [ + ['orders', 'BASE TABLE'], + ['v_orders', 'VIEW'], + ['mv_orders', 'MATERIALIZED VIEW'], + ['lower_view', 'view'], + ['untyped', null] + ] + }); + + const res = await GET(mockEvent({ level: 'tables', catalog: 'tpch', schema: 'tiny' })); + + expect(await res.json()).toEqual([ + { name: 'orders', kind: 'table' }, + { name: 'v_orders', kind: 'view' }, + { name: 'mv_orders', kind: 'materialized_view' }, + { name: 'lower_view', kind: 'view' }, + { name: 'untyped', kind: 'table' } + ]); + }); + + it('describes columns with a quoted table name and catalog/schema context', async () => { + vi.mocked(trinoMetadataQuery).mockResolvedValue({ + columns: [], + rows: [ + ['orderkey', 'bigint', '', ''], + ['custkey', 'bigint', '', ''] + ] + }); + + const res = await GET( + mockEvent({ level: 'columns', catalog: 'tpch', schema: 'tiny', table: 'a"b' }) + ); + + expect(await res.json()).toEqual(['orderkey', 'custkey']); + expect(submitted()).toEqual({ + sql: 'DESCRIBE "a""b"', + opts: { user: 'alice', catalog: 'tpch', schema: 'tiny' } + }); + }); + + it('dedupes and sorts function overloads', async () => { + vi.mocked(trinoMetadataQuery).mockResolvedValue({ + columns: [], + rows: [['substr'], ['abs'], ['substr'], ['abs'], ['concat']] + }); + + const res = await GET(mockEvent({ level: 'functions' })); + + expect(await res.json()).toEqual(['abs', 'concat', 'substr']); + expect(submitted()).toEqual({ sql: 'SHOW FUNCTIONS', opts: { user: 'alice' } }); + }); + + it('throws 502 with the Trino error message when the query fails', async () => { + vi.mocked(trinoMetadataQuery).mockRejectedValue(new Error('Access Denied')); + + await expect(GET(mockEvent({ level: 'catalogs' }))).rejects.toThrow( + expect.objectContaining({ status: 502, body: { message: 'Access Denied' } }) + ); + }); +}); diff --git a/src/routes/(app)/api/trino/query/server.test.ts b/src/routes/(app)/api/trino/query/server.test.ts new file mode 100644 index 00000000..e60a3160 --- /dev/null +++ b/src/routes/(app)/api/trino/query/server.test.ts @@ -0,0 +1,197 @@ +import { describe, it, expect, vi, beforeEach } from 'vitest'; + +vi.mock('$lib/server/trino/client.js', () => ({ + resolveTrinoClient: vi.fn() +})); + +vi.mock('$lib/server/trino/queries.js', () => ({ + startScript: vi.fn(), + getQuerySnapshots: vi.fn(), + cancelQuery: vi.fn(), + removeTabQuery: vi.fn() +})); + +import { POST, GET, DELETE } from './+server.js'; +import { resolveTrinoClient } from '$lib/server/trino/client.js'; +import { + startScript, + getQuerySnapshots, + cancelQuery, + removeTabQuery +} from '$lib/server/trino/queries.js'; +import type { TrinoClient } from '$lib/server/trino/client.js'; + +const TAB_ID = '6f1c2a4e-8b3d-4c5e-9f7a-1b2c3d4e5f60'; +const fakeClient = { fake: true } as unknown as TrinoClient; + +function mockLocals( + user: { id: string; username?: string } | null = { id: 'u1', username: 'alice' } +) { + return { + logger: { trace: vi.fn(), debug: vi.fn(), info: vi.fn(), warn: vi.fn(), error: vi.fn() }, + user + }; +} + +function postEvent(body: unknown, locals = mockLocals()) { + const request = new Request('http://localhost/api/trino/query', { + method: 'POST', + body: typeof body === 'string' ? body : JSON.stringify(body) + }); + return { request, locals } as unknown as Parameters[0]; +} + +function urlEvent(params: string, locals = mockLocals()) { + const url = new URL(`http://localhost/api/trino/query?${params}`); + return { url, locals } as unknown as Parameters[0]; +} + +describe('POST /api/trino/query', () => { + beforeEach(() => { + vi.clearAllMocks(); + vi.mocked(resolveTrinoClient).mockReturnValue(fakeClient); + vi.mocked(startScript).mockResolvedValue(); + }); + + it('returns 400 when no Trino connection is configured', async () => { + vi.mocked(resolveTrinoClient).mockReturnValue(null); + + const res = await POST(postEvent({ statements: ['SELECT 1'], tabId: TAB_ID })); + + expect(res.status).toBe(400); + expect(await res.json()).toEqual({ error: 'No Trino connection configured' }); + expect(startScript).not.toHaveBeenCalled(); + }); + + it('returns 400 for a body that is not valid JSON', async () => { + const res = await POST(postEvent('{not json')); + + expect(res.status).toBe(400); + expect(await res.json()).toEqual({ error: 'Invalid JSON body' }); + }); + + it.each([ + ['no statements', { statements: [], tabId: TAB_ID }], + ['an empty statement', { statements: [''], tabId: TAB_ID }], + ['a non-UUID tabId', { statements: ['SELECT 1'], tabId: 'tab-1' }], + ['a missing tabId', { statements: ['SELECT 1'] }] + ])('returns 400 for %s', async (_label, body) => { + const res = await POST(postEvent(body)); + + expect(res.status).toBe(400); + expect((await res.json()).error).toEqual(expect.any(String)); + expect(startScript).not.toHaveBeenCalled(); + }); + + it('starts the script and returns 204', async () => { + const res = await POST( + postEvent({ + statements: ['SELECT 1', 'SELECT 2'], + tabId: TAB_ID, + catalog: 'tpch', + schema: 'tiny' + }) + ); + + expect(res.status).toBe(204); + expect(resolveTrinoClient).toHaveBeenCalledWith('u1'); + expect(startScript).toHaveBeenCalledWith(fakeClient, 'u1', TAB_ID, ['SELECT 1', 'SELECT 2'], { + user: 'alice', + catalog: 'tpch', + schema: 'tiny' + }); + }); + + it('falls back to the anonymous user without a session user', async () => { + const res = await POST( + postEvent({ statements: ['SELECT 1'], tabId: TAB_ID }, mockLocals(null)) + ); + + expect(res.status).toBe(204); + expect(startScript).toHaveBeenCalledWith(fakeClient, 'anonymous', TAB_ID, ['SELECT 1'], { + user: 'anonymous', + catalog: undefined, + schema: undefined + }); + }); + + it('logs but does not fail the request when the script orchestrator rejects', async () => { + const err = new Error('boom'); + vi.mocked(startScript).mockRejectedValue(err); + const locals = mockLocals(); + + const res = await POST(postEvent({ statements: ['SELECT 1'], tabId: TAB_ID }, locals)); + + expect(res.status).toBe(204); + await vi.waitFor(() => + expect(locals.logger.error).toHaveBeenCalledWith( + expect.objectContaining({ err, tab_id: TAB_ID }), + 'script orchestrator crashed' + ) + ); + }); +}); + +describe('GET /api/trino/query', () => { + beforeEach(() => { + vi.clearAllMocks(); + vi.mocked(getQuerySnapshots).mockReturnValue([]); + }); + + it.each([ + ['missing', ''], + ['invalid', 'tabId=not-a-uuid'] + ])('returns 400 when tabId is %s', async (_label, params) => { + const res = await GET(urlEvent(params)); + + expect(res.status).toBe(400); + expect(getQuerySnapshots).not.toHaveBeenCalled(); + }); + + it('returns lightweight snapshots by default', async () => { + const snapshot = { state: 'RUNNING', sql: 'SELECT 1' }; + vi.mocked(getQuerySnapshots).mockReturnValue([snapshot] as never); + + const res = await GET(urlEvent(`tabId=${TAB_ID}`)); + + expect(res.status).toBe(200); + expect(await res.json()).toEqual([snapshot]); + expect(getQuerySnapshots).toHaveBeenCalledWith('u1', TAB_ID, true); + }); + + it('returns full snapshots when lightweight=false', async () => { + await GET(urlEvent(`tabId=${TAB_ID}&lightweight=false`)); + + expect(getQuerySnapshots).toHaveBeenCalledWith('u1', TAB_ID, false); + }); +}); + +describe('DELETE /api/trino/query', () => { + beforeEach(() => { + vi.clearAllMocks(); + vi.mocked(cancelQuery).mockResolvedValue(true); + }); + + it('returns 400 when tabId is invalid', async () => { + const res = await DELETE(urlEvent('tabId=nope')); + + expect(res.status).toBe(400); + expect(cancelQuery).not.toHaveBeenCalled(); + }); + + it('cancels the active query without removing tab state', async () => { + const res = await DELETE(urlEvent(`tabId=${TAB_ID}`)); + + expect(res.status).toBe(204); + expect(cancelQuery).toHaveBeenCalledWith('u1', TAB_ID); + expect(removeTabQuery).not.toHaveBeenCalled(); + }); + + it('cancels and removes tab state when cleanup=true', async () => { + const res = await DELETE(urlEvent(`tabId=${TAB_ID}&cleanup=true`)); + + expect(res.status).toBe(204); + expect(cancelQuery).toHaveBeenCalledWith('u1', TAB_ID); + expect(removeTabQuery).toHaveBeenCalledWith('u1', TAB_ID); + }); +}); From cf3a744a5d554a4f053c8538265e5591e34b4185 Mon Sep 17 00:00:00 2001 From: Johannes Fleck Date: Thu, 8 Oct 2026 13:28:50 +0200 Subject: [PATCH 2/7] test: cover Trino query store and script execution (#345) Covers sequential script execution and stop conditions, cancellation, snapshots, tab removal and the TTL-based eviction sweep. --- src/lib/server/trino/queries.test.ts | 493 +++++++++++++++++++++++++++ 1 file changed, 493 insertions(+) create mode 100644 src/lib/server/trino/queries.test.ts diff --git a/src/lib/server/trino/queries.test.ts b/src/lib/server/trino/queries.test.ts new file mode 100644 index 00000000..14a0722e --- /dev/null +++ b/src/lib/server/trino/queries.test.ts @@ -0,0 +1,493 @@ +import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'; +import type { TrinoClient, TrinoQueryResult } from './client.js'; +import type { TrinoQuery } from './queries.js'; + +// Stable mock functions shared across module re-imports (each test imports a +// fresh queries.ts so its in-memory store starts empty). +const mocks = vi.hoisted(() => ({ + env: {} as Record, + collectResults: vi.fn(), + resolveTrinoPublicUrl: vi.fn(), + queryTotalInc: vi.fn(), + activeInc: vi.fn(), + activeDec: vi.fn() +})); + +vi.mock('$env/dynamic/private', () => ({ env: mocks.env })); +vi.mock('$lib/server/logging', () => ({ + logger: { child: () => ({ info: vi.fn(), warn: vi.fn(), debug: vi.fn(), error: vi.fn() }) } +})); +vi.mock('$lib/server/metrics.js', () => ({ + trinoQueryTotal: { inc: mocks.queryTotalInc }, + trinoActiveQueries: { inc: mocks.activeInc, dec: mocks.activeDec } +})); +vi.mock('./client.js', () => ({ resolveTrinoPublicUrl: mocks.resolveTrinoPublicUrl })); +vi.mock('./result-collector.js', () => ({ collectResults: mocks.collectResults })); + +type Queries = typeof import('./queries.js'); + +const USER = 'u1'; +const TAB = 'tab-1'; +const OPTS = { user: 'alice', catalog: 'tpch', schema: 'tiny' }; +const SWEEP_INTERVAL_MS = 60_000; + +let q: Queries; +let client: { + submit: ReturnType; + cancel: ReturnType; + cancelViaUri: ReturnType; +}; + +function trinoClient(): TrinoClient { + return client as unknown as TrinoClient; +} + +/** A submit response that hands off to collectResults via nextUri. */ +function pending(id: string): TrinoQueryResult { + return { id, nextUri: `http://trino/v1/statement/queued/${id}/1`, stats: { state: 'QUEUED' } }; +} + +/** Make collectResults drive each polled query to the given terminal state. */ +function collectTo(state: 'FINISHED' | 'FAILED' | 'CANCELLED', error: string | null = null) { + return async (query: TrinoQuery) => { + query.columns = [{ name: 'x', type: 'integer' }]; + query.rows = [[1]]; + query.error = error; + q.terminateQuery(query, state); + }; +} + +/** Release functions of hung collectResults calls, flushed after each test. */ +const pendingReleases: (() => void)[] = []; + +/** Make collectResults hang until the returned release function is called. */ +function hangCollect(mutate?: (query: TrinoQuery) => void) { + let release!: () => void; + mocks.collectResults.mockImplementationOnce((query: TrinoQuery) => { + mutate?.(query); + return new Promise((resolve) => { + release = resolve; + pendingReleases.push(resolve); + }); + }); + return () => release(); +} + +async function loadQueries(env: Record = {}) { + for (const key of Object.keys(mocks.env)) delete mocks.env[key]; + Object.assign(mocks.env, env); + vi.resetModules(); + q = await import('./queries.js'); +} + +beforeEach(async () => { + vi.clearAllMocks(); + vi.useFakeTimers(); + mocks.resolveTrinoPublicUrl.mockReturnValue('https://trino.example.com'); + // mockReset also drops queued *Once implementations a failed test may have left behind. + mocks.collectResults.mockReset().mockImplementation(collectTo('FINISHED')); + client = { + submit: vi.fn(async (sql: string) => pending(`q-${sql}`)), + cancel: vi.fn(async () => {}), + cancelViaUri: vi.fn(async () => {}) + }; + await loadQueries(); +}); + +afterEach(() => { + // Unblock scripts left hanging by a test that failed before calling release(). + pendingReleases.splice(0).forEach((release) => release()); + vi.useRealTimers(); +}); + +describe('mapTrinoState', () => { + it.each(['QUEUED', 'PLANNING', 'RUNNING', 'FINISHING', 'FINISHED', 'FAILED'])( + 'passes %s through', + (state) => { + expect(q.mapTrinoState(state)).toBe(state); + } + ); + + it.each([['STARTING'], ['BLOCKED'], [undefined]])('maps %s to RUNNING', (state) => { + expect(q.mapTrinoState(state)).toBe('RUNNING'); + }); +}); + +describe('toQueryProgress', () => { + it('returns zeroed progress without stats', () => { + expect(q.toQueryProgress(undefined)).toEqual({ + progressPercentage: 0, + processedRows: 0, + elapsedTimeMillis: 0 + }); + }); + + it('defaults missing stat fields to 0', () => { + expect(q.toQueryProgress({ state: 'RUNNING', processedRows: 42 })).toEqual({ + progressPercentage: 0, + processedRows: 42, + elapsedTimeMillis: 0 + }); + }); +}); + +describe('terminateQuery', () => { + function runningQuery(): TrinoQuery { + return { state: 'RUNNING', completedAt: null } as TrinoQuery; + } + + it('applies only the first termination and decrements the gauge once', () => { + const query = runningQuery(); + + q.terminateQuery(query, 'CANCELLED'); + q.terminateQuery(query, 'FAILED'); + + expect(query.state).toBe('CANCELLED'); + expect(query.completedAt).toBe(Date.now()); + expect(mocks.activeDec).toHaveBeenCalledTimes(1); + }); + + it('leaves the gauge alone when decrementGauge is false', () => { + q.terminateQuery(runningQuery(), 'FAILED', { decrementGauge: false }); + + expect(mocks.activeDec).not.toHaveBeenCalled(); + }); +}); + +describe('startScript', () => { + it('runs statements in order and stores a result per statement', async () => { + await q.startScript(trinoClient(), USER, TAB, ['SELECT 1', 'SELECT 2'], OPTS); + + expect(client.submit.mock.calls).toEqual([ + ['SELECT 1', OPTS], + ['SELECT 2', OPTS] + ]); + expect(mocks.collectResults).toHaveBeenCalledTimes(2); + expect(mocks.activeInc).toHaveBeenCalledTimes(2); + expect(mocks.activeDec).toHaveBeenCalledTimes(2); + expect(mocks.queryTotalInc).toHaveBeenCalledWith({ outcome: 'submitted' }); + + const snapshots = q.getQuerySnapshots(USER, TAB); + expect(snapshots.map((s) => [s.sql, s.state])).toEqual([ + ['SELECT 1', 'FINISHED'], + ['SELECT 2', 'FINISHED'] + ]); + }); + + it('stops at the first statement that does not finish', async () => { + mocks.collectResults.mockImplementationOnce(collectTo('FAILED', 'Table not found')); + + await q.startScript(trinoClient(), USER, TAB, ['SELECT 1', 'SELECT 2'], OPTS); + + expect(client.submit).toHaveBeenCalledTimes(1); + const [snapshot] = q.getQuerySnapshots(USER, TAB); + expect(snapshot).toMatchObject({ state: 'FAILED', error: 'Table not found' }); + }); + + it('marks a statement without nextUri as finished without polling', async () => { + client.submit.mockResolvedValueOnce({ id: 'q1', stats: { state: 'RUNNING' } }); + + await q.startScript(trinoClient(), USER, TAB, ['SELECT 1', 'SELECT 2'], OPTS); + + expect(mocks.collectResults).toHaveBeenCalledTimes(1); // only the second statement + expect(mocks.activeInc).toHaveBeenCalledTimes(1); + expect(mocks.queryTotalInc).toHaveBeenCalledWith({ outcome: 'completed' }); + expect(q.getQuerySnapshots(USER, TAB).map((s) => s.state)).toEqual(['FINISHED', 'FINISHED']); + }); + + it('keeps results returned directly in the submit response', async () => { + client.submit.mockResolvedValueOnce({ + id: 'q1', + columns: [{ name: 'n', type: 'bigint' }], + data: [[7]], + stats: { state: 'FINISHED', processedRows: 1 } + }); + + await q.startScript(trinoClient(), USER, TAB, ['SELECT 7'], OPTS); + + expect(mocks.collectResults).not.toHaveBeenCalled(); + const [snapshot] = q.getQuerySnapshots(USER, TAB); + expect(snapshot).toMatchObject({ + state: 'FINISHED', + columns: [{ name: 'n', type: 'bigint' }], + rows: [[7]], + progress: { processedRows: 1 } + }); + }); + + it('fails and stops on an error in the submit response', async () => { + client.submit.mockResolvedValueOnce({ + id: 'q1', + error: { message: 'line 1:1: mismatched input' } + }); + + await q.startScript(trinoClient(), USER, TAB, ['SELEC 1', 'SELECT 2'], OPTS); + + expect(client.submit).toHaveBeenCalledTimes(1); + expect(mocks.collectResults).not.toHaveBeenCalled(); + expect(mocks.activeDec).not.toHaveBeenCalled(); + expect(mocks.queryTotalInc).toHaveBeenCalledWith({ outcome: 'failed' }); + const [snapshot] = q.getQuerySnapshots(USER, TAB); + expect(snapshot).toMatchObject({ state: 'FAILED', error: 'line 1:1: mismatched input' }); + }); + + it('falls back to a generic message when the submit error has none', async () => { + client.submit.mockResolvedValueOnce({ id: 'q1', error: {} }); + + await q.startScript(trinoClient(), USER, TAB, ['SELECT 1'], OPTS); + + expect(q.getQuerySnapshots(USER, TAB)[0].error).toBe('Query failed'); + }); + + it('stops without storing a result when submit throws', async () => { + client.submit.mockRejectedValueOnce(new Error('ECONNREFUSED')); + + await expect( + q.startScript(trinoClient(), USER, TAB, ['SELECT 1', 'SELECT 2'], OPTS) + ).resolves.toBeUndefined(); + + expect(client.submit).toHaveBeenCalledTimes(1); + expect(mocks.queryTotalInc).toHaveBeenCalledWith({ outcome: 'failed' }); + expect(q.getQuerySnapshots(USER, TAB)).toEqual([]); + }); + + it('stops when Trino returns no query ID', async () => { + client.submit.mockResolvedValueOnce({ id: '' }); + + await q.startScript(trinoClient(), USER, TAB, ['SELECT 1', 'SELECT 2'], OPTS); + + expect(client.submit).toHaveBeenCalledTimes(1); + expect(q.getQuerySnapshots(USER, TAB)).toEqual([]); + }); + + it('fails the query when the poll loop crashes', async () => { + mocks.collectResults.mockRejectedValueOnce(new Error('unexpected')); + + await q.startScript(trinoClient(), USER, TAB, ['SELECT 1', 'SELECT 2'], OPTS); + + expect(client.submit).toHaveBeenCalledTimes(1); + expect(mocks.activeDec).toHaveBeenCalledTimes(1); + expect(mocks.queryTotalInc).toHaveBeenCalledWith({ outcome: 'failed' }); + expect(q.getQuerySnapshots(USER, TAB)[0]).toMatchObject({ + state: 'FAILED', + error: 'Internal poll error' + }); + }); + + it('replaces earlier results for the same tab', async () => { + await q.startScript(trinoClient(), USER, TAB, ['SELECT 1', 'SELECT 2'], OPTS); + await q.startScript(trinoClient(), USER, TAB, ['SELECT 3'], OPTS); + + expect(q.getQuerySnapshots(USER, TAB).map((s) => s.sql)).toEqual(['SELECT 3']); + }); + + it('cancels a still-running query before starting a new script on the tab', async () => { + const release = hangCollect(); + const first = q.startScript(trinoClient(), USER, TAB, ['SELECT slow'], OPTS); + await vi.waitFor(() => expect(mocks.collectResults).toHaveBeenCalled()); + + await q.startScript(trinoClient(), USER, TAB, ['SELECT fast'], OPTS); + release(); + await first; + + expect(client.cancelViaUri).toHaveBeenCalledWith(pending('q-SELECT slow').nextUri, 'alice'); + expect(q.getQuerySnapshots(USER, TAB).map((s) => s.sql)).toEqual(['SELECT fast']); + }); + + it('keeps tabs and users isolated from each other', async () => { + await q.startScript(trinoClient(), USER, TAB, ['SELECT 1'], OPTS); + await q.startScript(trinoClient(), USER, 'tab-2', ['SELECT 2'], OPTS); + await q.startScript(trinoClient(), 'u2', TAB, ['SELECT 3'], OPTS); + + expect(q.getQuerySnapshots(USER, TAB).map((s) => s.sql)).toEqual(['SELECT 1']); + expect(q.getQuerySnapshots(USER, 'tab-2').map((s) => s.sql)).toEqual(['SELECT 2']); + expect(q.getQuerySnapshots('u2', TAB).map((s) => s.sql)).toEqual(['SELECT 3']); + }); +}); + +describe('getQuerySnapshots', () => { + it('returns an empty list for an unknown tab', () => { + expect(q.getQuerySnapshots(USER, 'unknown')).toEqual([]); + }); + + it('includes rows and columns unless lightweight', async () => { + await q.startScript(trinoClient(), USER, TAB, ['SELECT 1'], OPTS); + + expect(q.getQuerySnapshots(USER, TAB)[0]).toMatchObject({ + columns: [{ name: 'x', type: 'integer' }], + rows: [[1]] + }); + expect(q.getQuerySnapshots(USER, TAB, true)[0]).toMatchObject({ columns: [], rows: [] }); + }); + + it('links to the Trino UI via the public URL', async () => { + await q.startScript(trinoClient(), USER, TAB, ['SELECT 1'], OPTS); + + expect(q.getQuerySnapshots(USER, TAB)[0].trinoQueryUrl).toBe( + 'https://trino.example.com/ui/query.html?q-SELECT 1' + ); + expect(mocks.resolveTrinoPublicUrl).toHaveBeenCalledWith(USER); + }); + + it('omits the Trino UI link without a public URL', async () => { + mocks.resolveTrinoPublicUrl.mockReturnValue(null); + await q.startScript(trinoClient(), USER, TAB, ['SELECT 1'], OPTS); + + expect(q.getQuerySnapshots(USER, TAB)[0].trinoQueryUrl).toBeNull(); + }); +}); + +describe('getAllQuerySummaries', () => { + it('returns an empty object for a user without queries', () => { + expect(q.getAllQuerySummaries(USER)).toEqual({}); + }); + + it('groups lightweight snapshots by tab', async () => { + await q.startScript(trinoClient(), USER, TAB, ['SELECT 1', 'SELECT 2'], OPTS); + await q.startScript(trinoClient(), USER, 'tab-2', ['SELECT 3'], OPTS); + + const summaries = q.getAllQuerySummaries(USER); + + expect(Object.keys(summaries)).toEqual([TAB, 'tab-2']); + expect(summaries[TAB].map((s) => s.sql)).toEqual(['SELECT 1', 'SELECT 2']); + expect(summaries[TAB][0]).toMatchObject({ columns: [], rows: [] }); + }); +}); + +describe('removeTabQuery', () => { + it('drops the tab and leaves other tabs intact', async () => { + await q.startScript(trinoClient(), USER, TAB, ['SELECT 1'], OPTS); + await q.startScript(trinoClient(), USER, 'tab-2', ['SELECT 2'], OPTS); + + q.removeTabQuery(USER, TAB); + + expect(Object.keys(q.getAllQuerySummaries(USER))).toEqual(['tab-2']); + }); + + it('drops the user entry once the last tab is removed', async () => { + await q.startScript(trinoClient(), USER, TAB, ['SELECT 1'], OPTS); + + q.removeTabQuery(USER, TAB); + + expect(q.getAllQuerySummaries(USER)).toEqual({}); + }); + + it('is a no-op for an unknown user', () => { + expect(() => q.removeTabQuery('nobody', TAB)).not.toThrow(); + }); +}); + +describe('cancelQuery', () => { + it('returns false when the tab has no queries', async () => { + expect(await q.cancelQuery(USER, TAB)).toBe(false); + }); + + it('returns false when the last query is already terminal', async () => { + await q.startScript(trinoClient(), USER, TAB, ['SELECT 1'], OPTS); + + expect(await q.cancelQuery(USER, TAB)).toBe(false); + expect(client.cancelViaUri).not.toHaveBeenCalled(); + expect(client.cancel).not.toHaveBeenCalled(); + }); + + it('cancels via nextUri while results are still pending', async () => { + const release = hangCollect(); + const script = q.startScript(trinoClient(), USER, TAB, ['SELECT 1', 'SELECT 2'], OPTS); + await vi.waitFor(() => expect(mocks.collectResults).toHaveBeenCalled()); + + expect(await q.cancelQuery(USER, TAB)).toBe(true); + release(); + await script; + + expect(client.cancelViaUri).toHaveBeenCalledWith(pending('q-SELECT 1').nextUri, 'alice'); + expect(client.cancel).not.toHaveBeenCalled(); + expect(mocks.activeDec).toHaveBeenCalledTimes(1); + expect(mocks.queryTotalInc).toHaveBeenCalledWith({ outcome: 'cancelled' }); + // The script does not continue with the next statement after cancellation. + expect(client.submit).toHaveBeenCalledTimes(1); + expect(q.getQuerySnapshots(USER, TAB).map((s) => s.state)).toEqual(['CANCELLED']); + }); + + it('cancels by query ID when there is no nextUri', async () => { + const release = hangCollect((query) => (query.nextUri = undefined)); + const script = q.startScript(trinoClient(), USER, TAB, ['SELECT 1'], OPTS); + await vi.waitFor(() => expect(mocks.collectResults).toHaveBeenCalled()); + + await q.cancelQuery(USER, TAB); + release(); + await script; + + expect(client.cancel).toHaveBeenCalledWith('q-SELECT 1', 'alice'); + expect(client.cancelViaUri).not.toHaveBeenCalled(); + }); + + it('still marks the query cancelled when Trino rejects the cancel', async () => { + client.cancelViaUri.mockRejectedValueOnce(new Error('503')); + const release = hangCollect(); + const script = q.startScript(trinoClient(), USER, TAB, ['SELECT 1'], OPTS); + await vi.waitFor(() => expect(mocks.collectResults).toHaveBeenCalled()); + + expect(await q.cancelQuery(USER, TAB)).toBe(true); + release(); + await script; + + expect(q.getQuerySnapshots(USER, TAB)[0].state).toBe('CANCELLED'); + }); +}); + +describe('expired query sweep', () => { + const DEFAULT_TTL_MS = 1800 * 1000; + + it('evicts finished tabs that have not been accessed within the TTL', async () => { + await q.startScript(trinoClient(), USER, TAB, ['SELECT 1'], OPTS); + + vi.advanceTimersByTime(DEFAULT_TTL_MS - SWEEP_INTERVAL_MS); + expect(Object.keys(q.getAllQuerySummaries(USER))).toEqual([TAB]); + + vi.advanceTimersByTime(SWEEP_INTERVAL_MS); + expect(q.getAllQuerySummaries(USER)).toEqual({}); + }); + + it('keeps tabs that were polled recently', async () => { + await q.startScript(trinoClient(), USER, TAB, ['SELECT 1'], OPTS); + + vi.advanceTimersByTime(DEFAULT_TTL_MS / 2); + q.getQuerySnapshots(USER, TAB); // a client poll refreshes the access time + vi.advanceTimersByTime(DEFAULT_TTL_MS / 2 + SWEEP_INTERVAL_MS); + + expect(Object.keys(q.getAllQuerySummaries(USER))).toEqual([TAB]); + }); + + it('never evicts a tab with an active query', async () => { + const release = hangCollect(); + const script = q.startScript(trinoClient(), USER, TAB, ['SELECT 1'], OPTS); + await vi.waitFor(() => expect(mocks.collectResults).toHaveBeenCalled()); + + vi.advanceTimersByTime(DEFAULT_TTL_MS * 2); + + expect(Object.keys(q.getAllQuerySummaries(USER))).toEqual([TAB]); + release(); + await script; + }); + + it('only evicts the expired tabs of a user', async () => { + await q.startScript(trinoClient(), USER, TAB, ['SELECT 1'], OPTS); + vi.advanceTimersByTime(DEFAULT_TTL_MS / 2); + await q.startScript(trinoClient(), USER, 'tab-2', ['SELECT 2'], OPTS); + + vi.advanceTimersByTime(DEFAULT_TTL_MS / 2 + SWEEP_INTERVAL_MS); + + expect(Object.keys(q.getAllQuerySummaries(USER))).toEqual(['tab-2']); + }); + + it('honours STACKABLE_COCKPIT_QUERY_TTL', async () => { + await loadQueries({ STACKABLE_COCKPIT_QUERY_TTL: '120' }); + await q.startScript(trinoClient(), USER, TAB, ['SELECT 1'], OPTS); + + vi.advanceTimersByTime(60_000); + expect(Object.keys(q.getAllQuerySummaries(USER))).toEqual([TAB]); + + vi.advanceTimersByTime(120_000); + expect(q.getAllQuerySummaries(USER)).toEqual({}); + }); +}); From a280fc3b8102c3a3ce687b0d3224cdbd835ec014 Mon Sep 17 00:00:00 2001 From: Johannes Fleck Date: Thu, 8 Oct 2026 13:28:50 +0200 Subject: [PATCH 3/7] test: cover per-user Trino connections (#345) --- src/lib/server/trino/user-clients.test.ts | 87 +++++++++++++++++++++++ 1 file changed, 87 insertions(+) create mode 100644 src/lib/server/trino/user-clients.test.ts diff --git a/src/lib/server/trino/user-clients.test.ts b/src/lib/server/trino/user-clients.test.ts new file mode 100644 index 00000000..a48b00e3 --- /dev/null +++ b/src/lib/server/trino/user-clients.test.ts @@ -0,0 +1,87 @@ +import { describe, it, expect, vi, beforeEach } from 'vitest'; + +const { TrinoClientMock } = vi.hoisted(() => ({ + TrinoClientMock: vi.fn(function (this: Record, options: unknown) { + this.options = options; + }) +})); + +vi.mock('$lib/server/logging', () => ({ + logger: { child: () => ({ info: vi.fn(), warn: vi.fn(), debug: vi.fn() }) } +})); +vi.mock('./client.js', () => ({ + TrinoClient: TrinoClientMock, + buildBasicAuthHeader: (user: string, pass: string) => `Basic ${user}:${pass}`, + trinoUserImpersonation: true +})); + +type UserClients = typeof import('./user-clients.js'); +let uc: UserClients; + +const BASE = { url: 'http://trino:8080', authType: 'none', username: '', password: '' } as const; + +beforeEach(async () => { + vi.clearAllMocks(); + vi.resetModules(); + uc = await import('./user-clients.js'); +}); + +describe('buildUserTrinoClient', () => { + it('builds an unauthenticated client for authType none', () => { + uc.buildUserTrinoClient({ ...BASE, username: 'alice', password: 'secret' }); + + expect(TrinoClientMock).toHaveBeenCalledWith({ + serverUrl: 'http://trino:8080', + authorization: undefined, + impersonate: true + }); + }); + + it('sends Basic auth when both credentials are set', () => { + uc.buildUserTrinoClient({ ...BASE, authType: 'basic', username: 'alice', password: 'secret' }); + + expect(TrinoClientMock).toHaveBeenCalledWith( + expect.objectContaining({ authorization: 'Basic alice:secret' }) + ); + }); + + it.each([ + ['username', { username: '', password: 'secret' }], + ['password', { username: 'alice', password: '' }] + ])('omits Basic auth when the %s is missing', (_label, creds) => { + uc.buildUserTrinoClient({ ...BASE, authType: 'basic', ...creds }); + + expect(TrinoClientMock).toHaveBeenCalledWith( + expect.objectContaining({ authorization: undefined }) + ); + }); + + it('does not store the client', () => { + uc.buildUserTrinoClient(BASE); + + expect(uc.getUserTrinoClient('u1')).toBeNull(); + }); +}); + +describe('per-user connections', () => { + it('returns null for a user without a connection', () => { + expect(uc.getUserTrinoClient('u1')).toBeNull(); + expect(uc.getUserTrinoUrl('u1')).toBeNull(); + }); + + it('stores the client and URL per user', () => { + uc.createUserTrinoClient('u1', BASE); + + expect(uc.getUserTrinoClient('u1')).toBe(TrinoClientMock.mock.instances[0]); + expect(uc.getUserTrinoUrl('u1')).toBe('http://trino:8080'); + expect(uc.getUserTrinoClient('u2')).toBeNull(); + }); + + it('replaces an existing connection', () => { + uc.createUserTrinoClient('u1', BASE); + uc.createUserTrinoClient('u1', { ...BASE, url: 'http://other:8443' }); + + expect(uc.getUserTrinoClient('u1')).toBe(TrinoClientMock.mock.instances[1]); + expect(uc.getUserTrinoUrl('u1')).toBe('http://other:8443'); + }); +}); From 915fe714a3ff014b491594211e1ab8cb46007ded Mon Sep 17 00:00:00 2001 From: Johannes Fleck Date: Thu, 8 Oct 2026 13:35:48 +0200 Subject: [PATCH 4/7] fix: handle edge cases in Trino submit responses (#345) - Count a submit without query ID as failed, so trino_query_total stays balanced. - Don't expose a terminal state straight from the submit response. The query previously skipped draining any remaining pages and never got a completion time or outcome metric. --- src/lib/server/trino/queries.test.ts | 26 ++++++++++++++++++++++++++ src/lib/server/trino/queries.ts | 7 ++++++- 2 files changed, 32 insertions(+), 1 deletion(-) diff --git a/src/lib/server/trino/queries.test.ts b/src/lib/server/trino/queries.test.ts index 14a0722e..0462197f 100644 --- a/src/lib/server/trino/queries.test.ts +++ b/src/lib/server/trino/queries.test.ts @@ -206,6 +206,7 @@ describe('startScript', () => { await q.startScript(trinoClient(), USER, TAB, ['SELECT 7'], OPTS); expect(mocks.collectResults).not.toHaveBeenCalled(); + expect(mocks.queryTotalInc).toHaveBeenCalledWith({ outcome: 'completed' }); const [snapshot] = q.getQuerySnapshots(USER, TAB); expect(snapshot).toMatchObject({ state: 'FINISHED', @@ -215,6 +216,30 @@ describe('startScript', () => { }); }); + it('still drains remaining pages when the submit response is already finished', async () => { + client.submit.mockResolvedValueOnce({ + id: 'q1', + nextUri: 'http://trino/v1/statement/executing/q1/1', + data: [[1]], + stats: { state: 'FINISHED' } + }); + let stateWhileDraining: string | undefined; + mocks.collectResults.mockImplementationOnce(async (query: TrinoQuery) => { + stateWhileDraining = q.getQuerySnapshots(USER, TAB)[0].state; + query.rows.push([2]); + q.terminateQuery(query, 'FINISHED'); + }); + + await q.startScript(trinoClient(), USER, TAB, ['SELECT 1'], OPTS); + + expect(mocks.collectResults).toHaveBeenCalledTimes(1); + expect(stateWhileDraining).toBe('RUNNING'); + expect(q.getQuerySnapshots(USER, TAB)[0]).toMatchObject({ + state: 'FINISHED', + rows: [[1], [2]] + }); + }); + it('fails and stops on an error in the submit response', async () => { client.submit.mockResolvedValueOnce({ id: 'q1', @@ -257,6 +282,7 @@ describe('startScript', () => { await q.startScript(trinoClient(), USER, TAB, ['SELECT 1', 'SELECT 2'], OPTS); expect(client.submit).toHaveBeenCalledTimes(1); + expect(mocks.queryTotalInc).toHaveBeenCalledWith({ outcome: 'failed' }); expect(q.getQuerySnapshots(USER, TAB)).toEqual([]); }); diff --git a/src/lib/server/trino/queries.ts b/src/lib/server/trino/queries.ts index 935bebcc..835ec5dc 100644 --- a/src/lib/server/trino/queries.ts +++ b/src/lib/server/trino/queries.ts @@ -155,12 +155,17 @@ async function submitStatement( const trinoQueryId = submitResult.id; if (!trinoQueryId) { + trinoQueryTotal.inc({ outcome: 'failed' }); throw new Error('Trino did not return a query ID'); } + const initialState = submitResult.stats ? mapTrinoState(submitResult.stats.state) : 'QUEUED'; + const query: TrinoQuery = { trinoQueryId, - state: submitResult.stats ? mapTrinoState(submitResult.stats.state) : 'QUEUED', + // As in collectResults: don't expose a terminal state before all result pages + // are drained. startScript terminates the query once that has happened. + state: isTerminal(initialState) ? 'RUNNING' : initialState, progress: toQueryProgress(submitResult.stats), columns: submitResult.columns ?? [], rows: submitResult.data ?? [], From d2b9c722e160ccd3e58538d70829cb724980bc29 Mon Sep 17 00:00:00 2001 From: Johannes Fleck Date: Thu, 8 Oct 2026 13:47:09 +0200 Subject: [PATCH 5/7] fix: stop superseded and cancelled Trino scripts between statements (#345) While a script waited for Trino to accept its next statement, the tab had no active query: cancelling did nothing, and a re-run's results were mixed with those of the old script, which kept executing. Track the running script per tab with an AbortController. cancelQuery, resetTabQueries and a new script abort it; a statement accepted after the abort is cancelled in Trino instead of stored. Cancellation now updates state synchronously before calling Trino, so concurrent callers cannot cancel or count a query twice and a reset cannot wipe a newer script. --- src/lib/server/trino/queries.test.ts | 136 +++++++++++++++++++++++++- src/lib/server/trino/queries.ts | 138 +++++++++++++++++++++++---- 2 files changed, 252 insertions(+), 22 deletions(-) diff --git a/src/lib/server/trino/queries.test.ts b/src/lib/server/trino/queries.test.ts index 0462197f..f821fcb2 100644 --- a/src/lib/server/trino/queries.test.ts +++ b/src/lib/server/trino/queries.test.ts @@ -57,9 +57,26 @@ function collectTo(state: 'FINISHED' | 'FAILED' | 'CANCELLED', error: string | n }; } -/** Release functions of hung collectResults calls, flushed after each test. */ +/** Release functions of hung mock calls, flushed after each test. */ const pendingReleases: (() => void)[] = []; +function hang(): { promise: Promise; release: () => void } { + let release!: () => void; + const promise = new Promise((resolve) => (release = resolve)); + pendingReleases.push(release); + return { promise, release }; +} + +/** Make client.submit hang for the given SQL until the returned release function is called. */ +function hangSubmit(target: string) { + const { promise, release } = hang(); + client.submit.mockImplementation(async (sql: string) => { + if (sql === target) await promise; + return pending(`q-${sql}`); + }); + return release; +} + /** Make collectResults hang until the returned release function is called. */ function hangCollect(mutate?: (query: TrinoQuery) => void) { let release!: () => void; @@ -331,6 +348,123 @@ describe('startScript', () => { }); }); +describe('overlapping scripts in one tab', () => { + it('stops a script superseded while its next statement is being submitted', async () => { + const release = hangSubmit('SELECT 2'); + const first = q.startScript( + trinoClient(), + USER, + TAB, + ['SELECT 1', 'SELECT 2', 'SELECT 4'], + OPTS + ); + await vi.waitFor(() => expect(client.submit).toHaveBeenCalledTimes(2)); + + await q.startScript(trinoClient(), USER, TAB, ['SELECT 3'], OPTS); + release(); + await first; + + expect(q.getQuerySnapshots(USER, TAB).map((s) => s.sql)).toEqual(['SELECT 3']); + // The statement Trino accepted for the stale script is cancelled, and nothing after it runs. + expect(client.cancelViaUri).toHaveBeenCalledWith(pending('q-SELECT 2').nextUri, 'alice'); + expect(client.submit.mock.calls.map(([sql]) => sql)).toEqual([ + 'SELECT 1', + 'SELECT 2', + 'SELECT 3' + ]); + expect(mocks.queryTotalInc).toHaveBeenCalledWith({ outcome: 'cancelled' }); + expect(mocks.activeInc).toHaveBeenCalledTimes(2); + expect(mocks.activeDec).toHaveBeenCalledTimes(2); + }); + + it('cancels a script between statements', async () => { + const release = hangSubmit('SELECT 2'); + const script = q.startScript( + trinoClient(), + USER, + TAB, + ['SELECT 1', 'SELECT 2', 'SELECT 3'], + OPTS + ); + await vi.waitFor(() => expect(client.submit).toHaveBeenCalledTimes(2)); + + expect(await q.cancelQuery(USER, TAB)).toBe(true); + release(); + await script; + + expect(client.cancelViaUri).toHaveBeenCalledWith(pending('q-SELECT 2').nextUri, 'alice'); + expect(client.submit).toHaveBeenCalledTimes(2); + expect(q.getQuerySnapshots(USER, TAB).map((s) => [s.sql, s.state])).toEqual([ + ['SELECT 1', 'FINISHED'] + ]); + expect(await q.cancelQuery(USER, TAB)).toBe(false); + }); + + it('keeps only the latest script when two start while a query is still being cancelled', async () => { + hangCollect(); + const running = q.startScript(trinoClient(), USER, TAB, ['SELECT slow'], OPTS); + await vi.waitFor(() => expect(mocks.collectResults).toHaveBeenCalled()); + const cancel = hang(); + client.cancelViaUri.mockImplementationOnce(() => cancel.promise); + + const first = q.startScript(trinoClient(), USER, TAB, ['SELECT first'], OPTS); + const second = q.startScript(trinoClient(), USER, TAB, ['SELECT second'], OPTS); + await second; + cancel.release(); + await first; + pendingReleases.splice(0).forEach((r) => r()); + await running; + + expect(client.submit.mock.calls.map(([sql]) => sql)).toEqual(['SELECT slow', 'SELECT second']); + expect(q.getQuerySnapshots(USER, TAB).map((s) => s.sql)).toEqual(['SELECT second']); + // The slow query is cancelled and counted exactly once. + expect(client.cancelViaUri).toHaveBeenCalledTimes(1); + expect(mocks.queryTotalInc.mock.calls.filter(([l]) => l.outcome === 'cancelled')).toHaveLength( + 1 + ); + }); + + it('does not run the next statement when the current one finishes during cancellation', async () => { + const poll = hang(); + mocks.collectResults.mockImplementationOnce(async (query: TrinoQuery) => { + await poll.promise; + q.terminateQuery(query, 'FINISHED'); // last page arrives after the cancel + }); + const script = q.startScript(trinoClient(), USER, TAB, ['SELECT 1', 'DROP TABLE t'], OPTS); + await vi.waitFor(() => expect(mocks.collectResults).toHaveBeenCalled()); + const cancel = hang(); + client.cancelViaUri.mockImplementationOnce(() => cancel.promise); + + const cancelling = q.cancelQuery(USER, TAB); + poll.release(); + await script; + cancel.release(); + await cancelling; + + expect(client.submit).toHaveBeenCalledTimes(1); + expect(q.getQuerySnapshots(USER, TAB).map((s) => s.state)).toEqual(['CANCELLED']); + }); + + it('keeps results of a script started while the tab is being reset', async () => { + hangCollect(); + const running = q.startScript(trinoClient(), USER, TAB, ['SELECT slow'], OPTS); + await vi.waitFor(() => expect(mocks.collectResults).toHaveBeenCalled()); + const cancel = hang(); + client.cancelViaUri.mockImplementationOnce(() => cancel.promise); + + const reset = q.resetTabQueries(USER, TAB); + await q.startScript(trinoClient(), USER, TAB, ['SELECT new'], OPTS); + cancel.release(); + await reset; + pendingReleases.splice(0).forEach((r) => r()); + await running; + + expect(q.getQuerySnapshots(USER, TAB).map((s) => [s.sql, s.state])).toEqual([ + ['SELECT new', 'FINISHED'] + ]); + }); +}); + describe('getQuerySnapshots', () => { it('returns an empty list for an unknown tab', () => { expect(q.getQuerySnapshots(USER, 'unknown')).toEqual([]); diff --git a/src/lib/server/trino/queries.ts b/src/lib/server/trino/queries.ts index 835ec5dc..b634dc54 100644 --- a/src/lib/server/trino/queries.ts +++ b/src/lib/server/trino/queries.ts @@ -79,6 +79,39 @@ const userQueries = new Map>(); /** Tracks when each tab was last accessed (userId → tabId → timestamp). */ const tabLastAccessed = new Map>(); +/** Abort handle of the script currently executing in each tab (userId → tabId → controller). */ +const runningScripts = new Map>(); + +/** Register a new script run for a tab, superseding any run still registered. */ +function beginScriptRun(userId: string, tabId: string): AbortController { + let userMap = runningScripts.get(userId); + if (!userMap) { + userMap = new Map(); + runningScripts.set(userId, userMap); + } + userMap.get(tabId)?.abort(); + const run = new AbortController(); + userMap.set(tabId, run); + return run; +} + +/** Unregister a script run, unless a newer run has already replaced it. */ +function endScriptRun(userId: string, tabId: string, run: AbortController): void { + const userMap = runningScripts.get(userId); + if (userMap?.get(tabId) !== run) return; + userMap.delete(tabId); + if (userMap.size === 0) runningScripts.delete(userId); +} + +/** Abort the script running in a tab, if any. Returns whether one was running. */ +function abortScriptRun(userId: string, tabId: string): boolean { + const run = runningScripts.get(userId)?.get(tabId); + if (!run) return false; + run.abort(); + endScriptRun(userId, tabId, run); + return true; +} + function touchTab(userId: string, tabId: string): void { let userMap = tabLastAccessed.get(userId); if (!userMap) { @@ -130,12 +163,14 @@ function buildSnapshot( /** Clear all stored results for a tab and cancel any active query. */ export async function resetTabQueries(userId: string, tabId: string): Promise { - await cancelQuery(userId, tabId); - const tabMap = getUserTabMap(userId); - tabMap.set(tabId, []); + // Clear before awaiting Trino, so a script started meanwhile keeps its results. + abortScriptRun(userId, tabId); + const cancelled = markActiveQueryCancelled(userId, tabId); + getUserTabMap(userId).set(tabId, []); + if (cancelled) await cancelInTrino(cancelled); } -/** Submit a single SQL statement to Trino and store the TrinoQuery. */ +/** Submit a single SQL statement to Trino. The caller decides whether to store the TrinoQuery. */ async function submitStatement( client: TrinoClient, userId: string, @@ -185,19 +220,35 @@ async function submitStatement( trinoQueryTotal.inc({ outcome: 'failed' }); } + log.info({ trino_query_id: trinoQueryId, user_id: userId, tab_id: tabId }, 'query started'); + return query; +} + +function storeQuery(userId: string, tabId: string, query: TrinoQuery): void { const tabMap = getUserTabMap(userId); const queries = tabMap.get(tabId) ?? []; queries.push(query); tabMap.set(tabId, queries); +} - log.info({ trino_query_id: trinoQueryId, user_id: userId, tab_id: tabId }, 'query started'); - return query; +/** Best-effort cancellation of a query on the Trino side. */ +async function cancelInTrino(query: TrinoQuery): Promise { + try { + if (query.nextUri) { + await query.client.cancelViaUri(query.nextUri, query.trinoUser); + } else { + await query.client.cancel(query.trinoQueryId, query.trinoUser); + } + } catch (err) { + log.warn({ err, trino_query_id: query.trinoQueryId }, 'failed to cancel query in Trino'); + } } /** * Submit and sequentially execute an array of SQL statements. * Runs in the background (fire-and-forget from the POST handler). - * Stops on first failure, cancellation, or submit error. + * Stops on first failure, cancellation, or submit error. A later script in the + * same tab, or cancelQuery, aborts the run so it submits no further statements. */ export async function startScript( client: TrinoClient, @@ -206,7 +257,11 @@ export async function startScript( statements: string[], options: { user: string; catalog?: string; schema?: string } ): Promise { - await resetTabQueries(userId, tabId); + // Take over the tab synchronously, so the most recently started script wins + // even while an earlier one is still waiting for Trino. + const run = beginScriptRun(userId, tabId); + const cancelled = markActiveQueryCancelled(userId, tabId); + getUserTabMap(userId).set(tabId, []); touchTab(userId, tabId); log.info( @@ -214,7 +269,27 @@ export async function startScript( 'starting script execution' ); + try { + if (cancelled) await cancelInTrino(cancelled); + await runStatements(client, userId, tabId, statements, options, run.signal); + } finally { + endScriptRun(userId, tabId, run); + } +} + +async function runStatements( + client: TrinoClient, + userId: string, + tabId: string, + statements: string[], + options: { user: string; catalog?: string; schema?: string }, + signal: AbortSignal +): Promise { for (const sql of statements) { + // Cancelled or superseded during an earlier await (e.g. while cancelling the + // predecessor's query in Trino). + if (signal.aborted) break; + let query: TrinoQuery; try { query = await submitStatement(client, userId, tabId, sql, options); @@ -223,6 +298,22 @@ export async function startScript( break; } + if (signal.aborted) { + // Cancelled or superseded while Trino was accepting the statement: the + // result belongs to no one, so cancel it instead of storing it. + log.info( + { trino_query_id: query.trinoQueryId, user_id: userId, tab_id: tabId }, + 'script aborted during submit' + ); + if (!isTerminal(query.state)) { + terminateQuery(query, 'CANCELLED', { decrementGauge: false }); + trinoQueryTotal.inc({ outcome: 'cancelled' }); + await cancelInTrino(query); + } + break; + } + storeQuery(userId, tabId, query); + if (!isTerminal(query.state) && query.nextUri) { trinoActiveQueries.inc(); try { @@ -280,28 +371,33 @@ export function removeTabQuery(userId: string, tabId: string): void { } } +/** Stop the tab's running script and cancel its active query. Returns whether anything was running. */ export async function cancelQuery(userId: string, tabId: string): Promise { + // Abort first so the script submits no further statements, including while + // it is between statements and has no active query. + const scriptAborted = abortScriptRun(userId, tabId); + if (scriptAborted) log.info({ user_id: userId, tab_id: tabId }, 'script aborted'); + const cancelled = markActiveQueryCancelled(userId, tabId); + if (cancelled) await cancelInTrino(cancelled); + return scriptAborted || cancelled !== undefined; +} + +/** + * Mark the tab's active query as cancelled and return it, so the caller can + * cancel it in Trino. The state changes synchronously: concurrent callers + * never cancel or count the same query twice. + */ +function markActiveQueryCancelled(userId: string, tabId: string): TrinoQuery | undefined { const query = getActiveQuery(userId, tabId); - if (!query) return false; + if (!query) return undefined; log.info( { trino_query_id: query.trinoQueryId, user_id: userId, tab_id: tabId }, 'cancelling query' ); - - try { - if (query.nextUri) { - await query.client.cancelViaUri(query.nextUri, query.trinoUser); - } else { - await query.client.cancel(query.trinoQueryId, query.trinoUser); - } - } catch (err) { - log.warn({ err, trino_query_id: query.trinoQueryId }, 'failed to cancel query in Trino'); - } - terminateQuery(query, 'CANCELLED'); trinoQueryTotal.inc({ outcome: 'cancelled' }); - return true; + return query; } // --- Periodic eviction --- From 2f22617f3b478b079eee4b084df8740f1a1a6224 Mon Sep 17 00:00:00 2001 From: Johannes Fleck Date: Thu, 8 Oct 2026 13:47:19 +0200 Subject: [PATCH 6/7] fix: treat a FAILED first Trino response without error as failure (#345) Mapping terminal submit states to RUNNING turned a FAILED response without an error object into a completed query, and the script went on to the next statement. --- src/lib/server/trino/queries.test.ts | 13 +++++++++++++ src/lib/server/trino/queries.ts | 4 ++-- 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/src/lib/server/trino/queries.test.ts b/src/lib/server/trino/queries.test.ts index f821fcb2..f42f923a 100644 --- a/src/lib/server/trino/queries.test.ts +++ b/src/lib/server/trino/queries.test.ts @@ -273,6 +273,19 @@ describe('startScript', () => { expect(snapshot).toMatchObject({ state: 'FAILED', error: 'line 1:1: mismatched input' }); }); + it('treats a FAILED submit response without error details as a failure', async () => { + client.submit.mockResolvedValueOnce({ id: 'q1', stats: { state: 'FAILED' } }); + + await q.startScript(trinoClient(), USER, TAB, ['SELECT 1', 'SELECT 2'], OPTS); + + expect(client.submit).toHaveBeenCalledTimes(1); + expect(mocks.queryTotalInc).not.toHaveBeenCalledWith({ outcome: 'completed' }); + expect(q.getQuerySnapshots(USER, TAB)[0]).toMatchObject({ + state: 'FAILED', + error: 'Query failed' + }); + }); + it('falls back to a generic message when the submit error has none', async () => { client.submit.mockResolvedValueOnce({ id: 'q1', error: {} }); diff --git a/src/lib/server/trino/queries.ts b/src/lib/server/trino/queries.ts index b634dc54..ef5fce1a 100644 --- a/src/lib/server/trino/queries.ts +++ b/src/lib/server/trino/queries.ts @@ -214,8 +214,8 @@ async function submitStatement( completedAt: null }; - if (submitResult.error) { - query.error = submitResult.error.message ?? 'Query failed'; + if (submitResult.error || initialState === 'FAILED') { + query.error = submitResult.error?.message ?? 'Query failed'; terminateQuery(query, 'FAILED', { decrementGauge: false }); trinoQueryTotal.inc({ outcome: 'failed' }); } From 9e9cdd0b426e8d032963964f9c16fca80c6a2c6c Mon Sep 17 00:00:00 2001 From: Johannes Fleck Date: Thu, 8 Oct 2026 13:55:12 +0200 Subject: [PATCH 7/7] fix: fail polled Trino queries reporting FAILED without error (#345) collectResults only failed a query when the page carried an error object. A FAILED state without one was ignored and the query ended as FINISHED. --- src/lib/server/trino/result-collector.test.ts | 34 +++++++++++++++++++ src/lib/server/trino/result-collector.ts | 4 +-- 2 files changed, 36 insertions(+), 2 deletions(-) diff --git a/src/lib/server/trino/result-collector.test.ts b/src/lib/server/trino/result-collector.test.ts index 4fbf8131..025d8dc2 100644 --- a/src/lib/server/trino/result-collector.test.ts +++ b/src/lib/server/trino/result-collector.test.ts @@ -81,3 +81,37 @@ describe('collectResults row limit', () => { expect(query.state).toBe('FINISHED'); }); }); + +describe('collectResults failure', () => { + it('fails with the Trino error message', async () => { + const poll = vi.fn(async () => ({ + id: 'q1', + stats: { state: 'FAILED' }, + error: { message: 'Division by zero' } + })); + + const query = makeQuery({ poll }); + await collectResults(query); + + expect(query.state).toBe('FAILED'); + expect(query.error).toBe('Division by zero'); + }); + + it('fails on a FAILED state without error details', async () => { + const poll = vi + .fn() + .mockResolvedValueOnce({ + id: 'q1', + nextUri: 'http://trino/next-2', + stats: { state: 'FAILED' } + }) + .mockResolvedValue({ id: 'q1', stats: { state: 'FAILED' } }); + + const query = makeQuery({ poll }); + await collectResults(query); + + expect(poll).toHaveBeenCalledTimes(1); + expect(query.state).toBe('FAILED'); + expect(query.error).toBe('Query failed'); + }); +}); diff --git a/src/lib/server/trino/result-collector.ts b/src/lib/server/trino/result-collector.ts index 841b7cb4..7e40f632 100644 --- a/src/lib/server/trino/result-collector.ts +++ b/src/lib/server/trino/result-collector.ts @@ -60,8 +60,8 @@ export async function collectResults(query: TrinoQuery): Promise { query.progress = toQueryProgress(result.stats); } - if (result.error) { - query.error = result.error.message ?? 'Query failed'; + if (result.error || result.stats?.state === 'FAILED') { + query.error = result.error?.message ?? 'Query failed'; terminateQuery(query, 'FAILED'); trinoQueryTotal.inc({ outcome: 'failed' }); return;