Amazon Athena + Google Colabによる分析環境

1日目にも登場した謎の事業本部のデータエンジニアv0.1.0の @mashiike です。
こちらは Tech KAYAC Advent Calendar 2020 11日目の記事です。

今回は、いつものRedashの話ではなくAmazon AthenaとGoogle Colaboratoryの話をしたいと思います。

背景

弊社はWebアプリケーションをAWS上に構築することが多いです。そのため分析用の環境もAWS上に構築することが多くあります。社内でよく耳にするのは、ECSを用いてRedashを構築・運用、backupのデータベースや調査用のデータベースとつなげてグラフ化・分析などを行う環境です。複雑な分析はなどは、前年のTech KAYAC Advent Calendarでお話した、 redash + Google Colabの分析環境なども出てきたりします。

techblog.kayac.com

一方で、世の中の分析界隈では Google BigQuery と Cloud Composer の組み合わせが強力で、こちらの２つのサービスを使って分析基盤を構築する例はよく耳にすると思います。（中には、ヘビーユースする関係で自分たちで非マネージドなApatch Airflowを構築・運用する方もいらっしゃいます。）
このようなBigQueryを用いた環境では、認証周りの関係でGoogle Colabと連携しやすく非常にスムーズに分析できます。

より具体的には、公式ライブラリが提供されており以下のようなPythonコードでアクセスすることが可能です。

from google.colab import auth
from google.cloud import bigquery

auth.authenticate_user()
client = bigquery.Client(project='<プロジェクトID>')
query =  "<実行するクエリ>"
df = client.query(query).to_dataframe()

非常に魅力的です。
しかし、個人的な分析環境を構築するなら良いのですが、仕事してとなると話は変わります。
弊社は、どちらかと言うとAWS上での構築・運用の知見が集まっているため、分析向けのDBの管理やシステムの管理はAWS上に構築したほうが社内のサポートを受けやすいという状況です。ですので、AWS上で分析環境を構築しがちになります。

もっとGoogle Colabを使いたい！

社内ではGoogle SpreadsheetやGoogle Docsなども使用しております。分析の結果を必要としている人たちは、Google Workspaceと親和性が高いのでGoogle Colabをもっと活用したいと日々思っています。
しかしながら、先程の背景のとおり分析対象のデータはAWS上にあり非常にアクセスがしづらい状況が続いておりました。
（redashをバイパスすることでデータへのアクセス手段を手に入れてたのが去年です。）
技術は進化するものです。新たなデータアクセス手段を手に入れました。それが今回のGoogle Colab上からAmazon Athenaに直接クエリする方法です。

Web Identityを信頼したIAM ロール

Google ColabはマネージドJupter Notebookであり、お好きなpythonライブラリを活用することができます。もちろん boto3 も例外なく。

boto3 を使ってAWS上のリソースにアクセスすることは可能ですが、問題となるのが認証です。すこしインターネット上を検索してみると、専用のIAM ユーザーのアクセスキーをGoogle Driveにおいておくという方法などが見つかります。
しかし、専用のIAMユーザーとはいえ、Google Drive上にアクセスキーを置いておくのは心情的によろしくないので回避したいです。

そこで、登場するのが Web Identityを信頼したIAM ロールです。皆様、以下の画面はご存知でしょうか？ IAM ロールの作成画面なのですが、実はIAM ロールはWeb Identityを信頼されたエンティティとして選択することできます。そして、デフォルトで用意されている選択肢の中には、GoogleのWeb Identityも存在します。

f:id:ikeda-masashi:20201210115013p:plain — AWS コンソール - IAM ロール作成画面

今回は、一例として以下のような信頼関係の設定を紹介します。

{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Sid": "",
      "Effect": "Allow",
      "Principal": {
        "Federated": "accounts.google.com"
      },
      "Action": "sts:AssumeRoleWithWebIdentity",
      "Condition": {
        "StringEquals": {
          "accounts.google.com:email": "<分析者のGoogle Workspace上のemail address>",
          "accounts.google.com:aud": "32555940559.apps.googleusercontent.com"
        }
      }
    }
  ]
}

32555940559.apps.googleusercontent.com というのは Google Cloud SDKのaudです。
より強固なセキュリティにしたい場合は、専用のGCP Projectにて設定した値を用いるのが良いでしょう。
この信頼関係を設定したIAM ロールにAthenaへのクエリのためのポリシーをアタッチしておきます。
とりあえずは AmazonAthenaFullAccess をアタッチしても良いですが、本来なら色々と権限を絞ったポリシーを用意することをおすすめします。

Google Colab上でWeb Identityを取得する。

Google Colab上でWeb Identity(ID Token)を取得することはそれほど難しくなく、様々な手段があります。
一例として、以下のようなコードで取得することが可能です。

!pip install PyJWT
import google.colab.auth as _colab_auth
import google.auth as _google_auth
import google.auth.transport.requests as _google_auth_requests 
import jwt

_colab_auth.authenticate_user()
creds, _ = _google_auth.default()
transport = _google_auth_requests.Request()
creds.refresh(transport)
jwt.decode(creds.id_token, verify=False)

これでGoogleのID Tokenが取得できました。あとは、このID Tokenを用いてAssumeRoleするだけですね。

Google Colab上からAmazon Athenaへクエリする

boto3のドキュメントを確認すると、stsのクライアントには assume_role_with_web_identity という関数があります。この関数と先程のID Tokenを用いてsessionを取得し、Athenaクライアントを生成します。

サンプルコードとしては以下のような感じです。

import uuid
import boto3
from boto3.session import Session
session_id = str(uuid.uuid4())[-12:]

client = boto3.client('sts')
IAM_ROLE_ARN = '<Google Web Identityを信頼したIAM ロールのARN>'
IAM_ROLE_SESSION_NAME_PREFIX = 'gcolab'
REGION_NAME = 'ap-northeast-1'

response = client.assume_role_with_web_identity(
    RoleArn=IAM_ROLE_ARN,
    RoleSessionName=IAM_ROLE_SESSION_NAME_PREFIX+"-"+session_id,
    WebIdentityToken=creds.id_token,
    DurationSeconds=900
)
session = Session(aws_access_key_id=response['Credentials']['AccessKeyId'],
aws_secret_access_key=response['Credentials']['SecretAccessKey'],
aws_session_token=response['Credentials']['SessionToken'],
region_name=REGION_NAME)
athena_client = session.client('athena')

これで、boto3のathenaクライアントが取得できたので、あとはクエリするだけです。クエリの実行に関しては、Classmethod様のこちらのブログ記事が参考になるかと思います。

dev.classmethod.jp

もちろん、AthenaだけでなくIAM ロールにアタッチする権限を変更すれば、様々なAWSリソースにアクセスできます。 Google Colab上からAmazon Comprehendにアクセスして、Google Docsに書かれている内容の感情分析なども面白いかもしれません。

おわりに

先日、AWSの新たなサービスとして Amazon Managed Workflows for Apache Airflow (MWAA) が発表されました。

aws.amazon.com

Google BigQuery + Cloud Composerと同様に、Athena + MWAA という分析環境の構築もいかがでしょうか？　AWSには様々なAI/MLサービスもありますので、今後のAWS上で構築する分析環境のノウハウが活発になることを期待したいです。その際には、今回のAmazon Athena + Google Colabによる分析もご活躍できたら幸いです。

明日は、弊社のフロントエンドエンジニア @iemon がなにかを話してくれるそうです。