migrate(W-14179971): pg-v5: Upgrade pg:bloat

packages/cli/package.json

@@ -39,6 +39,7 @@
     "@opentelemetry/sdk-trace-node": "^1.15.1",
     "@opentelemetry/semantic-conventions": "^1.15.1",
     "@types/js-yaml": "^3.12.5",
+    "@types/tunnel-ssh": "4.1.1",
     "ansi-escapes": "3.2.0",
     "async-file": "^2.0.2",
     "bytes": "^3.1.2",
@@ -77,9 +78,9 @@
     "strftime": "^0.10.0",
     "strip-ansi": "^6",
     "term-img": "^4.1.0",
-    "tmp": "^0.0.33",
     "true-myth": "2.2.3",
     "tslib": "1.14.1",
+    "tunnel-ssh": "4.1.6",
     "urijs": "^1.19.11",
     "uuid": "3.3.2",
     "valid-url": "^1.0.9",
@@ -113,6 +114,7 @@
     "@types/std-mocks": "^1.0.4",
     "@types/strftime": "^0.9.8",
     "@types/supports-color": "^5.3.0",
+    "@types/tmp": "^0.2.6",
     "@types/urijs": "^1.19.4",
     "@types/uuid": "^8.3.0",
     "@types/validator": "^10.9.0",
@@ -136,6 +138,7 @@
     "sinon": "^7.2.4",
     "std-mocks": "^2.0.0",
     "strip-ansi": "6.0.1",
+    "tmp": "^0.2.3",
     "ts-node": "^10.9.1",
     "tsheredoc": "^1.0.1",
     "typescript": "4.8.4"

packages/cli/src/commands/addons/index.ts

@@ -141,7 +141,7 @@ function formatAttachment(attachment: Heroku.AddOnAttachment, showApp = true) {
   return output.join(' ')
 }
 
