tech

Drizzle + Vitest + PostgreSQL でトランザクションを使うテストの分離方法

はじめに

Drizzleを使って新しいアプリのAPIを作っていて、テストを書く時にtransaction周りの処理が上手くいかなかったのでメモです。

最初は SAVEPOINT 方式と TRUNCATE 方式の使い分けで解決したと思っていたのですが、Vitest はデフォルトでファイル並列実行(fileParallelism: true)が有効なため、そこで新たな問題が発生しました。最終的にはワーカー毎のスキーマ分離という方法に落ち着いたので、その経緯と実装を紹介します。

問題: SAVEPOINT と db.transaction() の衝突

テストのデータ分離には、「外側で BEGIN を張って、各テストは SAVEPOINT に戻す」方式を採用しようとしていました。

beforeAll(async () => {
  await testDb.execute(sql`BEGIN`)
})

beforeEach(async () => {
  await testDb.execute(sql.raw(`SAVEPOINT sp_${counter}`))
})

afterEach(async () => {
  await testDb.execute(sql.raw(`ROLLBACK TO SAVEPOINT sp_${counter}`))
})

afterAll(async () => {
  await testDb.execute(sql`ROLLBACK`)
})

この方式自体は最初のシンプルな実装では問題なく動いていたのですが、アプリコード内で db.transaction() を使うエンドポイントのテストを書き始めたところ、以下のエラーで落ちるようになりました。

ROLLBACK TO SAVEPOINT can only be used in transaction blocks

GET/DELETE など db.transaction() を使わないエンドポイントは通るのに、POST/PUT など db.transaction() を使うエンドポイントだけ落ちる、という状況でした。

原因

原因は「テストが beforeAll で張った外側のトランザクションを、アプリの db.transaction()COMMIT で閉じてしまう」ことでした。

流れにするとこんな感じです。

  1. テストが BEGIN(外側トランザクション開始)
  2. アプリコードが db.transaction() を実行(内部で BEGIN → 処理 → COMMIT
  3. アプリが db.transaction() 完了によって COMMIT が発行されると、外側のトランザクションも一緒に終了してしまう
  4. テストの afterEachROLLBACK TO SAVEPOINT ... を実行
  5. すでにトランザクション外になっているため失敗してしまう

PostgreSQL(と主要なRDBMS)ではトランザクション内で BEGIN しても新しいネストは作られません。ネストさせたい場合は ORM 側が SAVEPOINT に動的に変更する必要がありますが、Drizzle の transaction() は自動でやってくれるわけではないため、衝突しました。

関連する議論:

対策案の比較

論点は「テスト分離のための外側 BEGIN/SAVEPOINT」と「アプリの db.transaction()」をどう共存させるか、です。

案A: テスト環境で db.transaction() を SAVEPOINT 実装に差し替える

  • 長所: アプリコード無改修 / 既存の SAVEPOINT 方式を維持
  • 短所: テスト時の transaction() の意味が本番と微妙に変わる(コミット境界に依存する処理があると差分になり得る)

案B: 問題が起きるテストだけTRUNCATEでDBをリセットする方式にする

  • 長所: db.transaction() が本番どおり動く(実コミット前提で検証できる)
  • 短所: SAVEPOINT より遅くなりがち / ID の連番などに注意点が出る

案C: ワーカー毎にスキーマを分離する

  • 長所: 最も強い分離 / 並列実行でも安全
  • 短所: セットアップが少し複雑

結論

最終的には 案B + 案C の組み合わせ に落ち着きました。

  • ファイル単位で SAVEPOINT か TRUNCATE かを選べるようにする(案B)
  • ワーカー毎にスキーマを分離して並列実行でも安全にする(案C)

最初は案Bだけで十分だと思っていたのですが、 並列実行中に TRUNCATE が他のワーカーのデータを消してしまい競合してしまう問題が発生し、案Cも必要になりました。

実装

1) ファイル単位で分離方式を切り替える(案B)

テスト分離の方式を helper として用意する

テスト用のヘルパーファイルに 2 つの API を用意しました。

  • useSavepointIsolation()

    • beforeAll: BEGIN
    • beforeEach: SAVEPOINT
    • afterEach: ROLLBACK TO SAVEPOINT
    • afterAll: ROLLBACK
  • useResetIsolation()

    • afterEach: TRUNCATE ... CASCADE(ワーカースキーマ内の全テーブル)

使う側はテストファイルで呼び出すだけです。

import { useSavepointIsolation } from '../test/isolation'

useSavepointIsolation()

db.transaction() を使うテストファイルは、こちらにします。

import { useResetIsolation } from '../test/isolation'

useResetIsolation()

2) ワーカー毎にスキーマを分離する(案C)

問題: 並列実行時の競合

Vitest はデフォルトで fileParallelism: true(ファイル並列実行)なので、案Bだけだと以下のエラーが発生しました。

deadlock detected

useResetIsolation() は TRUNCATE を使うのですが、これがファイル並列実行時に他のワーカーのデータも消してしまったり、ロック競合を起こしていました。

解決: PostgreSQL の search_path を活用

各ワーカーが独自のスキーマを使うようにすれば、TRUNCATE は自分のスキーマ内だけに影響するため、他のワーカーには影響しません。

  1. テスト実行前に test_w0, test_w1, ... というスキーマを事前作成
  2. 各ワーカーは VITEST_WORKER_ID 環境変数を見て自分のスキーマを決定
  3. DB接続時に SET search_path TO test_w{id} を実行

