parallel-restore.rb

#!/usr/bin/env ruby

require 'fileutils'
require 'optparse'

ARGV << '-h' unless [6,8].include? ARGV.length

$options = {
  sql_host: 'localhost'
}

OptionParser.new do |opts|
  opts.banner = "Usage: #{$PROGRAM_NAME} [options]"
  opts.on('-d', '--backup-dir PATH', '/path/to/backup/directory') { |v| $options[:backup_dir] = v }
  opts.on('-s', '--batch-size SIZE', 'Number of FrontEnds to generate scripts for') { |v| $options[:batch_size] = v }
  opts.on('-i', '--batch-id ID', 'A number from 1 to Batch Size') { |v| $options[:batch_id] = v }
  opts.on('-S', '--sql-host HOST', String, 'The hostname of the postgresql server. (default: localhost)') { |v|
    $options[:sql_host] = v
  }
  opts.on_tail('-h', '--help', 'Show this message') do
    puts opts
    exit
  end
end.parse!

def org_path(org)
  File.expand_path(File.join($options[:backup_dir], 'organizations', org))
end

def tree_size(org_path)
  Dir[File.join(org_path, '**/*')].size
end

def organizations_path
  File.expand_path(File.join($options[:backup_dir], 'organizations'))
end

def populate_org_hash
  orgs = {}
  Dir[File.join(organizations_path, '**')].each do |org_path|
    org = File.basename(org_path)
    orgs[org] = {}
    orgs[org]['path'] = org_path
    orgs[org]['size'] = tree_size(org_path)
  end
  orgs
end

def batch_orgs_by_size(orgs)
  sorted_list = Hash[orgs.sort_by { |_k, v| v['size'] }.reverse]
  index = 0
  h = {}
  sorted_list.each do |s|
    batch = (index % $options[:batch_size].to_i) + 1
    h[batch] ||= []
    h[batch] << s
    index += 1
  end
  h
end

def write_script(batches)
  full_org_list_by_size = []
  batches.each_key do |key|
    full_org_list_by_size << batches[key].map {|k, v| k}.join(' ')
  end

  first_org = full_org_list_by_size.join(' ').split(' ')[0]
  first_time = ''
  this_org_list = batches[$options[:batch_id].to_i].map { |k, _v| k }

  if $options[:batch_id].to_i == 1
    this_org_list.shift
    first_time = "( time knife ec restore #{$options[:backup_dir]} --with-user-sql --sql-host #{$options[:sql_host]} -c /etc/opscode/pivotal.rb --only-org #{first_org} ) 2>&1 | tee -a restore.log"
  end

  fname = "chef_restore-#{$options[:batch_id].to_i}of#{$options[:batch_size].to_i}.sh"

  script = <<-EOS
#!/bin/bash

# run this script with sudo -E ie. sudo -E ./#{fname}

SECONDS=0

export PATH=/opt/opscode/embedded/bin:$PATH

# FULL ORG LIST
# #{full_org_list_by_size.join(' ')}
#
# THIS WAS GENERATED FOR SERVER #{$options[:batch_id].to_i} OF #{$options[:batch_size].to_i}
ORGS='#{this_org_list.join(" ")}'

check_for_org () {
 `knife raw -m GET "/organizations" -c /etc/opscode/pivotal.rb | grep -q "\\"#{first_org}\\""`
}

#{first_time}

for org in $ORGS; do
  check_for_org
  while [ $? -ne 0 ]; do
    echo "Checking for existence of Org #{first_org} before starting.."
    sleep 10
    check_for_org
  done
  ( time knife ec restore #{$options[:backup_dir]} -c /etc/opscode/pivotal.rb --only-org $org --skip-users --skip-useracl ) 2>&1 | tee -a restore.log
done

duration=$SECONDS
echo
echo "A total of $(($duration / 60)) minutes and $(($duration % 60)) seconds elapsed."
EOS

  puts "Writing: #{fname}"
  File.open(fname, 'w') do |f|
    f.write(script)
  end
  FileUtils.chmod "u=rwx", "#{fname}"
end

orgs = populate_org_hash
batches = batch_orgs_by_size(orgs)
write_script(batches)