Skip to content

Commit

Permalink
fix(pg-v5): return child process stdout (#1701)
Browse files Browse the repository at this point in the history 
https://heroku.support/945120

In [PR #1691][1], some refactoring was done to make the way psql is
executed more consistent. Unfortunately, this ended up in a regression
for commands that were expecting the result of `PSQL.exec` to return a
promise that resolved with the result of the psql child_process's stdout
stream. While the output is displayed, some commands, [like `ps` were
looking at the return result of `psql.exec`][2]. Because #1691 made this
return `undefined`, downstream commands like `ps` or others now error.

This changes the behavior in `packages/pg-v5/lib/psql.js` back to  what
it was doing before:

* Resolve the promise returned from `psql.exec` with the value of the
psql child_process's `stdout` stream
* lib/psql.js no longer pipes to stdout. This is once again the
responsibility of the caller.

Ideally, long term, we would return the child process (or a wrapper) and
let callers decide to do with output themselves. Node has a lot of great
utilities to make this easier. I tried going down this route but it
ended up being a larger change than intended due to tests. In the
interest of fixing this issue quickly, let's go with this approach until
we can revisit the return value of `PSQL.exec` in the future.

[1]: #1691
[2]: https://github.com/heroku/cli/blob/729acd6382b424649c51bff7f4afb02df99d46ac/packages/pg-v5/commands/ps.js#L27
  • Loading branch information
@fivetanley
fivetanley committed Dec 16, 2020
1 parent 729acd6 commit 7f86c49
Show file tree
Hide file tree
Showing 4 changed files with 97 additions and 32 deletions.
14 changes: 8 additions & 6 deletions packages/pg-v5/commands/psql.js
View file
Expand Up @@ -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,
Expand Down
41 changes: 32 additions & 9 deletions packages/pg-v5/lib/psql.js
View file
@@ -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')
Expand All @@ -18,8 +21,7 @@ function psqlQueryOptions (query, dbEnv) {
return {
dbEnv,
psqlArgs,
childProcessOptions,
pipeToStdout: true
childProcessOptions
}
}

Expand All @@ -36,8 +38,7 @@ function psqlFileOptions (file, dbEnv) {
return {
dbEnv,
psqlArgs,
childProcessOptions,
pipeToStdout: true
childProcessOptions
}
}

Expand Down Expand Up @@ -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 = {
Expand All @@ -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
}

Expand Down Expand Up @@ -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 {
Expand All @@ -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
Expand Down
34 changes: 26 additions & 8 deletions packages/pg-v5/test/commands/ps.js
View file
Expand Up @@ -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 = () => {
Expand All @@ -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)
}
}

Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -71,8 +84,13 @@ WHERE
FROM pg_stat_activity
WHERE
query <> '<insufficient privilege>'
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')
}
40 changes: 31 additions & 9 deletions packages/pg-v5/test/lib/psql.js
View file
Expand Up @@ -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
Expand Down Expand Up @@ -67,7 +74,7 @@ describe('psql', () => {

async function ensureFinished (promise) {
try {
await promise
return await promise
} finally {
if (fakeTunnel) {
if (!fakeTunnel.exited) {
Expand Down Expand Up @@ -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 () => {
Expand Down Expand Up @@ -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')
Expand All @@ -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',
Expand Down Expand Up @@ -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',
Expand Down Expand Up @@ -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',
Expand Down Expand Up @@ -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',
Expand Down Expand Up @@ -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', '')
})
})

Expand Down Expand Up @@ -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', '')
})
})

Expand Down Expand Up @@ -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) => {
Expand All @@ -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()
Expand All @@ -563,6 +584,7 @@ class TunnelStub extends EventEmitter {
super(...args)
this.exited = false
}

close () {
this.exited = true
process.nextTick(() => {
Expand Down

0 comments on commit 7f86c49

Please sign in to comment.