maybe daily dev notes

私の開発日誌

トップ > AWS > Aurora DSQLをPrismaで使う

Aurora DSQLをPrismaで使う

AWS Backend 開発 Lambda

次世代のサーバーレスRDBであるAurora DSQLを、TypeScript用ORMのPrismaと一緒に使う話です。DSQLドンドン使っていきたい！

DSQL x Prismaは実現できるのか

そもそも、今のPrismaはDSQLを扱えるでしょうか？

幸いにもDSQLは多くの面でPostgres互換を実現しているため、PrismaからもPostgresデータベースとして接続可能です。schema.prisma の datasource.provider を postgresql にすればOKです。

// schema.prisma
datasource db {
  provider = "postgresql"
  url      = env("DATABASE_URL")
}

ちょっとした model を追加してみましたが、正常にマイグレーションやデータ読み書きを実行できることが確認できます。

model User {
  id String @id
}

しかしながら、ひとつ考慮事項がありました。それは、上記の 環境変数 DATABASE_URL をどう渡すか？ です。この記事では、主にこの点についてフォーカスします。

DSQLと環境変数 DATABASE_URL の課題

DSQLでは、データベース接続で利用するパスワード文字列として、動的に生成することを前提とした authentication token (以下トークンと呼ぶ) を利用します。以下は動的にトークンを取得してデータベースURLを生成する例です:

import { DsqlSigner } from '@aws-sdk/dsql-signer';

const hostname = 'example.dsql.us-east-1.on.aws';
const signer = new DsqlSigner({
  hostname,
  expiresIn: 24 * 3600 * 7, // 期限は最長1週間
});
const token = await signer.getDbConnectAdminAuthToken();
const url = `postgres://admin:${encodeURIComponent(token)}@${hostname}:5432/postgres`;

一方Prismaでは、パスワード文字列は静的に環境変数として渡されることを(暗に)前提としています。参考: Connection URLs

この2つの食い違いが、DSQL x Prismaの利用体験を損ねる可能性があります。以下では、どのように実装すれば良い感じになるかを見ていきましょう。

実装例の紹介

いくつか考えた・見つけた実装のアイデアを紹介します。

実装例1. 動的にPrismaClientを生成する

動的に環境変数 DATABASE_URL を設定した後、PrismaClientを初期化する方法です。

// prisma.ts
import { DsqlSigner } from '@aws-sdk/dsql-signer';
import { PrismaClient } from '@prisma/client';

const hostname = 'example.dsql.us-east-1.on.aws';

async function generateToken() {
  const signer = new DsqlSigner({
    hostname,
  });
  return await signer.getDbConnectAdminAuthToken();
}

export const getClient = async () => {
  const token = await generateToken();
  process.env.DATABASE_URL = `postgres://admin:${encodeURIComponent(token)}@${hostname}:5432/postgres`; 
  // PrismaClientのコンストラクタ内で、上記環境変数が参照される
  return new PrismaClient();
};

// 呼び出し側
import { getClient } from './prisma';

const main = async () => {
  const prisma = await getClient();
  await prisma.user.findMany();
}

これは最もstraightforwardな方法だと思います。欠点としては、従来の使い方と比べて、以下が追加で必要になる点です:

  1. PrismaClientを非同期関数越しに取得する必要がある
    • 静的な環境変数を使っていたならただの変数として取得できるので、やや使い勝手は変わります
    • 例: await prisma.user.findMany → await (await getClient()).user.findMany など
  2. 再接続の処理が必要になる
    • トークンの期限が切れるとデータベースに新たに接続できなくなります (既存の接続は利用できるようです)
    • このため、期限切れ前にPrismaClientを再度初期化する必要があります
    • なおトークン期限は最長1週間まで指定できるので、AWS Lambdaなど、比較的短命なランタイムでは気にする必要がないかもしれません

参考までに、再接続も意識したコード (getClient 関数のみ抜粋) は以下のようになるでしょう:

let client: PrismaClient | undefined = undefined;
let lastEstablieshedAt = 0;
export const getClient = async () => {
  if (client) {
    // トークンの期限切れより前に更新する (以下は例として1時間)
    if (Date.now() - lastEstablieshedAt < 1 * 3600 * 1000) {
      return client;
    } else {
      await client.$disconnect();
    }
  }
  lastEstablieshedAt = Date.now();
  const token = await generateToken();
  process.env.DATABASE_URL = `postgres://admin:${encodeURIComponent(token)}@${hostname}:5432/postgres`;
  client = new PrismaClient();
  return client;
};

ちなみに、top-level awaitを使えばPrismaClient自体をexportできるため、もう少し使いやすくなります。しかし再接続の実装は難しくなるかもしれません。

