Dynamodb throughput scaling rake task

aws_dynamodb.rake

# OVERWRITE_AWS_CONFIG is used to create a dynamo_db client which use production environment when using rake task in development

namespace :aws do
  namespace :dynamodb do

    desc "DynamoDB scale by table_name"
    task :scale, [:table_name, :read, :write] => :environment do |t, args|
      args.with_defaults(read: 1200, write: 1200)
      requested_read = args[:read].to_i
      requested_write = args[:write].to_i

      dynamo_db = AWS::DynamoDB.new(OVERWRITE_AWS_CONFIG)
      
      puts "Loading #{args[:table_name]} schema..."
      table = dynamo_db.tables[args[:table_name]].load_schema
      sleep 1 while table.status == :updating

      actual_read = table.read_capacity_units
      actual_write = table.write_capacity_units

      # Very usefull when you only want to change one value without knowledge of the other one
      requested_read = actual_read if requested_read < 0
      requested_write = actual_write if requested_write < 0

      if actual_read != requested_read or actual_write != requested_write
        if requested_read > actual_read or requested_write > actual_write # SCALE UP
          while actual_read < requested_read or actual_write < requested_write
            puts "Provisioning throughput on #{args[:table_name]} : read: #{[actual_read * 2, requested_read].min}, write: #{[actual_write * 2, requested_write].min}"
            table.provision_throughput read_capacity_units: [actual_read * 2, requested_read].min, write_capacity_units: [actual_write * 2, requested_write].min
            sleep 1 while table.status == :updating
            actual_read = table.read_capacity_units
            actual_write = table.write_capacity_units
            puts "Provisioned throughput on #{args[:table_name]} : read: #{actual_read}, write: #{actual_write}"
          end
        else # SCALE DOWN
          puts "Provisioning throughput on #{args[:table_name]} : read: #{requested_read}, write: #{requested_write}"
          table.provision_throughput read_capacity_units: requested_read, write_capacity_units: requested_write
          sleep 1 while table.status == :updating
          actual_read = table.read_capacity_units
          actual_write = table.write_capacity_units
          puts "Provisioned throughput on #{args[:table_name]} : read: #{actual_read}, write: #{actual_write}"
        end
      end
    end

  end
end

Usage : aws:dynamodb:scale[table_name,500,500]