❄️

dbt Snowflakeでウェアハウスをモデルごとに切り替える

2023/03/12に公開

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

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("Warehouse size not one of " ~ valid_warehouse_sizes) }}
    {% endif %}
    {% if target.name in ('prod') %}
        {% do return('DBT_' ~ size + '_WH') %}
    {% elif target.name in ('ci') %}
        {% do return('DEV_XS_WH') %}
    {% elif target.name == 'dev_d' %}
        {# ローカルでの動的ウェアハウス切り替えチェック用 #}
        {% do return('DEV_' ~ size + '_WH') %}
    {% else %}
        {% do return(None) %}
    {% endif %}
{% endmacro %}

モデルでの利用方法:

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

with source as (
  ...
),

マクロの挙動の説明について軽く説明すると、dbt run する時に指定する --target xxx で使うウェアハウスを決定します。

具体的には、dbt run --target prod で、get_warehouse('XS') と設定されているなら、 DBT_XS_WH が使われ、get_warehouse('L') なら、DBT_L_WH が使われます。
dbt run --target dev または、dbt run --target ci の場合、設定に関係なく全て DEV_XS_WH で実行されます。
dbt run --target dev_d とすると、動的にウェアハウスを切り替えて実行することができます。

このマクロを使うことで、モデルごとにウェアハウス切り替えができますが、注意点もあるので、それについても言及しておきます。

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 で作るようにしたことで、この問題は解決されたので、事なきを得ました。

現状も残る課題

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

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

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