dump.py で取得したdumpファイルからDynamoDBを復元するプログラム

restore.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
Restore DynamoDB Tables from Dump File.
This program does not cover all DynamoDB format,
but It's good enough for me.

dump.py can dump Tables to JSON file.

"""
import sys
import json
import time
import codecs
import decimal
import logging
import argparse

import boto.dynamodb2

from boto.dynamodb2.fields import HashKey, RangeKey
from boto.dynamodb2.table import Table
from boto.dynamodb2.types import STRING, NUMBER

# --------------------------------
# Overwrite These Parameters
# --------------------------------

REGION='ap-northeast-1'
AWS_ACCESS_KEY_ID=''
AWS_SECRET_ACCESS_KEY=''

# --------------------------------

def _bool(s):
  return s in ['True', 'true', 'TRUE']


def _create_table(table, conn):
  exists_table = conn.list_tables()['TableNames']
  tablename = table['TableName']
  if tablename in exists_table:
    logging.info('There is "{}" Table Already.'.format(tablename))
    return
  
  schema = []
  if table.has_key('HashKey'):
    schema.append(HashKey(table['HashKey']['Name'], data_type=table['HashKey']['Type']))
  if table.has_key('RangeKey'):
    schema.append(RangeKey(table['RangeKey']['Name'], data_type=table['RangeKey']['Type']))
  
  Table.create(tablename,
    schema=schema,
    throughput={
      'read': table['ReadCapacityUnits'],
      'write': table['WriteCapacityUnits'],
    },
    connection=conn)
  
  time.sleep(1)
  return

def _group(xs, n):
  result = []
  group = []
  for x in xs:
    group.append(x)
    if len(group) >= n:
      result.append(group)
      group = []
  if len(group) > 0:
    result.append(group)
  return result


def _to_value(type, value):
  if type == NUMBER:
    return decimal.Decimal(value)
  return value


def _to_data(item):
  data = {}
  for (k,v) in item.items():
    data[k] = _to_value(v['Type'], v['Value'])
  return data


def _write_all_items(table_info, conn, force=False):
  table = Table(table_info['TableName'], connection=conn)
  groups = _group(table_info['Items'], 10)
  for group in groups:
    with table.batch_write() as batch:
      for item in group:
        batch.put_item(data=_to_data(item), overwrite=force)
  return


def restore(dumpfile):
  '''
  dumpしたJSONデータから環境を復元します
  
  '''
  conn = boto.dynamodb2.connect_to_region(
    REGION,
    aws_access_key_id=AWS_ACCESS_KEY_ID,
    aws_secret_access_key=AWS_SECRET_ACCESS_KEY
  )
  
  with codecs.open(dumpfile, 'r', 'utf_8') as f:
    tables = json.load(f)
  
  for table in tables:
    _create_table(table, conn)
    _write_all_items(table, conn, True)
  

if __name__ == '__main__':
  parser = argparse.ArgumentParser(description=u'DynamoDB Restore from dumpfile')
  parser.add_argument('--dumpfile', type=str, help=u'Dump file from dump.py output', required=True)
  
  args = parser.parse_args(sys.argv[1:])
  restore(args.dumpfile)