LayerX エンジニアブログ

LayerX の エンジニアブログです。

dbt-snowflakeでSingle Sign-Onを使用する際にウェブブラウザのタブが大量に開かれる問題を解決する #ベッテク月間

こんにちは。バクラク事業部 機械学習・データ部 データグループの@civitaspoです。LayerXの新オフィスは東銀座に位置しているのですが、築地に近いこともあり、お魚が美味しいです。今日はマグロのカマ焼きを食べました。脂が乗っていて非常に美味しかったです。

さて、先週に引き続き、Snowflakeに関する記事を書こうと思います。先週はSnowpipe StreamingとAmazon Data Firehoseを使用してSnowflakeにストリームデータをロードする方法を紹介しました。

tech.layerx.co.jp

今回の記事は、dbt-snowflakeでSingle Sign-Onを使用する際にウェブブラウザのタブが大量に開かれる問題と解決策を提示します。最初に、前提知識となる、dbt-snowflakeとdbt-snowflakeで使用するSingle Sign-Onについて説明します。その後、課題と解決策について説明します。

なお、7月はLayerXのエンジニアブログがたくさん出る#ベッテク月間です。LayerXの行動指針の一つである「Bet Technology」を略して「ベッテク」と呼んでいます。今後もベッテクな記事がたくさん出ますので、どんな記事がでるのかこちらのカレンダーからチェックしてみてください!

dbt-snowflakeとは

まず初めに、dbt-snowflakeの説明をします。dbtは、各データプラットフォームへの接続や専用機能を利用するために、アダプターと呼ばれるプラガブルな仕組みを提供しています。dbt-snowflakeはSnowflake向けに実装されたアダプターの一種です。

docs.getdbt.com github.com

dbt-snowflakeでSnowflakeの接続に利用可能な認証方式は、以下の4種類です。

  • パスワード認証
  • パスワード認証 + 他要素認証
  • Key-Pair認証
  • SSO(Single Sign-On)認証

弊社では、以前公開したブログに記載したとおり、認証方式に関して、以下のようなルールで運用しています。

・IdP上のユーザーを利用する場合はSAML認証、またはOAuth認証しか認めない。
・IdP上に存在しないユーザーはシステム間連携で必要な場合を除き、原則作成しない。作成する場合は暗号化された鍵を用いたKey-Pair認証のみ許可する。

ref. Don’t Use Passwords in Your Snowflake Account - LayerX エンジニアブログ

そのため、CI/CDシステムからのdbt実行に関してはKey-Pair認証を使用しています。特に、本番環境におけるデータの変更は、dbtを用いたCI/CDシステムからの変更のみを許可する運用としています。これは、コードレビューを経たデータの変更のみを許可し、本番環境における意図しないデータの変更を防ぐためです。

一方で、開発環境や試験環境*1のデータ変更に関しては、開発生産性を担保するため、開発者の端末からSSO認証を用いたdbt実行が可能な状態を担保しています。開発者は個人に割り当てられたIdP上のアカウントを使用してSSOを行いますが、Snowflake上での権限差異が発生しないよう、共通のロールを使用してdbtを実行しています。

本記事は、このdbt実行時に使用するSSO認証で発生する課題に焦点を当てます。

dbt-snowflakeにおけるSingle Sign-On

課題の説明に入る前に、dbt-snowflakeにおけるSSOを用いた認証について、より詳細に説明します。

dbt-snowflakeにおけるSSOでは、ウェブブラウザを介して認証を行います。dbt実行時にウェブブラウザが起動し、IdPによる認証が完了すると、Snowflakeとの接続が確立されます。Oktaを使用している場合と、それ以外のIdPでは設定方法等が異なるのですが、本記事では後者を扱っている前提で読んでいただきたいです。

docs.snowflake.com

この実装はSnowflake社が提供しているPython SDKを利用して実現されています。

github.com

また、認証時に取得したIDトークンは、以下の条件を満たす場合において、開発者の端末にキャッシュすることができます。

  • 接続先のSnowflakeアカウントで ALLOW_ID_TOKEN*2true に設定されている*3
  • snowflake-connector-python[secure-local-storage]*4がインストールされている
  • 開発者の端末で、MacOS KeyChain*5または Windows Credential Locker*6を使用できる*7 *8

キャッシュされたIDトークンは、MacOS KeyChainであれば、以下のように「snowflake」と検索することで確認することができます。

一部、モザイク加工しています。

キャッシュ時に使用している名称は、以下の実装のとおり ${ホスト名}:${ユーザー名}:${ドライバー名}:ID_TOKEN となります。 github.com

