launch scratch: add retry for connection string (#2659)

Synced from monorepo@5d3685eae93d62ef913d8c0d52f003adf66b36d0

Author: divyenduz

.sync-commit

@@ -1 +1 @@
-5162fbd21b5849c51c0f806b7fe2276aefc6c014
+5d3685eae93d62ef913d8c0d52f003adf66b36d0

package.json

@@ -1,7 +1,7 @@
 {
   "name": "xata-cli",
   "description": "Xata CLI",
-  "version": "1.4.4",
+  "version": "1.4.5",
   "private": true,
   "module": "src/app.ts",
   "type": "module",
@@ -37,15 +37,15 @@
     "zod": "4.4.3",
     "zod-config": "1.4.0",
     "@xata.io/ai": "0.1.0",
+    "@xata.io/api": "0.1.1",
+    "@xata.io/config": "0.0.1",
     "@xata.io/lang": "0.0.1",
     "@xata.io/pgroll": "0.9.0",
     "@xata.io/pgstream": "0.2.0",
     "@xata.io/sql": "0.1.4",
     "@xata.io/test-utils": "0.0.1",
     "@xata.io/tsconfig": "0.0.1",
-    "@xata.io/config": "0.0.1",
-    "@xata.io/utils": "0.1.0",
-    "@xata.io/api": "0.1.1"
+    "@xata.io/utils": "0.1.0"
   },
   "scripts": {
     "dev": "bun src/bin/cli.ts",

src/app.ts

@@ -51,8 +51,7 @@ const routes = buildRouteMap({
     hideRoute: {
       onboard: true,
       ai: true,
-      stream: true,
-      scratch: true
+      stream: true
     }
   }
 });

src/commands/scratch.ts

@@ -18,6 +18,8 @@ type Flags = {
 };
 
 const SIGNAL_CLEANUP_TIMEOUT_MS = 10 * 1000;
+const CONNECTION_STRING_RETRY_COUNT = 2;
+const CONNECTION_STRING_RETRY_DELAY_MS = 100;
 const SCRATCH_SCALE_TO_ZERO = {
   enabled: true,
   inactivityPeriodMinutes: 10
@@ -57,6 +59,10 @@ async function withTimeout<T>(promise: Promise<T>, ms: number) {
   }
 }
 
+async function sleep(ms: number) {
+  await new Promise((resolve) => setTimeout(resolve, ms));
+}
+
 function getErrorMessage(error: unknown) {
   if (error instanceof Error) return error.message;
   if (typeof error === 'object' && error !== null && 'message' in error) return String(error.message);
@@ -202,6 +208,28 @@ async function deleteScratchBranchByName(
   }
 }
 
+async function resolveScratchBranchConnectionString(
+  context: LocalContext,
+  organizationId: string,
+  projectId: string,
+  branch: ScratchBranch
+) {
+  if (branch.connectionString) return branch;
+
+  let resolvedBranch = branch;
+  for (let retry = 0; retry < CONNECTION_STRING_RETRY_COUNT; retry++) {
+    await sleep(CONNECTION_STRING_RETRY_DELAY_MS);
+
+    resolvedBranch = await context.api.branches.describeBranch({
+      pathParams: { organizationID: organizationId, projectID: projectId, branchID: branch.id }
+    });
+
+    if (resolvedBranch.connectionString) return resolvedBranch;
+  }
+
+  return resolvedBranch;
+}
+
 export async function implementation(this: LocalContext, flags: Flags, ...command: string[]) {
   const hasQuery = Boolean(flags.execute);
   const hasBinary = command.length > 0;
@@ -337,9 +365,11 @@ export async function implementation(this: LocalContext, flags: Flags, ...comman
 
     this.process.stderr.write(chalk.green(`Created scratch branch ${createdScratchBranch.name}\n`));
 
-    invariant(createdScratchBranch.connectionString, 'Scratch branch should have a connection string.');
+    scratchBranch = await resolveScratchBranchConnectionString(this, organizationId, projectId, createdScratchBranch);
+
+    invariant(scratchBranch.connectionString, 'Scratch branch should have a connection string.');
 
-    const connectionString = buildConnectionString(createdScratchBranch.connectionString, {
+    const connectionString = buildConnectionString(scratchBranch.connectionString, {
       database,
       endpointType: 'rw'
     });

src/commands/scratch.unit.test.ts

@@ -27,14 +27,14 @@ function buildContext() {
     name: body.name,
     createdAt: '2026-06-09T00:00:00.000Z',
     parentID: 'source-branch-id',
-    connectionString: 'postgresql://user:pass@localhost:5432/postgres'
+    connectionString: 'postgresql://user:pass@localhost:5432/postgres' as string | null
   }));
   const deleteBranch = mock(async () => ({}));
   const listBranches = mock(async (): Promise<{ branches: { id: string; name: string }[] }> => ({ branches: [] }));
   const describeBranch = mock(async () => ({
     id: 'scratch-branch-id',
     name: 'scratch-branch',
-    connectionString: 'postgresql://user:pass@localhost:5432/postgres',
+    connectionString: 'postgresql://user:pass@localhost:5432/postgres' as string | null,
     status: { statusType: 'STATUS_TYPE_HEALTHY', status: 'healthy' }
   }));
 
@@ -95,6 +95,31 @@ describe('scratch command', () => {
     expect(stderr.join('')).toContain('Deleted scratch branch scratch-');
   });
 
+  test('retries describing the scratch branch when create does not return a connection string', async () => {
+    const { context, createBranch, describeBranch, unsafe } = buildContext();
+    createBranch.mockImplementationOnce(async ({ body }: { body: { name: string } }) => ({
+      id: 'scratch-branch-id',
+      name: body.name,
+      createdAt: '2026-06-09T00:00:00.000Z',
+      parentID: 'source-branch-id',
+      connectionString: null
+    }));
+    describeBranch.mockImplementationOnce(async () => ({
+      id: 'scratch-branch-id',
+      name: 'scratch-branch',
+      connectionString: null,
+      status: { statusType: 'STATUS_TYPE_HEALTHY', status: 'healthy' }
+    }));
+
+    await implementation.call(context, { execute: 'select 1', json: false });
+
+    expect(describeBranch).toHaveBeenCalledTimes(2);
+    expect(describeBranch).toHaveBeenCalledWith({
+      pathParams: { organizationID: 'org-id', projectID: 'project-id', branchID: 'scratch-branch-id' }
+    });
+    expect(unsafe).toHaveBeenCalledWith('select 1');
+  });
+
   test('prints SQL results as JSON when --json is passed', async () => {
     const { context, stdout } = buildContext();