maybe daily dev notes

私の開発日誌

Aurora Postgres Data APIをあらゆるORMから使う試み

あけましておめでとうございます。冬休みの自由工作レポートを提出します。

はじめに

最近Amazon RDS AuroraでData APIが使えるようになりました。Auroraインスタンスに対してHTTP APISQLクエリを発行できる便利なものです。

この記事では、Data APIをより使いやすくするための方法を検討します (ネタバレ: 目標未完です) 。

Data APIのおさらい

Data APIに関する知識を箇条書きでまとめます。

メリット

  • 踏み台なしにインターネットからクエリ可能 (IAM認証)
    • 不要ならData API自体を無効化できる (デフォルトで無効)
    • IAM認証なので、DB認証情報の管理が不要になるのも嬉しい
  • Data API側でコネクションプールされる
    • Lambdaでもコネクション枯渇の可能性が抑制できる
    • RDS Proxyが不要になり、コスト減の可能性
  • CloudTrailで監査ログを記録可能

デメリット

個人的には踏み台とRDS Proxy不要のメリットが大きいと考え、Data API活用の方法を探ることにしました。

モチベーション

Data APIは独自のHTTPインターフェスを介する必要があるため、通常のORMライブラリなどではそのまま利用できません。このため、一部のライブラリではData API専用のアダプターが開発されています ( 例: kysely-data-api, typeorm-aurora-data-api-driver)。

しかしながら、これらのアダプターは各ライブラリ専用のものであり、他のORMでは利用できません。例えばTypeScriptのORMであるPrismaは、3年以上前からData API対応のissueが存在するものの、未だに実現されていません。

github.com

そもそもData APIアダプタの事例は少ないため、他のORMでも似た状況にあるものは多いと推察されます。

そこで今回はより汎用的な解決策として、下図のアイデアを考えました。

図中のData API Proxyが今回作りたいモノです。これはAuroraクラスタとクライアントの間に存在し、クライアントからのSQLリクエストをData APIに変換し、Data APIとしてAuroraクラスタに送信します。また、Data APIのレスポンスをPostgres/MySQLネイティブの形式に変換し、クライアントに返します。

これにより、クライアント側から見るとあたかも通常のPostgres/MySQLと通信しているように、Aurora Data APIを利用できます。つまり、従来のORMでもアダプターなしにそのまま、AuroraクラスタとData APIで通信できるわけです。便利そうですね。*1

この記事では上記のようなプロキシの実現を目指しました。以下で「プロキシ」と書かれたものは、これを指すことにします。 なお、2023/12現在Data APIはPostgresでのみ利用可能 (serverless v1を除く) なため、以下はPostgresの話に限定します。

Postgresサーバーとして振る舞わせる方法

このプロキシは、クライアントから見ると通常のPostgresサーバーとして振る舞う必要があります。つまり認証やクエリのインターフェースを模擬しなければいけません。

これを実現する方法を調べると、うってつけのフレームワークがありました:

github.com

PgwireはPostgresのインターフェースを実装したフレームワークで、開発者はクエリを受信した際の処理だけ実装すれば、自分だけのPostgresのサーバーを作れるという代物です。例えばこれは、Postgresインターフェースで使えるSQLiteサーバーの実装です: sqlite.rs

これを使えばまさにやりたいことが実現できそうですね。

そう簡単ではなかった

Pgwireのおかげで、シンプルなクエリをプロキシする程度のものは簡単にできました。しかし!実装を進める中でいくつかの課題が見つかり、一部はクリティカルに思えたので、作業を止めています。

以下では見つかった課題を簡単に紹介します。私自身はPostgres素人なので誤りを含む可能性がありますが、ご参考までに。

Extended Queryどうする?

Postgresがクエリを処理するフローには2種類あります: Simple QueryとExtended Queryです (詳細)。

Simple Queryは名前の通り単純なもので、SQLクエリを受け取って結果を返すだけです。 一方、Extended Queryはクエリの処理をいくつかのステップに分解してリクエストすることができます。(Parse, Bind, Execute, Sync + Describe, Close)。

難しいのは、Data APIはSimple Queryのモデルに近い点です。つまり、SQLクエリを投げてリクエストを返すことしかできないため、Extended Queryのリクエストに対しては、何らかの方法で、プロキシ内でレスポンスをでっち上げる必要があります。

実はPgwireはよくできており、Pgwire内でクライアントごとのストアを管理し、Parse/Bindで作成されたSQLクエリをExecute時に取り出してくれる仕組みが存在します (このあたり)。これにより、Pgwire利用者はそれほどExtended queryのフローを意識する必要なく、ただクエリに対する処理を書けば良いのです。

