モデルごとにウェアハウスを切り替えたい

Snowflake上で、dbtを使った開発していると、モデルごとにウェアハウスを切り替えたくなる時があります。
筆者の扱っているモデルでは、1割くらいのモデルが非常にデータ量が多い一方、その他のモデルは相対的に少ない性質がありました。データ量が多いモデルは、L(large)だと1分以内に終了するのに対して、XS(x-small)だとクエリがいつ終わるか分からない状況でした。しかしながら、他のモデルはXSでも現実的な時間で終了するため、一部のモデルのためにLを使うことは非常にコストパフォーマンスが悪いという状況でした。

dbtには、snowflake_warehouse というコンフィグがあります。それを使って、今回の課題を解決するのですが、少し工夫が必要だったので、それについてまとめてみました。

ちなみに、snowflake_warehouse は、以下のような使い方ができます。

dbt_project.yml に設定するパターン:

name: my_project
version: 1.0.0

...

models:
  # このプロジェクトのデフォルトでDBT_XS_WHが使われる
  +snowflake_warehouse: "DBT_XS_WH"
  my_project:
    big_data:
      +snowflake_warehouse: "DBT_L_WH" # big_dataモデルは、DBT_L_WHが使われる
    small_data: # 特に設定がなければ、DBT_XS_WHが使われる


snapshots:
  +snowflake_warehouse: "DBT_XS_WH"

モデルのコードに直接設定するパターン:

{{
  config(
    materialized='table',
    snowflake_warehouse='DBT_XS_WH'
  )
}}

with source as (
  ...
),

これ使ったら終わりじゃん！とはなりませんでした。なぜなら開発環境、CI(Continuous Integration)、本番環境では使っているウェアハウスが違うからです。
対応表にすると、以下のような感じでした。

環境	target.name	ウェアハウス名
本番	prod	DBT_XX_WH
CI	ci	DEV_XX_WH
開発	dev	DEV_XX_WH

筆者に限らず、環境ごとにウェアハウスを変えてる方は多いと思います。そのため、ハードコードでウェアハウス設定すると、本当にやりたいことはできません。そこで、それ用のマクロを作成することにしました。

ウェアハウス切替マクロ

このアイデアは、select.dev に掲載されていたイケてるアイデアです。

マクロ

