AWS LambdaによるRDSの自動リストアスクリプト

AmazonLambda

はじめに

DI部のおおたきです。前回AWS Lambdaを使ってRDSの削除スクリプトを書いてみましたが、今回は同様にpythonでスナップショットからのリストアスクリプトを書いてみました。
前回同様pythonコードは初めて書いているので、色々突っ込みどころがあるかもしれませんがご了承ください!
リストアのスクリプトを書きはじめて気がついたのですが、AWS SDKを使ってリストアする際にRDSのパラメータグループやセキュリティグループの設定はできません。これはAWS CLIでも同様みたいです。
そのため、リストア後に設定を変更する必要があります。リストア後に設定を変更するにはRDSの起動完了を待つ必要があるため、今回はLambdaからLambdaを呼出し起動が完了するのを待つようにしました。

ロールの設定をする

今回設定しているロールは以下になります。

  • AmazonRDSFullAccess
  • AWSLambdaBasicExecutionRole
  • AWSLambdaRole

ロールの設定の詳細は前回のAWS LambdaによるRDSの自動削除スクリプトを参考にしてください。 前回はAmazonRDSFullAccessのみ付与しましたが、今回はLambdaからLambdaを呼出すためにAWSLambdaRoleを、またCloudWatchにログを残すためにAWSLambdaBasicExecutionRoleも付与しました。

Lambdaの設定

今回はリストアするLambda Functionのrestore-rdsと、RDS起動後に設定を変更するLambda Functionのmodify-rdsを作成しました。
実行の流れとしては最初にCloudWatchのイベントでrestore-rdsが実行され、restore-rds内でmodify-rdsが呼出されます。
modify-rdsはRDSのステータスを確認し起動が完了していれば設定を変更する、していなければ再度自身を呼出すという処理をするようにしました。

リストア用スクリプトの設定

restore-rdsのスクリプトです。

import boto3
import logging
import os
import json

logger = logging.getLogger()
logger.setLevel(logging.INFO)
client = boto3.client('rds')

def lambda_handler(event, context):
    snapshot_id = os.environ.get('SNAPSHOT_ID')
    if not snapshot_id:
        logger.info('START get_snapshot_id')
        snapshot_id = get_snapshot_id()
        logger.info('END get_snapshot_id')
        if (snapshot_id is None):
            logger.info('not found snapshot')
            return

    logger.info('snapshot_id:' + snapshot_id)

    if (not is_exist_instance()):
        logger.info('START restore')
        restore(snapshot_id)
        logger.info('END restore')
        logger.info('START call_modify_lambda')
        call_modify_lambda()
        logger.info('END call_modify_lambda')
    else:
        logger.info('exist instance')

def restore(snapshot_id):
    logger.info('restore snapshot id:' + snapshot_id)

    client.restore_db_instance_from_db_snapshot(
        DBInstanceIdentifier = os.environ.get('INSTANCE_ID'),
        DBSnapshotIdentifier = snapshot_id,
        DBInstanceClass='db.t2.micro',
        Port=5432,
        AvailabilityZone='ap-northeast-1a',
        DBSubnetGroupName='xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx',
        MultiAZ=False,
        PubliclyAccessible=True,
        AutoMinorVersionUpgrade=True,
        OptionGroupName='default:postgres-9-6',
        StorageType='standard',
        CopyTagsToSnapshot=False
    )

def is_exist_instance():
    db_instances = client.describe_db_instances();
    target_instance = \
        filter(lambda x : x['DBInstanceIdentifier'] == os.environ.get('INSTANCE_ID'), db_instances['DBInstances'])
    return len(target_instance) > 0

def get_snapshot_id():
    snapshot_list = client.describe_db_snapshots(
        DBInstanceIdentifier=os.environ.get('INSTANCE_ID'),
        SnapshotType='manual',
     )

    logger.info(snapshot_list)
    snapshot_id = \
        snapshot_list['DBSnapshots'][0]['DBSnapshotIdentifier'] if len(snapshot_list['DBSnapshots']) > 0 else None

    return snapshot_id

def call_modify_lambda():
    params = {}
    params['instance_id'] = os.environ.get('INSTANCE_ID')
    params['retry_count'] = 0
    params['modified_flag'] = False
    client_lambda = boto3.client('lambda')
    client_lambda.invoke(
        FunctionName='modify-rds',
        InvocationType='Event',
        Payload=json.dumps(params)
    )

コードの解説をします。最初に環境変数からスナップショットIDを取得しています。復元対象のスナップショットはここで取得したスナップショットIDのスナップショットになります。

snapshot_id = os.environ.get('SNAPSHOT_ID')

