diff --git a/packages/pg-v5/commands/psql.js b/packages/pg-v5/commands/psql.js index af2024a4fb..4cd29b930d 100644 --- a/packages/pg-v5/commands/psql.js +++ b/packages/pg-v5/commands/psql.js @@ -9,27 +9,29 @@ function * run (context, heroku) { const { app, args, flags } = context - let namespace = flags.credential ? `credential:${flags.credential}` : null + const namespace = flags.credential ? `credential:${flags.credential}` : null let db try { db = yield fetcher.database(app, args.database, namespace) } catch (err) { - if (namespace && err.message === `Couldn't find that addon.`) { - throw new Error(`Credential doesn't match, make sure credential is attached`) + if (namespace && err.message === 'Couldn\'t find that addon.') { + throw new Error('Credential doesn\'t match, make sure credential is attached') } throw err } cli.console.error(`--> Connecting to ${cli.color.addon(db.attachment.addon.name)}`) if (flags.command) { - yield psql.exec(db, flags.command) + const output = yield psql.exec(db, flags.command) + process.stdout.write(output) } else if (flags.file) { - yield psql.execFile(db, flags.file) + const output = yield psql.execFile(db, flags.file) + process.stdout.write(output) } else { yield psql.interactive(db) } } -let cmd = { +const cmd = { description: 'open a psql shell to the database', needsApp: true, needsAuth: true, diff --git a/packages/pg-v5/lib/psql.js b/packages/pg-v5/lib/psql.js index b2606ec662..14110c8f15 100644 --- a/packages/pg-v5/lib/psql.js +++ b/packages/pg-v5/lib/psql.js @@ -1,6 +1,9 @@ 'use strict' const { once, EventEmitter } = require('events') +const Stream = require('stream') +const util = require('util') +const finished = util.promisify(Stream.finished) const bastion = require('./bastion') const debug = require('./debug') @@ -18,8 +21,7 @@ function psqlQueryOptions (query, dbEnv) { return { dbEnv, psqlArgs, - childProcessOptions, - pipeToStdout: true + childProcessOptions } } @@ -36,8 +38,7 @@ function psqlFileOptions (file, dbEnv) { return { dbEnv, psqlArgs, - childProcessOptions, - pipeToStdout: true + childProcessOptions } } @@ -72,7 +73,7 @@ function psqlInteractiveOptions (prompt, dbEnv) { } } -function execPSQL ({ dbEnv, psqlArgs, childProcessOptions, pipeToStdout }) { +function execPSQL ({ dbEnv, psqlArgs, childProcessOptions}) { const { spawn } = require('child_process') const options = { @@ -84,10 +85,6 @@ function execPSQL ({ dbEnv, psqlArgs, childProcessOptions, pipeToStdout }) { const psql = spawn('psql', psqlArgs, options) psql.once('spawn', () => debug('psql process spawned')) - if (pipeToStdout) { - psql.stdout.pipe(process.stdout) - } - return psql } @@ -148,11 +145,36 @@ const trapAndForwardSignalsToChildProcess = (childProcess) => { return cleanup } +function consumeStream(inputStream) { + let result = '' + const throughStream = new Stream.PassThrough() + + const promise = new Promise(async (resolve, reject) => { + try { + await finished(throughStream) + resolve(result) + } catch (err) { + reject(err) + } + }) + + throughStream.on('data', (chunk) => result += chunk.toString()) + inputStream.pipe(throughStream) + return promise; +} + async function runWithTunnel (db, tunnelConfig, options) { const tunnel = await Tunnel.connect(db, tunnelConfig) debug('after create tunnel') const psql = execPSQL(options) + // interactive opens with stdio: 'inherit' + // which gives the child process the same stdin,stdout,stderr of the node process (global `process`) + // https://nodejs.org/api/child_process.html#child_process_options_stdio + // psql.stdout will be null in this case + // return a string for consistency but ideally we should return the child process from this function + // and let the caller decide what to do with stdin/stdout/stderr + const stdoutPromise = psql.stdout ? consumeStream(psql.stdout) : Promise.resolve('') const cleanupSignalTraps = trapAndForwardSignalsToChildProcess(psql) try { @@ -174,6 +196,7 @@ async function runWithTunnel (db, tunnelConfig, options) { kill(psql, 'SIGKILL') debug('end tunnel cleanup') } + return stdoutPromise } // a small wrapper around tunnel-ssh diff --git a/packages/pg-v5/test/commands/ps.js b/packages/pg-v5/test/commands/ps.js index e78f122be8..93d77d4e86 100644 --- a/packages/pg-v5/test/commands/ps.js +++ b/packages/pg-v5/test/commands/ps.js @@ -5,6 +5,7 @@ const cli = require('heroku-cli-util') const expect = require('unexpected') const nock = require('nock') const proxyquire = require('proxyquire') +const { stdout } = require('stdout-stderr') const db = {} const fetcher = () => { @@ -13,10 +14,18 @@ const fetcher = () => { } } +const FAKE_OUTPUT_TEXT = ` +pid | state | source | username | running_for | transaction_start | waiting | query +-------+--------+---------+----------+-------------+-------------------+---------+------- + 17496 | active | standby | postgres | | | t | +(1 row) + +` + const psql = { exec: function (db, query) { this._query = query - return Promise.resolve('') + return Promise.resolve(FAKE_OUTPUT_TEXT) } } @@ -29,18 +38,20 @@ describe('pg:ps', () => { let api beforeEach(() => { + stdout.start() api = nock('https://api.heroku.com:443') cli.mockConsole() }) afterEach(() => { + stdout.stop() nock.cleanAll() api.done() }) - it('runs query', () => { - return cmd.run({ app: 'myapp', args: {}, flags: {} }) - .then(() => expect(psql._query.trim(), 'to equal', `SELECT + it('runs query', async () => { + await cmd.run({ app: 'myapp', args: {}, flags: {} }) + expect(removeEmptyLines(psql._query.trim()), 'to equal', removeEmptyLines(`SELECT pid, state, application_name AS source, @@ -55,11 +66,13 @@ WHERE AND state <> 'idle' AND pid <> pg_backend_pid() ORDER BY query_start DESC`)) + + expect(stdout.output, 'to equal', FAKE_OUTPUT_TEXT) }) - it('runs verbose query', () => { - return cmd.run({ app: 'myapp', args: {}, flags: { verbose: true } }) - .then(() => expect(psql._query.trim(), 'to equal', `SELECT + it('runs verbose query', async () => { + await cmd.run({ app: 'myapp', args: {}, flags: { verbose: true } }) + expect(removeEmptyLines(psql._query.trim()), 'to equal', removeEmptyLines(`SELECT pid, state, application_name AS source, @@ -71,8 +84,13 @@ WHERE FROM pg_stat_activity WHERE query <> '' - AND pid <> pg_backend_pid() ORDER BY query_start DESC`)) + + expect(stdout.output, 'to equal', FAKE_OUTPUT_TEXT) }) }) + +function removeEmptyLines (string) { + return string.split('\n').filter(str => str.trim().length > 0).join('\n') +} diff --git a/packages/pg-v5/test/lib/psql.js b/packages/pg-v5/test/lib/psql.js index c991eb0fc2..bf78191d84 100644 --- a/packages/pg-v5/test/lib/psql.js +++ b/packages/pg-v5/test/lib/psql.js @@ -36,6 +36,13 @@ const bastionDb = { hostname: 'localhost' } +const NOW_OUTPUT = ` +now +------------------------------- + 2020-12-16 09:54:01.916894-08 +(1 row) +` + describe('psql', () => { let fakePsqlProcess, fakeTunnel, tunnelStub let sandbox @@ -67,7 +74,7 @@ describe('psql', () => { async function ensureFinished (promise) { try { - await promise + return await promise } finally { if (fakeTunnel) { if (!fakeTunnel.exited) { @@ -110,8 +117,10 @@ describe('psql', () => { const promise = psql.exec(db, 'SELECT NOW();') await fakePsqlProcess.waitForStart() mock.verify() + fakePsqlProcess.stdout.write(NOW_OUTPUT) await fakePsqlProcess.simulateExit(0) - await ensureFinished(promise) + const output = await ensureFinished(promise) + expect(output, 'to equal', NOW_OUTPUT) }) it('runs psql and throws an error if psql exits with exit code > 0', async () => { @@ -147,7 +156,7 @@ describe('psql', () => { try { expect(fakePsqlProcess.exited, 'to equal', false) await fakePsqlProcess.simulateExit(1) - await ensureFinished(promise); + await ensureFinished(promise) throw new Error('psql.exec should have thrown') } catch (err) { expect(err.message, 'to equal', 'psql exited with code 1') @@ -156,7 +165,7 @@ describe('psql', () => { describe('private databases (not shield)', () => { it('opens an SSH tunnel and runs psql for bastion databases', async () => { - let tunnelConf = { + const tunnelConf = { username: 'bastion', host: 'bastion-host', privateKey: 'super-private-key', @@ -184,7 +193,7 @@ describe('psql', () => { }) it('closes the tunnel manually if psql exits and the tunnel does not close on its own', async () => { - let tunnelConf = { + const tunnelConf = { username: 'bastion', host: 'bastion-host', privateKey: 'super-private-key', @@ -215,7 +224,7 @@ describe('psql', () => { }) it('closes psql manually if the tunnel exits and psql does not close on its own', async () => { - let tunnelConf = { + const tunnelConf = { username: 'bastion', host: 'bastion-host', privateKey: 'super-private-key', @@ -280,7 +289,7 @@ describe('psql', () => { await ensureFinished(promise) }) it('opens an SSH tunnel and runs psql for bastion databases', async () => { - let tunnelConf = { + const tunnelConf = { username: 'bastion', host: 'bastion-host', privateKey: 'super-private-key', @@ -383,7 +392,10 @@ describe('psql', () => { await fakePsqlProcess.waitForStart() await fakePsqlProcess.simulateExit(0) mock.verify() - await ensureFinished(promise) + const output = await ensureFinished(promise) + // psql interactive doesn't pipe output to the process + // ensure promise returned resolves with a promise anyway + expect(output, 'to equal', '') }) }) @@ -432,7 +444,10 @@ describe('psql', () => { await fakePsqlProcess.waitForStart() await fakePsqlProcess.simulateExit(0) mock.verify() - await ensureFinished(promise) + const output = await ensureFinished(promise) + // psql interactive doesn't pipe output to the process + // ensure promise returned resolves with a promise anyway + expect(output, 'to equal', '') }) }) @@ -519,15 +534,18 @@ class FakeChildProcess extends EventEmitter { this.killed = false this.stdout = new PassThrough() } + async waitForStart () { if (!this.ready) { await once(this, 'ready') } } + start () { this.ready = true this.emit('ready') } + simulateExit (code) { if (!this.exited) { return new Promise((resolve) => { @@ -543,15 +561,18 @@ class FakeChildProcess extends EventEmitter { }) } } + kill (signal) { this.killed = true this._killedWithSignal = signal const killedWithCode = signals[signal] this.simulateExit(killedWithCode) } + get killedWithSignal () { return this._killedWithSignal } + async teardown () { await this.simulateExit(0) this.removeAllListeners() @@ -563,6 +584,7 @@ class TunnelStub extends EventEmitter { super(...args) this.exited = false } + close () { this.exited = true process.nextTick(() => {