【AWS】DynamoDBをもっと簡単に扱えるライブラリ作った(Python)

目的

DynamoDBへのアクセスをもうちょっと簡単に書けるように、汎用的な関数作ってライブラリ化したいと思ってたんです。

完成品は GitHub へどうぞ。
https://github.com/umaxyon/dymdao

何がイヤなのか

通常の書き方はこんな感じ。

import boto3

db = boto3.resource("dynamodb")

# 1件取得
dat = db.Table('stock_report').get_item(Key={'code': '1383'})
print(dat['Item'])

# ソートキーありの複合キーテーブルから、パーティションキー(HASHキー)のみ指定して複数件取得
dat = db.Table('tbl_hogehoge').query(
    TableName='tbl_hogehoge',
    KeyConditionExpression="#a = :val",
    ExpressionAttributeNames={"#a": 'mykey'},
    ExpressionAttributeValues={":val": '3238'},
    ScanIndexForward=True
)
print(dat['Items'])

• 「Key=」とか書きたくない。というかキー名も書きたくない。覚えたくない。
• 毎回「Item」とか書きたくない。複数件の時は「Items」とか面倒くさい。常にリストで返してくれよ。
• get_itemはまぁ簡単だけど、queryになると指定項目が長くて多いし面倒。覚えたくない。でもコピペもしたくない。
• queryは何でTableNameが必要なの？Tableで指定済みでしょ？
• というか、get_itemとかqueryとか使いわけたくない。パーティションキー(HASHキー)だけ指定したらシュッと欲しいモノのリスト返してくれよ。

NoSQLなんだしもっとキーバリューストアぽく、チャラっと書いてサクッと取れるAPIが欲しいんです。

こう書きたい

当記事では以下のような、独自メソッドを追加しつつ、既存のメソッドも使えるライブラリを作成する手順を説明していきます。

    # 独自のクラス作って、お気軽に使えるオリジナルメソッド「find」を作成します。
    dao = DymDao()
    tbl = dao.table('tbl_hogehoge')

    # HASH、RANGE 両指定(get_item実行)
    dat = tbl.find('3238', '2021/02/03')

    # HASH のみ指定(query実行)
    dat = tbl.find('3238')

    # query(降順)
    dat = tbl.find('3238', asc=False)

    # 別のオプションも指定可能
    dat = tbl.find('3238', option={"Limit": 2})

    # リストで戻ってくる
    print(dat[0])

    # 既存のメソッドも使える
    dat = tbl.get_item(Key={'mykey': '3238', 'row': '2021/02/03'})

    # 「Item」とか「Items」とかいらない。
    print(dat)

• 独自のfindメソッド等を自由に追加できる場所を用意したい。
• テーブル定義からキー名を探したい。
• でも既存の dynamodb.Table のメソッドも、必要に応じて使いたい。

作戦

キー名を書かずに済ますために、テーブル定義をどうにかして入手する必要があります。

{
    'TableName': 'tbl_hogehoge',
    'KeySchema': [
        {"AttributeName": 'mykey', "KeyType": 'HASH'},    # ← コレがほしい
        {'AttributeName': 'row', 'KeyType': 'RANGE'}    # ← コレがほしい
    ],
    "AttributeDefinitions": [
        {"AttributeName": 'mykey', "AttributeType": 'S'},
        {'AttributeName': 'row', 'AttributeType': 'N'}
    ],
    "ProvisionedThroughput": {"ReadCapacityUnits": 1, "WriteCapacityUnits": 1}
}

テーブル定義の取得法は、以下のどっちかで考えることに。

1. 「describe_table」メソッドで動的に取得する

awsへのアクセスが1回増える事に目をつむるなら、これが簡単かも。

client = boto3.client("dynamodb")
ddl = client.describe_table(TableName="tbl_hogehoge")

2. テーブル定義をYAMLファイルに保存しておいて読み込む

CloudFormationでテーブル作成している場合、以下のようなYAML、またはJSONがすでに手元にあるかも。

AWSTemplateFormatVersion: "2010-09-09"
Resources:
  TblHogehoge:
    Type: AWS::DynamoDB::Table
    Properties:
      TableName: "tbl_hogehoge"
      AttributeDefinitions:
        - AttributeName: "mykey"
          AttributeType: "S"
        - AttributeName: "row"
          AttributeType: "N"
      KeySchema:
        - AttributeName: "mykey"
          KeyType: "HASH"
        - AttributeName: "row"
          KeyType: "RANGE"
      ProvisionedThroughput:
        ReadCapacityUnits: "1"
        WriteCapacityUnits: "1"