// prisma.mts
const getClient = async () => {
  // 省略
  return new PrismaClient();
};

// 呼び出し側はこちらを使う
import { prisma } from './prisma.mts';
export const prisma = await getClient();

実装例2. authentication tokenを環境変数として埋め込む

2つ目は、動的なトークンを静的なものとして扱う方法です。以下のような仕組みを作れば実現できるはずです (未実装)。

トークンを取得するLambdaを作成し、そのLambdaからアプリケーション本体 (これもLambdaにあるとする) の環境変数を書き換えます。 このLambdaをトークンが期限切れするよりも前に定期的に呼び出します。

アプリケーションがECSの場合は、Secrets Managerから動的に環境変数を埋め込めるため、以下のような構成もありでしょう (タスク定義を直接更新しない):

この方法の利点は、Prisma視点ではデータベースURLの文字列が環境変数から得られる静的な文字列となるため、従来と全く同じ使用感を実現できる点です。

一方、欠点は以下が考えられます:

  • Lambdaの環境変数に認証情報を直書きすることを推奨しない組織もある
  • トークンを更新するたびにコールドスタートが生じる
    • せいぜい数時間おきなので、大した影響はないだろうが
  • トークンを更新する仕組みの実装・管理が必要
    • 認証情報のローテーションの仕組みと似たようなもので、一度作ればほぼ管理不要だが、一定の面倒くささはあり
    • CDKコンストラクトなどが再利用可能モジュールがあると嬉しいかも？

とはいえアプリ側の実装が単純になるのはやはり嬉しいものです。

実装例3. PrismaのDBドライバーを node-postgres に差し替える

PostgresのPrismaでは、DBドライバーをPrisma独自のものでなく、node-postgres (pg) に差し替えることができます。node-postgresではパスワードとしてasync関数を指定できるため、素直に今回の処理を実装できます。以下はコード例です:

// schema.prisma
generator client {
  provider = "prisma-client-js"
  // feature flagを有効化
  previewFeatures = ["driverAdapters"]
}

// index.ts
import { PrismaPg } from '@prisma/adapter-pg';
import { PrismaClient } from '@prisma/client';
import { Pool } from 'pg';
import { DsqlSigner } from '@aws-sdk/dsql-signer';

const hostname = 'example.dsql.us-east-1.on.aws';

const pool = new Pool({
  host: hostname,
  user: 'admin',
  database: 'postgres',
  port: 5432,
  ssl: true,
  password: async () => {
    const signer = new DsqlSigner({
      hostname,
    });
    return await signer.getDbConnectAdminAuthToken();
  },
});
const adapter = new PrismaPg(pool);
const prisma = new PrismaClient({ adapter });
const users = await prisma.user.findMany();

個人的には、この方法が最も単純で美しいように思います。新規接続のたびにpassword関数が呼び出されるため、期限切れの問題も自然と解決されています。

注意点としては、DBドライバーそのものが置き換わるため、Prismaの動作に影響が生じうることです。私自身はPrismaとnode-postgresを合わせて利用したことがないため、どの程度挙動に差異があるのかは不明です。まだpreviewの機能でもあることから、漠然とした不安のある選択肢だとは思います。

実はDSQL以前からあった普遍的な問題なのだ

諸々見てきましたが、実はこの問題、DSQLに限った話ではありません。

例えば以下に挙げる場合において、開発者たちは以前から同じ問題に直面していたはずです:

  • RDSのIAM認証で接続する
  • DB認証情報を環境変数でなく外部ストアから取得する
  • DB認証情報をSecrets Managerなどでローテーションする

Prismaでは3年以上前からこの問題について指摘されていますが、まだ解決はしていないようです。

Prisma特有の困難として、クエリエンジンがRustで書かれているために、JavaScriptでpasswordを取得するasync関数を書かれても、それを実行しづらいという点が挙げられています。最近はPrismaのRust部分をTypeScriptに移行するという話もある (全く別の文脈ですが) ので、そこに少し期待したいですね。

We’re addressing this by migrating Prisma’s core logic from Rust to TypeScript and redesigning the ORM to make customization and extension easier.

まとめ

Aurora DSQLとPrismaを併用する上で、特に認証情報 (authentication token)をどう扱うかについて考えました。Prisma特有？の困難は見えましたが、接続さえできれば普通に使えるので、今後サーバーレスPostgresとして活用していければと思います。

なお、検証に用いたコードはこちらに公開しています: aurora-dsql-prisma-example

今月のもなちゃん

引っ越しで内装が様変わりする中、かつての安息の地をディスプレイ下に見出した様子です。

ではまた！