Skip to content

Commit

Permalink
migrate(W-14179971): pg-v5: Upgrade pg:bloat
Browse files Browse the repository at this point in the history
  • Loading branch information
justinwilaby committed Apr 12, 2024
1 parent af845d8 commit cfffcb5
Show file tree
Hide file tree
Showing 13 changed files with 1,469 additions and 31 deletions.
5 changes: 4 additions & 1 deletion packages/cli/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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",
Expand All @@ -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"
Expand Down
2 changes: 1 addition & 1 deletion packages/cli/src/commands/addons/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
Original file line number Diff line number Diff line change
@@ -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 (
Expand Down Expand Up @@ -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),
]
132 changes: 132 additions & 0 deletions packages/cli/src/lib/pg/bastion.ts
Original file line number Diff line number Diff line change
@@ -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>) {
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,
dstPort: Number.parseInt(db.port, 10),
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: ReturnType<typeof tunnelConfig>, 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(),
},
)
}
23 changes: 22 additions & 1 deletion packages/cli/src/lib/pg/fetcher.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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')

Expand Down Expand Up @@ -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 attachment(heroku, app, db))).addon
}

export async function database(heroku: APIClient, app: string, db?: string, namespace?: string) {
const attached = await attachment(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
}
Loading

0 comments on commit cfffcb5

Please sign in to comment.