Redshiftのストアドプロシージャで単体試験自動化をやってみた

データウェアハウス開発部の川島です。現在は健康診断データ（以下、健診データ）のAWS移行のプロジェクトに参画しています。健診データの詳細については、以前の記事をご覧ください。

techblog.jmdc.co.jp

今年、JMDCではアドベントカレンダーに参加しています。

本記事は、JMDC Advent Calendar 2024 15日目の記事です。

はじめに

健診データ基盤について、オンプレミスからAWSへのクラウド移行を進めており、データの変換やチェック等の基盤としてRedshiftを採用しています。 Redshiftには、ストアドプロシージャやユーザー定義関数が存在しており、プログラム中でこれらのデータベースオブジェクトを使用しています。

クラウド化にあたって、オンプレ時代に環境構築やリリース作業でコストがかかっていた点を見直すため、CodePipelineを利用した自動リリース機能を開発しました。
内容としては、Github内に存在するデータベースオブジェクトのCreate文が記載されたSQLファイルを、CodeBuild内でRedshift DataAPI経由して実行することにより、Githubにコミットすることで自動でデータベースオブジェクトが作成されるというものになります。

テストの自動化を考えた理由

クラウド移行を進める中で自動リリースを実現したものの、依然としてデータベースオブジェクトのテストは手動で実施していました。そのため、下記の課題を抱えていました。

データベース内の各テーブルのデータを試験内容に適したものに更新し、証跡として取得する手間が大きい。
制度改定等によるデータ仕様の変更やチェック基準の見直しにより、プログラム（SQL）を変更する際のデグレード検証に掛かるコストを削減したい。
Redshiftは自動でバージョンアップがなされるため、バージョンアップによるプログラムへの影響を容易に検証できるようにしたい。

そのため、この自動リリースの仕組みを流用し、Redshift DataAPI経由でプログラムだけではなくテストコードもCreateし、実行することでテストそのものも自動化できるのではないかと思いました。

やってみたこと

サンプルとして、テーブルAからBへデータを転送するストアドプロシージャが存在するとします。

CREATE OR REPLACE PROCEDURE  dev.SP_01(I_ID IN BIGINT)
AS $$
DECLARE
BEGIN 
    DELETE FROM dev.TABLE_B;
    INSERT INTO dev.TABLE_B (ID,COL1) SELECT I_ID,COL1 FROM dev.TABLE_A;
END;
$$ LANGUAGE plpgsql;

このストアドプロシージャに対するテストコードの例は、下記になります。

CREATE OR REPLACE PROCEDURE  dev.TEST_SP_01()
AS $$
DECLARE
    V_ACTUAL_CNT                         BIGINT;
BEGIN 
    --Arrange
    DELETE FROM dev.TABLE_A;
    DELETE FROM dev.TABLE_B;

    INSERT INTO dev.TABLE_A (COL1) VALUES ('AA'),('BB'),('CC');
    INSERT INTO dev.TABLE_B(ID,COL1) VALUES (1,'XX');

    --Act
    CALL dev.SP_01(100);

    --Assert
    WITH V_EXPECT AS (SELECT * FROM (SELECT 100 AS ID, 'AA' AS COL1) UNION ALL
                     (SELECT 100 AS ID, 'BB' AS COL1) UNION ALL
                     (SELECT 100 AS ID, 'CC' AS COL1)),
        V_ACTUAL AS (SELECT ID,COL1 FROM dev.TABLE_B)
    SELECT COUNT(*) INTO V_ACTUAL_CNT 
    FROM  (SELECT * FROM    (SELECT * FROM V_ACTUAL MINUS SELECT * FROM V_EXPECT)
                             UNION ALL
                            (SELECT * FROM V_EXPECT MINUS SELECT * FROM V_ACTUAL)
        );
    IF V_ACTUAL_CNT <> 0 THEN
        RAISE EXCEPTION 'ERROR';
    END IF;

END;
$$ LANGUAGE plpgsql;

Arrangeセクションでは、テーブルを初期化（全件Delete）した上で、TABLE_Aにはテストデータを、TABLE_BにはダミーデータをInsertします。
Actセクションでは、試験対象のストアドプロシージャをCallします。
Assertセクションでは、MINUSを用いて、想定される結果（V_EXPACTに記載）と実際にTABLE_Aに登録された結果を比較します。MINUSした結果が0件であれば両者が完全一致しているため、正常に動作していることになります。

今回のケースではデータパターンが少なく少数のデータで結果の検証できるために、Arrangeセクションは直接Insert文で済ませていますが、データ量が多い場合はテストコードと合わせてCSVファイルをGithubに配置し、AWS CLIでS3に転送してからCOPY等で対象テーブルにセットすることも可能です。

そして、このストアドプロシージャをRedshift DataAPI経由でCallします。
CodeBuild内でRedshift DataAPIを実行するには、いくつかの方法が考えられるかとは思いますが、今回は一番簡単なShellScriptで記載してみます。
なお、APIで指定しているDBへの接続情報については、CodeBuildの環境変数で渡しています。

#!/bin/bash
i_sql="CALL dev.TEST_SP_01();"

EXECUTE_EXEC_ID=$(aws redshift-data execute-statement \
    --region ${AWS_REGION} \
    --secret ${SECRETS_ARN} \
    --workgroup-name ${WORKGROUP_NAME} \
    --sql "${i_sql}" \
    --database ${DATABASE} \
    --query 'Id' \
    --output text)

EXEC_STATUS=PRE
while [ ${EXEC_STATUS} != "FINISHED" ]; do
    EXEC_STATUS=$(aws redshift-data describe-statement \
        --id ${EXECUTE_EXEC_ID} \
        --query 'Status' \
        --output text)

    ##エラー時は抜ける
    if [ ${EXEC_STATUS} = "FAILED" ]; then
        exit 1
    elif [ ${EXEC_STATUS} = "ABORTED" ]; then
        exit 1
    fi

    sleep 5

done
exit 0

execute-statementでCall文を実行し、describe-statementでSQL文の実行結果が正常終了していれば、このShellScriptも正常終了させます。
もし、異常系の検証をする場合であれば、describe-statementでqueryの指定を「Error」とすることで、例外の内容を取得できます。取得した例外の内容に対してgrepコマンドを用いるなどで、エラー内容が想定通りであるかを確認できます。
ShellScriptが異常終了しますと、CodeBuildの結果が異常になりますので、結果の確認は容易です。