このIDトークンの有効期限は4時間に設定されており、有効期限が切れたあとは再度、ウェブブラウザを介した認証が必要になります。

With ALLOW_ID_TOKEN set to true, the ID token will expire in 4 hours at which time the user will be prompted to re-authenticate.
The expiration time is set in the token upon creation and it is evaluated when the token is presented to Snowflake for authentication.

ref. https://community.snowflake.com/s/article/How-long-is-an-ID-token-cached-when-ALLOW-ID-TOKEN-is-set-to-true

dbt-snowflakeにおけるSingle Sign-Onではウェブブラウザのタブが大量に開かれることがある

さて、ここからdbt-snowflakeにおけるSSOで大量にタブが開かれる問題について書いていきます。この問題はdbt-snowflakeの以下のGitHub Issueで議論されています。

github.com

dbtはスレッドを使用した並行実行をサポートしています。この並行実行では、スレッドごとにデータプラットフォームに対する接続を行う仕様となっています。一方、dbt-snowflakeにおけるSSOを使用したSnowflakeへの接続では、有効なIDトークンのキャッシュが存在しなければ、ウェブブラウザを立ち上げて認証を完了させようとします。そのため、指定したスレッド数分のウェブブラウザ立ち上げが行われ、大量にタブが開かれてしまうのです。

先ほど掲載したGitHub Issueでは、具体的な解決策は提示されないまま、Closeされてしまっています。この問題を解決するのが、本記事の主題です。

dbt-snowflakeがスレッド内でSnowflakeへの接続を試みる前にIDトークンをキャッシュしておけば良い

ここまでの前提知識や課題から、dbt-snowflakeがスレッド内でSnowflakeへの接続を試みる前に、IDトークンを取得して、キャッシュしておけば良いことがわかります。加えて、キャッシュするIDトークンの名前空間は${ホスト名}:${ユーザー名}:${ドライバー名}:ID_TOKENであるため、同一のドライバーを用いてIDトークンを取得すればよいこともわかります。つまり、dbtを実行する前に、Snowflake社が提供するPython SDKであるsnowflake-connector-pythonを使用して、IDトークンを取得しておけば課題は解決できるはずです。

Snowflake CLIを使用して事前にIDトークンを取得する

Snowflake社は、Snowflake CLIというコマンドラインツールを提供しています。このツールはPythonで書かれており、Snowflakeへの接続のためにsnowflake-connector-pythonを使用しています。

docs.snowflake.com

Snowflake CLIには、Snowflakeとの接続を試行するためのサブコマンド snow connection test が実装されています。

$ snow connection test --help

 Usage: snow connection test [OPTIONS]

 Tests the connection to Snowflake.

╭─ Options ─────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╮
│ --help  -h        Show this message and exit.                                                                                         │
╰───────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯
╭─ Connection configuration ────────────────────────────────────────────────────────────────────────────────────────────────────────────╮
│ --connection,--environment  -c      TEXT  Name of the connection, as defined in your `config.toml`. Default: `default`.               │
│ --account,--accountname             TEXT  Name assigned to your Snowflake account. Overrides the value specified for the connection.  │
│ --user,--username                   TEXT  Username to connect to Snowflake. Overrides the value specified for the connection.         │
│ --password                          TEXT  Snowflake password. Overrides the value specified for the connection.                       │
│ --authenticator                     TEXT  Snowflake authenticator. Overrides the value specified for the connection.                  │
│ --private-key-path                  TEXT  Snowflake private key path. Overrides the value specified for the connection.               │
│ --database,--dbname                 TEXT  Database to use. Overrides the value specified for the connection.                          │
│ --schema,--schemaname               TEXT  Database schema to use. Overrides the value specified for the connection.                   │
│ --role,--rolename                   TEXT  Role to use. Overrides the value specified for the connection.                              │
│ --warehouse                         TEXT  Warehouse to use. Overrides the value specified for the connection.                         │
│ --temporary-connection      -x            Uses connection defined with command line parameters, instead of one defined in config      │
│ --mfa-passcode                      TEXT  Token to use for multi-factor authentication (MFA)                                          │
│ --enable-diag                             Run python connector diagnostic test                                                        │
│ --diag-log-path                     TEXT  Diagnostic report path                                                                      │
│ --diag-allowlist-path               TEXT  Diagnostic report path to optional allowlist                                                │
╰───────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯
╭─ Global configuration ────────────────────────────────────────────────────────────────────────────────────────────────────────────────╮
│ --format           [TABLE|JSON]  Specifies the output format. [default: TABLE]                                                        │
│ --verbose  -v                    Displays log entries for log levels `info` and higher.                                               │
│ --debug                          Displays log entries for log levels `debug` and higher; debug logs contains additional information.  │
│ --silent                         Turns off intermediate output to console.                                                            │
╰───────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯

