diff --git a/src/internal/database/pg-connection.test.ts b/src/internal/database/pg-connection.test.ts index 2abf9c9af..e1153bcb5 100644 --- a/src/internal/database/pg-connection.test.ts +++ b/src/internal/database/pg-connection.test.ts @@ -7,6 +7,7 @@ import PgConnection from 'pg/lib/connection' import { vi } from 'vitest' import { getPgCancelConnectionTarget, + type PgExecutor, PgPoolExecutor, PgPoolManager, PgPoolStrategy, @@ -56,6 +57,62 @@ function createDatabaseError(code: string | undefined, message = 'database error return error } +type TestPgBeginTransactionOptions = { + timeout?: number + statementTimeoutMs?: number +} + +function normalizeTestStatementTimeoutMs( + options?: TestPgBeginTransactionOptions +): number | undefined { + const timeoutMs = options?.statementTimeoutMs ?? options?.timeout + + if (timeoutMs === undefined || !Number.isFinite(timeoutMs) || timeoutMs <= 0) { + return undefined + } + + return timeoutMs +} + +function createMockTenantConnectionWithTransaction( + overrides: Partial = {}, + query = vi.fn().mockResolvedValue({ rows: [] }) +) { + const client = { + query, + release: vi.fn(), + } as unknown as PoolClient + let transaction: PgTransaction | undefined + const beginTransaction = vi.fn( + async (options?: TestPgBeginTransactionOptions): Promise => { + transaction = new PgTransaction(client, undefined, { + statementTimeoutMs: normalizeTestStatementTimeoutMs(options), + }) + return transaction + } + ) + const pool = { + acquire: vi.fn().mockReturnValue({ + beginTransaction, + }), + } as unknown as PgPoolStrategy + const connection = new PgTenantConnection(pool, createPoolStrategySettings(overrides)) + + return { + beginTransaction, + client, + connection, + query, + getTransaction() { + if (!transaction) { + throw new Error('Expected test transaction to be created') + } + + return transaction + }, + } +} + async function expectQueryErrorRelease(error: Error): Promise> { const release = vi.fn() const client = Object.assign(new EventEmitter(), { @@ -402,6 +459,42 @@ describe('PgPoolExecutor', () => { expect(client.listenerCount('error')).toBe(0) }) + it('uses transaction timeout with isolation and read-only begin modes', async () => { + const client = Object.assign(new EventEmitter(), { + query: vi.fn().mockResolvedValue({ rows: [] }), + release: vi.fn(), + }) as unknown as PoolClient & EventEmitter + const pool = { + connect: vi.fn().mockResolvedValue(client), + } as unknown as Pool + const executor = new PgPoolExecutor(pool) + + const transaction = await executor.beginTransaction({ + timeout: 4321, + isolation: 'repeatable read', + readOnly: true, + }) + + expect(client.query).toHaveBeenCalledTimes(1) + expect(client.query).toHaveBeenNthCalledWith( + 1, + 'BEGIN ISOLATION LEVEL REPEATABLE READ, READ ONLY', + undefined + ) + + await transaction.query('SELECT 1') + + expect(client.query).toHaveBeenCalledTimes(3) + expect(client.query).toHaveBeenNthCalledWith( + 2, + "SELECT set_config('statement_timeout', $1, true)", + ['4321ms'] + ) + expect(client.query).toHaveBeenNthCalledWith(3, 'SELECT 1', undefined) + + await transaction.rollback() + }) + it('returns clients to the pool after regular SQL errors', async () => { for (const code of ['42P01', '23505', '23503', '42501', '22P02', '42703']) { const error = createDatabaseError(code) @@ -553,6 +646,93 @@ describe('PgPoolExecutor', () => { }) describe('PgTransaction', () => { + it('applies a pending statement timeout before the first direct query', async () => { + const client = { + query: vi.fn().mockResolvedValue({ rows: [] }), + release: vi.fn(), + } as unknown as PoolClient + const transaction = new PgTransaction(client, undefined, { statementTimeoutMs: 4321 }) + + await transaction.query('SELECT 1') + await transaction.query('SELECT 2') + + expect(client.query).toHaveBeenCalledTimes(3) + expect(client.query).toHaveBeenNthCalledWith( + 1, + "SELECT set_config('statement_timeout', $1, true)", + ['4321ms'] + ) + expect(client.query).toHaveBeenNthCalledWith(2, 'SELECT 1', undefined) + expect(client.query).toHaveBeenNthCalledWith(3, 'SELECT 2', undefined) + }) + + it('normalizes invalid constructor statement timeouts as disabled', async () => { + for (const statementTimeoutMs of [-5, Number.NaN, Number.POSITIVE_INFINITY, 0]) { + const client = { + query: vi.fn().mockResolvedValue({ rows: [] }), + release: vi.fn(), + } as unknown as PoolClient + const transaction = new PgTransaction(client, undefined, { statementTimeoutMs }) + + await transaction.query('SELECT 1') + + expect(client.query).toHaveBeenCalledTimes(1) + expect(client.query).toHaveBeenNthCalledWith(1, 'SELECT 1', undefined) + } + }) + + it('rejects a pre-aborted direct query before applying a pending statement timeout', async () => { + const client = { + query: vi.fn().mockResolvedValue({ rows: [] }), + release: vi.fn(), + } as unknown as PoolClient + const transaction = new PgTransaction(client, undefined, { statementTimeoutMs: 4321 }) + + await expect( + transaction.query('SELECT 1', { signal: AbortSignal.abort() }) + ).rejects.toMatchObject({ + name: 'AbortError', + code: 'ABORT_ERR', + }) + expect(client.query).not.toHaveBeenCalled() + expect(client.release).toHaveBeenCalledWith(expect.objectContaining({ name: 'AbortError' })) + + await expect(transaction.query('SELECT 2')).rejects.toThrow( + 'Cannot query a completed transaction' + ) + expect(client.query).not.toHaveBeenCalled() + }) + + it('honors abort signals while applying a pending statement timeout', async () => { + const controller = new AbortController() + const client = Object.assign(new EventEmitter(), { + query: vi.fn(() => { + controller.abort() + return new Promise(() => undefined) + }), + release: vi.fn(), + }) as unknown as PoolClient & EventEmitter + const transaction = new PgTransaction(client, undefined, { statementTimeoutMs: 4321 }) + + await expect( + transaction.query('SELECT 1', { signal: controller.signal }) + ).rejects.toMatchObject({ + name: 'AbortError', + code: 'ABORT_ERR', + }) + + expect(client.query).toHaveBeenCalledTimes(1) + expect(client.query).toHaveBeenNthCalledWith( + 1, + "SELECT set_config('statement_timeout', $1, true)", + ['4321ms'] + ) + expect(client.release).toHaveBeenCalledWith(expect.objectContaining({ name: 'AbortError' })) + + await transaction.rollback() + expect(client.query).toHaveBeenCalledTimes(1) + }) + it('rejects queries after commit releases the client', async () => { const client = { query: vi.fn().mockResolvedValue({ rows: [] }), @@ -708,11 +888,335 @@ describe('PgTenantConnection', () => { }) }) + it('treats non-finite transaction timeouts as disabled', async () => { + for (const timeout of [Number.NaN, Number.POSITIVE_INFINITY]) { + const { beginTransaction, connection, query, getTransaction } = + createMockTenantConnectionWithTransaction({ + isExternalPool: false, + }) + + const transaction = await connection.transaction({ timeout }) + expect(transaction).toBe(getTransaction()) + expect(beginTransaction).toHaveBeenCalledWith({ timeout }) + expect(query).not.toHaveBeenCalled() + + await connection.setScope(getTransaction()) + + expect(query).toHaveBeenCalledTimes(1) + const [scopeStatement, scopeValues] = query.mock.calls[0] + expect(scopeStatement).toContain("set_config('role', $1, true)") + expect(scopeStatement).not.toContain("set_config('statement_timeout'") + expect(scopeValues).toHaveLength(9) + } + }) + + it('defers statement_timeout for low-level Postgres beginTransaction', async () => { + const { beginTransaction, connection, query, getTransaction } = + createMockTenantConnectionWithTransaction({ + isExternalPool: false, + }) + + await expect( + connection.beginTransaction({ + timeout: 4321, + isolation: 'serializable', + readOnly: true, + }) + ).resolves.toBe(getTransaction()) + expect(beginTransaction).toHaveBeenCalledWith({ + timeout: 4321, + isolation: 'serializable', + readOnly: true, + }) + expect(query).not.toHaveBeenCalled() + + await connection.setScope(getTransaction()) + + expect(query).toHaveBeenCalledTimes(1) + const [scopeStatement, scopeValues] = query.mock.calls[0] + expect(scopeStatement).toContain("set_config('role', $1, true)") + expect(scopeStatement).toContain("set_config('statement_timeout', $10, true)") + expect(scopeValues).toEqual([ + 'authenticated', + 'authenticated', + 'jwt', + '', + JSON.stringify({ role: 'authenticated' }), + '{}', + '', + '', + '', + '4321ms', + ]) + }) + + it('defers statement_timeout setup until the first scope application', async () => { + const query = vi.fn().mockResolvedValue({ rows: [] }) + const client = { + query, + release: vi.fn(), + } as unknown as PoolClient + let transaction: PgTransaction + const beginTransaction = vi.fn( + async (options?: TestPgBeginTransactionOptions): Promise => { + transaction = new PgTransaction(client, undefined, { + statementTimeoutMs: normalizeTestStatementTimeoutMs(options), + }) + return transaction + } + ) + const pool = { + acquire: vi.fn().mockReturnValue({ + beginTransaction, + }), + } as unknown as PgPoolStrategy + const connection = new PgTenantConnection( + pool, + createPoolStrategySettings({ + isExternalPool: false, + }) + ) + + await expect(connection.transaction({ timeout: 4321 })).resolves.toBe(transaction!) + expect(beginTransaction).toHaveBeenCalledWith({ timeout: 4321 }) + expect(query).not.toHaveBeenCalled() + + await connection.setScope(transaction!) + + expect(query).toHaveBeenCalledTimes(1) + const [statement, values] = query.mock.calls[0] + expect(statement).toContain("set_config('role', $1, true)") + expect(statement).toContain("set_config('statement_timeout', $10, true)") + expect(values).toEqual([ + 'authenticated', + 'authenticated', + 'jwt', + '', + JSON.stringify({ role: 'authenticated' }), + '{}', + '', + '', + '', + '4321ms', + ]) + }) + + it('keeps external-pool search_path setup before deferring statement_timeout', async () => { + const query = vi.fn().mockResolvedValue({ rows: [] }) + const client = { + query, + release: vi.fn(), + } as unknown as PoolClient + let transaction: PgTransaction + const beginTransaction = vi.fn( + async (options?: TestPgBeginTransactionOptions): Promise => { + transaction = new PgTransaction(client, undefined, { + statementTimeoutMs: normalizeTestStatementTimeoutMs(options), + }) + return transaction + } + ) + const pool = { + acquire: vi.fn().mockReturnValue({ + beginTransaction, + }), + } as unknown as PgPoolStrategy + const connection = new PgTenantConnection( + pool, + createPoolStrategySettings({ + isExternalPool: true, + }) + ) + + await expect(connection.transaction({ timeout: 4321 })).resolves.toBe(transaction!) + expect(beginTransaction).toHaveBeenCalledWith({ timeout: 4321 }) + + expect(query).toHaveBeenCalledTimes(1) + expect(query).toHaveBeenNthCalledWith( + 1, + "SELECT set_config('search_path', $1, true)", + expect.any(Array) + ) + + await connection.setScope(transaction!) + + expect(query).toHaveBeenCalledTimes(2) + const [scopeStatement, scopeValues] = query.mock.calls[1] + expect(scopeStatement).toContain("set_config('role', $1, true)") + expect(scopeStatement).toContain("set_config('statement_timeout', $10, true)") + expect(scopeValues).toEqual([ + 'authenticated', + 'authenticated', + 'jwt', + '', + JSON.stringify({ role: 'authenticated' }), + '{}', + '', + '', + '', + '4321ms', + ]) + }) + + it('folds deferred statement_timeout into scope setup', async () => { + const query = vi.fn().mockResolvedValue({ rows: [] }) + const client = { + query, + release: vi.fn(), + } as unknown as PoolClient + let transaction: PgTransaction + const beginTransaction = vi.fn( + async (options?: TestPgBeginTransactionOptions): Promise => { + transaction = new PgTransaction(client, undefined, { + statementTimeoutMs: normalizeTestStatementTimeoutMs(options), + }) + return transaction + } + ) + const pool = { + acquire: vi.fn().mockReturnValue({ + beginTransaction, + }), + } as unknown as PgPoolStrategy + const connection = new PgTenantConnection( + pool, + createPoolStrategySettings({ + isExternalPool: true, + }) + ) + + await expect(connection.transaction({ timeout: 4321 })).resolves.toBe(transaction!) + expect(beginTransaction).toHaveBeenCalledWith({ timeout: 4321 }) + + expect(query).toHaveBeenCalledTimes(1) + expect(query).toHaveBeenNthCalledWith( + 1, + "SELECT set_config('search_path', $1, true)", + expect.any(Array) + ) + + await connection.setScope(transaction!) + + expect(query).toHaveBeenCalledTimes(2) + const [scopeStatement, scopeValues] = query.mock.calls[1] + expect(scopeStatement).toContain("set_config('statement_timeout', $10, true)") + expect(scopeValues).toEqual([ + 'authenticated', + 'authenticated', + 'jwt', + '', + JSON.stringify({ role: 'authenticated' }), + '{}', + '', + '', + '', + '4321ms', + ]) + }) + + it('omits statement_timeout setup for low-level beginTransaction without a positive timeout', async () => { + const cases: Array<{ + options?: { timeout: number } + }> = [{}, { options: { timeout: 0 } }] + + for (const { options } of cases) { + const { beginTransaction, connection, query, getTransaction } = + createMockTenantConnectionWithTransaction({ + isExternalPool: false, + }) + + await connection.beginTransaction(options) + expect(beginTransaction).toHaveBeenCalledWith(options) + expect(query).not.toHaveBeenCalled() + + await connection.setScope(getTransaction()) + + expect(query).toHaveBeenCalledTimes(1) + const [scopeStatement, scopeValues] = query.mock.calls[0] + expect(scopeStatement).toContain("set_config('role', $1, true)") + expect(scopeStatement).not.toContain("set_config('statement_timeout'") + expect(scopeValues).toHaveLength(9) + } + }) + + it('does not re-apply statement_timeout after setScope consumes it', async () => { + const query = vi.fn().mockResolvedValue({ rows: [] }) + const client = { + query, + release: vi.fn(), + } as unknown as PoolClient + let transaction: PgTransaction + const beginTransaction = vi.fn( + async (options?: TestPgBeginTransactionOptions): Promise => { + transaction = new PgTransaction(client, undefined, { + statementTimeoutMs: normalizeTestStatementTimeoutMs(options), + }) + return transaction + } + ) + const pool = { + acquire: vi.fn().mockReturnValue({ + beginTransaction, + }), + } as unknown as PgPoolStrategy + const connection = new PgTenantConnection( + pool, + createPoolStrategySettings({ + isExternalPool: false, + }) + ) + + await connection.transaction({ timeout: 4321 }) + await connection.setScope(transaction!) + await transaction!.query('SELECT 1') + + expect(query).toHaveBeenCalledTimes(2) + expect(query).toHaveBeenNthCalledWith(2, 'SELECT 1', undefined) + expect( + query.mock.calls.filter(([statement]) => String(statement).includes('statement_timeout')) + ).toHaveLength(1) + }) + + it('reuses precomputed scope JSON payloads across repeated scope applications', async () => { + const pool = { + acquire: vi.fn(), + } as unknown as PgPoolStrategy + const connection = new PgTenantConnection( + pool, + createPoolStrategySettings({ + headers: { + 'x-test-header': 'test-value', + }, + user: { + jwt: 'jwt', + payload: { + role: 'authenticated', + sub: 'user-id', + }, + }, + }) + ) + const executor = { + query: vi.fn().mockResolvedValue({ rows: [] }), + } as unknown as PgExecutor + const stringifySpy = vi.spyOn(JSON, 'stringify') + + try { + await connection.setScope(executor) + await connection.setScope(executor) + + expect(stringifySpy).not.toHaveBeenCalled() + } finally { + stringifySpy.mockRestore() + } + }) + it('preserves setup errors when external-pool rollback fails', async () => { const setupError = new Error('search_path setup failed') const rollbackError = new Error('rollback failed') const transaction = { - query: vi.fn().mockRejectedValue(setupError), + query: vi.fn(), + runSetupQuery: vi.fn().mockRejectedValue(setupError), rollback: vi.fn().mockRejectedValue(rollbackError), } as unknown as PgTransaction const pool = { diff --git a/src/internal/database/pg-connection.ts b/src/internal/database/pg-connection.ts index 643f6a6b4..331330766 100644 --- a/src/internal/database/pg-connection.ts +++ b/src/internal/database/pg-connection.ts @@ -45,6 +45,42 @@ export interface PgQueryOptions { type PgQueryArgument = PgQueryOptions | unknown[] +interface PgTransactionOptions { + statementTimeoutMs?: number +} + +type PgBeginTransactionOptions = TransactionOptions & PgTransactionOptions + +const defaultStatementTimeoutMs = normalizeStatementTimeoutMs(databaseStatementTimeout) + +// Shared frozen options for the common transaction() call without caller options. +const defaultBeginTransactionOptions: PgBeginTransactionOptions | undefined = + defaultStatementTimeoutMs !== undefined + ? Object.freeze({ statementTimeoutMs: defaultStatementTimeoutMs }) + : undefined + +const scopeConfigSetters = `set_config('role', $1, true), + set_config('request.jwt.claim.role', $2, true), + set_config('request.jwt', $3, true), + set_config('request.jwt.claim.sub', $4, true), + set_config('request.jwt.claims', $5, true), + set_config('request.headers', $6, true), + set_config('request.method', $7, true), + set_config('request.path', $8, true), + set_config('storage.operation', $9, true), + set_config('storage.allow_delete_query', 'true', true)` + +const scopeConfigSql = ` + SELECT + ${scopeConfigSetters}; + ` + +const scopeConfigSqlWithStatementTimeout = ` + SELECT + ${scopeConfigSetters}, + set_config('statement_timeout', $10, true); + ` + export interface PgExecutor { query( statement: string | PgStatement, @@ -53,7 +89,7 @@ export interface PgExecutor { } export interface PgTransactionalExecutor extends PgExecutor { - beginTransaction(options?: TransactionOptions): Promise + beginTransaction(options?: PgBeginTransactionOptions): Promise } interface PgPoolErrorContext { @@ -398,7 +434,7 @@ export class PgPoolExecutor implements PgTransactionalExecutor { } } - async beginTransaction(options?: TransactionOptions): Promise { + async beginTransaction(options?: PgBeginTransactionOptions): Promise { let client: PoolClient try { client = await this.pool.connect() @@ -411,10 +447,12 @@ export class PgPoolExecutor implements PgTransactionalExecutor { } const clientErrorTracker = new PgClientErrorTracker(client) - const transaction = new PgTransaction(client, clientErrorTracker) + const transaction = new PgTransaction(client, clientErrorTracker, { + statementTimeoutMs: options?.statementTimeoutMs ?? options?.timeout, + }) try { - await transaction.query(buildBeginStatement(options)) + await transaction.runSetupQuery(buildBeginStatement(options)) return transaction } catch (e) { if (!transaction.isCompleted()) { @@ -432,11 +470,15 @@ export class PgPoolExecutor implements PgTransactionalExecutor { export class PgTransaction implements PgExecutor { private completed = false + private statementTimeoutMs?: number constructor( private readonly client: PoolClient, - private readonly clientErrorTracker?: PgClientErrorTracker - ) {} + private readonly clientErrorTracker?: PgClientErrorTracker, + options: PgTransactionOptions = {} + ) { + this.statementTimeoutMs = normalizeStatementTimeoutMs(options.statementTimeoutMs) + } isCompleted(): boolean { return this.completed @@ -445,12 +487,44 @@ export class PgTransaction implements PgExecutor { async query( statement: string | PgStatement, options?: PgQueryArgument + ): Promise> { + return this.runQuery(statement, options, true) + } + + async runSetupQuery( + statement: string | PgStatement, + options?: PgQueryArgument + ): Promise> { + // Setup statements must not consume a deferred timeout before scope can fold it in. + return this.runQuery(statement, options, false) + } + + // Hands the deferred timeout to the caller (setScope) so it can be folded + // into its next statement instead of costing a separate set_config round trip. + takePendingStatementTimeoutMs(): number | undefined { + const timeoutMs = this.statementTimeoutMs + this.statementTimeoutMs = undefined + return timeoutMs + } + + private async runQuery( + statement: string | PgStatement, + options: PgQueryArgument | undefined, + applyPendingStatementTimeout: boolean ): Promise> { if (this.completed) { throw new Error('Cannot query a completed transaction') } + const signal = getQuerySignal(options) + try { + assertValidSignal(signal) + + if (applyPendingStatementTimeout) { + await this.applyPendingStatementTimeoutBeforeQuery(signal) + } + const result = await runPgQuery(this.client, statement, options) this.clientErrorTracker?.throwIfErrored() return result @@ -515,6 +589,22 @@ export class PgTransaction implements PgExecutor { this.clientErrorTracker?.detach() } } + + private async applyPendingStatementTimeoutBeforeQuery(signal?: AbortSignal): Promise { + const timeoutMs = this.takePendingStatementTimeoutMs() + if (!timeoutMs) { + return + } + + await runPgQuery( + this.client, + { + text: `SELECT set_config('statement_timeout', $1, true)`, + values: [`${timeoutMs}ms`], + }, + { signal } + ) + } } export class PgTenantConnection { @@ -522,12 +612,16 @@ export class PgTenantConnection { public readonly role: string private abortSignal?: AbortSignal private disposed = false + private readonly headersPayload: string + private readonly userPayload: string constructor( public readonly pool: PgPoolStrategy, protected readonly options: TenantConnectionOptions ) { this.role = options.user.payload.role || 'anon' + this.headersPayload = JSON.stringify(options.headers || {}) + this.userPayload = JSON.stringify(options.user.payload) } static stop() { @@ -562,6 +656,7 @@ export class PgTenantConnection { async beginTransaction(options?: TransactionOptions): Promise { this.assertNotDisposed() + // PgPoolExecutor derives the deferred statement_timeout from options.timeout. return this.pool.acquire().beginTransaction(options) } @@ -584,6 +679,7 @@ export class PgTenantConnection { this.assertNotDisposed() try { + const beginOptions = withDefaultStatementTimeout(opts) const transaction = await retry( async (bail) => { if (this.disposed) { @@ -592,7 +688,7 @@ export class PgTenantConnection { } try { - return await this.pool.acquire().beginTransaction(opts) + return await this.pool.acquire().beginTransaction(beginOptions) } catch (e) { if (isConnectionLimitError(e)) { throw e @@ -613,7 +709,7 @@ export class PgTenantConnection { if (this.options.isExternalPool) { try { - await transaction.query({ + await transaction.runSetupQuery({ text: `SELECT set_config('search_path', $1, true)`, values: [searchPath.join(',')], }) @@ -623,19 +719,6 @@ export class PgTenantConnection { } } - const statementTimeout = opts?.timeout ?? databaseStatementTimeout - if (statementTimeout > 0) { - try { - await transaction.query({ - text: `SELECT set_config('statement_timeout', $1, true)`, - values: [`${statementTimeout}ms`], - }) - } catch (e) { - await this.rollbackTransactionSafely(transaction, e, 'statement_timeout setup') - throw e - } - } - return transaction } catch (e) { if (isConnectionTimeoutError(e)) { @@ -674,36 +757,56 @@ export class PgTenantConnection { } async setScope(tnx: PgExecutor) { - const headers = JSON.stringify(this.options.headers || {}) + const statementTimeoutMs = + tnx instanceof PgTransaction ? tnx.takePendingStatementTimeoutMs() : undefined + + const values: unknown[] = [ + this.role, + this.role, + this.options.user.jwt || '', + this.options.user.payload.sub || '', + this.userPayload, + this.headersPayload, + this.options.method || '', + this.options.path || '', + this.options.operation?.() || '', + ] + + if (statementTimeoutMs) { + values.push(`${statementTimeoutMs}ms`) + } + await tnx.query({ - text: ` - SELECT - set_config('role', $1, true), - set_config('request.jwt.claim.role', $2, true), - set_config('request.jwt', $3, true), - set_config('request.jwt.claim.sub', $4, true), - set_config('request.jwt.claims', $5, true), - set_config('request.headers', $6, true), - set_config('request.method', $7, true), - set_config('request.path', $8, true), - set_config('storage.operation', $9, true), - set_config('storage.allow_delete_query', 'true', true); - `, - values: [ - this.role, - this.role, - this.options.user.jwt || '', - this.options.user.payload.sub || '', - JSON.stringify(this.options.user.payload), - headers, - this.options.method || '', - this.options.path || '', - this.options.operation?.() || '', - ], + text: statementTimeoutMs ? scopeConfigSqlWithStatementTimeout : scopeConfigSql, + values, }) } } +// PgPoolExecutor already derives the deferred statement_timeout from options.timeout, +// so options only need rewriting when the global default has to be injected. +function withDefaultStatementTimeout( + options: TransactionOptions | undefined +): PgBeginTransactionOptions | undefined { + if (!defaultBeginTransactionOptions || options?.timeout !== undefined) { + return options + } + + if (!options) { + return defaultBeginTransactionOptions + } + + return { ...options, statementTimeoutMs: defaultStatementTimeoutMs } +} + +function normalizeStatementTimeoutMs(timeoutMs: number | undefined): number | undefined { + if (timeoutMs === undefined || !Number.isFinite(timeoutMs) || timeoutMs <= 0) { + return undefined + } + + return timeoutMs +} + function createDisposedTenantConnectionError(): Error { return new Error('Cannot use a disposed PgTenantConnection') } diff --git a/src/test/storage-pg-db.test.ts b/src/test/storage-pg-db.test.ts index 90f8591df..e149a630c 100644 --- a/src/test/storage-pg-db.test.ts +++ b/src/test/storage-pg-db.test.ts @@ -352,28 +352,7 @@ describe('StoragePgDB bucket metadata', () => { }) it('restores parent transaction scope after failed super-user queries', async () => { - const authenticatedUser = { - jwt: 'storage-pg-db-authenticated-jwt', - payload: { - role: 'authenticated', - sub: randomUUID(), - }, - } - const authenticatedSettings = { - ...connectionSettings, - user: authenticatedUser, - superUser, - } - const authenticatedPool = new PgPoolStrategy(authenticatedSettings) - const authenticatedDb = new StoragePgDB( - new PgTenantConnection(authenticatedPool, authenticatedSettings), - { - tenantId, - host: 'localhost', - } - ) - - try { + await withAuthenticatedDb('storage-pg-db-authenticated-jwt', async (authenticatedDb) => { await authenticatedDb.withTransaction(async (tx) => { await expect(readCurrentRole(tx)).resolves.toBe('authenticated') @@ -409,9 +388,62 @@ describe('StoragePgDB bucket metadata', () => { ) expect(result.rows[0].n).toBe(1) }) - } finally { - await authenticatedPool.destroy() - } + }) + }) + + it('preserves custom statement timeout after nested super-user scope restores', async () => { + await withAuthenticatedDb('storage-pg-db-timeout-jwt', async (authenticatedDb) => { + await authenticatedDb.withTransaction( + async (tx) => { + await expect(readCurrentRole(tx)).resolves.toBe('authenticated') + await expect(readCurrentStatementTimeoutMs(tx)).resolves.toBe(4321) + + await expect( + runStorageQuery( + tx.asSuperUser() as StoragePgDB, + 'ReadSuperUserStatementTimeout', + async (pg) => { + await expect(readCurrentRoleFromExecutor(pg)).resolves.toBe(superUser.payload.role) + await expect(readCurrentStatementTimeoutMsFromExecutor(pg)).resolves.toBe(4321) + } + ) + ).resolves.toBeUndefined() + + await expect(readCurrentRole(tx)).resolves.toBe('authenticated') + await expect(readCurrentStatementTimeoutMs(tx)).resolves.toBe(4321) + }, + { timeout: 4321 } + ) + }) + }) + + it('preserves custom statement timeout after failed nested super-user scope rolls back', async () => { + await withAuthenticatedDb('storage-pg-db-timeout-rollback-jwt', async (authenticatedDb) => { + await authenticatedDb.withTransaction( + async (tx) => { + await expect(readCurrentRole(tx)).resolves.toBe('authenticated') + await expect(readCurrentStatementTimeoutMs(tx)).resolves.toBe(4321) + + await expect( + runStorageQuery( + tx.asSuperUser() as StoragePgDB, + 'FailSuperUserStatementTimeout', + async (pg) => { + await expect(readCurrentRoleFromExecutor(pg)).resolves.toBe(superUser.payload.role) + await expect(readCurrentStatementTimeoutMsFromExecutor(pg)).resolves.toBe(4321) + await pg.query("SELECT set_config('statement_timeout', '30s', true)") + await expect(readCurrentStatementTimeoutMsFromExecutor(pg)).resolves.toBe(30000) + throw new Error('failed nested super-user timeout query') + } + ) + ).rejects.toThrow('failed nested super-user timeout query') + + await expect(readCurrentRole(tx)).resolves.toBe('authenticated') + await expect(readCurrentStatementTimeoutMs(tx)).resolves.toBe(4321) + }, + { timeout: 4321 } + ) + }) }) it('preserves original errors when best-effort parent scope restoration fails', async () => { @@ -1748,6 +1780,74 @@ describe('StoragePgDB bucket metadata', () => { return result.rows[0].role } + function readCurrentStatementTimeoutMs(storage: StoragePgDB): Promise { + return runStorageQuery(storage, 'ReadCurrentStatementTimeout', (pg) => + readCurrentStatementTimeoutMsFromExecutor(pg) + ) + } + + async function readCurrentStatementTimeoutMsFromExecutor(pg: PgExecutor): Promise { + const result = await pg.query<{ statement_timeout: string }>('SHOW statement_timeout') + + return parseStatementTimeoutMs(result.rows[0].statement_timeout) + } + + function parseStatementTimeoutMs(statementTimeout: string): number { + const value = statementTimeout.trim() + if (value === '0') { + return 0 + } + + const match = /^(\d+(?:\.\d+)?)\s*(ms|s|min|h|d)?$/.exec(value) + if (!match) { + throw new Error(`Unexpected statement_timeout value: ${statementTimeout}`) + } + + const multipliers = { + ms: 1, + s: 1000, + min: 60 * 1000, + h: 60 * 60 * 1000, + d: 24 * 60 * 60 * 1000, + } + const amount = Number(match[1]) + const unit = (match[2] ?? 'ms') as keyof typeof multipliers + + return Math.round(amount * multipliers[unit]) + } + + async function withAuthenticatedDb( + jwt: string, + fn: (authenticatedDb: StoragePgDB) => Promise + ): Promise { + const authenticatedUser = { + jwt, + payload: { + role: 'authenticated', + sub: randomUUID(), + }, + } + const authenticatedSettings = { + ...connectionSettings, + user: authenticatedUser, + superUser, + } + const authenticatedPool = new PgPoolStrategy(authenticatedSettings) + const authenticatedDb = new StoragePgDB( + new PgTenantConnection(authenticatedPool, authenticatedSettings), + { + tenantId, + host: 'localhost', + } + ) + + try { + await fn(authenticatedDb) + } finally { + await authenticatedPool.destroy() + } + } + async function tableExists(tableName: string): Promise { const result = await pool.acquire().query<{ exists: boolean }>({ text: `SELECT to_regclass($1) IS NOT NULL AS "exists"`,