Wait for completion on pgroll migration push (#1434)

Author: andrew-farries

.changeset/light-cycles-repair.md

@@ -1,5 +1,5 @@
 ---
-"@xata.io/client": major
+'@xata.io/client': major
 ---
 
 Make XataApiClient to use ES Proxies

cli/src/commands/push/index.ts

@@ -1,5 +1,6 @@
 import { Args, Flags } from '@oclif/core';
 import { Schemas } from '@xata.io/client';
+import { PgRollMigrationDefinition } from '@xata.io/pgroll';
 import { BaseCommand } from '../../base.js';
 import {
   LocalMigrationFile,
@@ -11,10 +12,10 @@ import {
   allMigrationsPgRollFormat,
   getBranchDetailsWithPgRoll,
   isBranchPgRollEnabled,
-  isMigrationPgRollFormat
+  isMigrationPgRollFormat,
+  waitForMigrationToFinish
 } from '../../migrations/pgroll.js';
 import { MigrationFilePgroll } from '../../migrations/schema.js';
-import { PgRollMigrationDefinition } from '@xata.io/pgroll';
 
 export default class Push extends BaseCommand<typeof Push> {
   static description = 'Push local changes to a remote Xata branch';
@@ -103,10 +104,12 @@ export default class Push extends BaseCommand<typeof Push> {
         .flatMap((migration) => PgRollMigrationDefinition.parse(migration));
       for (const migration of migrationsToPush) {
         try {
-          await xata.api.migrations.applyMigration({
+          const { jobID } = await xata.api.migrations.applyMigration({
             pathParams: { workspace, region, dbBranchName: `${database}:${branch}` },
             body: migration
           });
+
+          await waitForMigrationToFinish(xata.api, workspace, region, database, branch, jobID);
         } catch (e) {
           this.log(`Failed to push ${migration} with ${e}. Stopping.`);
           this.exit(1);

cli/src/commands/push/push.test.ts

@@ -188,6 +188,14 @@ const baseFetch = (url: string, request: any) => {
         }
       })
     };
+  } else if (
+    url === `https://test-1234.us-east-1.xata.sh/db/db1:main/migrations/jobs/1234` &&
+    request.method === 'GET'
+  ) {
+    return {
+      ok: true,
+      json: async () => ({ status: 'completed' })
+    };
   }
 
   throw new Error(`Unexpected fetch request: ${url} ${request.method}`);

cli/src/migrations/pgroll.ts

@@ -1,9 +1,9 @@
-import { Schemas } from '@xata.io/client';
-import { migrationsDir, readMigrationsDir } from './files.js';
+import { Schemas, XataApiClient } from '@xata.io/client';
 import path from 'path';
-import { safeJSONParse, safeReadFile } from '../utils/files.js';
-import { migrationFilePgroll, MigrationFilePgroll } from './schema.js';
 import { XataClient } from '../base.js';
+import { safeJSONParse, safeReadFile } from '../utils/files.js';
+import { migrationsDir, readMigrationsDir } from './files.js';
+import { MigrationFilePgroll, migrationFilePgroll } from './schema.js';
 
 export const isBranchPgRollEnabled = (details: Schemas.DBBranch) => {
   // @ts-expect-error TODO: Fix this when api is finalized
@@ -121,3 +121,26 @@ export async function getBranchDetailsWithPgRoll(
 
   return details;
 }
+
+export async function waitForMigrationToFinish(
+  api: XataApiClient,
+  workspace: string,
+  region: string,
+  database: string,
+  branch: string,
+  jobId: string
+) {
+  const { status, error } = await api.migrations.getMigrationJobStatus({
+    pathParams: { workspace, region, dbBranchName: `${database}:${branch}`, jobId }
+  });
+  if (status === 'failed') {
+    throw new Error(`Migration failed, ${error}`);
+  }
+
+  if (status === 'completed') {
+    return;
+  }
+
+  await new Promise((resolve) => setTimeout(resolve, 1000));
+  return await waitForMigrationToFinish(api, workspace, region, database, branch, jobId);
+}

test/integration/sql.test.ts

@@ -373,7 +373,7 @@ describe('SQL proxy', () => {
     expect(record2).toBeDefined();
     expect(record2?.[4]).toBe('[C] Planes');
   });
-  
+
   test('xata.sql has a connection string', async () => {
     expect(xata.sql.connectionString).toBeDefined();
     expect(xata.sql.connectionString).toMatch(