背景
ユビーではNode.jsのアプリケーションにおいて、DBのORMにPrismaを採用しています。Prismaに限らず、実際のデータベースを使ったテストにおいてはテストを実行するごとにデータを削除することでテストごとのデータの干渉を防いでテストの安定性を保つという手法が用いられます。今回はPrismaでテスト時にデータ削除するときのパフォーマンス改善の事例について紹介します。
今開発しているアプリケーションでは以下のように、DBを用いるテストを実行する前に全てのテーブルを truncateするという素朴な方法でテストDBの掃除を行っていました。
import { Prisma, PrismaClient } from "@prisma/client";
const prisma = new PrismaClient();
export async function cleanupDatabase() {
const tableNames = Prisma.dmmf.datamodel.models.map((model) => model.dbName);
for (const table of tableNames) {
await prisma.$queryRawUnsafe(`TRUNCATE TABLE "${table}" CASCADE`);
}
}
describe("MyService", () => {
beforeEach(async() => {
await cleanupDatabase();
});
// DBを利用したテストコード
});
しかしテストコードが多くなってくるにつれ、このcleanupの処理のオーバーヘッドが気になるようになってきて、実際に計測してみるとcleanupの処理だけで1500msくらいかかっていました。今の時点ではテスト全体を実行して60s程度なので全然遅くはないんですが、既存機能を移行中のアプリケーションなので、今後もテストコードが増えていくのは自明であり、このままだとテストの増加にともなって実行時間が線形に伸びていくことが明らかだったため、今のうちになんとかしておこうと思い改善を試みました。
改善の方針と実装
最初に結果から言うと、一回のcleanup処理が1500msから100ms弱くらいになりテスト全体では60sから10s程度まで高速化しました。
基本的な方針としては全てのテーブルをデータ削除の対象にせず、データが作成されたテーブルだけを記録しておいてtruncateするテーブルを絞るというものです。PrismaにはClient Extensionsという拡張のための機能があって、これを使うことでPrismaのAPI呼び出しに対して処理をフックできます。これを使ってデータ作成のイベントを拾って記録し、必要なテーブルだけ削除処理をおこないます。
const prisma = new PrismaClient().$extends({
query: {
async $allOperations({ operation, model, args, query }) {
// ここで operation が create だったら model を保存しておいて
// テスト終了後に model から table を引いて trucate する
},
},
});
これいい感じにやってくれるライブラリを作りました。なお現時点だとPostgreSQLにしか対応していません。
こんな感じの機能を提供します。
import { Prisma, PrismaClient } from "@prisma/client";
import { PrismaCleaner } from "@ubie/prisma-cleaner";
const cleaner = new PrismaCleaner({
prisma: new PrismaClient(),
models: Prisma.dmmf.datamodel.models,
});
const prisma = new PrismaClient().$extends(cleaner.withCleaner());
(async () => {
await prisma.user.create({ ... });
// Userテーブルだけ削除される
await cleaner.cleanup(); // => TRUNCATE TABLE "public"."User" CASCADE
await prisma.post.create({ ... });
await prisma.comment.create({ ... });
// PostテーブルとCommentテーブルが削除される。Userテーブルは含まれない
await cleaner.cleanup(); // => TRUNCATE TABLE "public"."User", "public"."Comment" CASCADE
})();
jestなどに組み込む場合はセットアップ時にPrismaCleanerの初期化をしたらafterEachでcleanup()を呼ぶだけでいいようにしています。例えば以下のように書けば、テストを書くときはデータ削除については何も気にせずにテストを書けばテスト終了後に勝手にデータ削除が走ります。
export default {
// ...
setupFilesAfterEnv: ["<rootDir>/setup.ts"],
globalSetup: "<rootDir>/global-setup.ts",
}
import { Prisma, PrismaClient } from "@prisma/client";
import { PrismaCleaner } from "@ubie/prisma-cleaner";
export const cleaner = new PrismaCleaner({
prisma: new PrismaClient(),
models: Prisma.dmmf.datamodel.models,
});
import { PrismaClient } from "@prisma/client";
import { cleaner } from "./cleaner";
afterEach(async () => {
// 各テストケース終了後に毎回これが呼ばれてテストケース中に作成されたデータだけ消す
await cleaner.cleanup();
});
import { cleaner } from "./cleaner";
// この処理はテスト起動時の最初の一回だけ呼ばれる。
// 必須ではないけど、テスト完了後に何かの間違いでデータが残った場合などにテストが壊れることがあるので
// 最初に全テーブルきれいにしておくとよい。
export default async function setup() {
await cleaner.cleanupAllTables();
}
import { PrismaClient } from "@prisma/client";
import { UserService } from "./UserService";
import { cleaner } from "../test/cleaner";
describe("UserService", () => {
// テストで利用するprisma clientはcleaner.withCleaner()のextensionを有効にする
const prisma = new PrismaClient().$extends(cleaner.withCleaner()) as PrismaClient;
const userService = new UserService(prisma);
it("should create a new user", async () => {
// このレコードはsetup.tsのafterEachのcleanupで削除される
const user = await userService.createUser("xxx");
expect(user.name).toEqual("xxx");
expect(await prisma.user.count()).toEqual(1);
});
it("should be cleanup user table by cleaner", async () => {
// 勝手に削除されている
expect(await prisma.user.count()).toEqual(0);
});
});
テストのときだけ$extendsでcleanerのextensionを有効にするのがちょっと面倒で、上記のコードはテスト対象がPrismaClientを受け取るので比較的シンプルですが、シングルトンになる場合やNestJSでDIする場合などはひと工夫必要です。いくつかのケースをexampleにおいているので参照してみてください。
その他の手法
ここにたどり着くまでにいくつかの方法を試したので紹介しておきます。
jest-prisma
まず一番最初に試したのは jest-prisma です。
これはテストケースごとにtransactionを発行して、テストが終わるとrollbackするというやつです。結果としては上記の必要なものだけtruncateする場合と同じくらいの高速化はできたのですが、transaction内で実行されることが起因で一部のテストでエラーが発生してアプリケーションかテストのコードを直す必要がありました。
コードを直して対応してもよかったのですが、transaction起因で今後も同じような問題に遭遇する可能性を懸念して別の方法を探った結果、同じくらいの速度になったので今回はjest-prismaの採用を見送りました。なお、jest-prismaはユビー内の他のプロジェクトでは普通に使われていて、状況によっては採用するのは全然アリだと思っています。すごく便利です。Quramyさんいつもありがとうございます。
余談ですが、jest-prismaを導入したいモチベーションのもう一つに、テストを並列化したいというのもありました。DBを利用したテストでは並列にテストケースを実行するとデータが競合してテストが失敗するので、基本的にはjestのmaxWorkersを1にするなどして直列で実行します。
jest-prismaを導入するとテスト実行時にtransactionでデータの変更が分離されるので並列実行可能になりテストが高速化できないかと思っていたのですが、実行してみるとdeadlock祭りで死亡したので諦めました。もしこれが実現可能だったらメンテナスコストを払ってでもjest-prismaを採用するメリットのほうが大きいかなと思っていたのですがダメでした。テストの並列化についてはworkerごとにDBを分けるしかないかなと思っていて、いずれやるかもしれません。
まとめてtruncate
次に試したのがtruncateをまとめる方法です。PostgreSQLでは、以下のように複数テーブルをまとめてtruncateに指定できます。
TRUNCATE TABLE foo, bar, baz;
これを使って最初のtruncateのコードを以下のように書き換えました。
export async function cleanupDatabase() {
const tables = Prisma.dmmf.datamodel.models.map((model) => `"${model.dbName}"`).join(", ");
await prisma.$queryRawUnsafe(`TRUNCATE TABLE ${tables}`);
}
まとめて指定する場合だと、外部キー制約をたどってtruncateするためのCASCADEも不要になって効率がよくなり、クエリの発行も一回で済むので高速化できます。ちなみに複数のtruncateを並列で実行するというのも試してみましたがこちらはdeadlockで死にました。
このまとめて指定する方法だと、truncateの処理が1500msから300msくらいになり、かなり高速化できました。そんなに規模の大きくないアプリケーションであれば、無駄に複雑なものを持ち込むよりもこれぐらい素朴な実装でよさそうです。
ただ、この方法もテーブル数の増加に伴って遅くなるはずで、今回はモジュラモノリスなアプリケーションで今後もテーブル数が増加することがわかっているので、テーブルが増えても遅くならないようなアプローチを選択しました。
まとめ
Prismaのテストでデータ削除を高速化するための方法や実装について紹介しました。ユビーではNode.jsを使ったアプリケーション開発が活発に行われていますのでもし興味がある人がいたらぜひ一度お話しましょう！
Discussion