このsnow connection testのコマンドラインオプションに対して、dbt実行時に使用する接続設定と同じ値を指定することで、dbt-snowflakeでも利用可能なIDトークンがキャッシュされます。

$ cat profiles.yml
# 以下のような設定でdbtを実行するとする
sample_dbt_project:
  target: dev-via-sso
  outputs:
    base: &default-base
      type: snowflake
      account: sample-dev
      role: ROLE_DBT_BUILDER
      database: WORKSPACE
      schema: PUBLIC
      warehouse: WH_XS_DATA_LOADER
      threads: "{{ env_var('DBT_THREADS', 20) | int }}"
      connect_retries: 2

    dev-via-sso: &default-via-sso
      <<: *default-base
      user: sample@example.com
      authenticator: externalbrowser  # SSO
    stg-via-sso:
      <<: *default-via-sso
      account: layerx-bakuraku-stg
    prd-via-sso:
      <<: *default-via-sso
      account: layerx-bakuraku-prd

# profiles.yml に指定されている値と同じ値をコマンドラインオプションで指定する
$ snow connection test \
        --account sample-dev \
        --user 'sample@example.com' \
        --authenticator externalbrowser \
        --role ROLE_DBT_BUILDER \
        --database WORKSPACE \
        --warehouse WH_XS_DATA_LOADER \
        --temporary-connection

+-----------------------------------------------------+
| key             | value                             |
|-----------------+-----------------------------------|
| Connection name | None                              |
| Status          | OK                                |
| Host            | sample-dev.snowflakecomputing.com |
| Account         | sample-dev                        |
| User            | sample@example.com                |
| Role            | ROLE_DBT_BUILDER                  |
| Database        | WORKSPACE                         |
| Warehouse       | WH_XS_DATA_LOADER                 |
+-----------------------------------------------------+

このコマンドをdbt実行の直前に打つことで、課題だった大量にウェブブラウザのタブが開かれる問題が解消されます。毎回dbt実行前に、別途このコマンドを入力させるのは負担が大きいので、Wrapperを作るなど、その会社にあった形で運用に組み込むのが良いでしょう。

Future Work

ここまで説明しておきながら、本来はdbt-snowflakeにコミットして、根本的な解決を目指すべきだと考えています。今回の試行錯誤によって実装すべき内容は明らかになったので、根本的な解決についてはdbtのコミュニティメンバーとも話して気長に進めてみます💪

おわりに

この記事では、dbt-snowflakeでSSOを使用する際にウェブブラウザのタブが大量に開かれる問題と解決策を提示しました。非常にニッチな話題ですが、SSOを使用してdbtを実行している開発者にとっては有益な情報だったと信じています。

バクラク事業部では現在、データ基盤で利用するデータウェアハウスソリューションをGoogle BigQueryからSnowflakeへ移管するプロジェクトを進めています。そして、移管後にどういったデータ基盤を構築していくか、ワクワクしながら考えているフェーズです。まだまだ伸びしろが多く、やりたいこともたくさんあるので、一緒にワクワクしながらデータ基盤を作っていける仲間を募集しています。まずは、カジュアル面談からご応募ください!お話しましょう!お待ちしています!

jobs.layerx.co.jp

データ関連の求人票

open.talentio.com open.talentio.com

*1:本番データを複製して、dbtによるデータ変更、およびテストのみを実行可能な環境のことを指す。

*2:https://docs.snowflake.com/en/sql-reference/parameters#label-allow-id-token

*3:SSOに加えて多要素認証の設定も行っている場合、 ALLOW_CLIENT_MFA_CACHINGもtrueに設定しておく必要がある。https://docs.snowflake.com/en/sql-reference/parameters#label-allow-client-mfa-caching

*4:https://github.com/snowflakedb/snowflake-connector-python/blob/01b2c17e40527172702d028587d991eb4b4b6d0f/setup.cfg#L98-L99

*5:https://support.apple.com/en-US/guide/keychain-access/kyca1083/mac

*6:https://learn.microsoft.com/en-us/windows/uwp/security/credential-locker

*7:keyringというPythonのパッケージを使用している。そのため、MacOSまたはWindowsを使用している限りにおいて他のセキュアストレージサービスを利用できる可能性がある

*8:ドキュメント上はMacOSとWindowsのみ対応している旨が記載されているが、キャッシュの実装を読むとLinuxでも動作するように読める。しかし、ファイルに平文で保存される実装なので使用するべきではないだろう。