Ever ran a database migration only for it to fail halfway through, and then you're stuck with a partially migrated database? Then you try to run a rollback, which fails because the database is not on the previous version, nor is it on the next version, it's somewhere in between, so you have to cleanup manually. I've had this problem way too many times, so I finally got bothered by it enough to create a solution for myself.

From minor inconvenience to major f!@#ing problem

Ever ran a database migration only for it to fail halfway through, and then you're stuck with a partially migrated database? Then you try to run a rollback, which fails because the database is not on the previous version, nor is it on the next version, it's somewhere in between, so you have to cleanup manually before trying again.

This is a minor inconvenience when it happens in your local or in dev/test/staging, but sometimes you're lucky enough to get all the way to prod before your migration unexpectedly fails due to some fun edge case with the data, then you're in for a real treat.

I've had the pleasure of dealing with this problem enough times to finally get bothered enough by it to create a solution.

The idea

It might not be ideal in all scenarios, but my idea was to create a full copy of the database, then run the migration on the copy first. If that succeeds, there should be no reason for the migration to fail on the actual database. If the test migration fails, your real database is left untouched and you can just change your migration and try again.

Pros and cons

Pros:

  • If the migration fails on the copy, your real database is completely untouched. Just fix the migration and try again.
  • Uses PostgreSQL's native CREATE DATABASE ... WITH TEMPLATE which is a fast, server-side copy.
  • Works with Drizzle's migration system.
  • Easy to add to an existing project.

Cons:

  • CREATE DATABASE ... WITH TEMPLATE requires terminating all active connections to the source database. No new connections are allowed until the copy finishes, and naturally the copy time scales with database size.
  • The copy temporarily doubles your disk usage.

Most of these cons can be worked around though! If you already have regular database backups (which you should), you could restore a backup to a separate database instead of using CREATE DATABASE ... WITH TEMPLATE. This avoids the connection termination and downtime. The one thing you can't really work around is a slow migration, a worthwhile trade-off in my mind.

The script

Disclaimer: I'm no expert on Drizzle or Postgres, and got help from AI to set this up, however I've ran 150+ migrations over the course of a year for a decently large application, and so far it has yet to fail me!

First, the setup. We validate the environment and create some helpers to extract the database name and username from a connection string, which can be neatly done with JS's URL class:

import { z } from 'zod'
import { Pool } from 'pg'
import { sql } from 'drizzle-orm'
import { drizzle } from 'drizzle-orm/node-postgres'
import { migrate } from 'drizzle-orm/node-postgres/migrator'

const Env = z
  .object({
    DATABASE_URL: z.string().url(),
    MIGRATIONS_FOLDER: z.string().default('./drizzle'),
  })
  .safeParse(process.env)

if (!Env.success) {
  console.error('Invalid env:', Env.error.flatten())
  process.exit(1)
}

const { DATABASE_URL, MIGRATIONS_FOLDER } = Env.data

function getDbName(url: string) {
  return new URL(url).pathname.slice(1)
}

function getDbUsername(url: string) {
  return new URL(url).username
}

const DB_NAME = getDbName(DATABASE_URL)
const DB_USERNAME = getDbUsername(DATABASE_URL)
const TEST_DB_NAME = `${DB_NAME}_migrate_test`

Next, the main function. We create a copy of the database, run migrations on both, and cleanup. Since CREATE DATABASE ... WITH TEMPLATE requires that no other sessions are connected to the source database, we terminate them first.