しかし復元対象のスナップショットは常に最新にから復元したい場合が多いと思いますので、環境変数に値を設定していない場合はget_snapshot_idメソッドを呼出し最新のスナップショットIDを取得するようにしています。

    if not snapshot_id:
        logger.info('START get_snapshot_id')
        snapshot_id = get_snapshot_id()

次にリストア対象のインスタンスの存在確認をし、存在していなければリストアをします。リストアをしたら設定変更用のLambdaを呼出します。

    if (not is_exist_instance()):
        logger.info('START restore')
        restore(snapshot_id)
        logger.info('END restore')
        logger.info('START call_modify_lambda')
        call_modify_lambda()
        logger.info('END call_modify_lambda')
    else:
        logger.info('exist instance')
def restore(snapshot_id):
    logger.info('restore snapshot id:' + snapshot_id)

    client.restore_db_instance_from_db_snapshot(
        DBInstanceIdentifier = os.environ.get('INSTANCE_ID'),
        DBSnapshotIdentifier = snapshot_id,
        DBInstanceClass='db.t2.micro',
        Port=5432,
        AvailabilityZone='ap-northeast-1a',
        DBSubnetGroupName='xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx',
        MultiAZ=False,
        PubliclyAccessible=True,
        AutoMinorVersionUpgrade=True,
        OptionGroupName='default:postgres-9-6',
        StorageType='standard',
        CopyTagsToSnapshot=False
    )
def call_modify_lambda():
    params = {}
    params['instance_id'] = os.environ.get('INSTANCE_ID')
    params['retry_count'] = 0
    params['modified_flag'] = False
    client_lambda = boto3.client('lambda')
    client_lambda.invoke(
        FunctionName='modify-rds',
        InvocationType='Event',
        Payload=json.dumps(params)
    )

※restoreメソッド内のリストアするRDSの設定やcall_modify_lambdaメソッド内のFunctionNameはご使用する環境や作成するLambda Function名に合わせて変更して下さい。

環境変数の設定

使用する環境変数は以下の2つになります。

  • SNAPSHOT_ID:リストアする対象のスナップショットID。ただし値を設定していない場合は最新のスナップショットから復元する。
  • INSTANCE_ID:リストア対象のインタンスID

rds_restore_lambda01

設定変更用スクリプトの設定

modify-rdsのスクリプトです。

import boto3
import logging
import os
import json
from time import sleep

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

client = boto3.client('rds')

def lambda_handler(event, context):

    if (event.has_key('instance_id') and event['instance_id']):
        instance_id = event['instance_id']
    else:
        logger.info('undefined instance id')
        return

    if (event.has_key('retry_count') and isinstance(event['retry_count'],(int))):
        retry_count = event['retry_count']
    else:
        logger.info('undefined retry count')
        return

    if (event.has_key('modified_flag') and isinstance(event['modified_flag'],(bool))):
        modified_flag = event['modified_flag']
    else:
        logger.info('undefined modified flag')
        return

    logger.info('instance id:' + instance_id)
    logger.info('retry_count:' + str(retry_count))
    logger.info('modified_flag:' + str(modified_flag))

    if (retry_count > os.environ.get('RETRY_MAX_COUNT')):
        logger.info('over retry count')
        return

    response = client.describe_db_instances(
        DBInstanceIdentifier=instance_id
    )
    logger.info(response)
    instance_status = response['DBInstances'][0]['DBInstanceStatus']
    logger.info('instance status:' + instance_status)
    if (modified_flag and instance_status=='available'):
        logger.info("START reboot")
        reboot(instance_id)
        logger.info("END reboot")
    elif (modified_flag and instance_status=='modifying'):
        logger.info("waiting...")
        sleep(10)
        logger.info("START call_modify_lambda")
        call_modify_lambda(instance_id, retry_count,modified_flag)
        logger.info("END call_modify_lambda")
    elif (not modified_flag and instance_status=='available'):
        logger.info("START modify")
        modify(instance_id)
        logger.info("END modify")
        logger.info("START call_modify_lambda")
        call_modify_lambda(instance_id, retry_count,True)
        logger.info("END call_modify_lambda")
    elif (not modified_flag and (instance_status=='creating' or instance_status=='backing-up' or instance_status=='modifying')):
        logger.info("waiting...")
        sleep(120)
        logger.info("START call_modify_lambda")
        call_modify_lambda(instance_id, retry_count, modified_flag)
        logger.info("END call_modify_lambda")


def call_modify_lambda(instance_id, retry_count, modified_flag):
    params = {}
    params['instance_id'] = instance_id
    params['retry_count'] = retry_count + 1
    params['modified_flag'] = modified_flag
    client_lambda = boto3.client("lambda")
    client_lambda.invoke(
        FunctionName="modify-rds",
        InvocationType="Event",
        Payload=json.dumps(params)
    )

