Amazon Athena を自動テストするためのメモ

2021年11月24日

概要

技術よりのメモ

Athena を仕事で使っているのですが、AデータベースとBデータベースにクエリを投げて、その Count とか抜粋したデータが一致しているというようなテストをする場面がありました。
一度だけだったら手動でいいのですが、データやテーブル定義の変更などが発生する度にこの確認を行うのは面倒です。

そこで、定期的に Athena を自動実行できるテストができないかということで調べてみました。
ひとまずローカルからプログラムでテストできるようにすれば、定期実行もできるようになるはずです。

ということで、ローカルから Athena のテストをしたいわけで調べてみたのですが、aws credential を設定済みであればできそうであることがわかりました。

ローカルから Athena の実行

Athena をローカルから実行するためには、以下のクラスメソッドさんの記事が大変参考になりました。

こちらに載っているコード(athena.py)を流用することで、実行できると思います。

retrying を使っているので、pip install する必要があります。

実行した結果の比較

結果は csv として出力されるので、pytest で csv の比較を行いたいです。

先程の athena.py を import して、query_to_athena を実行してあげれば良さそうです。

単純な csv の比較は csv-diff というプラグインでできそうでした。

以下のようなコードで比較ができます。

from csv_diff import load_csv, compare
diff = compare(
    load_csv(open("one.csv"), key="id"),
    load_csv(open("two.csv"), key="id")
)

diffの各キーが空であれば、差分がないということになるので、そのようにpytestを書きます。

課題としては、カラムやデータの整形が必要な場合です。

カラム名に関しては、Athenaで実行するSQL側に AS句を入れるなどして対応ができそうですが、データの整形が必要になってくるような場合には、pandasを用いる等を検討したほうが良いでしょう。

pandasを使った csvの比較は以下の記事が参考になりそうです。

まとめ

ローカルからpytest を通じて Athenaを実行し、実行結果が期待したcsvと一致しているかどうかを調べることは可能でした。

Athenaにおいても、複数テーブルに同じようなSQLを投げるような確認が何度か必要なケースでは、プログラム化を検討したほうが良さそうです。

テストの実行時間は結構かかるようでしたが、恩恵は受けられるような気がします。