問題はDescribeです。これはクエリを実行する前に、クエリが返す行の型情報を取得するためのメッセージです。現状Data APIではこれに相当するAPIが用意されておらず、Describeメッセージに対する妥当なレスポンスをプロキシ内で作るのが難しいのです。適当な型情報が返せば良いかと試したところ、それでエラーになるクライアントもいるようでした。クエリを一回実行すれば型情報を得られるのですが、冪等でないクエリやINSERTのRETURNINGなどを考えると、そう簡単には実装できなそうです。

こうなるとExtended queryのサポートを諦めるのが現実的に思えますが、今回最も狙っていたORMのPrismaはExtended queryをガッツリ利用しているようでした (実際の挙動を見る限り)。他のORMがどうなのか調べられていないのですが、私が一番使うのがPrismaなのでモチベを失っています。

データ型の変換どうする?

Data APIで返されるレコードは、Data API独自のルールでデータ型が変換されています ( 詳細はこちら )。 例えば、DecimalやTimestampが文字列になるなどです。

プロキシでは、これをPostgresネイティブの型に復元してレスポンスを作る必要があります。元のデータ型はメタデータとして取得できるため、この情報を使えば問題なく復元できそうです。既存のORMアダプター実装 ( rds-data-api-client-library など) が参考になるはず。

全てのデータ型に対応するのは大変そうなので、プロキシの制約として妥協することになるかもしれません。

Transactionどうする?

Data APIは独自のフローでトランザクションを管理するため、透過的に扱うにはプロキシ側で一工夫必要です。基本的にはこのようなもので実現できるでしょう:

  1. BEGIN クエリを受け取ると、 beginTransaction APIを叩き、得られたTransaction IDを保存する
  2. 以降同じクライアントからクエリが来たら、保存したTransaction IDを付加して executeStatement APIを叩く
    • クライアントの識別は、socket_addr を見るのが良さそうです (参考)。ポート番号付きなので、コネクションを識別できる。ただし、ポート番号は使い回されるので、コネクション切断時にうまくクリーンアップする必要あり。
  3. COMMIT クエリを受け取ると、commitTransaction API を叩く
  4. ROLLBACK クエリを受け取ると、 rollbackTransaction API を叩く
    • ROLLBACK クエリ以外にロールバックが必要なシナリオはあるのか、要確認

まだ絵に描いた餅なので、いざ実装すると考慮漏れもあるかもしれません。

Prepared statementどうする?

Prepared statementはパースされたクエリをサーバー側に保持し、同一コネクション内で使い回すことで、同様のクエリの呼び出しを効率化する機能です。副産物としてSQLインジェクション対策にもなるため、例えばPrismaではデフォルトでPrepared statementを利用します。

しかしながら、試した限り、今のところData APIではPrepared statementをうまく使えないようです。(ドキュメントには明記されていないが、DBコネクションの概念がないので、それはそうだと思われる。)

このため、Prepared statementをプロキシで非対応とするのが一つの妥当な選択肢です。

あるいは、PgwireがExtended queryフローでサポートしているように、擬似的にアダプタ内でサポートすることもできるかもしれません。 PREPARE を受け取るとプロキシ内でステートメントと名前とクライアントIDを保持し、EXECUTE を受け取るとステートメントに引数を入れてData APIに投げるようなものです。……と書いてて思いましたが、これをやるためにはSQLクエリのパースが必要なため、Extended queryの方法より難易度が高そうです。あまり現実的ではないかもしれません。

まとめ

まとまりませんが、まとめです。Aurora Data APIを巷のORMで利用可能にするための汎用的な解決策として、Data API Proxyを考えました。しかしながら、いくつかの問題があり、真に汎用的で実用的なものにすることは難しそうです。

Extended queryの話はPostgres限定のため、MySQLだとまた話は少し変わるかもしれません (それでもPrepared statementの話は残るのですが)。Data APIMySQLにも対応したら、改めて見直してみたいと思います。*2

ちなみに作りかけたものはコチラに貼っています

*1:プロキシのインフラ管理が手間ではと思われる方もいるかもしれませんが、心配無用です。プロキシはクライアント間で共有する必要がないため、クライアント側にサイドカーなどとして配置でき、追加のインフラは不要です。

*2:MySQLにも似たようなフレームワークがあることだけは調査済みです https://github.com/jonhoo/msql-srv