fix(db): push the schema before the Project migration reads it

The bootstrap ran 0032_project_membership before Drizzle pushed anything, so a
development database from before a column the migration reads (such as the
fork parent) failed after committing the expansion, on every retry.

--prepare now only replays the expansion. push.ts then runs one Drizzle push
with workspace.project_id still nullable (SIM_DB_PUSH_PROJECT_BOOTSTRAP=1),
runs the migration through reconcile-project-membership.ts --bootstrap, and
finishes with the ordinary push. The connector and rollout marker are kept out
of Drizzle's diff until the migration retires them. A fresh bootstrap also
requires that no public table carries a project_id column.

Author: waleedlatif1

packages/db/drizzle.config.ts

@@ -15,9 +15,14 @@ export default {
   dbCredentials: {
     url: process.env.DATABASE_URL!,
   },
-  /** Runner-owned journals and resumable cleanup cursors must survive a development schema push. */
+  /**
+   * Runner-owned journals, resumable cleanup cursors, and a Project bootstrap's connector and
+   * rollout marker (until its migration retires them) must survive a development schema push.
+   */
   tablesFilter: [
     '!script_migrations',
+    '!project_workspace',
+    '!project_membership_rollout',
     '!project_backfill_archive_repairs',
     '!search_embedding_cleanup_progress',
     '!search_embedding_cleanup_targets',

packages/db/schema.ts

@@ -2171,6 +2171,20 @@ export const workspaceModeEnum = pgEnum('workspace_mode', [
 
 export type WorkspaceMode = (typeof workspaceModeEnum.enumValues)[number]
 
+/**
+ * A development push that bootstraps Projects first pushes this column nullable, so the columns
+ * `0032_project_membership` reads exist before it fills the column and enforces NOT NULL; the
+ * push that follows uses the required column. Every other push and generation reads it required.
+ */
+function workspaceProjectIdColumn() {
+  const column = text('project_id').references((): AnyPgColumn => project.id, {
+    onDelete: 'restrict',
+  })
+  return process.env.SIM_DB_PUSH_PROJECT_BOOTSTRAP === '1'
+    ? (column as ReturnType<typeof column.notNull>)
+    : column.notNull()
+}
+
 export const workspace = pgTable(
   'workspace',
   {
@@ -2198,9 +2212,7 @@ export const workspace = pgTable(
     organizationScopeKey: text('organization_scope_key').generatedAlwaysAs(
       sql`CASE WHEN organization_id IS NULL THEN 'personal' ELSE 'organization:' || organization_id END`
     ),
-    projectId: text('project_id')
-      .notNull()
-      .references((): AnyPgColumn => project.id, { onDelete: 'restrict' }),
+    projectId: workspaceProjectIdColumn(),
     workspaceMode: workspaceModeEnum('workspace_mode').notNull().default('grandfathered_shared'),
     billedAccountUserId: text('billed_account_user_id')
       .notNull()

packages/db/scripts/project-push-bootstrap.integration.ts

@@ -185,6 +185,30 @@ describe('development schema push of a database that predates Projects', () => {
     })
   }, 240_000)
 
+  it('pushes the columns the Project migration reads before running it', async () => {
+    await preProjectDatabase(async (sql, url) => {
+      await sql`ALTER TABLE workspace DROP COLUMN forked_from_workspace_id CASCADE`
+      const pushed = runPush(url, true)
+      expect(pushed.status, pushed.output).toBe(0)
+      expect(await projectSchema(sql)).toEqual({
+        project: true,
+        connector: false,
+        rollout: false,
+        required: true,
+      })
+      expect(
+        await sql`SELECT count(*)::int AS unassigned FROM workspace WHERE project_id IS NULL`
+      ).toEqual([{ unassigned: 0 }])
+      expect(
+        await sql`SELECT count(*)::int AS restored FROM pg_attribute WHERE attrelid = 'workspace'::regclass
+          AND attname = 'forked_from_workspace_id' AND NOT attisdropped`
+      ).toEqual([{ restored: 1 }])
+      expect(
+        await sql`SELECT name FROM script_migrations WHERE name = '0032_project_membership'`
+      ).toEqual([{ name: '0032_project_membership' }])
+    })
+  }, 240_000)
+
   it('resumes after the Project migration reports conflicts', async () => {
     await preProjectDatabase(async (sql, url) => {
       await sql`INSERT INTO workspace (id, name, owner_id, billed_account_user_id, forked_from_workspace_id)
@@ -298,6 +322,13 @@ describe('development schema push of a database that predates Projects', () => {
         await sql`ALTER TABLE workspace ADD COLUMN project_id text`
       },
     },
+    {
+      state: 'folder Project ownership without Projects',
+      async build(sql: Sql) {
+        await sql`ALTER TABLE folder ADD COLUMN project_id text`
+        await sql`UPDATE folder SET project_id = 'orphaned'`
+      },
+    },
     {
       state: 'a versioned expansion awaiting its script migration',
       async build(sql: Sql) {

packages/db/scripts/push.ts

@@ -17,8 +17,13 @@ const RECONCILIATION_COMMANDS = [
   ['bun', '--env-file=.env', 'run', './script-migrations/0031_drop_keyword_tin.ts'],
 ]
 
+/** Reads, on the fence session, whether `--prepare` left a Project bootstrap for this push to finish. */
+type ProjectBootstrapProbe = () => Promise<boolean>
+
 /** Fences preparation, Drizzle and reconcilers from a concurrent migration run. */
-async function withMigrationFence(run: (signal: AbortSignal) => Promise<number>): Promise<number> {
+async function withMigrationFence(
+  run: (signal: AbortSignal, projectBootstrapPending: ProjectBootstrapProbe) => Promise<number>
+): Promise<number> {
   const rawUrl = process.env.DATABASE_URL
   if (!rawUrl) {
     logger.error('DATABASE_URL is required for schema push')
@@ -48,7 +53,11 @@ async function withMigrationFence(run: (signal: AbortSignal) => Promise<number>)
         logger.error('A migration is running; retry schema push later')
         return 1
       }
-      return run(abort.signal)
+      return run(abort.signal, async () => {
+        const [state] = await tx<{ pending: boolean }[]>`
+          SELECT to_regclass('public.project_membership_rollout') IS NOT NULL AS pending`.simple()
+        return state.pending
+      })
     })) as number
   } catch (error) {
     logger.error(
@@ -82,11 +91,21 @@ export async function runPush(args: string[]): Promise<number> {
     return 1
   }
 
-  async function runCommands(signal?: AbortSignal): Promise<number> {
-    async function runCommand(command: string[]): Promise<number> {
+  async function runCommands(
+    signal?: AbortSignal,
+    projectBootstrapPending?: ProjectBootstrapProbe
+  ): Promise<number> {
+    async function runCommand(
+      command: string[],
+      env: Record<string, string> = {}
+    ): Promise<number> {
       signal?.throwIfAborted()
       const child = Bun.spawn(command, {
-        env: { ...process.env, SIM_DB_PUSH_RENAME_MODE: interactiveRenames ? 'prompt' : 'create' },
+        env: {
+          ...process.env,
+          SIM_DB_PUSH_RENAME_MODE: interactiveRenames ? 'prompt' : 'create',
+          ...env,
+        },
         stdin: 'inherit',
         stdout: 'inherit',
         stderr: 'inherit',
@@ -119,13 +138,21 @@ export async function runPush(args: string[]): Promise<number> {
     const pushArgs = args.filter((arg) => arg !== '--interactive-renames')
     // Launch the installed CLI directly so losing the fence also terminates its database connection.
     const cli = fileURLToPath(new URL('bin.cjs', import.meta.resolve('drizzle-kit')))
-    const exitCode = await runCommand([
-      'bun',
-      cli,
-      'push',
-      '--config=./drizzle.config.ts',
-      ...pushArgs,
-    ])
+    const drizzlePush = ['bun', cli, 'push', '--config=./drizzle.config.ts', ...pushArgs]
+    if (!help && process.env.SIM_DEV_DB_PUSH === '1' && (await projectBootstrapPending?.())) {
+      /** The Project migration reads columns an older database gains only from this push. */
+      const code = await runCommand(drizzlePush, { SIM_DB_PUSH_PROJECT_BOOTSTRAP: '1' })
+      if (code !== 0) return code
+      const bootstrapCode = await runCommand([
+        'bun',
+        '--env-file=.env',
+        'run',
+        './scripts/reconcile-project-membership.ts',
+        '--bootstrap',
+      ])
+      if (bootstrapCode !== 0) return bootstrapCode
+    }
+    const exitCode = await runCommand(drizzlePush)
     if (exitCode !== 0 || help) return exitCode
 
     for (const command of RECONCILIATION_COMMANDS) {

packages/db/scripts/reconcile-project-membership.ts

@@ -12,6 +12,9 @@ const url = process.env.MIGRATION_DATABASE_URL || process.env.DATABASE_URL
 if (!url) throw new Error('A database URL is required for Project schema push reconciliation')
 const sql = postgres(url, { max: 1, prepare: false, connect_timeout: 10, max_lifetime: null })
 
+const REFUSAL =
+  'Existing Project databases must complete the Project script migration 0032_project_membership before schema push. Use the phased migration path for retained data; schema push cannot bypass the writer-drain requirement'
+
 interface ProjectSchemaState {
   workspace: boolean
   project: boolean
@@ -20,6 +23,8 @@ interface ProjectSchemaState {
   journal: boolean
   /** Null when `workspace.project_id` does not exist; otherwise whether it is NOT NULL. */
   projectIdNotNull: boolean | null
+  /** Whether any public table carries a `project_id` column. */
+  projectColumns: boolean
 }
 
 async function readProjectSchemaState(): Promise<ProjectSchemaState> {
@@ -30,25 +35,29 @@ async function readProjectSchemaState(): Promise<ProjectSchemaState> {
       to_regclass('public.project_membership_rollout') IS NOT NULL AS rollout,
       to_regclass('drizzle.__drizzle_migrations') IS NOT NULL AS journal,
       (SELECT attnotnull FROM pg_attribute WHERE attrelid = to_regclass('public.workspace')
-        AND attname = 'project_id' AND NOT attisdropped) AS "projectIdNotNull"
+        AND attname = 'project_id' AND NOT attisdropped) AS "projectIdNotNull",
+      EXISTS (SELECT 1 FROM pg_attribute a JOIN pg_class c ON c.oid = a.attrelid
+        WHERE c.relnamespace = 'public'::regnamespace AND c.relkind IN ('r', 'p')
+          AND a.attname = 'project_id' AND NOT a.attisdropped) AS "projectColumns"
   `
   return state
 }
 
 /**
  * Classifies a development push database that has no connector writer to drain: either it
- * predates Projects entirely, or an earlier bootstrap committed the expansion and stopped before
- * enforcement. Only the versioned 0406 creates the rollout marker, so a marker that the migration
- * journal does not account for, beside a connector nothing has written, is a bootstrap to resume.
- * Every other shape returns null and keeps the writer-drain refusal.
+ * predates Projects entirely (no Project table, connector, marker, or `project_id` column
+ * anywhere), or an earlier bootstrap committed the expansion and has not finished. Only the
+ * versioned 0406 creates the rollout marker, so a marker that the migration journal does not
+ * account for, beside a connector nothing has written, is a bootstrap in progress. Every other
+ * shape returns null and keeps the writer-drain refusal.
  */
 async function classifyProjectBootstrap(
   state: ProjectSchemaState
 ): Promise<'fresh' | 'resume' | null> {
-  if (!state.workspace || state.projectIdNotNull) return null
-  const parts = [state.project, state.connector, state.rollout, state.projectIdNotNull !== null]
-  if (parts.every((present) => !present)) return 'fresh'
-  if (!parts.every(Boolean)) return null
+  if (process.env.SIM_DEV_DB_PUSH !== '1' || !state.workspace) return null
+  if (!state.project && !state.connector && !state.rollout && !state.projectColumns) return 'fresh'
+  if (!state.project || !state.connector || !state.rollout || state.projectIdNotNull !== false)
+    return null
   const [connector] = await sql<{ populated: boolean }[]>`
     SELECT EXISTS (SELECT 1 FROM public.project_workspace) AS populated`
   if (connector.populated) return null
@@ -70,13 +79,12 @@ async function migrationStatements(tag: string): Promise<string[]> {
 }
 
 /**
- * Applies the versioned Project expansion, then the registered script migration on this
- * single-session connection, recording it as the migration runner does. 0406 opens by committing
- * the migrator's batch; leaving that out commits the 0394 foundation together with the rollout
- * marker, so an interrupted bootstrap leaves either no Project schema or one it can resume.
+ * Replays the versioned Project expansion. 0406 opens by committing the migrator's batch; leaving
+ * that out commits the 0394 foundation together with the rollout marker, so an interrupted
+ * expansion leaves either no Project schema or one the next push resumes.
  */
-async function bootstrapProjects(mode: 'fresh' | 'resume'): Promise<void> {
-  logger.info('Bootstrapping Projects for a development push database without them', { mode })
+async function expandProjects(mode: 'fresh' | 'resume'): Promise<void> {
+  logger.info('Expanding Projects for a development push database without them', { mode })
   const foundation = mode === 'fresh' ? await migrationStatements('0394_project_foundation') : []
   const [batchCommit, ...expansion] = await migrationStatements('0406_workspace_project_column')
   if (batchCommit.trim() !== 'COMMIT;') {
@@ -89,28 +97,41 @@ async function bootstrapProjects(mode: 'fresh' | 'resume'): Promise<void> {
     await sql.unsafe('ROLLBACK').catch(() => undefined)
     throw error
   }
-  await runScriptMigrations(sql, [projectMembershipMigration])
 }
 
-const prepare = process.argv.includes('--prepare')
+function assertProjectMembershipComplete(state: ProjectSchemaState): void {
+  if (state.workspace && (state.connector || !state.projectIdNotNull || !state.project)) {
+    throw new Error(REFUSAL)
+  }
+}
+
+/**
+ * `--prepare` runs before the first Drizzle push and expands an eligible bootstrap. `push.ts`
+ * then pushes the schema with `workspace.project_id` still nullable, so the columns the Project
+ * migration reads exist, and calls `--bootstrap` to run the registered migration on this
+ * single-session connection, recorded as the migration runner records it. Without a flag, the
+ * post-push pass installs the native constraints.
+ */
 try {
   let state = await readProjectSchemaState()
-  if (prepare && process.env.SIM_DEV_DB_PUSH === '1') {
+  if (process.argv.includes('--prepare')) {
     const bootstrap = await classifyProjectBootstrap(state)
     if (bootstrap) {
-      await bootstrapProjects(bootstrap)
+      await expandProjects(bootstrap)
       state = await readProjectSchemaState()
     }
-  }
-  if (state.workspace && (state.connector || !state.projectIdNotNull || !state.project)) {
-    throw new Error(
-      'Existing Project databases must complete the Project script migration 0032_project_membership before schema push. Use the phased migration path for retained data; schema push cannot bypass the writer-drain requirement'
-    )
-  }
-  if (!prepare && state.workspace && state.project) {
-    /** Drizzle cannot express deferred FK timing; fresh push shares the migration finalizer. */
-    await enforceProjectMembership(sql)
-    logger.info('Project membership validation and native constraints completed')
+    if ((await classifyProjectBootstrap(state)) !== 'resume') assertProjectMembershipComplete(state)
+  } else if (process.argv.includes('--bootstrap')) {
+    if ((await classifyProjectBootstrap(state)) !== 'resume') throw new Error(REFUSAL)
+    await runScriptMigrations(sql, [projectMembershipMigration])
+    assertProjectMembershipComplete(await readProjectSchemaState())
+  } else {
+    assertProjectMembershipComplete(state)
+    if (state.workspace && state.project) {
+      /** Drizzle cannot express deferred FK timing; fresh push shares the migration finalizer. */
+      await enforceProjectMembership(sql)
+      logger.info('Project membership validation and native constraints completed')
+    }
   }
 } catch (error) {
   logger.error('Project schema push stopped; resolve the error and retry', {