pyyamlをpipインストールしてロードすると良いですね。
でも今回の説明は、テーブル定義をファイル管理してない人のために、1.で行きます。

ということで作る

DynamoDao 。。って名前にしようかと思ったけど、何か公式っぽくて恐れ多いので、DymDaoという名前にします。(Data Access Object で DAO)

1. とりあえずガワを作る

まずはクラス作って、元と同じく Tableメソッドでテーブル名指定できるようにします。
※でもPEP8に乗っ取り、 メソッド名 Table を lowercase に変えます。

import boto3

class DymDao:
    def __init__(self):
        self.db = boto3.resource('dynamodb')
        self.client = boto3.client('dynamodb')   # テーブル定義のdescribe用にclientも用意しとく

    def table(self, table_name):
        return self.db.Table(table_name)  # このTableに独自のメソッドを追加したい！

tableメソッドはとりあえず dynamodb.Table をそのままreturnしてますが、独自の find メソッド等を追加したいので、Tableをラップする 「WrapTable」クラスを作成することにします。

class WrapTable:
    def __init__(self, dao, table_name):
        self.dao = dao
        self.table_name = table_name
        self.table = dao.db.Table(table_name)  # Table取得を DymDao からここに移動

        # テーブル定義をdescribe
        self.ddl = dao.client.describe_table(TableName=table_name)
        self.key_schema = self.ddl['Table']['KeySchema']  

    def find(self, key):  # とりあえずfindメソッド追加
        pass

Table の取得・保持は WrapTable で行う事にして、DymDao は WrapTable を返すようにします。

class DymDao:
    def __init__(self):
        self.db = boto3.resource('dynamodb')
        self.client = boto3.client('dynamodb')

    def table(self, table_name):
        return WrapTable(self, table_name)   # 変更

ここまでで、以下のように書けるようになりました。

dao = DymDao()
dao.table('tbl_hogehoge').find('1234')  # まだfindメソッド内は未実装

でも WrapTable を返しているので、これでは既存の dynamodb.Table のメソッドが使えません。

2. WrapTableにDynamoDBのメソッドを移植する

WrapTable に「既存table の全メソッドに対する移譲メソッド」を作っても良いんですが、面倒なので python の inspectモジュールと setattr を使用して、WrapTable に既存メソッドを移植してしまいます。¹

WrapTableに以下のPrivateなメソッドを作り、コンストラクタから呼び出します。

import inspect
...

    def __init__(self, db, client, table_name):
        .... 略 ...
        self.__register_method()  # 呼び出し

    def __register_method(self):
        # dynamodb.Tableからメソッドを抽出
        methods = inspect.getmembers(self.table, inspect.ismethod)

        for name, func in methods:
            setattr(self, name, func)  # 自分(self)に、全メソッドをsetattr。

これおもしろいですね。メソッドだけ移植してもちゃんと動くんです。メソッドそのものが移動するんじゃなく、メソッドへの参照が setattr される感じなのかな。これにより WrapTable は dynamodb.Table の代理でメソッド呼び出しを受け付ける事が出来るようになります。

ここまでで、既存のメソッドをそのまま利用できて、かつ、好きなようにメソッドを追加可能な WrapTable という場所を手に入れることが出来ました。

dao = DymDao()

# WrapTable から 既存の get_item を呼び出せる！
dat = dao.table('tbl_hogehoge').get_item(Key={'mykey': '1383'})
print(dat['Item'])

3. 「Item」「Items」を何とかする

まず、戻り値からItemを取り除くメソッドを WrapTable に作成します。頑張って力業で取り出します。

    @staticmethod
    def pick_out_item(obj):
        if type(obj) == dict:  # 辞書以外は無視
            if 'Item' in obj:
                return obj['Item']
            if 'Items' in obj:
                items = []
                for row in obj['Items']:
                    items.append(row['Item'] if type(row) == dict and 'Item' in row else row)
                return items
        return obj

さらに、この Item除去処理を「後処理として追加した関数を作る」メソッドを作成します。デコレーター関数ですね。

    @staticmethod
    def __intercept(method):
        def _m(*args, **kwargs):
            ret = method(*args, **kwargs)   # 受け取ったメソッドを実行して、、
            return WrapTable.pick_out_item(ret)  # 結果retを Item除去メソッドにかけてから返す、、
        return _m  # 、、というクロージャ「_m」を作って返す

