From 3fb5a5b44de1f436ff5c37580ed2a2fbfaacb278 Mon Sep 17 00:00:00 2001
From: Andrew Farries <andrew.farries@xata.io>
Date: Tue, 9 Apr 2024 14:54:28 +0100
Subject: [PATCH] Wait for completion on `pgroll` migration push (#1434)

---
 .changeset/light-cycles-repair.md  |  2 +-
 cli/src/commands/push/index.ts     |  9 ++++++---
 cli/src/commands/push/push.test.ts |  8 ++++++++
 cli/src/migrations/pgroll.ts       | 31 ++++++++++++++++++++++++++----
 test/integration/sql.test.ts       |  2 +-
 5 files changed, 43 insertions(+), 9 deletions(-)

diff --git a/.changeset/light-cycles-repair.md b/.changeset/light-cycles-repair.md
index 61c123457..8ed66fc50 100644
--- a/.changeset/light-cycles-repair.md
+++ b/.changeset/light-cycles-repair.md
@@ -1,5 +1,5 @@
 ---
-"@xata.io/client": major
+'@xata.io/client': major
 ---
 
 Make XataApiClient to use ES Proxies
diff --git a/cli/src/commands/push/index.ts b/cli/src/commands/push/index.ts
index 596b88655..247c48458 100644
--- a/cli/src/commands/push/index.ts
+++ b/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);
diff --git a/cli/src/commands/push/push.test.ts b/cli/src/commands/push/push.test.ts
index 94fa78b23..0cfca5376 100644
--- a/cli/src/commands/push/push.test.ts
+++ b/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}`);
diff --git a/cli/src/migrations/pgroll.ts b/cli/src/migrations/pgroll.ts
index 4028d220c..37ea8130a 100644
--- a/cli/src/migrations/pgroll.ts
+++ b/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);
+}
diff --git a/test/integration/sql.test.ts b/test/integration/sql.test.ts
index 60f6deb07..e01782e82 100644
--- a/test/integration/sql.test.ts
+++ b/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(