Dataformワークフローをシェルスクリプトで動かす

2024.06.12

Dataformを何回も起動して検証データ作成するということをする際に、シェルスクリプトでサクッと自動化できないかなと思いやってみました。

やりたいこと

  • コンパイル変数をシェルスクリプト内で動的に設定し、コンパイルしてDataformのワークフロー実行

今回は例として、コンパイル変数を受け取りテーブルにInsertするという簡単なDataformワークフローを作成して、そのワークフローに動的にコンパイル変数を渡してワークフローを実行するということをしてみました。

前準備

Dataform側を設定してしまいます。

test.sqlx

config {
  type: "operations"
}

INSERT INTO `データセットID.shelltest_table` (yyyymmdd_str)
VALUES
  (${dataform.projectConfig.vars.compile_val});

shelltest_tableというテーブルをあらかじめ作成しております。構造としては、列名yyyymmdd_strをString型で持つテーブルとなります。
また、コンパイル変数を用いるのでdataform.jsonにコンパイル変数の設定をします。

dataform.json

{
  "defaultSchema": "データセットID",
  "assertionSchema": "dataform_assertions",
  "warehouse": "bigquery",
  "defaultDatabase": "プロジェクトID",
  "defaultLocation": "asia-northeast1",
  "vars": {"compile_var":"'default_val'"}
}

コンパイル変数としてcompile_varを定義して、デフォルト値もdefault_valと設定しています。
上記ワークフローをそのまま実行すると、shelltest_tableにはdefault_valという値がInsertされます。 このテーブルに対して、シェルスクリプトを通して指定期間の日付の値をInsertしてみたいと思います。

シェルスクリプトの解説

以下がシェルスクリプト全文です。

#!/bin/bash
# 開始日と終了日を指定(yyyy-mm-dd形式)
START_DATE="2023-01-01"
END_DATE="2023-01-10"

# プロジェクトIDとリポジトリ名を指定
PROJECT_ID="プロジェクトID"
REPO_NAME="リポジトリ名"

# 日付フォーマット変換関数
format_date() {
  date -d "$1" +"%Y%m%d"
}

# 開始日と終了日をyyyymmdd形式に変換
start_date=$(format_date "$START_DATE")
end_date=$(format_date "$END_DATE")

# 開始日から終了日までの期間をループ
current_date="$START_DATE"
while [ "$(format_date "$current_date")" -le "$end_date" ]; do
  # 日付をyyyymmdd形式で出力
  yyyymmdd_val=$(format_date "$current_date")
  yyyymmdd="\'${yyyymmdd_val}\'"
  echo "Processing date: $yyyymmdd"
  comp_result_response=$(curl -X POST "https://dataform.googleapis.com/v1beta1/projects/$PROJECT_ID/locations/asia-northeast1/repositories/$REPO_NAME/compilationResults" \
    --header "Authorization: Bearer $(gcloud auth print-access-token)" \
    --header "Content-Type: application/json" \
    --data "{\"gitCommitish\":\"main\", \"codeCompilationConfig\":{\"vars\":{\"compile_val\":\"$yyyymmdd\"}}}")

  # レスポンスからcompilationResultsのname値を抽出
  comp_result_name=$(echo "$comp_result_response" | jq -r '.name')
  echo "Compilation Result Name: $comp_result_name"

  curl "https://dataform.googleapis.com/v1beta1/projects/$PROJECT_ID/locations/asia-northeast1/repositories/$REPO_NAME/workflowInvocations" \
    --header "Authorization: Bearer $(gcloud auth print-access-token)" \
    --header "Content-Type: application/json" \
    --data "{\"compilationResult\":\"$comp_result_name\"}"

  current_date=$(date -I -d "$current_date + 1 day")
done

さらっとコードの解説をします。
1. START_DATEとEND_DATE変数に開始日と終了日をyyyy-mm-dd形式で指定します。
2. PROJECT_IDとREPO_NAME変数にプロジェクトIDとリポジトリ名を指定します。
3. format_date関数を定義し、指定された日付をyyyymmdd形式に変換します。
4. format_date関数を使用して、開始日と終了日をyyyymmdd形式に変換します。
5. current_date変数に開始日を設定し、終了日までの期間をwhileループで処理します。
6. 各日付をyyyymmdd形式で生成し、DataformAPI(コンパイル用)を実行します。
7. curlコマンドのレスポンスからcompilationResultsのname値をjqを使用して抽出します。
8. 抽出したname値を次のcurlコマンドに渡してDataformのワークフローを実行します。
9. 次の日付を計算します。

上記スクリプトをコマンドラインで実行すると、指定日付の日数分(2023年1月1日-10日までの日付)Dataformワークフローが起動されます。
また、コンパイル時にコンパイル変数に$yyyymmdd変数を渡しているので、日付ごとにDataformワークフローが実行されます。
コンパイルされて実行されるSQLとしては以下のようになります。

INSERT INTO `test_dataform.shelltest_table` (yyyymmdd_str)
VALUES
  ('20230101');

BigQueryに格納されるデータとしては以下のようになります。

このスクリプトでできないこと

Dataformのワークフロー実行APIは非同期実行となり、ワークフロー実行APIを呼び出した戻り値はワークフローの実行状態となります。
よって、ワークフロー終了後に後続の処理をさせるというような場合やワークフロー終了後にさらに別のワークフローを呼び出したいという場合はポーリングする実装を追加する必要があります。
また機会があればポーリングの実装を追加したシェルスクリプトを作ってみたいと思います。

参考

Dataform API