Event Systemを使ってBoto3をもっと便利に使おう!!

2020.01.27

CX事業本部@大阪の岩田です。

Lambda等で利用することの多いAWS SDK for Python (Boto3) はEvent Systemと呼ばれる機構が存在し、Event Systemを利用することで独自メソッドを追加したり、既存のメソッドを拡張することができます。このブログではEvent Systemの概要と利用例についてご紹介します。

Event Systemの概要

Boto3の各種イベントに対して関数を登録するための機構です。登録した関数は該当のイベント発生時に呼び出され、イベントに関連付けられたキーワード引数が渡されます。登録した関数からキーワード引数を操作することで、Boto3の各クラスをカスタマイズすることが可能になります。

詳細についてはBoto3のドキュメントをご参照下さい。

Extensibility Guide

利用可能なイベント

Boto3には以下の3つのイベントが定義されており、これらのイベントに対して独自処理を行う関数が登録可能です。

  • creating-client-class
    • このイベントは、サービスのクライアントクラスの作成時に発行されます。このイベントに独自の関数を登録しておき、クライアントクラスに独自メソッドを追加することができます。
  • creating-resource-class
    • このイベントは、サービスのクリソースクラスの作成時に発行されます。このイベントに独自の関数を登録しておき、リソースクラスに独自メソッドを追加することができます。
  • provide-client-params
    • このイベントは、クライアントメソッドに渡されるパラメーターのバリデーション前に発行されます。 このイベントを独自の関数を登録しておき、クライアントメソッドに渡されるパラメータを加工することが可能です。

関数の登録方法

関数を登録する対象のイベントは3段階の指定が可能で、以下の優先順位を持ちます。

  1. general.specific.more_specific
  2. general.specific
  3. general

generalには利用可能なイベント(creating-client-class,creating-resource-class,provide-client-paramsの3つのうちどれか)を指定します。specificにはS3等の具体的なサービス名を、more_specificにはlist_objectsのように具体的なオペレーション名を指定します。このオペレーション名はクライアントクラスのメソッド名と必ずしもイコールにならないので注意が必要です。例えばS3クライアントのupload_fileメソッドを拡張したい場合、独自関数を登録すべきオペレーション名はPutObjectです。オペレーション名はbotocoreのdataディレクトリ配下に存在する各種JSONファイルで定義されているので、オペレーション名が不明な場合はbotocore内のJSONファイルを漁ってみましょう。

具体的なコードは以下のようになります。

def add_my_general_bucket(params, **kwargs):
    ....略

s3 = boto3.client('s3')
event_system = s3.meta.events
event_system.register('provide-client-params.s3', add_my_general_bucket)

上記の例ではs3クライアントの各メソッド呼び出しに対してadd_my_general_bucketという独自関数を登録しています。

複数の指定方法を併用してイベントを指定した場合、優先度の高い指定方法で登録した関数から順番に呼び出されます。

s3 = boto3.client('s3')
event_system = s3.meta.events
event_system.register('provide-client-params.s3.ListObjects', my_func1)
event_system.register('provide-client-params.s3', my_func2)
event_system.register('provide-client-params', my_func3)

この例だと、S3クライアントのlist_objectsを呼び出した際my_func1,my_func2,my_func3の順に独自関数が呼び出されます。

実装例

ここからは実装例を紹介させて頂きます。

独自メソッドを追加する

S3クライアントにupload_file_exという独自メソッドを追加してみましょう。基本的には通常のupload_file_exと同様の動作ですが、独自にRequestIdというキーワード引数を追加、RequestIdが指定されていた場合はアップロードするオブジェクトのメタデータにrequest_idというメタデータを付与します。

from boto3.session import Session

def add_custom_method(class_attributes, **kwargs):
    class_attributes['upload_file_ex'] = custom_method

def custom_method(self, **kwargs):

    if 'RequestId' in kwargs:
        meta_data = kwargs.get('ExtraArgs', {}).get('Metadata', {})      
        meta_data.update({
          'request_id': kwargs.pop('RequestId')
        })
        kwargs.update({
          'ExtraArgs': {
            'Metadata': meta_data
          }
        })

    return self.upload_file(**kwargs)

def main():
    with open('/tmp/hoge.txt', 'w') as f:
        f.write('hoge')

    session = Session()
    session.events.register('creating-client-class.s3', add_custom_method)
    client = session.client('s3')
    client.upload_file_ex(
        Filename='/tmp/hoge.txt',
        Bucket=<適当なバケット名>,
        Key='hoge.txt',
        RequestId='abcdefghijk',
    )

