SnowflakeのWorkload identity federationをAWS Lambdaから利用してみた(Python)
全データエンジニアが待望にしていた Snowflake の Workload identity federation がGAになりました!🎉
めでたい!^ ^
LambdaでPythonを利用してSnowflakeと接続して処理をしていく仕組みを直近リリースする予定だったこともあり、早速検証してみたので簡単に内容をまとめていきます。
そもそもWorkload indentity federationって何?
クラウドプロバイダーのネイティブなアイデンティティを利用して認証ができる仕組みです。つまりAWSのIAMなどのセキュリティ強度が最高に高いレベルで確立している認証を利用してSnowflakeに接続することが可能になります。
クラウドプロバイダーのネイティブな仕組みに相乗りできるので、わざわざ認証のための別の仕組みをつくったり、シークレットの管理をしたりなどの運用が不要になるというのがポイントです。
つまり、この仕組みに対応してくれているクライアントアプリケーションでは以下の煩わしいことから解放されます。
- キーペア認証のための秘密鍵の管理
- 接続元のIP制限のためのネットワークポリシーの管理
- OAuthのアプリなどの設定
- Programatic Access Tokenという名のローテーションが必要なパスワードの管理
ドキュメントはこちら
今回取り上げるパターン
今回はAWSのIAMロールとsnowflake-connector-pythonを利用したシンプルな実装を検証していきます。(ちょうどそれで利用しようとしているので)
検証の流れ
- PythonのLambdaレイヤーを作成する
- Lambdaを作成する
- IAMロールのARNを取得する
- Snowflakeでユーザーを作成する
- Lambdaをテスト実行する
Lambdaレイヤーを作成する
私はローカルの環境を汚したくないのでDevcontainerを利用します。
{
"name": "Python 3",
"image": "mcr.microsoft.com/devcontainers/python:1-3.12-bullseye",
"features": {
"ghcr.io/jajera/features/zip:1": {}
}
}
Devcontainerを立ち上げたら、次はレイヤーのzipファイルを作成します。
Workload identity federation をサポートしている snowflake-connector-python のバージョンが 3.17.0 なので注意しましょう。
それではrequirements.txtを作成します。
snowflake-connector-python >= 3.17.0
それでは続いて以下のコマンドを順番に実行してzipファイルを作成します。
# 一応pipのバージョンを上げておく
pip install --upgrade pip
# pythonフォルダにパッケージをinstallする
pip install -r requirements.txt -t python/
# zipにまとめる
zip -r layer.zip python/
マネジメントコンソールにアクセスしてLambdaのレイヤーを作ります。
手順はこちら
Lambdaを作成する
マネジメントコンソールでLambdaを作ります。
IAMロールはcloudwatch logsの権限があれば大丈夫です。なので、関数を作成するときに一緒にロールを作成しちゃいます。
- ランタイムはdevcontainerと同じ Python 3.12 を選びます。
- アーキテクチャは arm64 を選びます。
- デフォルトの実行ロールの変更は何もいじらないで作りましょう。(実はここは少し罠なのですが…)
作成したら先ほど追加したカスタムのレイヤーの追加をして、コードも雑に作っちゃいます。
import json
import snowflake.connector
def lambda_handler(event, context):
"""Workload identity federationのテスト"""
with snowflake.connector.connect(
account='<snowflake_account>', # アカウント識別子を更新する
authenticator='WORKLOAD_IDENTITY',
workload_identity_provider='AWS'
) as conn:
with conn.cursor() as cur:
connected_user = cur.execute("SELECT current_user();").fetchall()
return {
'statusCode': 200,
'body': json.dumps(connected_user)
}
あと念のためタイムアウトを15秒くらいに設定しておきましょう!(9秒くらいかかる時あったのでデフォルトの3秒だと短すぎた)
IAMロールのARNを取得する
マネジメントコンソールからLambdaのIAMロールのARNを取得します。
こんな感じでした(アカウントIDだけ隠しています。)
arn:aws:iam::<account>:role/service-role/test_snowflake_workload_identity_federation-role-lo7dzts3
AWSが作成したということもあり、ロール名のパス(service-role/)がついているのが特徴的です。
ちなみにこのIAMロールのパスって何ものなの?という解説はこちらが詳しいです。
Snowflakeでユーザーを作成する
雑にユーザーを作ります。本当はロールとか色々設定が必要ですが最小限の設定だけです。
-- Workload identity federation を設定したユーザを作成
-- accountだけ隠しています
create user test_workload_identity_federation
WORKLOAD_IDENTITY = (
TYPE = AWS
ARN = 'arn:aws:iam::<account>:role/service-role/test_snowflake_workload_identity_federation-role-lo7dzts3'
)
type = service
default_role = public;
これで準備は完了です。
一応desc userを実行すると「HAS_WORKLOAD_IDENTITY」がtrueになっている事がわかります。
(あれ、情報ってこれだけ?ってのがまた一つの罠なのです。そのうち改善されるのかもしれないのですが、一番気軽なdesc userではどんな値が設定されているのが簡単にわからないのが難しいところです。)
desc user test_workload_identity_federation;
「wokload」で検索している図
Lambdaをテスト実行する
それではドキドキとワクワクを感じながら意気揚々とテスト実行してみましょう!
ポチッとな!!
…
…
エラー!!!
errorMessage": "250001 (08001): Failed to connect to DB: <account>.snowflakecomputing.com:443. Unable to find an associated Snowflake user for the verified AWS Caller Identity.",
あれ、ちゃんと設定したつもりなのに…
(茶番ですが…)
何が悪かったのか?
結論から言うと、Snowflake側の設定をミスっていたのです。
(desc userで設定があるかどうかのフラグ情報しか出なかったから気づくのに時間がかかりました)
トラブルシューティング的に別のユーザー名でユーザーを作ってみましょう。
create user test_workload_identity_federation_2
WORKLOAD_IDENTITY = (
TYPE = AWS
ARN = 'arn:aws:iam::<account>:role/service-role/test_snowflake_workload_identity_federation-role-lo7dzts3'
)
type = service
default_role = public;
エラーが出ます!
AWS Authenticator for AWS Account '<account>' and IAM role 'service-role' already exists in this Snowflake account.
あれ?
IAM role名が service-role になっているぞ…?と気づくのです。
それでは service-role のパス部分を除いて設定し直してみましょう。
create user test_workload_identity_federation_2
WORKLOAD_IDENTITY = (
TYPE = AWS
ARN = 'arn:aws:iam::<account>:role/test_snowflake_workload_identity_federation-role-lo7dzts3'
)
type = service
default_role = public;
今度こそ!とテスト実行をポチッとする!
成功!
{
"statusCode": 200,
"body": "[[\"TEST_WORKLOAD_IDENTITY_FEDERATION_2\"]]"
}
やったー!🎉
Workload identity federationの設定値の確認方法
私が確認している範囲にはなるので他の方法もあるかもしれませんが、snowflakeデータベースのaccount_usageの中にcredentialsというビューがあるので、これを利用して設定値のトラブルシューティングもできると思います。(最大2時間の遅延があるので注意)
use role accountadmin;
select
user_name,
type,
domain,
additional_details,
from snowflake.account_usage.credentials
where domain = 'WORKLOAD_IDENTITY_FEDERATION_METHOD';
CREDENTIALSビューをクエリしてみた
まとめ
SnowflakeのWorkload identity federationを検証してみました。
学びはこんなところでしょうか?
- IAMロールとSnowflakeのユーザーは1:1で紐づくからユーザー名は設定不要
- IAMロールのパスは除く必要がある
- セキュリティ面も堅牢になるし、運用の負荷も下がる
今後もっといろんなクライアントツールが対応してくれる事に期待したいですね。
おしまい(最後まで読んでいただきありがとうございました!)
Discussion