async function main() {
  const setupPool = new Pool({ connectionString: DATABASE_URL })
  const setupDb = drizzle(setupPool)

  const exists = await setupDb.execute(
    sql`SELECT 1 FROM pg_database WHERE datname = ${TEST_DB_NAME}`
  )

  if (exists.rowCount && exists.rowCount > 0) {
    await setupDb.execute(
      sql.raw(`DROP DATABASE "${TEST_DB_NAME}" WITH (FORCE)`)
    )
  }

  await setupDb.execute(sql.raw(
    `SELECT pg_terminate_backend(pg_stat_activity.pid)
     FROM pg_stat_activity
     WHERE pg_stat_activity.datname = '${DB_NAME}'
     AND pid <> pg_backend_pid()`
  ))

  await setupDb.execute(sql.raw(
    `CREATE DATABASE "${TEST_DB_NAME}"
     WITH TEMPLATE "${DB_NAME}" OWNER "${DB_USERNAME}"`
  ))

  await setupPool.end()

  const testUrl = new URL(DATABASE_URL)
  testUrl.pathname = `/${TEST_DB_NAME}`

  await runMigrations(testUrl.toString())
  await runMigrations(DATABASE_URL)

  const cleanupPool = new Pool({ connectionString: DATABASE_URL })
  const cleanupDb = drizzle(cleanupPool)
  await cleanupDb.execute(
    sql.raw(`DROP DATABASE IF EXISTS "${TEST_DB_NAME}" WITH (FORCE)`)
  )
  await cleanupPool.end()
}

main().catch((err) => {
  console.error(err)
  process.exit(1)
})

And then the runMigrations function. One challenge I ran into is that Drizzle's internal __drizzle_migrations table sometimes doesn't have a auto-incrementing sequence on the id column, which can cause issues when running migrations on a copied database. The ensureMigrationsTable function handles this. You could try to skip it and see if it works for your setup.

async function runMigrations(url: string) {
  const pool = new Pool({ connectionString: url })
  const db = drizzle(pool)
  await ensureMigrationsTable(db)
  await migrate(db, { migrationsFolder: MIGRATIONS_FOLDER })
  await pool.end()
}

async function ensureMigrationsTable(
  db: ReturnType<typeof drizzle>
) {
  await db.execute(sql.raw(`
    DO $$ BEGIN
      IF NOT EXISTS (
        SELECT 1 FROM pg_namespace WHERE nspname = 'drizzle'
      ) THEN
        EXECUTE 'CREATE SCHEMA drizzle';
      END IF;
    END$$
  `))

  await db.execute(sql.raw(`
    CREATE TABLE IF NOT EXISTS drizzle.__drizzle_migrations (
      id integer PRIMARY KEY,
      hash text NOT NULL,
      created_at bigint NOT NULL
    )
  `))

  await db.execute(sql.raw(`
    DO $$
    DECLARE seq text; def_expr text;
    BEGIN
      SELECT pg_get_serial_sequence(
        'drizzle.__drizzle_migrations', 'id'
      ) INTO seq;
      IF seq IS NULL THEN
        SELECT column_default INTO def_expr
        FROM information_schema.columns
        WHERE table_schema = 'drizzle'
          AND table_name = '__drizzle_migrations'
          AND column_name = 'id';
        IF def_expr IS NOT NULL THEN
          SELECT substring(
            def_expr from 'nextval\\(''([^'']+)''::regclass\\)'
          ) INTO seq;
        END IF;
        IF seq IS NULL THEN
          IF NOT EXISTS (
            SELECT 1 FROM pg_class c
            JOIN pg_namespace n ON n.oid = c.relnamespace
            WHERE c.relkind = 'S'
              AND n.nspname = 'drizzle'
              AND c.relname = '__drizzle_migrations_id_seq'
          ) THEN
            EXECUTE 'CREATE SEQUENCE drizzle.__drizzle_migrations_id_seq';
          END IF;
          EXECUTE 'ALTER SEQUENCE drizzle.__drizzle_migrations_id_seq
            OWNED BY drizzle.__drizzle_migrations.id';
          EXECUTE 'ALTER TABLE drizzle.__drizzle_migrations
            ALTER COLUMN id SET DEFAULT
            nextval(''drizzle.__drizzle_migrations_id_seq'')';
          seq := 'drizzle.__drizzle_migrations_id_seq';
        END IF;
      END IF;
      IF seq IS NOT NULL THEN
        EXECUTE format(
          'SELECT setval(%L, COALESCE(
            (SELECT MAX(id) FROM drizzle.__drizzle_migrations), 0
          ) + 1, false);', seq
        );
      END IF;
    END$$
  `))
}