fix(pg-v5): return child process stdout

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,

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

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 <> '<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')
+}

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(() => {