AuroraServerlessのDataAPIの返値をいい感じに整形する

以前の記事でAPI GatewayからAuroraServerlessにDataAPIを介してクエリを投げることは成功したが、戻り値がなんとも扱いにくい。。。

DataAPIでの戻り値

[
    {
    "stringValue": "001"
    },
    {
    "stringValue": "fuga"
    },
    {
    "longValue": 123
    }
]

そこで、DataAPIの返り値をいい感じに整形するlambdaサンプルコードを書いた。
カラム名は返り値のメタデータから取得する。

サンプルコード

import json

def lambda_handler(event, context):
    key_list = []
    value_list = []

    columns = event["columnMetadata"]
    records = event["records"]

    # create key list
    for column in columns:
        key_list.append(column["name"])

    # create value list
    for record in records:
        for item in record:
            value_list.extend(list(item.values()))

    # format to json
    tmp_result = ''
    result = ''
    for record_index, item in enumerate(value_list):
        column_index = record_index % len(key_list)
        tmp_result += '"{0}": "{1}",'.format(key_list[column_index], item)
        if column_index == getLastIndexOfArray(len(key_list)):
            result += '{' + tmp_result[:-1] + '},'
            tmp_result = ''
    result = '[' + result[:-1] + ']' 

    return json.loads(result)

formatした戻り値

  {
    "id": "001",
    "name": "fuga",
    "age": "123"
  }

また、eventに渡すのはrecordsのみではなくメタデータも入れるために以下のようにincludeResultMetadataにtrueを指定する。

dataApiResponse = rdsData.execute_statement(
            resourceArn = cluster_arn,
            secretArn = secret_arn,
            includeResultMetadata = True,  #←追加
            database = 'データベース名',
            sql = 'SQL文',
            )

本来は、ここでAPI Gateway→stepfunctionsとして上記のフォーマットした結果を返そうと企んだものの、難しいことがわかった。

• API GatewayからStartExecutionを指定した場合、stepfunctionsのレスポンスを返すことができない。
• 以下にあるように、DescribeExecutionを指定したAPIメソッドを別途用意し、ポーリングすることで取得できるというがなんか微妙だ。

→ lambdaからlambdaをキックすることで対応することで解決しよう！

DataAPIコールのlambdaに記述

    Payload = json.dumps(dataApiResponse)
    res = boto3.client('lambda').invoke(
        FunctionName='上記の整形用lambda名',
        InvocationType='RequestResponse', #←Eventを指定すると非同期。
        Payload=Payload
    )
    return json.loads(res['Payload'].read())

ただし、API Gatewayは最大でも29000ミリ秒でタイムアウトする。
AuroraServerlessの初回アクセス時にはインスタンスが起動するのに20秒近くかかるのでタイムアウトは必至。ここは考慮しておく必要がある。

追記

さらに不便なことにNullが入ったデータをDataAPIで取得するとこんなヤバイのが返ってくる。。

{
    "isNull": True
}

上のコードではbooleanの結果と区別がつかなくなって困るので修正

サンプル(修正後)

def lambda_handler(event, context):

    key_list = list()
    value_type_list = list() #←追記
    value_list = list()

    columns = event["columnMetadata"]
    records = event["records"]

    # create key list
    for column in columns:
        key_list.append(column["name"])

    # create value_type list & value list
    for record in records:
        for item in record:
            value_type_list.extend(list(item.keys())) #←追記
            value_list.extend(list(item.values()))

    print(value_type_list)

    # format to json
    tmp_result = ''
    result = ''
    for record_index, item in enumerate(value_list):
        column_index = record_index % len(key_list)
        if value_type_list[record_index] == "isNull":                   #←追記
            tmp_result += '"{}": null,'.format(key_list[column_index])  #←追記
        else:                                                           #←追記
            tmp_result += '"{0}": "{1}",'.format(key_list[column_index], item)
        if column_index == getLastIndexOfArray(len(key_list)):
            result += '{' + tmp_result[:-1] + '},'
            tmp_result = ''
    result = '[' + result[:-1] + ']' 

    return json.loads(result)