def reboot(instance_id):
    client.reboot_db_instance(
        DBInstanceIdentifier=instance_id,
        ForceFailover=False
     )

def modify(instance_id):
    client.modify_db_instance(
        DBInstanceIdentifier=instance_id,
        VpcSecurityGroupIds=[
            'sg-xxxxxxxx',
        ],
        DBParameterGroupName='xxxxxxxxxxxxxxxxxxx',
    )

コードの解説をします。最初に引数のevent変数にインスタンスID、リトライ数、モディファイフラグが設定されているかチェックします。 モディファイフラグはインスタンス起動後の設定変更が終わっているかのフラグになります。本スクリプトはインスタンスの起動が終わるまで自身のLambdaを繰り返し呼出しますが、インスタンスの設定変更後も再度自身のLambdaを繰り返し呼出しをするため、設定が完了しているかを保持するようにしています。

    if (event.has_key('instance_id') and event['instance_id']):
        instance_id = event['instance_id']
    else:
        logger.info('undefined instance id')
        return

    if (event.has_key('retry_count') and isinstance(event['retry_count'],(int))):
        retry_count = event['retry_count']
    else:
        logger.info('undefined retry count')
        return

    if (event.has_key('modified_flag') and isinstance(event['modified_flag'],(bool))):
        modified_flag = event['modified_flag']
    else:
        logger.info('undefined modified flag')
        return

次にリトライ回数をチェックします。これは何らかの原因でLambdaが無限ループに陥るのを防ぐ為に設定しており回数を超えたら処理は終了です。

    if (retry_count > os.environ.get('RETRY_MAX_COUNT')):
        logger.info('over retry count')
        return

環境変数に設定してある最大リトライ回数以下だったら次にインタンスのステータスを取得します。

    response = client.describe_db_instances(
        DBInstanceIdentifier=instance_id
    )
    logger.info(response)
    instance_status = response['DBInstances'][0]['DBInstanceStatus']

最後にモディファイフラグとインタンスのステータス状態によって呼出す処理を分けています。
if分の条件の上から順に以下の処理をしています。

条件 処理内容
設定変更済かつステータスがavailable インスタンスを再起動して処理を終了します。
設定変更済かつステータスがmodifying 10秒待機して再度自身のLambdaを呼出します。
設定未変更かつステータスがavailable 設定変更処理をして再度自身のLambdaを呼出します。
設定未変更かつステータスがcreatingまたはbacking-upまたはmodifying 120秒待機して再度自身のLambdaを呼出します。
    if (modified_flag and instance_status=='available'):
        logger.info("START reboot")
        reboot(instance_id)
        logger.info("END reboot")
    elif (modified_flag and instance_status=='modifying'):
        logger.info("waiting...")
        sleep(10)
        logger.info("START call_modify_lambda")
        call_modify_lambda(instance_id, retry_count,modified_flag)
        logger.info("END call_modify_lambda")
    elif (not modified_flag and instance_status=='available'):
        logger.info("START modify")
        modify(instance_id)
        logger.info("END modify")
        logger.info("START call_modify_lambda")
        call_modify_lambda(instance_id, retry_count,True)
        logger.info("END call_modify_lambda")
    elif (not modified_flag and (instance_status=='creating' or instance_status=='backing-up' or instance_status=='modifying')):
        logger.info("waiting...")
        sleep(120)
        logger.info("START call_modify_lambda")
        call_modify_lambda(instance_id, retry_count, modified_flag)
        logger.info("END call_modify_lambda")

modifyメソッド内ではパラメータグループとVPCセキュリティグループを変更しています。SDKのmodify_db_instanceではその他多くの項目を変更することが出来ますで、必要に応じて設定をしてください。
詳しくはboto3のマニュアルを参照下さい。

def modify(instance_id):
    client.modify_db_instance(
        DBInstanceIdentifier=instance_id,
        VpcSecurityGroupIds=[
            'sg-xxxxxxxx',
        ],
        DBParameterGroupName='xxxxxxxxxxxxxxxxxxx',
    )

環境変数の設定

使用する環境変数は以下になります。

  • RETRY_MAX_COUNT:最大リトライ数。

※今回はt2 microインタンスで試しましたが6回のリトライで処理は完了しています。 rds_restore_lambda02

まとめ

自動削除と比べ少し複雑な処理になってしまいましたが以上でRDSの自動リストアスクリプトの完成です。restore-rdsスクリプトをCloudWatchで時間指定して実行してあげれば自動起動するようになります。今回のLambdaのTimeOutはrestore-rdsが1分、restore-rdsが5分で設定して実行しました。
前回作成した自動削除スクリプトとあわせて参考にしていただければ幸いです。