feat: adding pg:upgrade:wait command

packages/cli/src/commands/pg/upgrade/wait.ts

@@ -0,0 +1,118 @@
+import color from '@heroku-cli/color'
+import {Command, flags} from '@heroku-cli/command'
+import {Args, ux} from '@oclif/core'
+import debug from 'debug'
+import heredoc from 'tsheredoc'
+import {getAddon} from '../../../lib/pg/fetcher'
+import pgHost from '../../../lib/pg/host'
+import notify from '../../../lib/notify'
+import {AddOnAttachmentWithConfigVarsAndPlan, AddOnWithRelatedData, PgUpgradeStatus} from '../../../lib/pg/types'
+import {HTTPError} from '@heroku/http-call'
+import {nls} from '../../../nls'
+import {formatResponseWithCommands} from '../../../lib/pg/util'
+
+const wait = (ms: number) => new Promise(resolve => {
+  setTimeout(resolve, ms)
+})
+
+export default class Wait extends Command {
+  static topic = 'pg'
+  static description = 'provides the status of an upgrade and blocks it until the operation is complete'
+  static flags = {
+    'wait-interval': flags.integer({description: 'how frequently to poll in seconds (to avoid rate limiting)'}),
+    'no-notify': flags.boolean({description: 'do not show OS notification'}),
+    app: flags.app({required: true}),
+    remote: flags.remote(),
+  }
+
+  static examples = [
+    heredoc(`
+      # Wait for upgrade to complete with default settings
+      $ heroku pg:upgrade:wait postgresql-curved-12345 --app myapp
+    `),
+    heredoc(`
+      # Wait with custom polling interval
+      $ heroku pg:upgrade:wait postgresql-curved-12345 --app myapp --wait-interval 10
+    `),
+    heredoc(`
+      # Wait without showing OS notifications
+      $ heroku pg:upgrade:wait postgresql-curved-12345 --app myapp --no-notify
+    `),
+  ]
+
+  static args = {
+    database: Args.string({description: `${nls('pg:database:arg:description')}`}),
+  }
+
+  public async run(): Promise<void> {
+    const {flags, args} = await this.parse(Wait)
+    const {app, 'wait-interval': waitInterval} = flags
+    const dbName = args.database
+    const pgDebug = debug('pg')
+
+    const waitFor = async (db: AddOnAttachmentWithConfigVarsAndPlan | AddOnWithRelatedData) => {
+      const interval = (!waitInterval || waitInterval < 0) ? 5 : waitInterval
+      let status
+      let waiting = false
+      let retries = 20
+      const notFoundMessage = 'Waiting to provision...'
+
+      while (true) {
+        try {
+          ({body: status} = await this.heroku.get<PgUpgradeStatus>(
+            `/client/v11/databases/${db.id}/upgrade/wait_status`,
+            {hostname: pgHost()},
+          ))
+        } catch (error) {
+          if (error instanceof HTTPError && (!retries || error.statusCode !== 404)) {
+            const httpError = error as HTTPError
+            pgDebug(httpError)
+            throw httpError
+          }
+
+          retries--
+          status = {'waiting?': true, message: notFoundMessage}
+        }
+
+        let message = formatResponseWithCommands(status.message)
+        if (status.step)
+          message = heredoc(`(${status.step}) ${message}`)
+
+        if (status['error?']) {
+          notify('error', `${db.name} ${message}`, false)
+          ux.error(message || '', {exit: 1})
+        }
+
+        if (!status['waiting?']) {
+          if (waiting) {
+            ux.action.stop(message)
+          } else {
+            ux.log(heredoc(`Waiting for database ${color.yellow(db.name)}... ${message}`))
+          }
+
+          return
+        }
+
+        if (!waiting) {
+          waiting = true
+          ux.action.start(`Waiting for database ${color.yellow(db.name)}`, message)
+        }
+
+        ux.action.status = message
+
+        await wait(interval * 1000)
+      }
+    }
+
+    let dbs: AddOnAttachmentWithConfigVarsAndPlan[] | AddOnWithRelatedData[] | [] = []
+    if (dbName) {
+      dbs = [await getAddon(this.heroku, app, dbName)]
+    } else {
+      ux.error(heredoc('You must provide a database. Run `--help` for more information on the command.'))
+    }
+
+    for (const db of dbs) {
+      await waitFor(db)
+    }
+  }
+}

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

@@ -131,6 +131,13 @@ export type PgStatus = {
   message: string
 }
 
+export type PgUpgradeStatus = {
+  'waiting?': boolean
+  'error?': boolean
+  message: string
+  step: string
+}
+
 type TenantInfoNames =
   'Plan'
   | 'Status'

