Zenn
🐈

dbt で出力先の dataset 名を開発と本番で切り替える

2022/08/07に公開
BigQuery
dbt
tech

やりたいこと

  • target = dev なら git のブランチ名をいい感じに入れてユニークにしたい
  • target = prod なら schema 名をそのまま使いたい
  • いずれの場合も カスタムスキーマ名を使いたい
    • デフォルトでは、 profiles.yml の dataset が接頭辞になってしまう
  • Github でブランチが削除されたら、対応する開発用データセットも削除したい

開発中はブランチ名の入ったデータセットへ出力

後でGithubActions に乗せる予定なので、dbt プロジェクトディレクトリの直下に profile.yml を置いて、 dev と prod を用意しておく。

macros/override/get_custom_schema.sql

{% macro generate_schema_name(custom_schema_name, node) -%}
  {%- set default_schema = target.schema -%}
  {%- if custom_schema_name is none -%}
    {{ env_var('DBT_SCHEMA_PREFIX', '') ~ default_schema | trim}}
  {%- else -%}
    {{ env_var('DBT_SCHEMA_PREFIX', '') ~ custom_schema_name | trim }}
  {%- endif -%}
{%- endmacro %}

環境変数 DBT_SCHEMA_PREFIX に使いたいプレフィックスを入れるだけ。

あとは、こんな感じのラッパースクリプトを dbtw とかで保存しておいて…

#!/bin/bash

# Required Parameter -------------------------------------------------
export DBT_PROFILES_DIR=$(cd "$(dirname "$0")" && pwd)
export DBT_TARGET=${TARGET:-"dev"}
CURRENT_BRANCH="$(cd "$DBT_PROFILE_DIR" && git branch --show-current)"

if [ $DBT_TARGET = 'dev' ]; then
  export DBT_SCHEMA_PREFIX="$(echo "__dbt__${DBT_TARGET}__${CURRENT_BRANCH}__" | sed -e 's/[\/-]/_/g')"
fi

# run dbt command ----------------------------------
env | grep "^DBT_"
echo ""

echo dbt "$@" --target "${DBT_TARGET}"
echo ""

dbt "$@" --target "${DBT_TARGET}"

./dbtw run

run すると __dbt__dev__<ブランチ名>__<スキーマ/データセット名> にモデルが出力される。

ちなみに、BigQuery の SQL ワークスペースでは、アンダースコアから始まるデータセットは表示されない。 bq ls -a とかすると見える。

Github 上でブランチが削除されたらデータセットも削除する

ここからはBigQuery向けです。他のDWH向けには、削除クエリと削除対象のデータセット抽出処理の変更が必要そう。

リストで渡した削除対象を削除するマクロ。
macros/drop_dev_schemas.sql

{% macro drop_dev_schemas(schemas) %}
{%- call statement('states', fetch_result=True) -%}
{% for schema in schemas %}
  {{ log("DROP SCHEMA IF EXISTS " ~ schema  ~ " CASCADE;", info=True) }}
  DROP SCHEMA IF EXISTS {{ schema }} CASCADE;
{% endfor %}
{%- endcall -%}
{% endmacro %}

本当はブランチのリストもクエリできればよかったけど _ はじまりのデータセットはINFORMATION_SCHEMA からは見えないっぽい。。
仕方ないので、 gcloud alpha bq datasets list -all で探す。

dbtw

#!/bin/bash

# Required Parameter -------------------------------------------------
export DBT_IMPERSONATE_SA='hoge@fuga.iam.gserviceaccount.com'

# -------------------------------------------------
export DBT_PROFILES_DIR=$(cd "$(dirname "$0")" && pwd)
export DBT_TARGET=${TARGET:-"dev"}
CURRENT_BRANCH="$(cd "$DBT_PROFILE_DIR" && git branch --show-current)"

if [ $DBT_TARGET = 'dev' ]; then
  export DBT_SCHEMA_PREFIX="$(echo "__dbt__${DBT_TARGET}__${CURRENT_BRANCH}__" | sed -e 's/[\/-]/_/g')"
fi

# -------------------------------------------------
if [ $1 = '--drop-dev-schemas' ]; then
  if [ $DBT_TARGET != 'dev' ]; then
    echo "Can't drop dataset on target=${DBT_TARGET}. It works target='dev' only."
    exit -1
  fi

  DROP_SCHEMA_PREFIX=${2:-$DBT_SCHEMA_PREFIX}
  if [ $(echo $DROP_SCHEMA_PREFIX | grep '^__dbt__dev__..*' | wc -l) -le 0 ]; then
    echo "Not matched."
    exit -1
  fi

  DATASETS=$(gcloud alpha bq datasets list --all --page-size=1024 --limit=1024 --impersonate-service-account="${DBT_IMPERSONATE_SA}" \
         | grep "${DROP_SCHEMA_PREFIX}" | sed -r 's/[: ]+/,/g' \
         | awk -F, '{print "\"`" $1 "." $2 "`\""}' | jq -s -c
  )

  if [[ $(echo ${DATASETS} | jq -r '.[]' | wc -l) -le 0 ]]; then
    echo "No dataset matched to '${DROP_SCHEMA_PREFIX}'. "
    exit
  fi

  echo "DROP schemas ..."
  echo ${DATASETS} | jq -r '.[]'
  echo ""

  env | grep "^DBT_"
  echo ""

  ARGS="{schemas: ${DATASETS}}"
  dbt run-operation drop_dev_schemas --arg "${ARGS}" --target "${DBT_TARGET}"
  echo "Completed."

  exit
fi

これで、 以下コマンドで当該データセットを削除できるようになった。

$ TARGET=dev ./dbtw --drop-dev-schemas
or
$ TARGET=dev ./dbtw --drop-dev-schemas "__dbt__dev__<branch_name>__<schema_name>"

Github Actions に乗せるならこんな感じに。ブランチが削除されたらそのブランチで使用されていたテスト用データセットを自動で削除する。
.github/workflows/dbt_clean_schemas.yml

name: dbt_cleanup

on:
  delete

permissions:
  contents: 'read'
  id-token: 'write'

jobs:
  celeanup:
    runs-on: ubuntu-latest
    timeout-minutes: 5

    steps:
      - uses: actions/checkout@v3
      - uses: actions/setup-python@v4
        with:
          python-version: '3.10'
          cache: 'pip'
      - run: pip install -r requirements.txt

      - id: 略！

      - name: install dbt packages
        run: TARGET=dev ./dbtw deps

      - name: dbt drop development datasets on this branch
        run:
          DBT_SCHEMA_PREFIX="$(echo "__dbt__dev__${{github.event.ref}}__" | sed -e 's/[\/-]/_/g')"
          TARGET=dev ./dbtw --drop-dev-schemas "$DBT_SCHEMA_PREFIX"

こっちにも sed の置換ロジックが出てきてしまっているのがイマイチ。

Discussion