diff --git a/packages/pg-cursor/index.js b/packages/pg-cursor/index.js index 5d566882b..5284e46bd 100644 --- a/packages/pg-cursor/index.js +++ b/packages/pg-cursor/index.js @@ -27,7 +27,10 @@ class Cursor extends EventEmitter { } _ifNoData() { - if (this.state !== 'done' && this.state !== 'error') { + // Only start a queued read when the portal is first ready. An in-flight + // execute (state === 'busy') must not be interrupted, and closed/failed + // cursors must stay closed. + if (this.state === 'submitted' || this.state === 'initialized') { this.state = 'idle' this._shiftQueue() } @@ -108,7 +111,9 @@ class Cursor extends EventEmitter { handleRowDescription(msg) { this._result.addFields(msg.fields) - if (this.state !== 'done' && this.state !== 'error') { + // If read() already sent Execute (state === 'busy'), keep that in-flight + // callback. Resetting to idle here would start a queued read and overwrite _cb. + if (this.state === 'submitted' || this.state === 'initialized') { this.state = 'idle' this._shiftQueue() } @@ -124,23 +129,40 @@ class Cursor extends EventEmitter { if (this.state !== 'done') { this.state = 'idle' } - setImmediate(() => { - const cb = this._cb - // remove callback before calling it - // because likely a new one will be added - // within the call to this callback - this._cb = null - if (cb) { - this._result.rows = this._rows - cb(null, this._rows, this._result) - } - this._rows = [] - }) + const cb = this._cb + const rows = this._rows + // Drop the callback synchronously so a later ReadyForQuery cannot + // deliver the same batch twice after CommandComplete already settled it. + this._cb = null + this._rows = [] + if (this.state === 'idle') { + this._shiftQueue() + } + if (cb) { + setImmediate(() => { + this._result.rows = rows + cb(null, rows, this._result) + }) + } + } + + _fulfillQueue(err, rows) { + const queue = this._queue.splice(0, this._queue.length) + for (let i = 0; i < queue.length; i++) { + const queuedCallback = queue[i][1] + setImmediate(() => queuedCallback.call(this, err, rows)) + } } handleCommandComplete(msg) { this._result.addCommandComplete(msg) this._closePortal() + // CommandComplete means the portal is exhausted. Settle the in-flight + // read with whatever rows remain (possibly fewer than requested, or none) + // instead of waiting for ReadyForQuery, which the client may never + // dispatch to this cursor. + this._sendRows() + this._fulfillQueue(null, []) } handlePortalSuspended() { @@ -150,6 +172,7 @@ class Cursor extends EventEmitter { handleReadyForQuery() { this._sendRows() this.state = 'done' + this._fulfillQueue(null, []) this.emit('end', this._result) } diff --git a/packages/pg-cursor/test/read-exhaustion.js b/packages/pg-cursor/test/read-exhaustion.js new file mode 100644 index 000000000..d4ca6e018 --- /dev/null +++ b/packages/pg-cursor/test/read-exhaustion.js @@ -0,0 +1,161 @@ +const assert = require('assert') +const EventEmitter = require('events') +const Cursor = require('../') +const pg = require('pg') + +class TestConnection extends EventEmitter { + constructor() { + super() + this.calls = { + close: 0, + execute: 0, + sync: 0, + } + } + + parse() {} + bind() {} + describe() {} + flush() {} + + execute() { + this.calls.execute++ + } + + close() { + this.calls.close++ + } + + sync() { + this.calls.sync++ + } +} + +const INT4 = { name: 'num', dataTypeID: 23 } + +const pushIntRows = (cursor, values) => { + for (const value of values) { + cursor.handleDataRow({ fields: [String(value)] }) + } +} + +const settleOrTimeout = (promise, ms = 150) => + Promise.race([ + promise.then((value) => ({ value })), + new Promise((resolve) => setTimeout(() => resolve({ timeout: true }), ms)), + ]) + +describe('read() when rowCount exceeds remaining rows (#2949)', function () { + it('settles as soon as the portal is exhausted', async function () { + const cursor = new Cursor('SELECT num FROM generate_series(1, 2) num') + const connection = new TestConnection() + cursor.submit(connection) + + const read = cursor.read(1000) + cursor.handleRowDescription({ fields: [INT4] }) + pushIntRows(cursor, [1, 2]) + cursor.handleCommandComplete({ text: 'SELECT 2' }) + + const outcome = await settleOrTimeout(read) + assert.ok(!outcome.timeout, 'read() hung after CommandComplete when rowCount exceeded remaining rows') + assert.deepStrictEqual(outcome.value, [{ num: 1 }, { num: 2 }]) + assert.strictEqual(cursor.state, 'done') + assert.deepStrictEqual(connection.calls, { close: 1, execute: 1, sync: 1 }) + }) + + it('returns a shorter batch when only some rows remain', async function () { + const cursor = new Cursor('SELECT num FROM generate_series(1, 5) num') + const connection = new TestConnection() + cursor.submit(connection) + cursor.handleRowDescription({ fields: [INT4] }) + + const first = cursor.read(3) + pushIntRows(cursor, [1, 2, 3]) + cursor.handlePortalSuspended() + assert.deepStrictEqual(await first, [{ num: 1 }, { num: 2 }, { num: 3 }]) + + const remaining = cursor.read(1000) + pushIntRows(cursor, [4, 5]) + cursor.handleCommandComplete({ text: 'SELECT 5' }) + + const outcome = await settleOrTimeout(remaining) + assert.ok(!outcome.timeout, 'read() hung when fetching more rows than remained') + assert.deepStrictEqual(outcome.value, [{ num: 4 }, { num: 5 }]) + }) + + it('settles queued reads with leftover or empty rows once the portal is exhausted', async function () { + const cursor = new Cursor('SELECT num FROM generate_series(1, 2) num') + const connection = new TestConnection() + cursor.submit(connection) + + const first = cursor.read(1000) + const queued = cursor.read(1000) + cursor.handleRowDescription({ fields: [INT4] }) + pushIntRows(cursor, [1, 2]) + cursor.handleCommandComplete({ text: 'SELECT 2' }) + + const firstOutcome = await settleOrTimeout(first) + const queuedOutcome = await settleOrTimeout(queued) + assert.ok(!firstOutcome.timeout, 'in-flight read hung after the portal was exhausted') + assert.ok(!queuedOutcome.timeout, 'queued read hung after the portal was exhausted') + assert.deepStrictEqual(firstOutcome.value, [{ num: 1 }, { num: 2 }]) + assert.deepStrictEqual(queuedOutcome.value, []) + }) + + it('keeps an in-flight oversized read when row description arrives late', async function () { + const cursor = new Cursor('SELECT num FROM generate_series(1, 2) num') + const connection = new TestConnection() + cursor.submit(connection) + + const first = cursor.read(1000) + const queued = cursor.read(5) + cursor.handleRowDescription({ fields: [INT4] }) + pushIntRows(cursor, [1, 2]) + cursor.handleCommandComplete({ text: 'SELECT 2' }) + + const firstOutcome = await settleOrTimeout(first) + const queuedOutcome = await settleOrTimeout(queued) + assert.ok(!firstOutcome.timeout, 'in-flight read was overwritten when row description arrived') + assert.ok(!queuedOutcome.timeout, 'queued read hung after a late row description') + assert.deepStrictEqual(firstOutcome.value, [{ num: 1 }, { num: 2 }]) + assert.deepStrictEqual(queuedOutcome.value, []) + assert.strictEqual(connection.calls.execute, 1) + }) +}) + +describe('read() against PostgreSQL when rowCount exceeds remaining rows (#2949)', function () { + beforeEach(function (done) { + const client = (this.client = new pg.Client()) + client.connect(done) + }) + + afterEach(function () { + this.client.end() + }) + + it('returns all 100 rows when asked for 1000', function (done) { + const cursor = this.client.query(new Cursor('SELECT generate_series as num FROM generate_series(1, 100)')) + cursor.read(1000, function (err, rows) { + assert.ifError(err) + assert.strictEqual(rows.length, 100) + assert.strictEqual(rows[0].num, 1) + assert.strictEqual(rows[99].num, 100) + cursor.read(1000, function (err, empty) { + assert.ifError(err) + assert.strictEqual(empty.length, 0) + done() + }) + }) + }) + + it('returns the leftover rows after a partial page', async function () { + const cursor = this.client.query(new Cursor('SELECT generate_series as num FROM generate_series(1, 100)')) + const first = await cursor.read(40) + assert.strictEqual(first.length, 40) + const remaining = await cursor.read(1000) + assert.strictEqual(remaining.length, 60) + assert.strictEqual(remaining[0].num, 41) + assert.strictEqual(remaining[59].num, 100) + assert.deepStrictEqual(await cursor.read(1000), []) + }) +})