dbt-privacyのAmazon Athena対応から学ぶdbtマクロ

この記事はヌーラバー真夏のブログリレー2024の17日目の記事です。

はじめに

dbt（data build tool）はデータパイプラインのELT処理のうち、T（Transform）を効率的に行うことができる便利なツールなのですが、悲しいかな、Amazon Athena（以下Athena）の公式のコネクタはありません。Communtyで開発されているAthena用のコネクタはあるので、dbtを利用すること自体は可能なのですが、dbtのpackageが対応していない場合があります。

そこで今回、dbt-privacyというデータ保護を行うpackageをAthenaに対応させてみたところ、対応する過程で、dbtのmacroの挙動を知ることができたので、まとめてみました。

dbt-privacyでのデータ保護

dbt-privacyは2024/8現在、PostgresSQL、Snowflake、BigQuery、Redshiftに対応しています。
細かい使い方については、公式ドキュメントを見ていただければと思いますが、下記のようなSQLを書くと、

select
    id
    ,{{ dbt_privacy.hash("email", pepper_scope="project") }} as email
from hoge

以下のように変換されます。これはdbtのConnectorがAmazon Redshiftの場合の例です。

select
    id
    ,sha2(
      ('' || lower(email) || sha2( ('pepperstr_jaffle_shop')::varchar,
      256))::varchar, 256) as email
from hoge

処理としては、

1. カラムの値を小文字化
2. ペッパー文字列を生成してSha256でハッシュ化（引数や環境変数の設定によってここは変わる）
3. カラムの値とペッパーを結合してSha256でハッシュ化

という流れになっています。
このようにdbtのpackageで呼び出されたメソッドは、コンパイル時にSQLの関数に書き換えられて実行されます。

dbtでのmacroの実行順序

dbt-privacyをAthenaに対応させる前に、dbtにおけるmacroの実行順について見ておきたいと思います。
公式ドキュメントによると以下のようになっていました。

1. global project - default
2. global project - plugin specific
3. imported package - default
4. imported package - plugin specific
5. local project - default
6. local project - plugin specific

まず1と2の global project ですが、これはdbt-coreの中に含まれているmacroになっています。
defaultというのはconnectorの種類と特定しない汎用的に定義されたmacro、plugin specific というのは特定のConnectionに対応したmacroのことを指しています。
global projectのmacroは、dbt coreの中に定義されているので、場所でいうとこのあたりにあります。

# global project
## default
dbt/core/dbt/include/global_project/macros/

## plugin specific (redshiftの場合)
dbt/adapters/redshift/dbt/include/redshift/macros/

次に3と4の imported package ですが、これは先程のdbt-privacyのように、あとからdbt projectに追加されたpackageに含まれるmacroを指しています。
default と plugin specific の区別はglobal projectの場合と同様です。
以下の場所に定義されています。

# imported package
## default
dbt_project_root/dbt_packages/[pacage_name]/macros/foobar.sql

## plugin specific (Redshiftの場合)
dbt_project_root/dbt_packages/[pacage_name]/macros/redshift_foobar.sql

最後の5、6のlocal projectは、dbt_project.ymlで定義されているmacro-pathsに格納されているもので、開発者が定義したmacroになります。。こちらも default と plugin specific の区別はglobal projectの場合と同様です。

# local project
## default
dbt_project_root/macros/foobar.sql

## plugin specific (Redshiftの場合)
dbt_project_root/macros/redshift_foobar.sql

この順番は「優先順」ではなく、「実行順」であることが重要です。
dbtでは、同名のmacroが複数定義されている場合、あとから読み込まれた内容で上書きされるので、
つまり、強さの順としては

1. local project
2. imported package
3. global project

になります。

dbt-privacyのmacroを解析してみる

前置きが長くなっちゃいましたが、dbt-privacyのmacroを見てみます。
dbt-privacyにはいろんなメソッドがあるので、冒頭でも例示した、一番基本的な hash というmacroについて細かく見てみたいと思います。
dbt-privacyの中身は下記のようなディレクトリ構成になっており、mask.sql というファイルにhashというmacroが定義されていました。

dbt_packages/dbt_privacy
└── macros
    ├── hash.sql
    ├── hash_unique.sql
    ├── mask.sql
    ├── mask_email.sql
    ├── redact_unique.sql
    ├── safe_mask.sql
    └── utils
        ├── arg_validation.sql
        ├── docs
        │   ├── generate_pepper.yml
        │   └── helpers.yml
        ├── generate_pepper.sql
        ├── simple_hash.sql
        ├── sql_lower.sql
        └── sql_quote_string.sql

mask.sql

{%- macro hash(
    expr,
    lowercase=true,
    salt_expr="''",
    pepper_scope="model",
    pepper_persistence="permanent",
    digest_size=256
) -%}

{%- if lowercase == true -%}{% set expr = dbt_privacy.sql_lower(expr) %}{%- endif -%}

{%- if salt_expr not in ("", "''") -%}
{% set salt = dbt_privacy.simple_hash(salt_expr, digest_size) %}
{%- else -%}
{% set salt = "''" %}
{%- endif -%}

{%- set pepper = dbt_privacy.generate_pepper(
    pepper_scope=pepper_scope, pepper_persistence=pepper_persistence
) -%}

{{- dbt_privacy.simple_hash(salt ~ ' || ' ~ expr ~ ' || ' ~ pepper, digest_size) -}}

