DataAPIで楽々、Lambda関数から別アカウントのRedshiftにクエリを投げてみた。

別アカウントのLambdaからでもRedshiftにクエリが叩けるぞ!
2020.11.04

先日Lambda関数からData API for Redshift(以降「DataAPI」)クエリの実行を行ってみました。

サブネットもセキュリティグループも考慮不要!Lambda関数からData API for Redshiftでクエリ実行してみた。

(以降、「前回の記事」)

今回はそこから少し発展させて、Lambda関数から別アカウントのVPCにあるRedshiftにクエリを投げてみます。 DataAPIはIAMによる認証でクエリが投げられるサービスですので、 AWSアカウントが違う場合でも、IAMによる設定を適切に行えばクエリを投げることが可能になります!

構成は簡単ですが、図にするとこんな感じです。

Redshiftがあるアカウントを「Redshiftアカウント」、 Lambda関数があるアカウントを「Lambdaアカウント」と呼称します。

なお、前回の記事がある前提で基本的な説明は省略していますので、 この記事の前に前回の記事をご参照頂ければと思います。

Lambda関数(暫定)を用意

アクセスを許可するLambdaアカウント側のIAMロールを確定させるためにLambda関数を作成します。 一旦前回の記事のLambda関数コードを貼り付けておきます。 もちろんこのまま実行しても失敗します。

RedshiftアカウントにIAMロールを作成する

Redshiftアカウントに、Redshiftにアクセス可能なIAMロールを作成します。 Lambda関数からこのロールにAssumeRoleを行うことによって、 別アカウントのRedshiftへの接続を可能にします。

「別のAWSアカウント」を選び、LambdaアカウントのIDを入力。

このロールがRedshiftにアクセスするので「AmazonRedshiftFullAccess」を選択する。

タグは任意で。

ロールの名前を入力。

ロールが作られました。このロールの引き受けができるリソースを上記で作成したLambda関数に限定したいので、 「信頼関係の編集」をクリック。

PrincipalAWSの部分に、Lambda関数にアタッチされているIAMロールを指定します。

RedshiftアカウントのIAMロールの設定は完了です。 続いて、接続元のLambda関数にアタッチされたロールに上記IAMロールへのAssumeRoleを許可します。

以上でIAM周りの設定は完了です。

Lambda関数の修正

Lambda関数のコードに、RedshiftアカウントのIAMロールにAssumeRoleする処理を入れなくてはいけません。 とは言っても、以下のようにsessionを作成するようにするだけです。

import json
import time
import boto3
from boto3.session import Session

# Redshift接続情報
CLUSTER_NAME='cluster_name'
DATABASE_NAME='sample_db'
DB_USER='dbuser'
ROLE_ARN='arn:aws:iam::10XXXXXXXXXX:role/AllowRedshiftDataApiCrossAccount'

# 実行するSQLを設定
sql = '''
    SELECT * FROM public.sample_table;
'''

def create_session():
    sts_client = boto3.client('sts', endpoint_url='https://sts.ap-northeast-1.amazonaws.com')
    credentials = sts_client.assume_role(RoleArn=ROLE_ARN,
                                             RoleSessionName='ResfhitDataAPI')
    accesskey = credentials['Credentials']['AccessKeyId']
    secretkey = credentials['Credentials']['SecretAccessKey']
    session_token = credentials['Credentials']['SessionToken']
    session = Session(aws_access_key_id=accesskey,
                      aws_secret_access_key=secretkey,
                      aws_session_token=session_token)
    return session


def lambda_handler(event, context):
    # Redshiftにクエリを投げる。非同期なのですぐ返ってくる
    session = create_session()
    data_client = session.client('redshift-data')
    result = data_client.execute_statement(
        ClusterIdentifier=CLUSTER_NAME,
        Database=DATABASE_NAME,
        DbUser=DB_USER,
        Sql=sql,
    )

    # 実行IDを取得
    id = result['Id']
    #print('id = {}'.format(id))

    # クエリが終わるのを待つ
    statement = ''
    status = ''
    while status != 'FINISHED' and status != 'FAILED' and status != 'ABORTED':
        statement = data_client.describe_statement(Id=id)
        #print(statement)
        status = statement['Status']
        time.sleep(1)

    # 結果の表示
    if status == 'FINISHED':
        if int(statement['ResultSize']) > 0:
            # select文等なら戻り値を表示
            statement = data_client.get_statement_result(Id=id)
            print(json.dumps(statement['Records']))
        else:
            # 戻り値がないものはFINISHだけ出力して終わり
            print('QUERY FINSHED')
    elif status == 'FAILED':
        # 失敗時
        print('QUERY FAILED\n{}'.format(statement))
    elif status == 'ABORTED':
        # ユーザによる停止時
        print('QUERY ABORTED: The query run was stopped by the user.')

sessionを作ってsession.client('redshift-data')することによって、 そこで作られるインスタンスはRedshiftアカウント内でのロール権限で動作します。 あとは前回同様、クエリを流して終了を待てば結果を得ることができます。

まとめ

DataAPIを使用することで、別アカウントのRedshiftにもLambdaからクエリを投げることができました!

アカウントAのRedshiftからデータを取り出してアカウントBのRedshiftにデータを入れる などの処理が単独のLambda関数で実行可能となり、構成がかなりシンプルになります。

参照リンク