🔐

Prisma で RLS (Row Level Security) を使う

2023/05/05に公開

PostgreSQL

Prisma

RLS

tech

概要

Prisma で PostgreSQL の RLS(Row Level Security) を使う方法のまとめ。
Prisma では RLS は公式にサポートされていない^[1]ため、RLSを使う手順が少し煩雑だったためメモ。

バージョン情報

Prisma: 5.0.0

方法

拡張する

次のコードのように Prisma Client を拡張します

// 実行パラメータとしてテナントIDをセットするSQLを実行する関数
function setTenant(
  prisma: Pick<PrismaClient, "$executeRaw">,
  tenantId: string
): PrismaPromise<number> {
  return prisma.$executeRaw`SELECT set_config('app.current_tenant', ${tenantId}, TRUE)`;
}

// 全てのクエリの実行の前に、`$transaction()` を使用して、テナントIDをセットする拡張
// 注意： この拡張を使用する場合、`$transaction()` を含めた全てのクエリが、新しいトランザクションとしてラップされる。
// 　　　 そのため、`$transaction()` が正しく動作しない点に注意する。 `$transaction()` は別途`extendTx` にて拡張する。
// see: https://github.com/prisma/prisma-client-extensions/tree/900e0a6466fbbebc31205b314ecbbc8bd9dcff7b/row-level-security#caveats
function extend(tenantId: string) {
  return Prisma.defineExtension((prisma) =>
    prisma.$extends({
      query: {
        $allModels: {
          async $allOperations({ args, query }) {
            const [, result] = await prisma.$transaction([
              setTenant(prisma, tenantId),
              query(args),
            ]);
            return result;
          },
        },
      },
    })
  );
}

// トランザクションの実行時に、テナントIDをセットする拡張
// 上記の `extend` 拡張の問題を解決するために、別途 `$transaction()` を拡張する関数
// 注意: 上記の `extend` 拡張と結合して使用できない点に注意する
function extendTx(tenantId: string) {
  return Prisma.defineExtension((prisma) =>
    prisma.$extends({
      client: {
        $transaction: async (args: any) => {
          if (typeof args === "function") {
            return prisma.$transaction(async (tx) => {
              await setTenant(tx, tenantId);
              return args(tx);
            });
          }
          const [, ...results] = await prisma.$transaction([
            setTenant(prisma, uid),
            ...args,
          ]);
          return results;
        },
      },
    })
  );
}

// PrismaClientを拡張する
const prisma = new PrismaClient();
const tenantId = "org-1";
const client = prisma.$extends(extend(tenantId));
// `$transaction` 関数を強制的にオーバーライドして、テナントIDを付与する拡張を適用する。
client.$transaction = prisma.$extends(extendTx(tenantId)).$transaction;

簡単な解説

extend

Client extensions を使いクエリを拡張し、全てのクエリの実行前に、トランザクションの開始とテナントIDのパラメータを設定します。

function extend(tenantId: string) {
  return Prisma.defineExtension((prisma) =>
    prisma.$extends({
      query: {
        $allModels: {
          async $allOperations({ args, query }) {  // すべての操作を
            const [, result] = await prisma.$transaction([ 　// トランザクションでラップして
              setTenant(prisma, tenantId),  // 実行パラメーターをセットする
              query(args),
            ]);
            return result;
          },
        },
      },
    })
  );
}

extendTx

現状の Client extensions では $allOperations を使用したクエリ拡張時に、トランザクションが開始されているかどうかがわかりません。

そのため、$allOperations 内でクエリをトランザクションでラップすると、$transaction を含む全てのクエリにトランザクションが開始されてしまい、$transactionを使用したトランザクションが正しく動作しない問題があります。

そこで、$extends 関数だけを拡張する extension を定義します。

function extendTx(tenantId: string) {
  return Prisma.defineExtension((prisma) =>
    prisma.$extends({
      client: {
        $transaction: async (args: any) => {  // トランザクション関数で
          if (typeof args === "function") {  // Interactive transactionsの場合
            return prisma.$transaction(async (tx) => {
              await setTenant(tx, tenantId);
              return args(tx);
            });
          }
          const [, ...results] = await prisma.$transaction([
            setTenant(prisma, uid),
            ...args,
          ]);  　　// Sequential　operationsの場合
          return results;
        },
      },
    })
  );
}

拡張する

PrismaClient を拡張します。

この時、client.$transaction に extendTx 関数にて拡張した $transaction 関数を代入します。 (お行儀は悪いですが...)

こうすることで、$allOperations を使用した拡張とはコンフリクトせずに、$transaction を拡張します。

const prisma = new PrismaClient();
const tenantId = "org-1";
const client = prisma.$extends(extend(tenantId));
// `$transaction` 関数を強制的にオーバーライドして、テナントIDを付与する拡張を適用する。
client.$transaction = prisma.$extends(extendTx(tenantId)).$transaction;

まとめ

この記事ではPrisma ClientのExtensionsを使用して、RLSを実現する方法を紹介しました。

ポイントは次の通りです

全てのクエリの実行前に、テナントIDをトランザクションにて付与する拡張を追加する
上記の拡張のため、$transaction 関数内では2重にトランザクションが実行されてしまう。そこで、$transaction メソッドをオーバーライドし、拡張した関数内でトランザクションの開始とテナントIDの付与をする必要がある

備考

この記事では、$transaction 関数を強制的にオーバーライドする方法で、$transaction が正しく動作しない問題を回避しました。

ですが、公式例の注意書き^[2]にあるように、今後のバージョンでは、トランザクションの状態が判別できるようになる可能性があるため、このようなオーバーライドは不要になる可能性が高そうです。

NOTE: Because this example extension wraps every query in a new batch transaction, explicitly running transactions with companyPrisma.$transaction() may not work as intended. In a future version of Prisma Client, query extensions will have access to information about whether they are run inside a transaction, similar to the runInTransaction parameter provided to Prisma middleware. When this is available, this example will be updated to work for queries run inside explicit transactions

脚注

概要

バージョン情報

方法

拡張する

簡単な解説

extend

extendTx

拡張する

まとめ

備考

Discussion