📑

DataformのProject Configでカスタム変数を使う（Python版）

2022/12/02に公開

GCP

dataform

tech

Version

Python v3.11.0
google-cloud-dataform v0.2.3

はじめに

DataformではProject Configでカスタム変数を使用することができ、compile時に任意の値をセットできます。
この記事ではCLIとPythonライブラリの場合のカスタム変数の使用方法について紹介します。

準備

あらかじめDataformのコンソールからrepository, workspaceを作成しておきます。

DataformのService Account（service-${GCP_PROJECT_ID}@gcp-sa-dataform.iam.gserviceaccount.com）に bigquery.dataEditor と bigquery.jobUser の権限を付与しておきます。

以下のようにdataform.jsonにカスタム変数（"vars": { "sample": "0" }）を追加します。

dataform.json

{
  "defaultSchema": "dataform",
  "assertionSchema": "dataform_assertions",
  "warehouse": "bigquery",
  "defaultDatabase": "your-gcp-project-id",
  "defaultLocation": "US",
  "vars": {
    "sample": "0"
  }
}

Project Configの変数を取得する単純なSQLXを定義します。

definitions/first_view.sqlx

config {
  type: "table",
  schema: "sample_dataset",
}

SELECT ${ dataform.projectConfig.vars.sample } AS var

CLIの場合

dataform compile --varsまたはdataform run --varsでカスタム変数をセットできます。

これについては特に困ることはないので、本題のPythonライブラリの場合のサンプルコードを以下に示します。

Pythonの場合

準備

python -m venv dataform-sample
. dataform-sample/bin/activate
dataform-sample/bin/pip install google-cloud-dataform

サンプルコード

dataform.py

import argparse
import json
from google.cloud import dataform_v1beta1

def run_workflow(project, region, repository, workspace, config_vars):
    client = dataform_v1beta1.DataformClient()

    parent = f"projects/{project}/locations/{region}/repositories/{repository}"
    workspace = f"{parent}/workspaces/{workspace}"

    compilation_result_name = compile(client, parent, workspace, config_vars)
    invoke(client, parent, compilation_result_name)

def compile(client, parent, workspace, config_vars):
    code_compilication_config =  dataform_v1beta1.CompilationResult.CodeCompilationConfig()
    code_compilication_config.vars = config_vars

    compilation_result = dataform_v1beta1.CompilationResult()
    compilation_result.git_commitish = "main"
    compilation_result.workspace = workspace
    compilation_result.code_compilation_config = code_compilication_config

    request = dataform_v1beta1.CreateCompilationResultRequest(
        parent=parent,
        compilation_result=compilation_result,
    )

    response = client.create_compilation_result(request=request)

    print(response)
    return response.name

def invoke(client, parent, compilation_result_name):
    workflow_invocation = dataform_v1beta1.WorkflowInvocation()
    workflow_invocation.compilation_result=compilation_result_name

    request = dataform_v1beta1.CreateWorkflowInvocationRequest(
        parent=parent,
        workflow_invocation=workflow_invocation
    )

    response = client.create_workflow_invocation(request=request)

    print(response)
    return response.name

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Run a Dataform workflow.')
    parser.add_argument('--workspace', dest='workspace', required=True, help='Dataform workspace.')
    parser.add_argument('--project', dest='project', required=True, help='Google Cloud Project ID.')
    parser.add_argument('--region', dest='region', required=True, help='Region of Google Cloud.')
    parser.add_argument('--repository', dest='repository', required=True, help='Dataform repository.')
    parser.add_argument('--config-vars', type=json.loads, dest='config_vars', help='Custom variables of config of Dataform workspace. type: MapField<String, String>')
    args = parser.parse_args()

    run_workflow(args.project, args.region, args.repository, args.workspace, args.config_vars)

適宜環境変数をセットして以下のコマンドを実行するとBigQueryにテーブルが出力されます。

python dataform.py --workspace $WORKSPACE \
    --project $PROJECT \
    --region $REGION \
    --repository $REPOSITORY \
    --config-vars '{ "sample": "1" }'

result_1

--config-varsの引数を'{ "sample": "2" }' に変更して実行するとテーブルに反映されます。

result_2

注意点としてはCompilationResult.Types.CodeCompilationConfig.VarsのtypeがMapField<String, String> であるため、valueをString以外で渡すとエラーになります。そのため、ネストしたJSONもセットできません。

個人的にはテーブルごとにSQLXでカスタム変数を定義したいのですが、現状ではそのような方法はないようです。また、dataform.jsonはProjectごとにひとつのみで、カスタム変数にネストしたJSONもセットできないため、もっとカスタム変数が柔軟に使えるようになってほしいと願っています。

まとめ

Dataformでのカスタム変数の使用方法について紹介しました。
compileやworkflow実行時にSQLの値を動的に変えたい場合に有用なため、そのようなユースケースで活用していきたいです。