-function renderAttachment(attachment: Heroku.AddOnAttachment, app: string, isFirst = false) {
+export function renderAttachment(attachment: Heroku.AddOnAttachment, app: string, isFirst = false): string {
   const line = isFirst ? '\u2514\u2500' : '\u251C\u2500'
   const attName = formatAttachment(attachment, attachment.app?.name !== app)
   return printf(' %s %s', color.dim(line), attName)

packages/pg-v5/commands/bloat.js → packages/cli/src/commands/pg/bloat.ts

@@ -1,14 +1,9 @@
-'use strict'
+import {Command, flags} from '@heroku-cli/command'
+import {Args} from '@oclif/core'
+import {database} from '../../lib/pg/fetcher'
+import {exec} from '../../lib/pg/psql'
 
-const cli = require('heroku-cli-util')
-
-async function run(context, heroku) {
-  const fetcher = require('../lib/fetcher')(heroku)
-  const psql = require('../lib/psql')
-
-  let db = await fetcher.database(context.app, context.args.database)
-
-  let query = `
+const query = `
 WITH constants AS (
   SELECT current_setting('block_size')::numeric AS bs, 23 AS hdr, 4 AS ma
 ), bloat_info AS (
@@ -71,19 +66,24 @@ FROM
 ORDER BY raw_waste DESC, bloat DESC
 `
 
-  let output = await psql.exec(db, query)
-  process.stdout.write(output)
-}
+export default class Bloat extends Command {
+  static topic = 'pg';
+  static description = 'show table and index bloat in your database ordered by most wasteful';
+  static flags = {
+    app: flags.app({required: true}),
+    remote: flags.remote(),
+  };
+
+  static args = {
+    database: Args.string(),
+  };
 
-const cmd = {
-  topic: 'pg',
-  description: 'show table and index bloat in your database ordered by most wasteful',
-  needsApp: true,
-  needsAuth: true,
-  args: [{name: 'database', optional: true}],
-  run: cli.command({preauth: true}, run),
+  public async run(): Promise<void> {
+    const {flags, args} = await this.parse(Bloat)
+    const {app} = flags
+    const db = await database(this.heroku, app, args.database)
+    const output = await exec(db, query)
+    process.stdout.write(output)
+  }
 }
 
-module.exports = [
-  Object.assign({command: 'bloat'}, cmd),
-]

packages/cli/src/commands/pg/reset.ts

@@ -1,9 +1,9 @@
 import color from '@heroku-cli/color'
 import {Command, flags} from '@heroku-cli/command'
 import {Args, ux} from '@oclif/core'
+import confirmCommand from '../../lib/confirmCommand'
 import pgHost from '../../lib/pg/host'
 import {getAddon} from '../../lib/pg/fetcher'
-import confirmApp from '../../lib/apps/confirm-app'
 import heredoc from 'tsheredoc'
 
 export default class Reset extends Command {
@@ -32,7 +32,7 @@ export default class Reset extends Command {
         .sort()
     }
 
-    await confirmApp(app, confirm, heredoc(`
+    await confirmCommand(app, confirm, heredoc(`
       Destructive action
       ${color.addon(db.name)} will lose all of its data
     `))

packages/cli/src/commands/spaces/destroy.ts

@@ -3,7 +3,7 @@ import color from '@heroku-cli/color'
 import {Command, flags} from '@heroku-cli/command'
 import * as Heroku from '@heroku-cli/schema'
 import heredoc from 'tsheredoc'
-import confirmApp from '../../lib/apps/confirm-app'
+import confirmCommand from '../../lib/confirmCommand'
 import {displayNat} from '../../lib/spaces/spaces'
 
 type RequiredSpaceWithNat = Required<Heroku.Space> & {outbound_ips?: Required<Heroku.SpaceNetworkAddressTranslation>}
@@ -47,7 +47,7 @@ export default class Destroy extends Command {
       }
     }
 
-    await confirmApp(
+    await confirmCommand(
       spaceName as string,
       confirm,
       `Destructive Action\nThis command will destroy the space ${color.bold.red(spaceName as string)}\n${natWarning}\n`,

packages/cli/src/lib/pg/bastion.ts

@@ -0,0 +1,132 @@
+const debug = require('debug')('pg')
+import {APIClient} from '@heroku-cli/command'
+import * as EventEmitter from 'node:events'
+import * as createTunnel from 'tunnel-ssh'
+import {promisify} from 'util'
+import host from './host'
+import {getConnectionDetails} from './util'
+
+export const getBastion = function (config:Record<string, string>, baseName: string) {
+  // If there are bastions, extract a host and a key
+  // otherwise, return an empty Object
+
+  // If there are bastions:
+  // * there should be one *_BASTION_KEY
+  // * pick one host from the comma-separated list in *_BASTIONS
+  // We assert that _BASTIONS and _BASTION_KEY always exist together
+  // If either is falsy, pretend neither exist
+
+  const bastionKey = config[`${baseName}_BASTION_KEY`]
+  const bastions = (config[`${baseName}_BASTIONS`] || '').split(',')
+  const bastionHost = bastions[Math.floor(Math.random() * bastions.length)]
+  return (bastionKey && bastionHost) ? {bastionHost, bastionKey} : {}
+}
+
+export const env = (db: ReturnType<typeof getConnectionDetails>) => {
+  const baseEnv = Object.assign({
+    PGAPPNAME: 'psql non-interactive',
+    PGSSLMODE: (!db.host || db.host === 'localhost') ? 'prefer' : 'require',
+  }, process.env)
+  const mapping:Record<string, keyof Omit<typeof db, 'attachment'>> = {
+    PGUSER: 'user',
+    PGPASSWORD: 'password',
+    PGDATABASE: 'database',
+    PGPORT: 'port',
+    PGHOST: 'host',
+  }
+  Object.keys(mapping).forEach(envVar => {
+    const val = db[mapping[envVar]]
+    if (val) {
+      baseEnv[envVar] = val as string
+    }
+  })
+  return baseEnv
+}
+
+export function tunnelConfig(db: ReturnType<typeof getConnectionDetails>): createTunnel.Config {
+  const localHost = '127.0.0.1'
+  // eslint-disable-next-line no-mixed-operators
+  const localPort = Math.floor(Math.random() * (65535 - 49152) + 49152)
+  return {
+    username: 'bastion',
+    host: db.bastionHost,
+    privateKey: db.bastionKey,
+    dstHost: db.host || undefined,
+    dstPort: (db.port && Number.parseInt(db.port as string, 10)) || undefined,
+    localHost: localHost,
+    localPort: localPort,
+  }
+}
+
+export function getConfigs(db: ReturnType<typeof getConnectionDetails>) {
+  const dbEnv: NodeJS.ProcessEnv = env(db)
+  const dbTunnelConfig = tunnelConfig(db)
+  if (db.bastionKey) {
+    Object.assign(dbEnv, {
+      PGPORT: dbTunnelConfig.localPort,
+      PGHOST: dbTunnelConfig.localHost,
+    })
+  }
+
+  return {
+    dbEnv,
+    dbTunnelConfig,
+  }
+}
+
+class Timeout {
+  private readonly timeout: number
+  private readonly message: string
+  private readonly events = new EventEmitter()
+  private timer: NodeJS.Timeout | undefined
+
+  constructor(timeout: number, message: string) {
+    this.timeout = timeout
+    this.message = message
+  }
+
+  async promise() {
+    this.timer = setTimeout(() => {
+      this.events.emit('error', new Error(this.message))
+    }, this.timeout)
+
+    try {
+      await EventEmitter.once(this.events, 'cancelled')
+    } finally {
+      clearTimeout(this.timer)
+    }
+  }
+
+  cancel() {
+    this.events.emit('cancelled')
+  }
+}
+
+export async function sshTunnel(db: ReturnType<typeof getConnectionDetails>, dbTunnelConfig: createTunnel.Config, timeout = 10000) {
+  if (!db.bastionKey) {
+    return null
+  }
+
+  const timeoutInstance = new Timeout(timeout, 'Establishing a secure tunnel timed out')
+  const createSSHTunnel = promisify(createTunnel)
+  try {
+    return await Promise.race([
+      timeoutInstance.promise(),
+      createSSHTunnel(dbTunnelConfig),
+    ])
+  } catch (error) {
+    debug(error)
+    throw new Error('Unable to establish a secure tunnel to your database.')
+  } finally {
+    timeoutInstance.cancel()
+  }
+}
+
+export async function fetchConfig(heroku:APIClient, db: {id: string}) {
+  return heroku.get<{host: string, private_key:string}>(
+    `/client/v11/databases/${encodeURIComponent(db.id)}/bastion`,
+    {
+      hostname: host(),
+    },
+  )
+}

packages/cli/src/lib/pg/fetcher.ts

@@ -3,10 +3,11 @@ import type {AddOnAttachment} from '@heroku-cli/schema'
 import * as Heroku from '@heroku-cli/schema'
 import debug from 'debug'
 import {AmbiguousError, appAttachment, NotFound} from '../addons/resolve'
+import {fetchConfig} from './bastion'
 import {getConfig} from './config'
 import color from '@heroku-cli/color'
 import type {AddOnAttachmentWithConfigVarsAndPlan} from './types'
-import {getConfigVarName} from './util'
+import {bastionKeyPlan, getConfigVarName, getConnectionDetails} from './util'
 
 const pgDebug = debug('pg')
 
@@ -132,3 +133,23 @@ async function allAttachments(heroku: APIClient, app_id: string) {
 export async function getAddon(heroku: APIClient, app: string, db = 'DATABASE_URL') {
   return ((await getAttachment(heroku, app, db))).addon
 }
+
+export async function database(heroku: APIClient, app: string, db?: string, namespace?: string) {
+  const attached = await getAttachment(heroku, app, db, namespace)
+
+  // would inline this as well but in some cases attachment pulls down config
+  // as well, and we would request twice at the same time but I did not want
+  // to push this down into attachment because we do not always need config
+  const config = await getConfig(heroku, attached.app.name as string)
+
+  const database = getConnectionDetails(attached, config)
+  if (bastionKeyPlan(attached.addon) && !database.bastionKey) {
+    const {body: bastionConfig} = await fetchConfig(heroku, attached.addon)
+    const bastionHost = bastionConfig.host
+    const bastionKey = bastionConfig.private_key
+
+    Object.assign(database, {bastionHost, bastionKey})
+  }
+
+  return database
+}