diff --git a/docs/pages/apis/client.mdx b/docs/pages/apis/client.mdx index 973b1a832..c466e82ca 100644 --- a/docs/pages/apis/client.mdx +++ b/docs/pages/apis/client.mdx @@ -92,6 +92,9 @@ type QueryConfig { // custom type parsers just for this query result types?: Types; + // cancel queued or active work when this signal aborts + signal?: AbortSignal; + // TODO: document queryMode?: string; } @@ -151,6 +154,53 @@ console.log(result.rows) // ['brianc'] await client.end() ``` +### Query cancellation + +Pass an `AbortSignal` in a query config object to cancel queued or active work: + +```js +const controller = new AbortController() +const query = client.query({ + text: 'SELECT pg_sleep(30)', + signal: controller.signal, +}) + +controller.abort() +await query // rejects with PostgreSQL error code 57014 if cancellation wins +``` + +A query aborted before submission rejects with `signal.reason` and sends nothing. Once submitted, cancellation races with normal completion: PostgreSQL may return error `57014`, or the query may finish normally first. An aborted transaction remains failed until you issue `ROLLBACK`. + +For explicit cancellation, use the exported helper with the client that owns the query: + +```js +import { cancelQuery } from 'pg' + +const query = client.query('SELECT pg_sleep(30)') +const cancellation = cancelQuery(client) + +await Promise.allSettled([query, cancellation]) +``` + +`cancelQuery(client)` resolves `true` after the cancellation connection finishes and the target client reaches `ReadyForQuery`. It resolves `false` when there is no active query and rejects if the cancel request cannot be completed. A `true` result does not guarantee that cancellation won the race. + +See [Query cancellation](/features/query-cancellation) for the separate-connection design, completion barrier, and race outcomes. + +Cancellation uses a separate connection to the selected PostgreSQL endpoint. A custom `stream` must therefore be a factory; a concrete custom stream instance cannot support cancellation. AbortSignal and explicit cancellation are not supported in pipeline mode. + +With a pool, check out a client so the query and cancellation target are unambiguous. Do not use a `PoolClient` after `release()`: + +```js +const client = await pool.connect() +try { + const query = client.query('SELECT pg_sleep(30)') + const cancellation = cancelQuery(client) + await Promise.allSettled([query, cancellation]) +} finally { + client.release() +} +``` + **client.query with a `Submittable`** If you pass an object to `client.query` and the object has a `.submit` function on it, the client will pass it's PostgreSQL server connection to the object and delegate query dispatching to the supplied object. This is an advanced feature mostly intended for library authors. It is incidentally also currently how the callback and promise based queries above are handled internally, but this is subject to change. It is also how [pg-cursor](https://github.com/brianc/node-pg-cursor) and [pg-query-stream](https://github.com/brianc/node-pg-query-stream) work. diff --git a/docs/pages/features/_meta.js b/docs/pages/features/_meta.js index 7ddd35a5c..f364cb8fe 100644 --- a/docs/pages/features/_meta.js +++ b/docs/pages/features/_meta.js @@ -1,6 +1,7 @@ export default { connecting: 'Connecting', queries: 'Queries', + 'query-cancellation': 'Query cancellation', pipelining: 'Pipelining', pooling: 'Pooling', transactions: 'Transactions', diff --git a/docs/pages/features/queries.mdx b/docs/pages/features/queries.mdx index 63ecdde1e..b57ce1846 100644 --- a/docs/pages/features/queries.mdx +++ b/docs/pages/features/queries.mdx @@ -79,6 +79,8 @@ const res = await client.query(query) console.log(res.rows[0]) ``` +The config object can also include an [`AbortSignal`](/apis/client#query-cancellation). This is available for `client.query` and `pool.query`; use a checked-out client when you need the explicit `cancelQuery(client)` helper. See [Query cancellation](/features/query-cancellation) for the lifecycle and race semantics. + The query config object allows for a few more advanced scenarios: ### Prepared statements diff --git a/docs/pages/features/query-cancellation.mdx b/docs/pages/features/query-cancellation.mdx new file mode 100644 index 000000000..ff0c0956c --- /dev/null +++ b/docs/pages/features/query-cancellation.mdx @@ -0,0 +1,46 @@ +--- +title: Query cancellation +--- + +## Why cancellation uses another connection + +PostgreSQL cannot receive a cancellation command on a connection while that connection is busy running a query. node-postgres therefore opens a short-lived connection to the same selected endpoint and sends a `CancelRequest` containing the active backend's process ID and secret key. + +This does not check out another `pg.Pool` client or create an authenticated database session. Concurrent cancellation calls for one client share the same operation, but different clients can still open short-lived cancel sockets at the same time and put pressure on PostgreSQL or proxy connection acceptance. + +The cancellation connection reuses the selected TCP address or Unix socket, SSL mode, and logical TLS server name. This matters when DNS returns multiple addresses or a connection uses failover: sending the request to a different server would not cancel the active backend. + +## Why cancellation has two completion conditions + +A cancellation request and the original query complete independently: + +1. The cancellation connection must finish writing the `CancelRequest` and reach EOF. +2. The original query connection must reach `ReadyForQuery`. + +node-postgres pauses that client's query queue until both conditions are met. Without this barrier, a delayed cancellation request could arrive after the next query starts and cancel the wrong work. + +`cancelQuery(client)` resolving `true` means both sides completed safely. It does not mean cancellation won the race. + +## Race outcomes + +PostgreSQL decides the outcome once a query has been submitted: + +- If the query finishes first, its normal result is preserved. +- If cancellation wins, the query rejects with PostgreSQL error code `57014`. +- If an `AbortSignal` fires before submission, the query rejects with `signal.reason` and sends nothing. + +When delivery may have started but cannot be confirmed, node-postgres closes the original connection. This fail-closed behavior prevents an uncertain late cancellation request from affecting later work. A manual cancellation that fails before any write leaves the original query running and rejects only the `cancelQuery` promise. + +## Pools and transactions + +Explicit cancellation needs the client that owns the active query, so use `pool.connect()` rather than `pool.query()`. Keep the client checked out until the query and cancellation operation settle, then release it. A released `PoolClient` reference must not be reused. + +Cancellation does not roll back a transaction. When PostgreSQL cancels a statement inside a transaction, the transaction remains failed until the application issues `ROLLBACK`. + +## Limitations + +- Query cancellation is not supported in pipeline mode because multiple queries may already be in flight. +- A concrete custom `stream` instance cannot create the second connection. Supply a stream factory when cancellation is required. +- Successful dispatch cannot guarantee server-side cancellation because normal query completion may win the race. + +See the [`client.query` cancellation API](/apis/client#query-cancellation) for signatures and examples. diff --git a/packages/pg/esm/index.mjs b/packages/pg/esm/index.mjs index 587d80c1e..6e3f5e57a 100644 --- a/packages/pg/esm/index.mjs +++ b/packages/pg/esm/index.mjs @@ -10,6 +10,7 @@ export const Query = pg.Query export const DatabaseError = pg.DatabaseError export const escapeIdentifier = pg.escapeIdentifier export const escapeLiteral = pg.escapeLiteral +export const cancelQuery = pg.cancelQuery export const Result = pg.Result export const TypeOverrides = pg.TypeOverrides diff --git a/packages/pg/lib/abort.js b/packages/pg/lib/abort.js new file mode 100644 index 000000000..a6d6ea292 --- /dev/null +++ b/packages/pg/lib/abort.js @@ -0,0 +1,23 @@ +'use strict' + +function isAbortSignal(signal) { + return ( + signal && + typeof signal === 'object' && + typeof signal.aborted === 'boolean' && + typeof signal.addEventListener === 'function' && + typeof signal.removeEventListener === 'function' + ) +} + +function getAbortReason(signal) { + if (signal.reason !== undefined) { + return signal.reason + } + const error = new Error('This operation was aborted') + error.name = 'AbortError' + error.code = 'ABORT_ERR' + return error +} + +module.exports = { getAbortReason, isAbortSignal } diff --git a/packages/pg/lib/cancel-connection.js b/packages/pg/lib/cancel-connection.js new file mode 100644 index 000000000..48b7bbad5 --- /dev/null +++ b/packages/pg/lib/cancel-connection.js @@ -0,0 +1,129 @@ +'use strict' + +module.exports = function cancelConnection(Connection, connection, processID, secretKey, timeoutMillis, signal) { + const rejectBeforeConnect = (message) => { + const error = new Error(message) + Object.defineProperty(error, 'cancelDispatchMayHaveStarted', { value: false }) + return Promise.reject(error) + } + + if (connection._hasCustomStream && !connection._streamFactory) { + return rejectBeforeConnect('Cannot cancel a query using a concrete custom stream instance') + } + + const port = connection._cancelPort || connection._connectPort + const host = connection._cancelHost || connection._connectHost + if (port === undefined || port === null) { + return rejectBeforeConnect('Cannot cancel a query before the connection endpoint is known') + } + + const transport = new Connection({ + stream: connection._streamFactory || undefined, + ssl: connection.ssl, + sslNegotiation: connection.sslNegotiation, + sslServername: connection._sslServername || connection._connectHost, + }) + const timeout = timeoutMillis > 0 ? timeoutMillis : 5000 + + return new Promise((resolve, reject) => { + let settled = false + let writeAttempted = false + let writeCompleted = false + + const cleanup = (keepErrorListener) => { + clearTimeout(timer) + transport.removeListener('connect', onConnect) + transport.removeListener('sslconnect', sendCancel) + if (!keepErrorListener) { + transport.removeListener('error', fail) + } + transport.removeListener('end', onEnd) + signal?.removeEventListener?.('abort', onAbort) + } + + const finish = (error) => { + if (settled) { + return + } + settled = true + if (error) { + cleanup(true) + Object.defineProperty(error, 'cancelDispatchMayHaveStarted', { value: writeAttempted }) + transport.stream.destroy?.() + reject(error) + } else { + cleanup(false) + resolve() + } + } + + const fail = (error) => finish(error instanceof Error ? error : new Error(String(error))) + + const onAbort = () => { + const reason = signal.reason + fail(reason instanceof Error ? reason : new Error('Cancel request aborted')) + } + + const onEnd = () => { + if (!writeCompleted) { + fail(new Error('Cancel connection ended before the request was written')) + return + } + finish() + } + + const sendCancel = () => { + if (settled || writeAttempted) { + return + } + writeAttempted = true + try { + const accepted = transport.cancel(processID, secretKey, (error) => { + if (error) { + fail(error) + return + } + writeCompleted = true + }) + if (accepted === false) { + writeAttempted = false + fail(new Error('Cancel connection is not writable')) + } + } catch (error) { + fail(error) + } + } + + const onConnect = () => { + if (!transport.ssl) { + sendCancel() + } else if (transport.sslNegotiation !== 'direct') { + transport.requestSsl() + } + } + + const timer = setTimeout(() => { + const error = new Error('Cancel request timeout') + error.code = 'PG_CANCEL_TIMEOUT' + fail(error) + }, timeout) + timer.unref?.() + + transport.on('connect', onConnect) + transport.on('sslconnect', sendCancel) + transport.on('error', fail) + transport.on('end', onEnd) + signal?.addEventListener?.('abort', onAbort, { once: true }) + + if (signal?.aborted) { + onAbort() + return + } + + try { + transport.connect(port, host) + } catch (error) { + fail(error) + } + }) +} diff --git a/packages/pg/lib/cancel-query.js b/packages/pg/lib/cancel-query.js new file mode 100644 index 000000000..134df7745 --- /dev/null +++ b/packages/pg/lib/cancel-query.js @@ -0,0 +1,8 @@ +'use strict' + +module.exports = function cancelQuery(client) { + if (!client || typeof client._cancelQuery !== 'function') { + return Promise.reject(new TypeError('cancelQuery requires a connected pg Client')) + } + return client._cancelQuery(false) +} diff --git a/packages/pg/lib/client.js b/packages/pg/lib/client.js index 2b13c1de7..2cac2f068 100644 --- a/packages/pg/lib/client.js +++ b/packages/pg/lib/client.js @@ -9,6 +9,7 @@ const Query = require('./query') const defaults = require('./defaults') const Connection = require('./connection') const crypto = require('./crypto/utils') +const cancellation = require('./query-cancellation').js const activeQueryDeprecationNotice = nodeUtils.deprecate( () => {}, @@ -98,6 +99,7 @@ class Client extends EventEmitter { }) this._queryQueue = [] this._sentQueryQueue = [] + this._cancelState = null this.pipeline = Boolean(c.pipeline) this.binary = c.binary || defaults.binary this.processID = null @@ -386,12 +388,16 @@ class Client extends EventEmitter { this.emit('connect') } const activeQuery = this._getActiveQuery() + if (activeQuery) { + activeQuery._pgSubmitState = 'settled' + } this._activeQuery = null this._txStatus = msg?.status ?? null this.readyForQuery = true if (activeQuery) { activeQuery.handleReadyForQuery(this.connection) } + cancellation.queryDone(this, activeQuery) this._pulseQueryQueue() } @@ -418,6 +424,7 @@ class Client extends EventEmitter { return this._handleErrorWhileConnecting(err) } this._queryable = false + cancellation.fail(this, err) this._errorAllQueries(err) this.emit('error', err) } @@ -571,27 +578,12 @@ class Client extends EventEmitter { return data } - cancel(client, query) { - if (client.activeQuery === query) { - const con = this.connection - - if (this.host && this.host.indexOf('/') === 0) { - con.connect(this.host + '/.s.PGSQL.' + this.port) - } else { - con.connect(this.port, this.host) - } + _cancelQuery(failClosed) { + return cancellation.cancel(this, failClosed) + } - // once connection is established send cancel message - con.on('connect', function () { - con.cancel(client.processID, client.secretKey) - }) - } else if (client._queryQueue.indexOf(query) !== -1) { - client._queryQueue.splice(client._queryQueue.indexOf(query), 1) - } else if (client._sentQueryQueue.indexOf(query) !== -1) { - // Query already sent on wire — can't remove it without corrupting the - // pipeline. No-op the callback so the result is silently discarded. - query.callback = () => {} - } + cancel(client, query) { + cancellation.cancelLegacy(client, query) } setTypeParser(oid, format, parseFn) { @@ -614,6 +606,12 @@ class Client extends EventEmitter { } _pulseQueryQueue() { + if (!this._queryable || this._ended) { + return + } + if (this._cancelState) { + return + } if (this.pipeline) { this._pulsePipelinedQueryQueue() return @@ -624,8 +622,10 @@ class Client extends EventEmitter { if (activeQuery) { this.readyForQuery = false this.hasExecuted = true + cancellation.submitStart(activeQuery) const queryError = activeQuery.submit(this.connection) + cancellation.submitEnd(this, activeQuery, !queryError) if (queryError) { process.nextTick(() => { activeQuery.handleError(queryError, this.connection) @@ -691,7 +691,9 @@ class Client extends EventEmitter { }).catch((err) => { // replace the stack trace that leads to `TCP.onStreamRead` with one that leads back to the // application that created the query - Error.captureStackTrace(err) + if (err instanceof Error) { + Error.captureStackTrace(err) + } throw err }) } else if (typeof query.callback !== 'function') { @@ -699,6 +701,10 @@ class Client extends EventEmitter { } } + if (query instanceof Query && !cancellation.prepareQuery(this, query)) { + return result + } + const readTimeout = config.query_timeout || this.connectionParameters.query_timeout if (readTimeout) { const queryCallback = query.callback || (() => {}) @@ -748,6 +754,9 @@ class Client extends EventEmitter { // queries written behind it are answered out of its portal, so rows land on the wrong query and // the reads that follow fail with 'portal does not exist'. Refuse it instead of corrupting. if (this.pipeline) { + if (cancellation.rejectPipelineSignal(this, query)) { + return result + } const portalQuery = typeof config.submit === 'function' && !(query instanceof Query) ? 'Custom query classes such as pg-cursor and pg-query-stream are' @@ -779,6 +788,7 @@ class Client extends EventEmitter { if (this._queryQueue.length > 0 && !this.pipeline) { queryQueueLengthDeprecationNotice() } + cancellation.queue(query) this._queryQueue.push(query) this._pulseQueryQueue() return result diff --git a/packages/pg/lib/connection.js b/packages/pg/lib/connection.js index 62d38fa69..b5d4d8e31 100644 --- a/packages/pg/lib/connection.js +++ b/packages/pg/lib/connection.js @@ -5,6 +5,7 @@ const EventEmitter = require('events').EventEmitter const { parse, serialize } = require('pg-protocol') const stream = require('./stream') const { getStream } = stream +const cancelConnection = require('./cancel-connection') const flushBuffer = serialize.flush() const syncBuffer = serialize.sync() @@ -16,6 +17,8 @@ class Connection extends EventEmitter { super() config = config || {} + this._streamFactory = typeof config.stream === 'function' ? config.stream : null + this._hasCustomStream = Boolean(config.stream) this.stream = config.stream || getStream(config.ssl) if (typeof this.stream === 'function') { this.stream = this.stream(config) @@ -27,6 +30,7 @@ class Connection extends EventEmitter { this.submittedNamedStatements = {} this.ssl = config.ssl || false this.sslNegotiation = config.sslNegotiation || 'postgres' + this._sslServername = config.sslServername this._ending = false this._emitMessage = false const self = this @@ -41,10 +45,14 @@ class Connection extends EventEmitter { const self = this this._connecting = true + this._connectPort = port + this._connectHost = host this.stream.setNoDelay(true) this.stream.connect(port, host) this.stream.once('connect', function () { + self._cancelPort = self.stream.remotePort || port + self._cancelHost = self.stream.remoteAddress || host if (self._keepAlive) { self.stream.setKeepAlive(true, self._keepAliveInitialDelayMillis) } @@ -114,8 +122,9 @@ class Connection extends EventEmitter { } const net = require('net') - if (net.isIP && net.isIP(host) === 0) { - options.servername = host + const servername = self._sslServername || host + if (servername && net.isIP && net.isIP(servername) === 0) { + options.servername = servername } try { self.stream = stream.getSecureStream(options) @@ -146,8 +155,12 @@ class Connection extends EventEmitter { this.stream.write(serialize.startup(config)) } - cancel(processID, secretKey) { - this._send(serialize.cancel(processID, secretKey)) + cancel(processID, secretKey, callback) { + return this._send(serialize.cancel(processID, secretKey), callback) + } + + cancelWithClone(processID, secretKey, timeoutMillis, signal) { + return cancelConnection(Connection, this, processID, secretKey, timeoutMillis, signal) } password(password) { @@ -162,11 +175,12 @@ class Connection extends EventEmitter { this._send(serialize.sendSCRAMClientFinalMessage(additionalData)) } - _send(buffer) { + _send(buffer, callback) { if (!this.stream.writable) { return false } - return this.stream.write(buffer) + const accepted = this.stream.write(buffer, callback) + return callback ? true : accepted } query(text) { diff --git a/packages/pg/lib/index.js b/packages/pg/lib/index.js index e8b746149..83f2a7fd9 100644 --- a/packages/pg/lib/index.js +++ b/packages/pg/lib/index.js @@ -9,6 +9,7 @@ const Pool = require('pg-pool') const TypeOverrides = require('./type-overrides') const { DatabaseError } = require('pg-protocol') const { escapeIdentifier, escapeLiteral } = require('./utils') +const cancelQuery = require('./cancel-query') const poolFactory = (Client) => { return class BoundPool extends Pool { @@ -30,6 +31,7 @@ const PG = function (clientConstructor) { this.TypeOverrides = TypeOverrides this.escapeIdentifier = escapeIdentifier this.escapeLiteral = escapeLiteral + this.cancelQuery = cancelQuery this.Result = Result this.utils = utils } diff --git a/packages/pg/lib/native/client.js b/packages/pg/lib/native/client.js index 9ec3c8c03..122a49d14 100644 --- a/packages/pg/lib/native/client.js +++ b/packages/pg/lib/native/client.js @@ -12,6 +12,7 @@ const TypeOverrides = require('../type-overrides') const EventEmitter = require('events').EventEmitter const util = require('util') const ConnectionParameters = require('../connection-parameters') +const cancellation = require('../query-cancellation').native const NativeQuery = require('./query') @@ -32,6 +33,7 @@ const Client = (module.exports = function (config) { }) this._queryQueue = [] + this._cancelState = null this._ending = false this._connecting = false this._connected = false @@ -82,6 +84,10 @@ Client.prototype._errorAllQueries = function (err) { this._queryQueue.length = 0 } +Client.prototype._cancelQuery = function (failClosed) { + return cancellation.cancel(this, failClosed) +} + // connect to the backend // pass an optional callback to be called once connected // or with an error if there was a connection error @@ -110,6 +116,7 @@ Client.prototype._connect = function (cb) { // handle connection errors from the native layer self.native.on('error', function (err) { self._queryable = false + cancellation.fail(self, err) self._errorAllQueries(err) self.emit('error', err) }) @@ -182,13 +189,19 @@ Client.prototype.query = function (config, values, callback) { resolveOut = resolve rejectOut = reject }).catch((err) => { - Error.captureStackTrace(err) + if (err instanceof Error) { + Error.captureStackTrace(err) + } throw err }) query.callback = (err, res) => (err ? rejectOut(err) : resolveOut(res)) } } + if (query instanceof NativeQuery && !cancellation.prepareQuery(this, query)) { + return result + } + if (readTimeout) { queryCallback = query.callback || (() => {}) @@ -240,6 +253,10 @@ Client.prototype.query = function (config, values, callback) { queryQueueLengthDeprecationNotice() } + if (cancellation.rejectPipelineSignal(this, query)) { + return result + } + cancellation.queue(query) this._queryQueue.push(query) this._pulseQueryQueue() return result @@ -267,6 +284,8 @@ Client.prototype.end = function (cb) { self.native.end(function () { self._connected = false + cancellation.fail(self, new Error('Connection terminated during query cancellation')) + self._errorAllQueries(new Error('Connection terminated')) process.nextTick(() => { @@ -290,12 +309,15 @@ Client.prototype._hasActiveQuery = function () { } Client.prototype._pulseQueryQueue = function (initialConnection) { - if (!this._connected) { + if (!this._connected || !this._queryable) { return } if (this.pipeline && !initialConnection) { return this._pulsePipelinedQueryQueue() } + if (this._cancelState) { + return + } if (this._hasActiveQuery()) { return } @@ -307,11 +329,18 @@ Client.prototype._pulseQueryQueue = function (initialConnection) { return } this._activeQuery = query - query.submit(this) const self = this query.once('_done', function () { + query._pgSubmitState = 'settled' + if (self._activeQuery === query) { + self._activeQuery = null + } + cancellation.queryDone(self, query) self._pulseQueryQueue() }) + cancellation.submitStart(query) + query.submit(this) + cancellation.submitEnd(this, query, query.state !== 'error' && query.state !== 'end') } Client.prototype._pulsePipelinedQueryQueue = function () { @@ -401,11 +430,7 @@ Client.prototype._pulsePipelinedQueryQueue = function () { // attempt to cancel an in-progress query Client.prototype.cancel = function (query) { - if (this._activeQuery === query) { - this.native.cancel(function () {}) - } else if (this._queryQueue.indexOf(query) !== -1) { - this._queryQueue.splice(this._queryQueue.indexOf(query), 1) - } + cancellation.cancelLegacy(this, query) } Client.prototype.ref = function () {} diff --git a/packages/pg/lib/native/query.js b/packages/pg/lib/native/query.js index 8cb561979..3b675603b 100644 --- a/packages/pg/lib/native/query.js +++ b/packages/pg/lib/native/query.js @@ -11,6 +11,7 @@ const NativeQuery = (module.exports = function (config, values, callback) { this.values = config.values this.name = config.name this.queryMode = config.queryMode + this.signal = config.signal this.callback = config.callback this.state = 'new' this._arrayMode = config.rowMode === 'array' @@ -49,7 +50,7 @@ const errorFieldMap = { NativeQuery.prototype.handleError = function (err) { // copy pq error fields into the error object const fields = this.native && this.native.pq.resultErrorFields() - if (fields) { + if (fields && err && typeof err === 'object') { for (const key in fields) { const normalizedFieldName = errorFieldMap[key] || key err[normalizedFieldName] = fields[key] diff --git a/packages/pg/lib/query-cancellation.js b/packages/pg/lib/query-cancellation.js new file mode 100644 index 000000000..cf4f7376a --- /dev/null +++ b/packages/pg/lib/query-cancellation.js @@ -0,0 +1,311 @@ +'use strict' + +const { getAbortReason, isAbortSignal } = require('./abort') + +function create(native) { + const handleQueryError = (client, query, error, attachNative) => { + if (native && attachNative) { + query.native = client.native + } + if (native) { + query.handleError(error) + } else { + query.handleError(error, client.connection) + } + } + + const rejectQuery = (client, query, error, attachNative = false) => { + process.nextTick(() => handleQueryError(client, query, error, attachNative)) + } + + const settle = (client, state, error, value) => { + if (state.done) { + return + } + state.done = true + if (state.onEnd) { + client.removeListener('end', state.onEnd) + } + if (client._cancelState === state) { + client._cancelState = null + } + if (error) { + state.reject(error) + } else { + state.resolve(value) + } + client._pulseQueryQueue() + } + + const maybeComplete = (client, state) => { + if (state.queryDone && state.transportDone) { + settle(client, state, null, true) + } + } + + const handleJsFailure = (client, state, error) => { + if (state.done) { + return + } + if (state.failClosed || error.cancelDispatchMayHaveStarted) { + state.failure = error + client._queryable = false + client._ending = true + if (client._ended) { + settle(client, state, error) + } else { + client.connection.stream.destroy() + } + return + } + settle(client, state, error) + } + + const begin = (client, state) => { + if (state.done || state.started) { + return + } + state.started = true + + if (native) { + const complete = (error) => { + if (state.done) { + return + } + if (error) { + settle(client, state, error) + if (state.failClosed) { + client._queryable = false + client.end(() => {}) + } + return + } + state.transportDone = true + maybeComplete(client, state) + } + try { + client.native.cancel(complete) + } catch (error) { + complete(error) + } + return + } + + state.controller = new AbortController() + state.onEnd = () => { + const error = state.failure || new Error('Connection terminated during query cancellation') + state.controller.abort(error) + settle(client, state, error) + } + client.once('end', state.onEnd) + + let cancellation + try { + cancellation = client.connection.cancelWithClone( + client.processID, + client.secretKey, + client._connectionTimeoutMillis, + state.controller.signal + ) + } catch (error) { + handleJsFailure(client, state, error) + return + } + cancellation + .then(() => { + if (state.done) { + return + } + state.transportDone = true + maybeComplete(client, state) + }) + .catch((error) => handleJsFailure(client, state, error)) + } + + const cancel = (client, failClosed) => { + const reject = (error) => new client._Promise((resolve, reject) => reject(error)) + const resolve = (value) => new client._Promise((resolve) => resolve(value)) + + if (client.pipeline) { + return reject(new Error('Query cancellation is not supported in pipeline mode')) + } + if (!client._connected || !client._queryable || client._ending || (!native && client._ended)) { + return reject(new Error('Client is not connected and queryable')) + } + if (client._cancelState) { + client._cancelState.failClosed ||= failClosed + return client._cancelState.promise + } + + const target = native ? (client._hasActiveQuery() ? client._activeQuery : null) : client._getActiveQuery() + if (!target) { + return resolve(false) + } + + const state = { + target, + failClosed: Boolean(failClosed), + started: false, + queryDone: false, + transportDone: false, + done: false, + } + state.promise = new client._Promise((resolve, reject) => { + state.resolve = resolve + state.reject = reject + }) + client._cancelState = state + + if (target._pgSubmitState !== 'submitting') { + begin(client, state) + } + return state.promise + } + + const fail = (client, error) => { + const state = client._cancelState + if (!state) { + return + } + if (native) { + settle(client, state, error) + return + } + const cancelError = state.failure || error + state.controller?.abort(cancelError) + settle(client, state, cancelError) + } + + const queryDone = (client, query) => { + const state = client._cancelState + if (!state || state.done || (native ? state.target !== query : !state.started)) { + return + } + state.queryDone = true + maybeComplete(client, state) + } + + const submitStart = (query) => { + query._pgSubmitState = 'submitting' + if (query.signal != null) { + query._abortState = 'submitting' + } + } + + const submitEnd = (client, query, submitted) => { + const state = client._cancelState + if (!submitted) { + query._pgSubmitState = 'settled' + if (state && state.target === query && !state.started) { + settle(client, state, null, false) + } + return + } + + query._pgSubmitState = 'active' + if (query.signal != null) { + query._abortState = 'active' + } + if (state && state.target === query && !state.started) { + begin(client, state) + } + } + + const prepareQuery = (client, query) => { + const signal = query.signal + if (signal == null) { + return true + } + if (!isAbortSignal(signal)) { + rejectQuery(client, query, new TypeError('Query signal must be an AbortSignal'), native) + return false + } + + const state = { abortRequested: false, settled: false } + query._abortState = 'preparing' + const originalCallback = query.callback + const cleanup = () => signal.removeEventListener('abort', onAbort) + + query.callback = (error, result) => { + if (state.settled) { + return + } + state.settled = true + query._abortState = 'settled' + cleanup() + originalCallback(error, result) + } + + const onAbort = () => { + if (state.settled) { + return + } + const reason = getAbortReason(signal) + if (query._abortState === 'preparing') { + state.abortRequested = true + state.reason = reason + return + } + if (query._abortState === 'queued') { + const index = client._queryQueue.indexOf(query) + if (index !== -1) { + client._queryQueue.splice(index, 1) + query._abortState = 'aborting' + rejectQuery(client, query, reason) + client._pulseQueryQueue() + } + return + } + if (query._abortState === 'submitting' || query._abortState === 'active') { + cancel(client, true).catch(() => {}) + } + } + + signal.addEventListener('abort', onAbort, { once: true }) + if (signal.aborted) { + onAbort() + } + if (state.abortRequested) { + rejectQuery(client, query, state.reason, native) + return false + } + return true + } + + const rejectPipelineSignal = (client, query) => { + if (!client.pipeline || query.signal == null) { + return false + } + rejectQuery(client, query, new Error('AbortSignal is not supported in pipeline mode'), native) + return true + } + + const queue = (query) => { + if (query.signal != null) { + query._abortState = 'queued' + } + } + + const cancelLegacy = (client, query) => { + const activeQuery = native ? client._activeQuery : client._getActiveQuery() + if (activeQuery === query) { + cancel(client, false).catch(() => {}) + return + } + + const index = client._queryQueue.indexOf(query) + if (index !== -1) { + client._queryQueue.splice(index, 1) + rejectQuery(client, query, getAbortReason({})) + client._pulseQueryQueue() + } else if (!native && client._sentQueryQueue.indexOf(query) !== -1) { + // Query already sent on wire — can't remove it without corrupting the + // pipeline. No-op the callback so the result is silently discarded. + query.callback = () => {} + } + } + + return { cancel, cancelLegacy, fail, prepareQuery, queryDone, queue, rejectPipelineSignal, submitEnd, submitStart } +} + +module.exports = { js: create(false), native: create(true) } diff --git a/packages/pg/lib/query.js b/packages/pg/lib/query.js index 6b9214199..9b0e05c91 100644 --- a/packages/pg/lib/query.js +++ b/packages/pg/lib/query.js @@ -18,6 +18,7 @@ class Query extends EventEmitter { this.name = config.name this.queryMode = config.queryMode this.binary = config.binary + this.signal = config.signal // use unique portal name each time this.portal = config.portal || '' this.callback = config.callback diff --git a/packages/pg/test/integration/client/cancel-query-tests.js b/packages/pg/test/integration/client/cancel-query-tests.js new file mode 100644 index 000000000..f8b4be77e --- /dev/null +++ b/packages/pg/test/integration/client/cancel-query-tests.js @@ -0,0 +1,111 @@ +'use strict' + +const assert = require('assert') +const helper = require('../test-helper') + +const { Client } = helper +const { cancelQuery } = helper.pg +const suite = new helper.Suite() +let testId = 0 + +const delay = (milliseconds) => new Promise((resolve) => setTimeout(resolve, milliseconds)) + +function config(label) { + testId++ + return { ...helper.config, application_name: `pg-cancel-${label}-${process.pid}-${testId}` } +} + +async function waitUntilActive(observer, applicationName) { + for (let attempt = 0; attempt < 50; attempt++) { + const result = await observer.query( + "SELECT 1 FROM pg_stat_activity WHERE application_name = $1 AND state = 'active'", + [applicationName] + ) + if (result.rowCount === 1) { + return + } + await delay(20) + } + throw new Error('query did not become active') +} + +suite.test('cancels an active query and preserves the client for following work', async () => { + const clientConfig = { ...config('manual'), ssl: { rejectUnauthorized: false } } + const client = new Client(clientConfig) + const observer = new Client(helper.config) + await Promise.all([client.connect(), observer.connect()]) + + try { + const running = client.query('SELECT pg_sleep(30)') + await waitUntilActive(observer, clientConfig.application_name) + const cancellation = cancelQuery(client) + + await assert.rejects(running, (error) => error.code === '57014') + assert.strictEqual(await cancellation, true) + assert.strictEqual((await client.query('SELECT 42 AS answer')).rows[0].answer, 42) + assert.strictEqual(await cancelQuery(client), false) + } finally { + await Promise.all([client.end(), observer.end()]) + } +}) + +suite.test('AbortSignal cancels an active transaction until explicit rollback', async () => { + const clientConfig = config('signal') + const client = new Client(clientConfig) + const observer = new Client(helper.config) + await Promise.all([client.connect(), observer.connect()]) + + try { + await client.query('BEGIN') + const controller = new AbortController() + const running = client.query({ text: 'SELECT pg_sleep(30)', signal: controller.signal }) + await waitUntilActive(observer, clientConfig.application_name) + controller.abort() + const cancellation = cancelQuery(client) + + await assert.rejects(running, (error) => error.code === '57014') + assert.strictEqual(await cancellation, true) + assert.strictEqual(client.getTransactionStatus(), 'E') + await assert.rejects(client.query('SELECT 1'), (error) => error.code === '25P02') + await client.query('ROLLBACK') + assert.strictEqual(client.getTransactionStatus(), 'I') + assert.strictEqual((await client.query('SELECT 1 AS ok')).rows[0].ok, 1) + } finally { + await Promise.all([client.end(), observer.end()]) + } +}) + +suite.test('holds a released size-one pool client until cancellation completes', async () => { + const poolConfig = config('pool') + const pool = new helper.pg.Pool({ ...poolConfig, max: 1 }) + const observer = new Client(helper.config) + await observer.connect() + + try { + const checkedOut = await pool.connect() + const running = checkedOut.query('SELECT pg_sleep(30)') + await waitUntilActive(observer, poolConfig.application_name) + const cancellation = cancelQuery(checkedOut) + const timeline = [] + const cancelledQuery = assert.rejects(running, (error) => error.code === '57014').then(() => timeline.push('query')) + const cancelled = cancellation.then((result) => { + assert.strictEqual(result, true) + timeline.push('cancel') + }) + + checkedOut.release() + const reacquired = await pool.connect() + assert.strictEqual(reacquired, checkedOut) + const following = reacquired.query('SELECT 7 AS answer').then((result) => { + timeline.push('following') + return result + }) + + await Promise.all([cancelledQuery, cancelled]) + assert.strictEqual((await following).rows[0].answer, 7) + assert.strictEqual(timeline[timeline.length - 1], 'following') + reacquired.release() + } finally { + await Promise.all([pool.end(), observer.end()]) + } +}) diff --git a/packages/pg/test/unit/client/cancel-query-tests.js b/packages/pg/test/unit/client/cancel-query-tests.js new file mode 100644 index 000000000..a1c59ac81 --- /dev/null +++ b/packages/pg/test/unit/client/cancel-query-tests.js @@ -0,0 +1,268 @@ +'use strict' + +const assert = require('assert') +const helper = require('./test-helper') +const { cancelQuery } = require('../../../lib') + +const suite = new helper.Suite() +const test = suite.test.bind(suite) + +function connectedClient() { + const client = helper.client() + client.connection.emit('readyForQuery', { status: 'I' }) + return client +} + +test('returns false without an active query', async () => { + const client = connectedClient() + assert.strictEqual(await cancelQuery(client), false) +}) + +test('shares one cancellation and holds the next query until ready and cancel EOF', async () => { + const client = connectedClient() + let finishCancel + let cancelCalls = 0 + client.connection.cancelWithClone = () => { + cancelCalls++ + return new Promise((resolve) => { + finishCancel = resolve + }) + } + + const first = client.query('SELECT 1') + const firstCancel = cancelQuery(client) + const duplicateCancel = cancelQuery(client) + assert.strictEqual(firstCancel, duplicateCancel) + assert.strictEqual(cancelCalls, 1) + + const second = client.query('SELECT 2') + client.connection.emit('readyForQuery', { status: 'I' }) + await first + assert.deepStrictEqual(client.connection.queries, ['SELECT 1']) + + finishCancel() + assert.strictEqual(await firstCancel, true) + assert.deepStrictEqual(client.connection.queries, ['SELECT 1', 'SELECT 2']) + + client.connection.emit('readyForQuery', { status: 'I' }) + await second +}) + +test('rejects an already-aborted query without submitting it', async () => { + const client = connectedClient() + const controller = new AbortController() + const reason = new Error('stop before queue') + controller.abort(reason) + + await assert.rejects(client.query({ text: 'SELECT 1', signal: controller.signal }), (error) => error === reason) + assert.deepStrictEqual(client.connection.queries, []) +}) + +test('rejects invalid signals before submitting', async () => { + const client = connectedClient() + + await assert.rejects(client.query({ text: 'SELECT 1', signal: {} }), /must be an AbortSignal/) + assert.deepStrictEqual(client.connection.queries, []) +}) + +test('removes and rejects an aborted queued query', async () => { + const client = connectedClient() + const first = client.query('SELECT 1') + const controller = new AbortController() + const reason = new Error('stop in queue') + const queued = client.query({ text: 'SELECT 2', signal: controller.signal }) + + controller.abort(reason) + await assert.rejects(queued, (error) => error === reason) + assert.deepStrictEqual(client.connection.queries, ['SELECT 1']) + + client.connection.emit('readyForQuery', { status: 'I' }) + await first +}) + +test('resolves deferred manual cancellation false when submit fails', async () => { + const client = connectedClient() + let cancelPromise + let cancelCalls = 0 + client.connection.cancelWithClone = () => { + cancelCalls++ + return Promise.resolve() + } + const submitError = new Error('submit failed') + const queryError = new Promise((resolve) => { + client.query({ + callback: (error) => resolve(error), + submit() { + cancelPromise = cancelQuery(client) + return submitError + }, + handleError(error) { + this.callback(error) + }, + }) + }) + + assert.strictEqual(await cancelPromise, false) + assert.strictEqual(await queryError, submitError) + assert.strictEqual(cancelCalls, 0) +}) + +test('defers a signal fired synchronously during submit', async () => { + const client = connectedClient() + const controller = new AbortController() + let cancelCalls = 0 + let insideSubmit = false + client.connection.cancelWithClone = () => { + assert.strictEqual(insideSubmit, false) + cancelCalls++ + return Promise.resolve() + } + client.connection.query = function (text) { + insideSubmit = true + this.queries.push(text) + controller.abort() + assert.strictEqual(cancelCalls, 0) + insideSubmit = false + } + + const query = client.query({ text: 'SELECT 1', signal: controller.signal }) + assert.strictEqual(cancelCalls, 1) + client.connection.emit('readyForQuery', { status: 'I' }) + await query +}) + +test('holds the queue when the server error and ready arrive before cancel EOF', async () => { + const client = connectedClient() + let finishCancel + client.connection.cancelWithClone = () => + new Promise((resolve) => { + finishCancel = resolve + }) + + const first = client.query('SELECT pg_sleep(1)') + const cancellation = cancelQuery(client) + const second = client.query('SELECT 2') + const serverError = Object.assign(new Error('canceling statement due to user request'), { code: '57014' }) + client.connection.emit('errorMessage', serverError) + client.connection.emit('readyForQuery', { status: 'I' }) + + await assert.rejects(first, (error) => error === serverError) + assert.deepStrictEqual(client.connection.queries, ['SELECT pg_sleep(1)']) + finishCancel() + assert.strictEqual(await cancellation, true) + assert.deepStrictEqual(client.connection.queries, ['SELECT pg_sleep(1)', 'SELECT 2']) + + client.connection.emit('readyForQuery', { status: 'I' }) + await second +}) + +test('manual cancellation failure before dispatch leaves the client usable', async () => { + const client = connectedClient() + const failure = new Error('connect failed') + Object.defineProperty(failure, 'cancelDispatchMayHaveStarted', { value: false }) + client.connection.cancelWithClone = () => Promise.reject(failure) + + const first = client.query('SELECT 1') + await assert.rejects(cancelQuery(client), (error) => error === failure) + assert.strictEqual(client._queryable, true) + client.connection.emit('readyForQuery', { status: 'I' }) + await first + + const second = client.query('SELECT 2') + client.connection.emit('readyForQuery', { status: 'I' }) + await second +}) + +test('settles cancellation when the transport throws synchronously', async () => { + const client = connectedClient() + const failure = new Error('stream factory failed') + client.connection.cancelWithClone = () => { + throw failure + } + + const query = client.query('SELECT 1') + await assert.rejects(cancelQuery(client), (error) => error === failure) + assert.strictEqual(client._cancelState, null) + client.connection.emit('readyForQuery', { status: 'I' }) + await query +}) + +test('settles cancellation when the original stream errors without closing', async () => { + const client = connectedClient() + const failure = new Error('original stream failed') + client.connection.cancelWithClone = () => new Promise(() => {}) + client.on('error', () => {}) + + const query = client.query('SELECT 1') + const cancellation = cancelQuery(client) + client.connection.emit('error', failure) + + await assert.rejects(cancellation, (error) => error === failure) + await assert.rejects(query, (error) => error === failure) + assert.strictEqual(client._cancelState, null) +}) + +test('signal upgrades a shared manual cancellation failure to fail closed', async () => { + const client = connectedClient() + const controller = new AbortController() + const failure = new Error('cancel transport failed') + Object.defineProperty(failure, 'cancelDispatchMayHaveStarted', { value: false }) + let rejectCancel + let cancelCalls = 0 + client.connection.cancelWithClone = () => { + cancelCalls++ + return new Promise((resolve, reject) => { + rejectCancel = reject + }) + } + client.connection.stream = { + destroy() { + client.connection.emit('end') + }, + } + + const query = client.query({ text: 'SELECT 1', signal: controller.signal }) + const manual = cancelQuery(client) + controller.abort() + assert.strictEqual(cancelCalls, 1) + rejectCancel(failure) + + await assert.rejects(manual, (error) => error === failure) + await assert.rejects(query, /Connection terminated/) + assert.strictEqual(client._queryable, false) +}) + +test('removes the abort listener when a query settles', async () => { + const client = connectedClient() + const controller = new AbortController() + const signal = controller.signal + const add = signal.addEventListener.bind(signal) + const remove = signal.removeEventListener.bind(signal) + let listeners = 0 + signal.addEventListener = (...args) => { + listeners++ + return add(...args) + } + signal.removeEventListener = (...args) => { + listeners-- + return remove(...args) + } + + const query = client.query({ text: 'SELECT 1', signal }) + assert.strictEqual(listeners, 1) + client.connection.emit('readyForQuery', { status: 'I' }) + await query + assert.strictEqual(listeners, 0) +}) + +test('rejects signal queries in pipeline mode before submit', async () => { + const client = helper.client({ pipeline: true }) + client.connection.emit('readyForQuery', { status: 'I' }) + const controller = new AbortController() + + await assert.rejects( + client.query({ text: 'SELECT 1', signal: controller.signal }), + /AbortSignal is not supported in pipeline mode/ + ) + assert.deepStrictEqual(client.connection.queries, []) +}) diff --git a/packages/pg/test/unit/client/native-cancel-query-tests.js b/packages/pg/test/unit/client/native-cancel-query-tests.js new file mode 100644 index 000000000..81208420a --- /dev/null +++ b/packages/pg/test/unit/client/native-cancel-query-tests.js @@ -0,0 +1,149 @@ +'use strict' + +const assert = require('assert') +const EventEmitter = require('events') +const helper = require('./test-helper') +const { cancelQuery } = require('../../../lib') + +class FakeNative extends EventEmitter { + constructor() { + super() + this.pq = { resultErrorFields: () => null } + this.queries = [] + this.cancelCallbacks = [] + FakeNative.instance = this + } + + query(text, values, callback) { + if (typeof values === 'function') { + callback = values + } + this.queries.push({ text, callback }) + if (this.onQuery) { + this.onQuery() + } + } + + cancel(callback) { + this.cancelCallbacks.push(callback) + } + + end(callback) { + if (callback) { + setImmediate(callback) + } + } +} + +const pgNativePath = require.resolve('pg-native') +require.cache[pgNativePath] = { id: pgNativePath, filename: pgNativePath, loaded: true, exports: FakeNative } +delete require.cache[require.resolve('../../../lib/native/client')] +const NativeClient = require('../../../lib/native/client') + +const suite = new helper.Suite() +const test = suite.test.bind(suite) +const tick = () => new Promise((resolve) => setImmediate(resolve)) + +function connectedClient(config) { + const client = new NativeClient(config) + client._connected = true + return client +} + +test('shares one native cancellation and holds the queue until both sides finish', async () => { + const client = connectedClient() + const native = FakeNative.instance + const first = client.query('SELECT 1') + const firstCancel = cancelQuery(client) + const duplicateCancel = cancelQuery(client) + const second = client.query('SELECT 2') + + assert.strictEqual(firstCancel, duplicateCancel) + assert.strictEqual(native.cancelCallbacks.length, 1) + native.queries[0].callback(null, [], []) + await first + await tick() + assert.deepStrictEqual( + native.queries.map((query) => query.text), + ['SELECT 1'] + ) + + native.cancelCallbacks[0]() + assert.strictEqual(await firstCancel, true) + assert.deepStrictEqual( + native.queries.map((query) => query.text), + ['SELECT 1', 'SELECT 2'] + ) + + native.queries[1].callback(null, [], []) + await second +}) + +test('rejects pre-aborted native work without submitting it', async () => { + const client = connectedClient() + const controller = new AbortController() + const reason = 'native stop' + controller.abort(reason) + + await assert.rejects(client.query({ text: 'SELECT 1', signal: controller.signal }), (error) => error === reason) + assert.deepStrictEqual(FakeNative.instance.queries, []) +}) + +test('defers reentrant native cancellation until submit returns', async () => { + const client = connectedClient() + const native = FakeNative.instance + let cancelPromise + native.onQuery = () => { + cancelPromise = cancelQuery(client) + assert.strictEqual(native.cancelCallbacks.length, 0) + } + + const query = client.query('SELECT 1') + assert.strictEqual(native.cancelCallbacks.length, 1) + native.cancelCallbacks[0]() + native.queries[0].callback(null, [], []) + await query + assert.strictEqual(await cancelPromise, true) +}) + +test('settles native cancellation when libpq cancel throws synchronously', async () => { + const client = connectedClient() + const native = FakeNative.instance + const failure = new Error('native cancel threw') + native.cancel = () => { + throw failure + } + + const query = client.query('SELECT 1') + await assert.rejects(cancelQuery(client), (error) => error === failure) + assert.strictEqual(client._cancelState, null) + native.queries[0].callback(null, [], []) + await query +}) + +test('native signal upgrades a shared manual cancellation failure to fail closed', async () => { + const client = connectedClient() + const native = FakeNative.instance + const controller = new AbortController() + const query = client.query({ text: 'SELECT 1', signal: controller.signal }) + const manual = cancelQuery(client) + controller.abort() + assert.strictEqual(native.cancelCallbacks.length, 1) + const failure = new Error('native cancel failed') + native.cancelCallbacks[0](failure) + + await assert.rejects(manual, (error) => error === failure) + await assert.rejects(query, /Connection terminated/) + assert.strictEqual(client._queryable, false) +}) + +test('rejects native signal queries in pipeline mode before submit', async () => { + const client = connectedClient({ pipeline: true }) + const controller = new AbortController() + + await assert.rejects( + client.query({ text: 'SELECT 1', signal: controller.signal }), + /AbortSignal is not supported in pipeline mode/ + ) + assert.deepStrictEqual(FakeNative.instance.queries, []) +}) diff --git a/packages/pg/test/unit/connection/cancel-query-tests.js b/packages/pg/test/unit/connection/cancel-query-tests.js new file mode 100644 index 000000000..1de7f7db0 --- /dev/null +++ b/packages/pg/test/unit/connection/cancel-query-tests.js @@ -0,0 +1,190 @@ +'use strict' + +const assert = require('assert') +const EventEmitter = require('events') +const Connection = require('../../../lib/connection') +const stream = require('../../../lib/stream') +const helper = require('./test-helper') + +const suite = new helper.Suite() +const test = suite.test.bind(suite) + +class CancelStream extends EventEmitter { + constructor(onWrite) { + super() + this.onWrite = onWrite + this.writable = true + this.destroyed = false + } + + setNoDelay() {} + setKeepAlive() {} + + connect(port, host) { + this.port = port + this.host = host + process.nextTick(() => this.emit('connect')) + } + + write(packet, callback) { + this.packet = packet + const accepted = this.onWrite(this, callback) + return accepted === undefined ? true : accepted + } + + destroy() { + this.destroyed = true + } +} + +function cancellableConnection(onWrite) { + const streams = [] + const connection = new Connection({ + stream: () => { + const stream = new CancelStream(onWrite) + streams.push(stream) + return stream + }, + }) + connection._connectPort = connection._cancelPort = 5432 + connection._connectHost = 'db.example.test' + connection._cancelHost = '192.0.2.10' + return { connection, streams } +} + +test('uses the selected endpoint and resolves only after write completion then EOF', async () => { + const { connection, streams } = cancellableConnection((stream, callback) => { + callback() + process.nextTick(() => stream.emit('close')) + }) + + await connection.cancelWithClone(123, 456, 50) + const stream = streams[1] + assert.strictEqual(stream.host, '192.0.2.10') + assert.strictEqual(stream.port, 5432) + assert.strictEqual(stream.packet.toString('hex'), '0000001004d2162e0000007b000001c8') +}) + +test('rejects EOF before write completion as ambiguous', async () => { + const { connection } = cancellableConnection((stream) => { + stream.emit('close') + }) + + await assert.rejects(connection.cancelWithClone(123, 456, 50), (error) => { + assert.strictEqual(error.cancelDispatchMayHaveStarted, true) + return /before the request was written/.test(error.message) + }) +}) + +test('rejects concrete custom stream instances before connect', async () => { + const stream = new CancelStream(() => {}) + const connection = new Connection({ stream }) + connection._connectPort = 5432 + connection._connectHost = 'localhost' + + await assert.rejects(connection.cancelWithClone(123, 456, 50), (error) => { + assert.strictEqual(error.cancelDispatchMayHaveStarted, false) + return /concrete custom stream/.test(error.message) + }) +}) + +test('marks a rejected write as not dispatched', async () => { + const { connection, streams } = cancellableConnection(() => { + throw new Error('write must not be called') + }) + const cancellation = connection.cancelWithClone(123, 456, 50) + streams[1].writable = false + + await assert.rejects(cancellation, (error) => { + assert.strictEqual(error.cancelDispatchMayHaveStarted, false) + return /not writable/.test(error.message) + }) + streams[1].emit('error', new Error('late socket error')) +}) + +test('treats write backpressure as accepted dispatch', async () => { + const { connection } = cancellableConnection((socket, callback) => { + callback() + process.nextTick(() => socket.emit('close')) + return false + }) + + await connection.cancelWithClone(123, 456, 50) +}) + +test('reuses a selected Unix socket path', async () => { + const { connection, streams } = cancellableConnection((socket, callback) => { + callback() + process.nextTick(() => socket.emit('close')) + }) + connection._cancelPort = '/var/run/postgresql/.s.PGSQL.5432' + connection._cancelHost = undefined + connection._connectHost = undefined + + await connection.cancelWithClone(123, 456, 50) + assert.strictEqual(streams[1].port, '/var/run/postgresql/.s.PGSQL.5432') + assert.strictEqual(streams[1].host, undefined) +}) + +test('bounds a written request that never reaches EOF', async () => { + const { connection, streams } = cancellableConnection((socket, callback) => callback()) + + await assert.rejects(connection.cancelWithClone(123, 456, 5), (error) => { + assert.strictEqual(error.code, 'PG_CANCEL_TIMEOUT') + assert.strictEqual(error.cancelDispatchMayHaveStarted, true) + return true + }) + assert.strictEqual(streams[1].destroyed, true) +}) + +async function testTlsCancel(sslNegotiation) { + const rawStreams = [] + const secureStreams = [] + const tlsOptions = [] + const originalGetSecureStream = stream.getSecureStream + + stream.getSecureStream = (options) => { + tlsOptions.push(options) + const secure = new CancelStream((raw, callback) => { + callback() + process.nextTick(() => options.socket.emit('close')) + }) + secureStreams.push(secure) + return secure + } + + try { + const connection = new Connection({ + stream: () => { + const raw = new CancelStream((socket, callback) => { + assert.strictEqual(socket.packet.toString('hex'), '0000000804d2162f') + callback?.() + process.nextTick(() => socket.emit('data', Buffer.from('S'))) + return true + }) + rawStreams.push(raw) + return raw + }, + ssl: { rejectUnauthorized: false }, + sslNegotiation, + }) + connection._connectPort = connection._cancelPort = 5432 + connection._connectHost = 'db.example.test' + connection._cancelHost = '192.0.2.10' + + await connection.cancelWithClone(123, 456, 50) + assert.strictEqual(tlsOptions[0].servername, 'db.example.test') + assert.strictEqual(secureStreams[0].packet.toString('hex'), '0000001004d2162e0000007b000001c8') + if (sslNegotiation === 'direct') { + assert.deepStrictEqual(tlsOptions[0].ALPNProtocols, ['postgresql']) + assert.strictEqual(rawStreams[1].packet, undefined) + } + } finally { + stream.getSecureStream = originalGetSecureStream + } +} + +test('preserves standard and direct TLS cancellation', async () => { + await testTlsCancel('postgres') + await testTlsCancel('direct') +})