注意点として、VITEST_WORKER_ID は 0-indexed です。

最大ワーカー数を定数として定義

Vitest は最大ワーカー数を export していないため、vitest.config.ts と global-setup で共有する定数を用意します。

// constants.ts(テスト用定数ファイル)
export const MAX_TEST_WORKERS = 32

vitest.config.ts でワーカー数を制限

export default defineConfig({
  test: {
    // ...
    maxWorkers: MAX_TEST_WORKERS,
  },
})

Vitest の globalSetup でスキーマを事前作成

async function setup() {
  // マイグレーションを実行(public スキーマにテーブルを作成)
  execSync('npm run db:migrate', { ... })

  const client = new Client({ connectionString: TEST_DATABASE_URL })
  await client.connect()

  // public スキーマのテーブルを取得
  const tables = await client.query(`
    SELECT tablename FROM pg_tables WHERE schemaname = 'public'
  `)

  // 各ワーカー用のスキーマを作成(0-indexed)
  for (let i = 0; i < MAX_TEST_WORKERS; i++) {
    const schemaName = `test_w${i}`

    // スキーマを作成(既存の場合はドロップして再作成)
    await client.query(`DROP SCHEMA IF EXISTS ${schemaName} CASCADE`)
    await client.query(`CREATE SCHEMA ${schemaName}`)

    // テーブル構造をコピー
    for (const row of tables.rows) {
      const tableName = row.tablename
      await client.query(`
        CREATE TABLE ${schemaName}.${tableName}
        (LIKE public.${tableName} INCLUDING ALL)
      `)
    }
  }
}

DB接続モジュールでワーカー固有のスキーマを設定

// ワーカーIDに基づいてスキーマ名を決定
// VITEST_WORKER_ID は 0-indexed(未設定時は 0 にフォールバック)
function resolveWorkerId(): number {
  const rawWorkerId = process.env.VITEST_WORKER_ID
  if (!rawWorkerId) return 0
  const parsed = Number.parseInt(rawWorkerId, 10)
  if (!Number.isFinite(parsed) || parsed < 0) return 0
  return parsed
}

const workerId = resolveWorkerId()
export const testSchemaName = `test_w${workerId}`

export async function setWorkerSchema(): Promise<void> {
  await client.query(`SET search_path TO ${testSchemaName}, public`)
}

テスト分離ヘルパーでスキーマを指定した TRUNCATE

drizzle-seedreset() はスキーマ未指定時に常に public を参照するため、ワーカースキーマでは使えません。代わりにスキーマを明示的に指定して TRUNCATE を実行します。

async function truncateAllTables() {
  // ワーカースキーマ内の全テーブルを取得
  const result = await testDb.execute<{ tablename: string }>(
    sql.raw(`SELECT tablename FROM pg_tables WHERE schemaname = '${testSchemaName}'`),
  )
  const tableNames = result.rows.map((row) => row.tablename)

  if (tableNames.length === 0) return

  // 全テーブルを TRUNCATE(CASCADE で外部キー制約も処理)
  const truncateQuery = tableNames.map((t) => `"${testSchemaName}"."${t}"`).join(', ')
  await testDb.execute(sql.raw(`TRUNCATE ${truncateQuery} CASCADE`))
}

export function useResetIsolation() {
  afterEach(async () => {
    await truncateAllTables()
  })
}

テストセットアップで接続後にスキーマを切り替え

beforeAll(async () => {
  await pgClient.connect()
  await setWorkerSchema()
})

afterAll(async () => {
  await pgClient.end()
})

これで各ワーカーが完全に独立したスキーマを使うようになり、TRUNCATE の競合が解消されました。

運用メモ

  • TRUNCATE には RESTART IDENTITY を付けていないため、「ID が毎回 1 になる前提」のテストは作らない方が安全です(factory で作った ID を使う、など)。
  • describe.concurrent() / it.concurrent() は避ける必要があります。同一ファイル内のテストが並列実行されると、SAVEPOINT の順序が壊れます。
  • ワーカースキーマの数(MAX_TEST_WORKERS)は、CI の並列数に合わせて調整してください。足りないとワーカーがスキーマを共有してしまいます。

まとめ

  • SAVEPOINT 方式は速くて便利ですが、Drizzle の db.transaction() と同居させると衝突しやすいです。
  • ファイル単位で SAVEPOINT か TRUNCATE かを選べるようにすることで、速度と確実性を両立できます。
  • Vitest のデフォルト並列実行では、ワーカー毎にDBスキーマを分離することで TRUNCATE の競合を回避できます。

おわりに

SAVEPOINT 方式はテストが高速で気に入っていたのですが、Drizzle の transaction() との相性でハマりました。さらに Vitest のデフォルト並列実行で TRUNCATE の競合でもハマりました。

ワーカー毎のスキーマ分離は最初「重そう」と思って避けていましたが、実際にやってみると search_path を変えるだけなので思ったより軽量でした。テストセットアップ時にスキーマをコピーする処理が入りますが、CI の実行時間に大きな影響はありませんでした。

ただ「テストの並列実行」という目的に対して、ここまで複雑な仕組みが必要になるのは想定外でした。何かもっと良い方法があれば知りたいなという気持ちです。

同じ問題に当たった方の参考になれば幸いです。

関連リンク

share