PrismaでRLSを使う

はじめに

ここ一年ほどはDBクライアントを書くときにORMのPrismaを使っていますが、テーブル定義を書きやすかったりコード補完がバチバチに効いてくれたりと非常に気に入ってます。

そんなPrismaでPostgreSQLのRLS(Row Level Security)を使う時の設定について紹介します。
RLSのしっかりした説明はしていません🙇

以下のレポジトリにサンプルがあります。このサンプルに基づいて紹介していきます。

https://github.com/s-ysk/prisma-rls-example

この記事では、以下の手順でRLSの設定をしていきます。

1. RLSをかけるDBユーザの追加・設定
2. Prismaスキーマファイルの作成
3. RLS定義マイグレーションファイル作成
4. Prisma Extensionsの定義
5. RLSが有効になったPrisma Clientを使う

RLSをかけるDBユーザの追加・設定

まずは、RLSをかけるDBユーザを追加します。アプリケーションからは、ここで作成したユーザを使ってDBアクセスを行います。以下のシェルスクリプトのように新しくユーザを作って、必要最低限の権限を与えます。

HOST=localhost
PORT=5432
USER=app
DB=prisma-local
PASSWORD=mypaSsWd

psql "postgresql://root:password@$HOST:$PORT/$DB" <<-EOSQL
	CREATE USER $USER WITH PASSWORD '$PASSWORD';
	ALTER USER $USER CREATEDB;
	GRANT CREATE ON DATABASE "$DB" TO $USER;
	GRANT CREATE ON SCHEMA public TO $USER;
	GRANT DELETE, INSERT, SELECT, UPDATE ON ALL TABLES IN SCHEMA public TO $USER;
	GRANT SELECT, UPDATE ON ALL SEQUENCES IN SCHEMA public TO $USER;

※このスクリプトはレポジトリ内のscript/setup_db_client.shにあります。

Prismaスキーマファイルの作成

次にPrismaのスキーマファイルを作成します。RLSを試すためにテナントとプロジェクトという二つのテーブルを作成し、プロジェクトはテナントに所属すると定義します。組織ごとにプロジェクトを持っているというイメージです。

generator client {
  provider = "prisma-client-js"

  // ①エクステンションを有効にする
  previewFeatures = ["clientExtensions"]
}

datasource db {
  provider = "postgresql"
  url      = env("DATABASE_URL_WITH_ROOT")
}

// Tables
model Tenant {
  id        Int       @id @default(autoincrement())
  name      String
  createdAt DateTime  @default(now())
  updatedAt DateTime  @updatedAt
  projects  Project[]
}

model Project {
  id        Int      @id @default(autoincrement())
  name      String
  createdAt DateTime @default(now())
  updatedAt DateTime @updatedAt
  tenant   Tenant @relation(fields: [tenantId], references: [id])
  
  // ②プロジェクト作成時にトランザクションに設定されたテナントIDを自動的に設定する
  tenantId Int    @default(dbgenerated("(current_setting('app.tenant_id'::TEXT))::INT"))
}

ポイントとしては、以下の二点です。

1. エクステンションを明示的に有効化する。
2. ProjectテーブルのtenantIdカラムの初期値として、current_setting('app.tenant_id'::TEXT)::INTを設定する。

一点目は2023年3月31日時点でエクステンションはプレビューリリースだったためです。

二目点では、DBのcurret_settingからapp.tenant_idの値を読み込み、それをtenantIdの初期値としています。curret_settingが何の値を取得しているのかは後ほど説明しますが、このようにすることで「tenantIdにDBトランザクションに設定された値を入力」できます。

さらにPrismaに渡す環境変数を定義しておきます。Prisma CLIがDBの操作に使うためのRootユーザと、アプリケーションがDBアクセスに使うApp Userを定義しました。

# ROOT DB user. Use this for prisma operations.
DATABASE_URL_WITH_ROOT="postgresql://root:password@localhost:5432/prisma-local?schema=public"

# RSL enabled user. Use this to access DB from application.
DATABASE_URL_WITH_APP_USER="postgresql://app:mypaSsWd@localhost:5432/prisma-local?schema=public"

※この辺で使われているユーザ名やパスワードはレポジトリ内のdocker-compose.yamlファイルやscript/ディレクトリ内のスクリプトに定義されています。

ここまででPrismaの設定が完了したので、以下のコマンドを使ってDBに設定を反映させます。

