Prisma リードレプリカ拡張によるクエリの負荷分散

2024年6月24日 09:51

皆様こんにちは。
バックエンドエンジニアの栄です。

今回は、Node.js の ORM である Prisma の拡張機能 @prisma/extension-read-replicas について紹介します。

@prisma/extension-read-replicas とは

@prisma/extension-read-replicas とは、Prisma が発行するクエリの向き先を、プライマリデータベースとレプリカデータベースに分散できる Prisma 公式の拡張機能です。
この拡張機能を導入することで、Prisma を使用するコード上では向き先が制御されていることや複数の接続先が存在していることをほとんど意識しない実装をすることができます。

Prisma 本体に搭載されていてもおかしくなさそうな機能ですが、拡張機能として提供されているのには、本体のリリースとは別でリリースできるようにしておくことで機能改善しやすくするといった狙いがあるようです。

@prisma/extension-read-replicas を使ってみる

ここからはサンプルコードを通して使用方法を紹介していきます。

環境

node 20.10.0
@prisma/client
@prisma/extension-read-replicas
@types/node 20.14.2
typescript 5.4.5
prisma 5.15.0

クライアントの初期化

クライアントの初期化のコードは非常にシンプルです。
readReplicas を呼び出し、レプリカのデータベース URL を指定します。

import { PrismaClient } from "@prisma/client";
import { readReplicas } from "@prisma/extension-read-replicas";

const prisma = new PrismaClient().$extends(
  readReplicas({
    url: process.env.DATABASE_REPLICA_URL!,
  })
);

レプリカのデータベース URL は複数指定することができます。
複数の指定がある場合はレプリカの中からランダムで接続先が選択されます。

import { PrismaClient } from "@prisma/client";
import { readReplicas } from "@prisma/extension-read-replicas";

const prisma = new PrismaClient().$extends(
  readReplicas({
    url: [
      process.env.DATABASE_REPLICA_URL!,
      process.env.DATABASE_REPLICA_URL_2!,
      process.env.DATABASE_REPLICA_URL_3!,
    ],
  })
);

今回紹介した拡張機能に限った話ではないですが、$extends を使うとクライアントの型が @prisma/client の PrismaClient ではなくなってしまい、他のモジュールから利用し辛くなります。
クライアントの型を他のモジュールで利用したい場合には、ユーティリティ型を使ってインスタンスから型を抽出して定義しておくのが良いです。

// 拡張したクライアントの型
export type ExtendedPrismaClient = typeof prisma;

// 拡張したクライアントのトランザクションクライアントの型
export type ExtendedPrismaTransactionClient = Parameters<
  Parameters<typeof prisma.$transaction>[0]
>[0];

接続先の自動分離

拡張したクライアントを使って呼び出した操作については、内部で自動的に適切なデータベースに振り分けられます。
読み込み実行はレプリカに、書き込み、トランザクション、生クエリ実行はプライマリに振り分けられます。

// レプリカの接続が使用される
prisma.user.findUnique({
  where: {
    id: 1,
  },
});
prisma.user.findMany({
  where: {
    id: {
      in: [1, 2, 3],
    },
  },
});

// プライマリの接続が使用される
prisma.user.create({
  data: {
    name: "foo",
  },
});
prisma.user.update({
  data: {
    name: "bar",
  },
  where: {
    id: 1,
  },
});
prisma.$transaction(async (tx) => {
  await tx.user.update({
    data: {
      name: "foo",
    },
    where: {
      id: 1,
    },
  });

  await tx.user.update({
    data: {
      name: "bar",
    },
    where: {
      id: 2,
    },
  });
});
prisma.$queryRaw`INSERT INTO user (name) VALUES ('baz')`;

具体的にどのメソッドが読み込みと判定されるのかはドキュメントには記載がないですが、内部実装から把握することができます。
本記事執筆時点の実装では以下のメソッドが読み込みと判定されているようです。

findFirst
findFirstOrThrow
findMany
findUnique
findUniqueOrThrow
groupBy
aggregate
count
findRaw
aggregateRaw

参考: https://github.com/prisma/extension-read-replicas/blob/d00866cd26696e8cb6c2e5d8c14f90cd3cd335bc/src/extension.ts#L9

明示的な接続先指定

前項では自動的に接続先を制御してくれることを紹介しましたが、プライマリに対して読み込みを実行したいといったような、明示的に接続先を指定したいケースもあると思います。
接続先の指定については、$primary または $replica を呼び出すことで実現可能です。

prisma.$replica().$queryRaw`SELECT * FROM user`;

prisma.$primary().user.findUnique({
  where: {
    id: 1,
  },
});

$replica を指定しつつ書き込み操作を実行するようなコードを書くことも可能である点には注意が必要です。

// 🚨レプリカに対して書き込み操作を実行できる🚨
prisma.$replica().user.create({
  data: {
    name: "foo",
  },
});

最後に

@prisma/extension-read-replicas の概要と、簡単なサンプルコードを紹介しました。
今回紹介した機能の他にも様々な拡張機能が公開されており、公式ドキュメントでも紹介されています。
Prisma の機能を使いつつも細かい制御をしたい場合などは、ユースケースにあった拡張機能を探してみるのが良さそうです。

本記事が Prisma のクエリの分散方法を模索している方や、@prisma/extension-read-replicas の導入を検討している方の助けになれば幸いです。

参考

この記事が気に入ったらサポートをしてみませんか？