{%- endmacro -%}

細かなmacroの書き方についての説明は割愛しますが、冒頭で例示したhashの処理内容と比べてみると、

1. カラムの値を小文字化

{%- if lowercase == true -%}{% set expr = dbt_privacy.sql_lower(expr) %}{%- endif -%}

2. ペッパー文字列を生成してSha256でハッシュ化（引数や環境変数の設定によってここは変わる）

{%- if salt_expr not in ("", "''") -%}
{% set salt = dbt_privacy.simple_hash(salt_expr, digest_size) %}
{%- else -%}
{% set salt = "''" %}
{%- endif -%}

{%- set pepper = dbt_privacy.generate_pepper(
    pepper_scope=pepper_scope, pepper_persistence=pepper_persistence
) -%}

3. カラムの値とペッパーを結合してSha256でハッシュ化

{{- dbt_privacy.simple_hash(salt ~ ' || ' ~ expr ~ ' || ' ~ pepper, digest_size) -}}

というような対応になっていました。
2のペッパー文字列の生成と、3のカラム値のハッシュ化は、両方ともdbt_privacy内のsimple_hash というマクロを使っていますね。
このsimple_hashというmacroは、utils/simple_hash.sqlに定義されています。

simple_hash

{%- macro simple_hash(expr, digest_size=256) -%}
{{- dbt_privacy.raise_on_bad_digest_size(digest_size) -}}
{{- return(adapter.dispatch("simple_hash", "dbt_privacy")(expr, digest_size)) -}}
{%- endmacro -%}

{%- macro default__simple_hash(expr, digest_size=256) -%}
sha2( ({{ expr }})::varchar, {{ digest_size }})
{%- endmacro -%}

{%- macro postgres__simple_hash(expr, digest_size=256) -%}
encode(sha{{ digest_size }} ( ({{ expr }})::varchar::bytea), 'hex')
{%- endmacro -%}

{%- macro redshift__simple_hash(expr, digest_size=256) -%}
sha2( ({{ expr }})::varchar, {{ digest_size }})
{%- endmacro -%}

{%- macro bigquery__simple_hash(expr, digest_size=256) -%}
{{- dbt_privacy.raise_on_bad_digest_size(digest_size, ok_sizes=[256, 512]) -}}
to_hex(sha{{ digest_size }} (cast({{ expr }} as string)))
{%- endmacro -%}

#エラーハンドリングのmacroについては省略#

{%- endmacro -%}

simple_hashから別のmacroは呼ばれていないので、どうやらこのmacroが処理の実体のようです。

simple_hash.sql の中に adapter.dispatch というメソッドがあるのですが、これはdbt-coreで実装されているメソッドで、profiles.ymlに定義されているconnectorの情報から、適切なmacroを探してくる処理を担っています。第１引数に対象のmacro名、第2引数に検索する名前空間を指定します。ここではdbt_privacyが名前空間として指定されています。
ただ、第2引数に指定された名前空間は、最初に探す場所というだけで、指定された場所になければ他の場所も検索するようになっています。

adapter.dispatchで呼び出されるmacro名は [adapter名]__[macro名] という命名になります。たとえば、profiles.ymlに設定されているconnectorがAmazon Redshiftだった場合、dbtはredsfhit__simple_hashというmacroを探します。
探した結果、どこにも該当するものがなければadapterとしてdefaultが設定され、ここではdefault__simple_hashが使われることになります。

上記のコードを見ると分かるとおり、ここにはathena__simple_hashというmacroがありません。
そのため、default__simple_hashが使われるのですが、Athenaでは以下のような記法は対応していないためエラーになります。

sha2( ({{ expr }})::varchar, {{ digest_size }})

これがdbt-privacyがAthenaに対応していない理由になります。

Athena用のsimple_hashを追加する

前のセクションで書いたとおり、dbtはmacroを呼び出すときに指定した名前空間以外の場所も探します。
つまり、ローカルプロジェクトに定義されているmacroも探してくれるので、ローカルにAthena用のsimple_hashを追加すればOKです。

というわけで、AthenaのSQL文法に合わせて、以下のmacroを追加しました。
格納場所はdbt_project.ymlのmacro-pathの場所です。

macros/simple_hash.sql

{%- macro athena__simple_hash(expr, digest_size=256) -%}
to_hex(sha{{ digest_size }}(to_utf8(cast({{ expr }} as varchar))))
{%- endmacro -%}

これでAthenaでもdbt-privacyのメソッドを使うことができました。
あらためて処理順を書くと、以下のとおりになります。

1. dbt_privacy.hashメソッドが呼び出される
2. dbt_packages/dbt_privacy/hash.sqlが呼び出される
3. 2からdbt_packages/dbt_privacy/simple_hash.sqlが呼び出される
4. simple_hash内で、adapter.dispatchメソッドによって、connectorに合わせたmacroが検索される
5. ローカルプロジェクトで定義したathena__simple_hashが使われる

まとめ

dbtのpackageには高機能なものがたくさんあり、どんどん使っていきたい気持ちになります。
一方で、自社の環境に対応していなかったり、処理をカスタマイズする必要がある場面など出てくるかもしれません。そういうときに、完全にイチから作り直すのではなくて、既存のmacroの追加や上書きなどで対応することができれば、パッケージのソースコードを有効に使うこともできると思います。

それでは。