if __name__ == '__main__':
    main()

session.events.register('creating-client-class.s3', add_custom_method)の部分でS3クライアント生成時にadd_custom_methodという関数の呼び出しを追加、add_custom_methodの中ではupload_file_exという名前でcustom_methodというメソッドを追加しています。追加したcustom_methodの中ではキーワード引数RequestIdの有無をチェックし、RequestIdが指定されていた場合はオブジェクトのメタデータrequest_idを追加してからupload_fileを呼び出しています。

試しに上記のコードを実行後、AWS CLIから確認すると、メタデータが付与できていることが分かります。

$ aws s3api head-object --bucket <適当なバケット名> --key hoge.txt
{
    "AcceptRanges": "bytes",
    "LastModified": "Mon, 27 Jan 2020 03:35:54 GMT",
    "ContentLength": 4,
    "ETag": "\"ea703e7aa1efda0064eaa507d9e8ab7e\"",
    "ContentType": "binary/octet-stream",
    "Metadata": {
        "request_id": "abcdefghijk"
    }
}

各メソッドに渡されたパラメータを保存しておき、例外発生時にログに出力する

Lambdaを実行している場合、バグ等に起因してClientErrorが発生することがあります。ログを見ればClientErrorが発生したことは分かるのですが、ClientErrorを発生させた原因のパラメータを追いかけることができません。Lambdaのコード例です。

import json
from botocore.exceptions import ClientError
import boto3

client = boto3.client('s3')

def handler(event, context):

    try:
        res = client.list_objects(Bucket='non_exists_bucket_name')
    except ClientError as e:
        print(e)
        raise e
    
    return {
        'statusCode': 200,
        'body': json.dumps(res)
    }

list_objectsの引数に存在しないバケット名を指定した場合、このコードを実行するとCloudWatch LogsにはAn error occurred (NoSuchBucket) when calling the ListObjects operation: The specified bucket does not existというログが残ります。しかしログからはどんなバケット名を指定していたのかを追跡できません。Event Systemを使ってBoto3のクライアントクラスに渡されたパラメータを記憶しておき、例外発生時にログ出力するようにするとトレーサビリティを上げることができます。

Lambdaのコードを以下のように修正してみます。

import json
from botocore.exceptions import ClientError
import boto3

class MyClient:

    _clients = {}
    _last_params = None

    @classmethod
    def create_client(cls, service_name: str, **kwargs):

        def _provide_client_param_hook(*params, **kwargs):
            MyClient._last_params = kwargs.get('params')

        if service_name in cls._clients:
            return cls._clients[service_name]

        client = boto3.client(service_name, **kwargs)
        event_system = client.meta.events
        event_system.register('provide-client-params', _provide_client_param_hook)
        cls._clients[service_name] = client
        return client

    @classmethod
    def get_last_params(cls):
        return cls._last_params

client = MyClient.create_client('s3')

def handler(event, context):

    try:
        res = client.list_objects(Bucket='non_exists_bucket_name')
    except ClientError as e:
        print(e)
        print(f'Boto3 Client Last param:{MyClient.get_last_params()}')
        raise e
    
    return {
        'statusCode': 200,
        'body': json.dumps(res)
    }

クライアントの生成処理をMyClientというクラスのクラスメソッドcreate_clientに移譲します。このcreate_clientの処理では作成済みのクライアントクラスが存在しない場合にクライアントクラスを作成しつつ、provide-client-paramsを指定して_provide_client_param_hookという関数を登録しています。この_provide_client_param_hookという関数の中ではクライアントクラスのメソッド呼び出し時に渡されたパラメータを、クラス変数の_last_paramsにセットしています。これで例外発生時にget_last_params経由で_last_paramsを取得すれば例外を発生させた原因のパラメータが特定できるという流れです。

修正後のLambdaを実行するとCloudWatch Logsには以下のように出力されます。

An error occurred (NoSuchBucket) when calling the ListObjects operation: The specified bucket does not exist
Boto3 Client Last param:{'Bucket': 'non_exists_bucket_name', 'EncodingType': 'url'}

これでバケット名に意図せずnon_exists_bucket_nameという文字列がセットされていたことが分かります。ここまで分かればバグの修正も早そうですね。

まとめ

Boto3のEvent Systemについてご紹介しました。Boto3のクライアントクラス、リソースクラスは実行時に動的に生成されることから、これらのクラスを継承して独自メソッドを追加して...といった戦略が取れませんが、Event Systemをうまく活用することで、独自に拡張していくことは可能です。Boto3をより便利に利用するためにEvent Systemの利用を検討してみてはいかがでしょうか?