packages/cli/test/unit/commands/pg/upgrade/wait.unit.test.ts

@@ -0,0 +1,93 @@
+import {stdout, stderr} from 'stdout-stderr'
+import {expect} from 'chai'
+import * as nock from 'nock'
+import * as proxyquire from 'proxyquire'
+import heredoc from 'tsheredoc'
+import {CLIError} from '@oclif/core/lib/errors'
+import runCommand from '../../../../helpers/runCommand'
+import expectOutput from '../../../../helpers/utils/expectOutput'
+
+const all = [
+  {id: 1, name: 'postgres-1', plan: {name: 'heroku-postgresql:hobby-dev'}},
+  {id: 2, name: 'postgres-2', plan: {name: 'heroku-postgresql:hobby-dev'}},
+]
+const fetcher =  {
+  all: () => Promise.resolve(all),
+  getAddon: () => Promise.resolve(all[0]),
+}
+
+const {default: Cmd} = proxyquire('../../../../../src/commands/pg/upgrade/wait', {
+  '../../../lib/pg/fetcher': fetcher,
+})
+
+describe('pg:upgrade:wait', function () {
+  let pg: nock.Scope
+
+  beforeEach(function () {
+    pg = nock('https://api.data.heroku.com')
+  })
+
+  afterEach(function () {
+    pg.done()
+    nock.cleanAll()
+  })
+
+  it('waits till upgrade is finished', async function () {
+    pg
+      .get('/client/v11/databases/1/upgrade/wait_status').reply(200, {'waiting?': true, message: 'preparing upgrade service'})
+      .get('/client/v11/databases/1/upgrade/wait_status').reply(200, {'waiting?': false, message: 'recreating followers', step: '7/7'})
+
+    await runCommand(Cmd, [
+      '--app',
+      'myapp',
+      '--wait-interval',
+      '1',
+      'DATABASE_URL',
+    ])
+    expect(stdout.output).to.equal('')
+    expectOutput(stderr.output, heredoc(`
+      Waiting for database postgres-1... preparing upgrade service
+      Waiting for database postgres-1... (7/7) recreating followers
+    `))
+  })
+
+  it('displays when the upgrade has been scheduled', async function () {
+    pg
+      .get('/client/v11/databases/1/upgrade/wait_status').reply(200, {'waiting?': false, message: 'upgrade is scheduled on 2025-04-17 20:30:00 UTC. You could also run the upgrade immediately using `heroku pg:upgrade:run`.'})
+
+    await runCommand(Cmd, [
+      '--app',
+      'myapp',
+      '--wait-interval',
+      '1',
+      'DATABASE_URL',
+    ])
+    expect(stdout.output).to.equal('Waiting for database postgres-1... upgrade is scheduled on 2025-04-17 20:30:00 UTC. You could also run the upgrade immediately using heroku pg:upgrade:run.\n')
+  })
+
+  it('requires a database', async function () {
+    await runCommand(Cmd, [
+      '--app',
+      'myapp',
+    ]).catch(error => {
+      expectOutput(error.message, heredoc(`
+      You must provide a database. Run \`--help\` for more information on the command.
+    `))
+    })
+  })
+
+  it('displays errors', async function () {
+    pg
+      .get('/client/v11/databases/1/upgrade/wait_status').reply(200, {'error?': true, message: 'this is an error message'})
+
+    await runCommand(Cmd, [
+      '--app',
+      'myapp',
+      'DATABASE_URL',
+    ]).catch(error => {
+      const {message, oclif} = error as CLIError
+      expect(message).to.equal('this is an error message')
+      expect(oclif.exit).to.equal(1)
+    })
+  })
+})

packages/cli/test/unit/commands/pg/wait.unit.test.ts

@@ -77,4 +77,23 @@ describe('pg:wait', function () {
       expect(oclif.exit).to.equal(1)
     })
   })
+
+  it('receives steps but does not display them', async function () {
+    pg
+      .get('/client/v11/databases/1/wait_status').reply(200, {'waiting?': true, message: 'upgrading', step: '1/3'})
+      .get('/client/v11/databases/1/wait_status').reply(200, {'waiting?': false, message: 'available'})
+
+    await runCommand(Cmd, [
+      '--app',
+      'myapp',
+      '--wait-interval',
+      '1',
+      'DATABASE_URL',
+    ])
+    expect(stdout.output).to.equal('')
+    expectOutput(stderr.output, heredoc(`
+      Waiting for database postgres-1... upgrading
+      Waiting for database postgres-1... available
+    `))
+  })
 })