yarn prisma migrate dev

RLS定義マイグレーションファイル作成

今度はRLSを定義します。上のセクションでprismaコマンドの実行後に、./prisma/migrationsディレクトリ以下に下のようなマイグレーションフォルダ・ファイルが作られていると思います。

./prisma
├── migrations
│   ├── 20230401040420_add_tables
│   │   └── migration.sql
│   └── migration_lock.toml
└── schema.prisma

ここで、migrations/の下に新たなマイグレーションフォルダ・ファイルを作成し、以下の内容を記入します。

-- Enable Row Level Security
ALTER TABLE "Tenant" ENABLE ROW LEVEL SECURITY;
ALTER TABLE "Project" ENABLE ROW LEVEL SECURITY;

-- Force Row Level Security for table owners
ALTER TABLE "Tenant" FORCE ROW LEVEL SECURITY;
ALTER TABLE "Project" FORCE ROW LEVEL SECURITY;

-- Create row security policies
CREATE POLICY tenant_isolation_policy ON "Tenant" USING ("id" = current_setting('app.tenant_id', TRUE)::INT);
CREATE POLICY tenant_isolation_policy ON "Project" USING ("tenantId" = current_setting('app.tenant_id', TRUE)::INT);

-- Create row security policies
CREATE POLICY bypass_rls_policy ON "Tenant" USING (current_setting('app.bypass_rls', TRUE)::text = 'on');
CREATE POLICY bypass_rls_policy ON "Project" USING (current_setting('app.bypass_rls', TRUE)::text = 'on');

ここではTenantテーブルとProjectテーブルにて、RLSを有効化しポリシーを設定しています。ポリシーに設定した条件がTRUEのレコードに対してのみ読み書きが可能になります。

ここで設定しているポリシーは二つになります。

1. tenant_isolation_policy

トランザクションに設定されたapp.tenant_idの値を確認し、その値がテーブル内のtenantIdと一致するレコードについてのみ読み書き可能とするポリシーです。

2. bypass_rls_policy

トランザクションに設定されたapp.bypass_rlsの値を確認し、その値がONであればどのレコードに対しても読み書き可能とするポリシーです。新規テナント作成時など、tenantIdが決まる前に何かしら操作をしたいケースがあるのでRLSをバイパス(迂回)するポリシーがあると便利です。

ちなみにRLSポリシーは OR 条件で適用されるので、bypass_rls_policyがTRUEになっているトランザクションではtenant_isolation_policyによらず全てのレコードを読み書きできます。

上記のマイグレーションファイルを作成したら、以下のコマンドでDBに設定を反映させます。

yarn prisma migrate deploy

ここまででDBの設定は完了しました。

Prisma Extensionsの定義

次はクライアントサイドの定義です。Prisma Clientを拡張し、クエリが実行される前に「トランザクション内のみで有効な設定」を挿入するようにします。

import { PrismaClient } from "@prisma/client";

const prisma = new PrismaClient({
  datasources: { db: { url: process.env["DATABASE_URL_WITH_APP_USER"] } },
});

// クエリ実行時にbypass_rls_policyを有効にする拡張
export const tenantGuardedPrisma = (tenantId: number) => {
  return prisma.$extends({
    query: {
      // 全てのモデルに対して拡張機能を有効にする
      $allModels: {
        // 全ての操作(CREATE, UPDATEなど)に拡張機能を有効にする
        async $allOperations({ args, query }) {
	  // トランザクション定義
          const [, , result] = await prisma.$transaction([
	    // app.tenant_idにtenantIdを設定
            prisma.$executeRaw`SELECT set_config('app.tenant_id', ${tenantId.toString()}, TRUE)`,
	    // tenant_isolation_policyを使いたいので、app.bypass_rlsはオフ
            prisma.$executeRaw`SELECT set_config('app.bypass_rls', 'off', TRUE)`,
	    // もともと実行しようとしていたクエリ
            query(args),
          ]);
          return result;
        },
      },
    },
  });
};

// クエリ実行時にtenant_isolation_policyを有効にする拡張
export const bypassedPrisma = () => {
  return prisma.$extends({
    query: {
      $allModels: {
        async $allOperations({ args, query }) {
          const [, , result] = await prisma.$transaction([
	    // app.tenant_idに-1(サービスで使っていない値)を入れてtenant_isolation_policyをオフにする
            prisma.$executeRaw`SELECT set_config('app.tenant_id', ${(-1).toString()}, TRUE)`,
            prisma.$executeRaw`SELECT set_config('app.bypass_rls', 'on', TRUE)`,
            query(args),
          ]);
          return result;
        },
      },
    },
  });
};

