DataformでBigQuery予約・割り当てを自動化

あるお客様は、BigQuery上で研究開発(R&D)プロジェクトを運用しており、開発チームとデータサイエンスチームの双方が利用しています。両チームの利用パターンには次のような特徴があります。

1. インタラクティブクエリは、主に平日の業務時間中に集中する。
2. 同時実行クエリ数が多く、短いレスポンスタイムが求められる。
3. 特定のタスクは夜間や週末に実行されるようスケジュールされている。

これらの要件を満たすため、業務時間中はクエリのレイテンシーと高い同時実行性を確保できる「On-Demand」プランを利用し、業務時間外は最大100スロットの「Enterprise Edition」に切り替えます。

料金プランの切り替えは、自動化フローで実現します。

現時点で、BigQueryの料金プランを自動的に切り替えるGoogle純正のソリューションは提供されていません。ただし、CLIまたはSQLで予約を管理するインターフェースは用意されています。

本記事では、DataformとそのSQLX言語を活用し、料金プラン更新を効率化していきます。この方法であれば、予約や割り当ての変更をDataformワークフロー内でシームレスに管理でき、Gitによるバージョン管理も可能になります。

Dataformを使うもう一つの利点は、そのシンプルさです。すべての処理をBigQuery内のSQLだけで完結でき、外部サービスを必要としません。Dataformは無料で利用でき、追加コストも発生しません。

準備として、最大スロット数100の「Enterprise Edition」予約 test_rd を(手動で)作成しておきます。

続いて、予約構成と割り当ての変更を管理するために、BigQuery Reservation API(SQL DDL)を利用します。

以下は、USリージョンで ndv-rd という予約に対し、assignment_idを df_rd として割り当てを作成するSQLコードの例です。

CREATE ASSIGNMENT `ndv-playground-bq-mgmt.region-us.ndv-rd.df_rd`
OPTIONS( assignee="projects/ndv-playground", job_type="QUERY");

Dataformに馴染みがない方は、まずBuild SQL pipelines to BigQuery with Dataformのブログを読み、Quick startを試してみることをおすすめします。

本記事に関連するコードはすべてこちらのGitHubリポジトリからご確認いただけます。

自動化ワークフローの構築は、次の6ステップで進めます。

[1] リポジトリの作成とロールの付与

BigQuery管理プロジェクトでDataformリポジトリを作成するには、Dataformセクションに移動し、リポジトリ作成のフローを開始します(Dataformリポジトリの作成)。

ワークフローを正常に実行するには、開発者プリンシパル(ユーザー)とDataformのSA(サービスアカウント)の双方に、次のロールが必要です:bigquery.resourceAdmin, bigquery.user

SAにロールを付与する際は(以下のコード参照)、リポジトリ作成後に表示されるメッセージに記載されたSA IDを使用します。

DataformのSA

gcloud projects add-iam-policy-binding <PROJECT-ID> \
--member="<service-account-dataform>" \
--role=roles/bigquery.resourceAdmin --condition=None
gcloud projects add-iam-policy-binding <PROJECT-ID> \
--member="<service-account-dataform>" \
--role=roles/bigquery.user --condition=None

[2] 開発ワークスペースの作成

新しいDataform開発ワークスペースを作成するには、次の手順に沿って進めてください:Dataformワークスペースの作成。

初期化直後のワークスペースには、SQLXのサンプルメソッドと「dataform.json」という構成ファイルが用意されています。サンプルファイル(first_view.sqlx, second_view.sqlx)は削除して問題ありません。

Dataformメソッドの構成にあたっては、「vars」セクションを追加し、以下のパラメータを記述します。このセクションは既存の最後のパラメータの後ろに配置してください。ここで設定する各パラメータの既定値は、後段のワークフローで上書きできます。

1. 「vars」セクションの前には、必ずカンマ "," を追加してください。
2. "<>" で示されたパラメータは実値に置き換えてください("<>"自体は削除します)。

"vars": {
"assigned_project":"'<R&D project>'",
"Assignment_id":"`df_rd_assignment`",
"reservation_prefix":"<BQ admin project>.region-<reservation location>",
"reservation_name":"`<R&D reservation name>`",
"region":"<reservation location>"
}

今回の実装例は次のとおりです。

"vars": {
"assigned_project":"'ndv-playground'",
"Assignment_id":"`df_rd_assignment`",
"reservation_prefix":"`ndv-playground-bq-mgmt.region-us`",
"reservation_name":"`ndv-rd`",
"region":"region-us"
}

