aurora-dsql-samples

Aurora DSQL対応のテストを、本物のクラスタなしでCIで回す PostgreSQL用の宣言的スキーマ管理ツールpistachioを、Amazon Aurora DSQLに対応させました。--engine dsqlを付けると、plan、apply、dumpがDSQLを相手に動きます。 DSQLはPostgreSQL 16のプロトコルを話しますが、受け付けるDDLも、カタログが返す値もPostgreSQLとは少し違います。この違いをテストしたいものの、CIから本物のクラスタには接続したくありません。そこで、次の順にテストを組み立てました。 本物のクラスタの応答をプロキシで記録し、CIでリプレイする ローカルのPostgreSQLをDSQLのように振る舞わせるプロキシを作り、シナリオテストをCIで回す そのプロキシが本物からずれていないかを、記録とconformance testで確かめる この記事では、各段階で何が解決し、何が残ったのかを説明します。 DSQLはPostgreSQLとどう違うのか まず、pistachioの対応で問題になった違いを示します。すべて実際のクラスタで確かめたものです。 操作 PostgreSQL DSQL インデックスの作成 CREATE INDEX idx ON t USING btree (c) CREATE INDEX ASYNC idx ON t (c)のみ。job_idが返り、CALL sys.wait_for_job(...)で完了を待つ デフォルト値付きの列追加 ADD COLUMN c text DEFAULT 'x' 拒否される。ADD COLUMNとALTER COLUMN ... SET DEFAULTに分ける 既存テーブルへのUNIQUE制約の追加 ADD CONSTRAINT ... UNIQUE (c) CREATE UNIQUE INDEX ASYNCの完了を待ち、ADD CONSTRAINT ... UNIQUE USING INDEXを実行する 既存テーブルへのCHECK制約・外部キーの追加 ADD CONSTRAINT ... NOT VALIDで追加し、ALTER TABLE ASYNC ... VALIDATE CONSTRAINTで検証する トランザクション内のDDL 複数実行できる 1つのみ 文法の違いであれば、生成するDDLを書き換えれば済みます。厄介なのは、文法は通るのにカタログの読み取り結果が違う場合です。 たとえばPRIMARY KEY (id)で作ったテーブルについて、DSQLのカタログは次の定義を返します。 PRIMARY KEY (id) INCLUDE (name, email) 主キー以外の全列がINCLUDEに入るためです。そのまま比較すると、実際には差分がないのに、planを実行するたびに主キーを削除して追加するDDLが出力されます。しかもDSQLは主キーのDROP CONSTRAINTを拒否するので、そのDDLは適用もできません。インデックスのアクセスメソッドがbtreeではなくbtree_indexと報告されること、text列の圧縮方式がlz4になることも、同じ種類の問題です。 つまり、DDLがエラーにならないことを確かめるだけでは足りません。applyの後にもう一度planを実行し、差分が出ないことまで確かめる必要があり、それには本物に近い応答が要ります。 CIから本物のクラスタを使いたくない DSQLはAWSのマネージドサービスなので、Dockerで手元に立てることができません。テストのたびに本物へ接続することも考えられますが、CIで回すには次の問題があります。 費用: DSQLは使用量に応じて課金されます。CIはプッシュのたびに動くので、テストを増やすほど費用も増えます。 安全性: クラスタに接続するには、CIにAWSの認証情報を置く必要があります。公開リポジトリのCIに認証情報を置きたくありません。また、GitHub Actionsはフォークからのプルリクエストにシークレットを渡さないので、外部からのプルリクエストではテストが動きません。 同時実行: テストはpublicスキーマが空の状態から始めます。そのため、複数のCIジョブが1つのクラスタを同時に使うことができません。 応答を記録してリプレイする 最初に作ったのがpgstubです。PostgreSQLのワイヤプロトコルを話すスタブで、recordモードではプロキシとして動きます。 record: テスト -> pgstub -> DSQLクラスタ （応答をYAMLに保存） replay: テスト -> pgstub （YAMLから応答を返す） クラスタに接続して一度テストを実行すると、送ったSQLと応答がYAMLに保存されます。 - query: CREATE INDEX ASYNC users_name_idx ON public.users (name); columns: [job_id] rows: - [job-1] complete: CREATE INDEX - query: CALL sys.wait_for_job('job-1') columns: [succeeded] types: [bool] rows: - [t] complete: CALL このファイルをコミットしておけば、CIはクラスタもPostgreSQLもなしにTestDSQLをリプレイできます。現在の記録は232クエリ分、約12,500行です。カタログを読むクエリも応答も丸ごと入っているので、INCLUDEやbtree_indexを含む値を読んで差分が出ないことも確かめられます。 記録を運用するうえでは、工夫が2つ必要でした。 1つ目は、実行ごとに変わる値の扱いです。OIDとジョブIDは記録し直すたびに変わるため、何も変更していなくても差分が出ます。pgstubのNormalizeOIDsとNormalizeColumn("job_id", "job-")で、これらを出現順の番号に置き換えました。上のjob-1がその結果です。応答が同じならファイルを書き換えないので、記録ファイルの差分は本当に応答が変わった箇所だけになります。 2つ目は、記録ファイルをリポジトリに公開することへの対策です。make record-dsqlは記録の後でファイルを検査し、クラスタの名前やAWSの認証情報らしき文字列が含まれていれば失敗します。 記録だけでは足りなかったこと リプレイには、仕組み上の限界があります。 pgstubはSQLの中身を解釈せず、一致するSQLに記録した応答を返すだけです。pistachioが送るSQLを1つでも変えると、カタログのクエリであってもリプレイは失敗します。クラスタの権限がなければ記録し直せないので、DSQLのCIジョブは、記録を更新するまで失敗したままになります。 記録にない操作を試すこともできません。もう少し複雑なテストとして、aws-samples/aurora-dsql-samplesのスキーマを段階的に構築、変更、削除し、そのたびに差分が残らないことを確かめるシナリオテストを書きました。しかし、これは本物のクラスタがないと動かないので、CIでは実行できませんでした。 DSQLのように振る舞うプロキシを作る 既存のDSQLモックもいくつか試しましたが、使えませんでした。試したものはPostgreSQLのラッパーで、SQLは受け付けるものの、カタログがDSQLと異なるためです。pistachioはカタログを読んで差分を計算するので、INCLUDEやbtree_indexが再現されないと、DSQLでの動作を確かめられません。 そこで、SQLに応じて答えを返せるものとしてinternal/dsqlproxyを作りました。クライアントとローカルのPostgreSQL 16の間に入り、SQLを書き換えて中継します。 pista -> dsqlproxy -> PostgreSQL 16 テスト以外のコードは、setup.sqlを含めて約840行です。やっていることは、大きく3つです。 1つ目は、DSQLが拒否するSQLを同じエラーで拒否することです。SQLSTATEとメッセージをDSQLに揃えるため、PostgreSQL側の関数で例外を投げます。 -- クライアントが送ったSQL ALTER TABLE t ALTER COLUMN c SET NOT NULL -- PostgreSQLに送るSQL SELECT dsql_proxy.refuse('0A000', 'unsupported ALTER TABLE ALTER COLUMN ... SET NOT NULL statement') 2つ目は、ASYNCのジョブを再現することです。ASYNCはPostgreSQLの文法にないので、構文解析の前に正規表現で取り出し、関数の呼び出しに置き換えます。 -- クライアントが送ったSQL CREATE INDEX ASYNC idx ON t (c) -- PostgreSQLに送るSQL SELECT job_id FROM dsql_proxy.run_job('INDEX_BUILD', 'CREATE INDEX idx ON t (c)') run_jobはインデックスをその場で作りますが、pg_index.indisvalidをfalseにしておき、sys.wait_for_jobが呼ばれたときにtrueへ戻します。こうすると、ジョブを待たずに次へ進むバグがあった場合に、DSQLと同じく無効なインデックスが見えます。カタログを直接書き換えるので、プロキシが接続するロールにはスーパーユーザー権限が必要です。 3つ目は、カタログの値をDSQLに合わせることです。pg_get_constraintdefやpg_get_indexdefの呼び出しを、dsql_proxyスキーマに定義した同名の関数に置き換えます。これらの関数がINCLUDEやbtree_indexを付け足します。 -- PRIMARY KEYの定義に、主キー以外の全列をINCLUDEとして足す SELECT pg_catalog.pg_get_constraintdef($1, $2) || ' INCLUDE (' || inc.cols || ')' これで、シナリオテストをCIで実行できるようになりました。 プロキシが本物からずれていないか確かめる 自作のモックには、テストが通っても本物で動く保証がないという問題があります。プロキシが間違っていれば、そのプロキシで通ったテストに意味はありません。そこで、本物から得た結果と比較する仕組みを2つ用意しました。 記録を比較する 本物のクラスタで記録したTestDSQLを、プロキシ経由でも記録して比較します。 TestDSQL -> pgstub -> dsqlproxy -> PostgreSQL 16 できた記録がコミット済みのdsql.yamlと1バイトでも違えば、CIのdsql-recordingジョブが失敗します。232クエリ分の応答について、プロキシがクラスタと同じ答えを返していることになります。CIで実行できるので、クラスタは要りません。 拒否されるSQLを本物で確かめる この比較には抜けがあります。記録にはエラー応答が1件も入っていないので、プロキシが返す拒否メッセージを検証できません。 そこで、TestConformanceというテストを作りました。SQLごとに、受け付けられるかどうかと、拒否される場合のSQLSTATEとメッセージを、約60件のテーブルにまとめました。 {sql: "CREATE INDEX probe_x_idx ON probe_t (c)", code: refused, msg: "unsupported mode. please use CREATE INDEX ASYNC."}, {sql: "CREATE INDEX ASYNC probe_x_idx ON probe_t (c)", undo: "DROP INDEX probe_x_idx"}, 同じテーブルを、普段のテストではプロキシに対して、記録を更新するときは本物のクラスタに対して実行します。 初めて本物のクラスタで実行すると、プロキシの拒否メッセージが8件間違っていました。たとえば次のような間違いです。 SQL プロキシの応答 DSQLの応答 ALTER COLUMN ... SET NOT NULL unsupported ALTER TABLE SET NOT NULL statement unsupported ALTER TABLE ALTER COLUMN ... SET NOT NULL statement CREATE TYPE ... AS ENUM unsupported statement: CreateEnum CREATE TYPE not supported プロキシは手で書いたものなので、こうした間違いが入り込みます。本物に対して実行して、初めてそれが見つかりました。 現在のテスト構成 テスト 相手 実行する場所 確かめること make test-dsql pgstubによるリプレイ CI 本物の応答に対してpistachioが正しく動くか make test-dsql-proxy dsqlproxy + PostgreSQL 16 CI 実際のスキーマの構築、変更、削除で差分が残らないか make check-dsql-recording pgstub + dsqlproxy + PostgreSQL 16 CI プロキシの応答が本物の記録と一致するか TestConformance dsqlproxy、または本物 CI（プロキシ）、手元（本物） SQLの受け付けと拒否が想定どおりか make record-dsql 本物 手元 記録の更新とconformance test make test-dsql-scenario 本物 手元 シナリオテストが本物でも通るか プロキシでは分からないこと プロキシは、本物で確かめて分かったことしか再現できません。次の点は、本物でなければ確かめられません。 DSQL自体の変化: DSQLは機能が増え続けています。調査中にも、DROP COLUMNなどが新たに使えるようになりました。こうした変化はCIでは検出できず、make record-dsqlを実行したときに初めて分かります。 非同期の実行: プロキシのジョブはその場で完了します。ビルドに時間がかかる場合の挙動は再現しません。 接続まわり: プロキシはTLSもIAM認証も扱いません。5分のトランザクションの上限と60分の接続の上限も再現していません。 プロキシに期待しているのは、本物を完全に再現することではなく、現状のDSQLの仕様をテストで固定することです。DSQLの変化を取り込むときや送るSQLを変えたときには、本物に対してmake record-dsqlを実行します。結果が違っていれば、テストのテーブルとプロキシを直します。 まとめ Aurora DSQLは手元に立てられず、CIから本物のクラスタに接続したくもありませんでした。そこで、本物の応答を記録したリプレイと、DSQLのように振る舞うプロキシを用意し、本物のクラスタなしでCIのテストを通せるようにしました。本物のクラスタに接続するのは、記録を更新するときなど、手元で実行するときだけです。 実装はpistachioのinternal/dsqlproxyとpgstubのリポジトリにあります。