以下のようにset_configの第三引数にTRUEをセットすると、「現在のトランザクション内でのみ有効な設定」を定義することができます。そのため上記の例では「tenantIdを与えると、RLSによってそのトランザクション内ではtenantIdに関連付けされたデータしか読み書きできない」よう設定しているということです。

set_config('app.tenant_id', ${tenantId.toString()}, TRUE)

Prismaスキーマファイルの作成で出てきたcurrent_setting(app.tenant_id)はset_configで設定した値を読んでいます。

ここまで定義したら、あとは拡張されたPrisma Clientをビジネスロジックから使うだけです。

RLSが有効になったPrisma Clientを使う

以下のようにTenantとそれに関連するProjectを定義して、RLSポリシーごとにどのような結果を得られるか確認してみました。

import {
  bypassedPrisma,
  tenantGuardedPrisma,
} from "./db";

// DBのルートユーザとしてアクセス可能なPrisma Client
// DBの初期設定に利用する
export const rootPrisma = new PrismaClient({
  datasources: { db: { url: process.env["DATABASE_URL_WITH_ROOT"] } },
});

const main = async () => {
  console.debug("Start example.");
  
  const myTenantId = 1;
  const otherTenantId = 2;
  const bypassed = bypassedPrisma();
  const myTenantPrisma = tenantGuardedPrisma(myTenantId);
  const otherTenantPrisma = tenantGuardedPrisma(otherTenantId);

  // 1. データ初期設定
  // Root user can do anything.
  await rootPrisma.tenant.createMany({
    data: [
      // 二つのテナントを作成
      { name: "test-my-tenant", id: myTenantId },
      { name: "test-other-tenant", id: otherTenantId },
    ],
  });
  await rootPrisma.project.createMany({
    // tenantId == 1でプロジェクトを作っておく
    data: ["my-p1", "my-p2"].map((e) => ({ name: e, tenantId: myTenantId })),
  });
  await rootPrisma.project.createMany({
    // tenantId == 2でプロジェクトを作っておく
    data: ["other-p1"].map((e) => ({ name: e, tenantId: otherTenantId })),
  });


  // 2. tenant_isolation_policyを使ってみる
  // 2-1. My tenantPrisma.
  console.debug("### Check myTenantPrisma can get its own projects");
  const myProjects = await myTenantPrisma.project.findMany();
  console.debug("myProjects", myProjects);

  // 2-2. Other tenantPrisma.
  console.debug("### Check otherTenantPrisma can get its own projects");
  const otherTenantPrismaProjects = await otherTenantPrisma.project.findMany();
  console.debug("otherTenantPrismaProjects", otherTenantPrismaProjects);

  // 3. bypass_rls_policyを使ってみる
  console.debug("### Check bypassedPrisma can get the all projects independent of tenantId");
  const bypassedPrismaProjects = await bypassed.project.findMany();
  console.debug("bypassedPrismaProjects", bypassedPrismaProjects);
};

main();

出力結果は以下のようになると思います(見やすさのために、creatdAtやupdatedAtなどのカラムは削除しています)

Start example.

### Check myTenantPrisma can get its own projects
myProjects [
  {
    id: 1,
    name: 'my-p1',
    tenantId: 1
  },
  {
    id: 2,
    name: 'my-p2',
    tenantId: 1
  }
]

### Check otherTenantPrisma can get its own projects
otherTenantPrismaProjects [
  {
    id: 3,
    name: 'other-p1',
    tenantId: 2
  }
]

### Check bypassedPrisma can get the all projects independent of tenantId
bypassedPrismaProjects [
  {
    id: 1,
    name: 'my-p1',
    tenantId: 1
  },
  {
    id: 2,
    name: 'my-p2',
    tenantId: 1
  },
  {
    id: 3,
    name: 'other-p1',
    tenantId: 2
  }
]

おわりに

こんな感じでPrismaにRLSを導入することができました。
マルチテナントのサービスを作る際はtenantIdを全てのテーブルに持たせ、全てのクエリにRLSを効かせるよう設定すると安心です。

参考文献

1. Prisma Client extensions
2. Example: Row level security