PythonでDynamoDBにローカルからテストデータを投入する

June 6, 2022

#Python

#AWS DynamoDB

やりたいこと

開発環境でDynamoDBへテストデータを投入したい時、そのためにわざわざS3イベント使ってLambdaから投入するのは面倒

ローカルからAWS CLIのprofileを使ってDynamoDBへアクセスし、テストデータを投入する方法を備忘録として残しておきます

前準備

AWS CLIのprofileを設定

事前にIAM Userを作成してaws_access_key_id, aws_secret_access_keyを取得し、以下のようにinsert_dbという名前でprofileを作成します

$ aws configure --profile insert_db

AWS Access Key ID [****************XXXX]: 
AWS Secret Access Key [****************XXXX]: 
Default region name [ap-northeast-1]: 
Default output format [json]:

DynamoDBをCDKで構築

dynamodb_stack.py

from aws_cdk import (
    Stack,
    RemovalPolicy,
    aws_dynamodb as dynamodb
)
from constructs import Construct


class DynamoDBStack(Stack):

    def __init__(self, scope: Construct, construct_id: str, **kwargs) -> None:
        super().__init__(scope, construct_id, **kwargs)

        Table = dynamodb.Table(self, "Table",
                               partition_key=dynamodb.Attribute(
                                   name="id", type=dynamodb.AttributeType.STRING),
                               sort_key=dynamodb.Attribute(
                                   name="meta", type=dynamodb.AttributeType.STRING),
                               billing_mode=dynamodb.BillingMode.PAY_PER_REQUEST,
                               time_to_live_attribute="ttl",
                               removal_policy=RemovalPolicy.DESTROY
                               )
        Table.add_global_secondary_index(
            partition_key=dynamodb.Attribute(
                name="meta", type=dynamodb.AttributeType.STRING),
            sort_key=dynamodb.Attribute(
                name="id", type=dynamodb.AttributeType.STRING),
            index_name="meta-id-idx"
        )

大量件数入れた後、deleteコマンドで消す場合も課金かかってしまうのでtime_to_live_attributeを設定しています。スタックを消す際にDBリソースも消えるようremoval_policyも設定しました。

リソースを作ったらDynamoDBのリソース名を控えておいてください以下のようにCLIからも取得できます

$ aws dynamodb list-tables --profile insert_db
{
    "TableNames": [
        "DynamoDBStack-TableXXXXXXXXXXXXXXXX"
    ]
}

実行ファイルの作成

insert_users.py

import sys
from datetime import datetime as dt
import random
from boto3.session import Session
from logging import getLogger, StreamHandler, Formatter, INFO

logger = getLogger(__name__)
logger.setLevel(INFO)

handler = StreamHandler(sys.stdout)
logger.addHandler(handler)
fmt = Formatter('%(levelname)s %(asctime)s %(message)s')
handler.setFormatter(fmt)

profile = "insert_db"
table_name = "DYNAMODB_TABLE"

session = Session(profile_name=profile)
dynamodb = session.resource('dynamodb')
table = dynamodb.Table(table_name)


def batch_write(table, items):
    try:
        with table.batch_writer() as batch:
            for item in items:
                batch.put_item(Item=item)
    except Exception as e:
        raise e


def create_userview(
    user_id,
    meta,
    timestamp_now,
    ttl
):
    random_num = random.randrange(3)
    record = {
        "id": f"UserView:{user_id}",
        "meta": meta,
        "random_number": random_num,
        "created_at": timestamp_now,
        "ttl": ttl
    }
    return record


def main():
    start_user_num = 0
    number_of_user = 10000
    timestamp_now = int(dt.timestamp(dt.now()))

    ttl_day = 1
    ttl = timestamp_now + (ttl_day * 24 * 60 * 60)

    batch_write_records = []
    logger.info(f"start_user_num: {start_user_num}")

    for user_i in range(start_user_num, number_of_user + 1):
        if user_i % 1000 == 0:
            logger.info(f"now user_i: {user_i}")

        user_id = f"CSV-STRESS-TEST-{user_i}"

        userview_record = create_userview(
            user_id=user_id,
            meta="User",
            timestamp_now=timestamp_now,
            ttl=ttl
        )
        batch_write_records.append(userview_record)

    logger.info("start db isert")
    batch_write(
        table=table,
        items=batch_write_records
    )
    logger.info("db isert is done")


if __name__ == "__main__":
    main()

boto3の場合、大量投入用にbatch_writer()が使えます

CLIなどでは(boto3にもあります)batch_write_itemがありますがこちらは一度に25itemまでが上限ですが、batch_writer()では特に制限なく投入できます

CDKの方で設定したtime_to_live_attribute用にttlを項目に設けました

削除したい時間をUNIXタイムスタンプ(ミリ秒ではなく秒)で設定することで書き込み料金がかからず消すことが出来ます

profileからboto3でDynamoDBへリクエストする設定はこちら

profile = "insert_db"
table_name = "DYNAMODB_TABLE"

session = Session(profile_name=profile)
dynamodb = session.resource('dynamodb')
table = dynamodb.Table(table_name)

以上、準備が出来たら実行します

$ python3 insert_users.py

このようにログが出たら成功です

おつかれさまでした

INFO 2022-06-06 11:18:15,961 start_user_num: 0
INFO 2022-06-06 11:18:15,961 now user_i: 0
INFO 2022-06-06 11:18:15,972 now user_i: 1000
INFO 2022-06-06 11:18:15,981 now user_i: 2000
INFO 2022-06-06 11:18:16,001 now user_i: 3000
INFO 2022-06-06 11:18:16,009 now user_i: 4000
INFO 2022-06-06 11:18:16,017 now user_i: 5000
INFO 2022-06-06 11:18:16,024 now user_i: 6000
INFO 2022-06-06 11:18:16,032 now user_i: 7000
INFO 2022-06-06 11:18:16,040 now user_i: 8000
INFO 2022-06-06 11:18:16,048 now user_i: 9000
INFO 2022-06-06 11:18:16,079 now user_i: 10000
INFO 2022-06-06 11:18:16,079 start db isert
INFO 2022-06-06 11:18:16,079 db isert is done

← Back to home