{% macro get_warehouse(size) %}
    {% set available_sizes = ["XS", "S", "M", "L", "XL"] %}
    {% if size not in available_sizes %}
        {{
            exceptions.raise_compiler_error(
                "想定されていないウェアハウスサイズが指定されています: "
                ~ valid_warehouse_sizes
            )
        }}
    {% endif %}
    {% if target.name in ("prod") %}
        {# backfill時は、モデルに設定されたウェアハウスサイズだと処理しきれないため、強めのウェアハウスを設定する #}
        {%- if var("is_transform_backfill") -%} {% do return("DBT_BACKFILL_XL_WH") %}
        {% else %} {% do return("DBT_" ~ size + "_WH") %}
        {% endif %}
    {% elif target.name in ("ci") %} {% do return("DEV_XS_WH") %}
    {# devはXSサイズを強制することで、課金事故を防ぐ #}
    {% elif target.name == "dev" %} {% do return("DEV_XS_WH") %}
    {# 明示的にdev_dを使う場合は、設定されたウェアハウスサイズで動かす #}
    {% elif target.name == "dev_d" %} {% do return("DEV_" ~ size + "_WH") %}
    {% else %}
        {{
            exceptions.raise_compiler_error(
                "想定されていないtargetが指定されています: " ~ target.name
            )
        }}
    {% endif %}
{% endmacro %}

モデルでの利用方法

{{
  config(
    materialized='table',
    snowflake_warehouse=get_warehouse('M'),
  )
}}

with source as (
  ...
),

🚨モデル定義のyamlではマクロが使えない

以下はエラーになります。

models:
  +snowflake_warehouse: "{{ get_warehouse('M')}}"

# Compilation Error
#  Could not render {{ get_warehouse('M')}}: 'get_warehouse' is undefined

直接書く場合は動きます。

models:
  +snowflake_warehouse: DEV_M_WH

テストでの利用方法

モデルの方のyamlでは、マクロは使えませんが、テストの方では正しく動作します。

models:
  - name: hoge
    columns:
      - name: fuga
        data_type: VARCHAR
        data_tests:
          - not_null:
              config:
                snowflake_warehouse: get_warehouse("L")

マクロの振る舞い

本番環境

本番環境では2つの実行モードがあります。

通常実行時は、各モデルで指定したウェアハウスサイズが使用されます。

dbt run --target prod

再集計実行時は、専用の大型ウェアハウスを使用します。
モデルの設定に関係なくDBT_BACKFILL_XL_WH(XLサイズウェアハウス)が使用されます。
こうすることで、再集計用に独立したウェアハウスを用意することで、通常の処理とキューが競合せず、スムーズな再計算が可能になります。

dbt run --target prod --vars '{"is_transform_backfill": true}'

開発環境

コスト抑制のため、最小サイズのウェアハウス（DEV_XS_WH）を強制使用します。

dbt run --target dev

開発時に本番相当の実行が必要な場合は、以下のコマンドで各モデルの指定サイズのウェアハウスを使用できます。

dbt run --target dev_d

CI環境

コスト抑制のため、最小サイズのウェアハウス（DEV_XS_WH）を強制使用します。

dbt run --target ci

CI環境ではクエリの実行可能性の検証のみを目的としています。
これを実現するために、モデルの末尾にマクロでlimit 0を付与することで、最小限のコストで実行可能性を検証しています。

state:modifiedが常にmodifiedになってしまう問題

ウェアハウス切り替えマクロを使い始めてから、CIで使っている state:modified が効かなくなりました。このメソッドについて知らない方は、The "state" method に詳しい説明があるので、読んでみてください。

state:modified は、例えばPull Requestの内容と本番を比較して、変更のあったモデルだけ実行する時に使うメソッドです。
ちなみに、筆者の環境では、以下のようなコマンドをCIに実行させています。

$ aws s3 sync --exact-timestamps --delete s3://dbt-state/main/ remote_state/ --region ap-northeast-1
$ dbt run --full-refresh --target ci --select state:modified+ --state remote_state
$ dbt run --target ci --select state:modified+ --state remote_state
$ dbt test --target ci --select state:modified+ --state remote_state

起きていたこととしては、使っているマニフェストファイルが、--target prod で作られているマニフェストだったため、--target ci で state:modified を使うと、モデルが使うウェアハウスが全てズレるため、全てに変更がある判定をされていました。（prodでは、DBT_XX_WH だけど、 ciでは常に DEV_XS_WH になるため。）
この事象に対して、--target prod で作っていたマニフェストを単純に、--target ci で作るようにしたことで、この問題は解決されたので、事なきを得ました。

現状も残る課題

以前はテストごとのウェアハウス切り替えは難しかったですが、dbt v1.9.0b1 以降から回避策が用意されたので、テストごとにウェアハウスは切り替えれます。

これで全ての問題が解決された！かと思いきや、そうでもありません。
モデルに対しては、snowflake_warehouse が提供されていたため、対策できましたが、テストに関してはこれが用意されていません。
[CT-1123] Adapter pre_model_hook + post_model_hook for tests and compilation, too #5766 というissueで、議論はされていますが、それなりに変更が面倒そうな箇所なため、すぐには実装されなさそうです。
ちなみに、カスタムテスト作って、ウェアハウス切り替えようとしても、そもそも、そういう実装ができる作りではないので、現状素直な解決策はありません。あったら教えて下さいｗ

紹介した手法以外に検討した方法

Snowflakeには、クエリアクセラレータや、マルチクラスターウェアハウスがあるので、これをdbtが使うウェアハウスで適用すればいいんじゃね？とか考えたんですが、事前に重たいことが分かりきってるモデルに使うのはユースケースが違いそうなのと、実際に試したところ、劇的に改善されることはなかったので見送りました。

Discussion