これを、前節の setattr でメソッド移植している箇所に適用します。

    def __register_method(self):
        methods = inspect.getmembers(self.table, inspect.ismethod)
        for name, func in methods:
            # 既存メソッド func をデコレートして setattr する
            setattr(self, name, self.__intercept(func))   # 変更

これで、DymDaoを使用した既存のメソッド呼び出しの戻り値から、「Item」「Items」を消し去る事が出来ました。

dao = DymDao()
dat = dao.table('tbl_hogehoge').get_item(Key={'mykey': '1383'})

print(dat)  # Itemとか書かなくていい！

4. findメソッドを実装する

まず、最初に取得しておいたテーブル定義の KeySchema は、「List in Dict」で、こんな形。

# リストの中に辞書が入ってる。
[
    {"AttributeName": 'key', "KeyType": 'HASH'},
    {'AttributeName': 'row', 'KeyType': 'RANGE'}
]

ここからHASHキーとRANGEキーについて、それぞれ AttributeName を取り出す処理を作ります。 (ただし、RANGEキーは無い場合もあります。)

どこかで見かけた next とジェネレータを使ったワンライナーで取り出します。便利。
以下のメソッドを WrapTable に追加し、コンストラクタのKeySchema取得処理と置き換えます。

    def __get_key_name(self):
        schema = self.ddl['Table']['KeySchema']

        # next関数は第2引数でデフォルト値指定できるので、無ければNoneを返すように指定しておく。
        hash_name = next((r['AttributeName'] for r in schema if r['KeyType'] == "HASH"), None)
        range_name = next((r['AttributeName'] for r in schema if r['KeyType'] == "RANGE"), None)

        return hash_name, range_name  # taple でいいや

コンストラクタ側。

class WrapTable:
    def __init__(self, dao, table_name):
        self.dao = dao
        self.table_name = table_name

        self.table = self.dao.db.Table(table_name)
        self.ddl = self.dao.client.describe_table(TableName=table_name)

        self.key = self.__get_key_name()  # 呼び出し
        self.__register_method()

これで材料はそろったので、find メソッドを実装します。HASHとRANGEの両方を指定された場合は get_item を、それ以外の時は query を実行するメソッドにします。

    def find(self, hash_value, range_value=None, asc=True, option=None):
        hash_name, range_name = self.key
        opt = option if option is not None else {}

        if range_value is not None:
            # range_value指定在りの時は、get_itemで1件取得
            key_param = {hash_name: hash_value, range_name: range_value}
            ret = self.table.get_item(Key=key_param, **opt)
            ret = self.pick_out_item(ret)

            # query の実行結果と合わせるためにリストに入れて返す。無ければ空リスト。
            return [ret] if ret is not None else []
        else:
            # range_value指定なしの時は、queryで取得
            query_params = {
                "TableName": self.table_name,
                "KeyConditionExpression": "#a = :val",
                "ExpressionAttributeValues": {":val": hash_value},
                "ExpressionAttributeNames": {"#a": hash_name},
                "ScanIndexForward": asc,
            }
            query_params.update(opt)  # 追加オプションがあれば足す

            ret = self.table.query(**query_params)
            return self.pick_out_item(ret)

完成

ということで、完成
以下のように使います。

    dao = DymDao()
    tbl = dao.table('tbl_hogehoge')

    # HASH、RANGE 両指定(get_item実行)
    dat = tbl.find('3238', '2021/02/03')

    # HASH のみ指定(query実行)
    dat = tbl.find('3238')

    # query(降順)
    dat = tbl.find('3238', asc=False)

    # 別のオプションも指定可能
    dat = tbl.find('3238', option={"Limit": 2})

    # リストで戻ってくる
    print(dat[0])

    # 既存のメソッドも使える
    dat = tbl.get_item(Key={'mykey': '3238', 'row': '2021/02/03'})

    # 「Item」とか「Items」とかいらない。
    print(dat)

これで DynamoDB の検索は「とりあえずfind呼んどけ」って感じで使えるかな。足りない機能は WrapTable の中で頑張ればどうにでも出来ますね。煩雑な処理を閉じ込める共通処理を書く場所として使っていきたいです。
（全文は GitHub で公開してます。）

1. 「わざわざこんなことしなくても、Tableクラスを拡張すればよいのでは？」と思うかもしれないですが、Tableクラスはboto3内で「ResourceFactory」というクラスが動的に生成しており、拡張クラスを定義することができないのです。 ↩