[3] SQLXメソッドの開発

このステップでは、予約と割り当てを処理するSQLXコードを作成します。

• 'definitions' フォルダ配下に drop_and_create_assignment.sqlx という名前のファイルを作成します。
• すべてのプロジェクト割り当てを削除し、必要な割り当てを作成する以下のSQLXコードを記述します。
• ✅マークが表示され、コンパイルが成功していることを確認します。

drop_and_create_assignment.sqlx のSQLXコードは次のとおりです。

config {type: "operations"}
BEGIN

CREATE TEMP TABLE to_drop (drop_string STRING) AS (
SELECT
CONCAT("DROP ASSIGNMENT IF EXISTS `",project_id,".",SPLIT(ddl,".")[OFFSET(1)],
       ".",reservation_name,".",assignment_id,"`")
FROM
${dataform.projectConfig.vars.region}.INFORMATION_SCHEMA.ASSIGNMENTS_BY_PROJECT
WHERE
assignee_id = ${dataform.projectConfig.vars.assigned_project}
AND job_type = 'QUERY' );

FOR drop_statement IN (
SELECT drop_string FROM to_drop)
DO EXECUTE IMMEDIATE drop_statement.drop_string;
END FOR;

END;

本構成で生成が想定されるコードは以下のとおりです。

BEGIN
CREATE TEMP TABLE to_drop (drop_string STRING) AS (
SELECT
CONCAT("DROP ASSIGNMENT IF EXISTS `",project_id,".",SPLIT(ddl,".")[OFFSET(1)],".",reservation_name,".",assignment_id,"`")
FROM
region-us.INFORMATION_SCHEMA.ASSIGNMENTS_BY_PROJECT
WHERE
assignee_id = 'ndv-playground' AND job_type = 'QUERY' );
FOR drop_statement IN (
SELECT drop_string FROM to_drop)
DO EXECUTE IMMEDIATE drop_statement.drop_string;
END FOR;
END;

CREATE ASSIGNMENT
`ndv-playground-bq-mgmt.region-us`.`ndv-rd`.`df_rd_assignment`
OPTIONS(
assignee=CONCAT("projects/",'ndv-playground'),
job_type="QUERY");

[4] SQLXメソッドのテストと検証

テスト実行の前に、割り当て対象プロジェクトの設定が dataform.json ファイルの値に基づいて変更される点にご注意ください。テスト用プロジェクトでの検証をおすすめします。

コードの実行とテストには、2通りの方法があります。ご自身の認証情報で「Run」ボタンを押す方法と、以下のメッセージにあるとおりDataformのSAで実行する「Start Execution」を選ぶ方法です。

すべてのアクションを実行するか、一部のアクションを選択して実行できます。サービスアカウント [email protected] が使用されます。

実行結果に This statement created a new assignment と表示されることを確認したうえで、「Capacity management」ビューで変更内容をチェックしてください。

予約がプロジェクトに割り当てられた状態

SAでの検証が成功したら、変更をコミットしてデフォルトブランチへプッシュします。これで開発は完了し、リリースに進む準備が整いました。

[5] Dataformリリース構成の作成

コードのスナップショットを作成し、上書きしたパラメータで実行するには release configuration を使用します。今回はシンプルにするため、設定するパラメータは1つに絞り、リフレッシュは「Never」にします。

変更を加えたら必ずコンパイルすることをお忘れなく。

R&Dプロジェクトでは、「On-Demand」用(reservation_name= `none`)と「Enterprise edition」用(reservation_name=`ndv-rd`)の2つのリリースを用意します(バッククォート ` の扱いにご注意ください)。

選択したアクションで「release configuration」をテストします。

ndv-playground-bq-mgmt.dataform.drop_and_create_assignment

Enterprise予約用のリリース

On-demand予約用のリリース

[6] Dataformワークフロー構成の作成

ワークフロー構成では、SQLXアクションを自動的にどう実行するかを定義します。各ワークフローは「release configuration」とそのパラメータをベースとし、選択したアクションが指定したトリガーに従って実行されます。

R&Dプロジェクトでは、異なるリリース構成を使い、別々の時間にトリガーされる2つのワークフローを用意します。開始時刻は 'cron' 形式で指定します。

On-demand用ワークフロー

Enterprise edition用ワークフロー