parallel-restorev3.rb

#!/usr/bin/env ruby

# Use this script if you have 3 or more Chef Organizations
# and wish to speed up restores by using parallelism.

require 'fileutils'
require 'optparse'

CONCURRENCY = 1
BACKUP_LOG = "restore_#{Time.now.strftime('%m%d%Y%H%M')}.log".freeze
SCRIPT_OUTPUT = 'chef-restore-by-org.sh'.freeze
KNIFE_CMD = "hab pkg exec chef/knife-ec-backup knife"

ARGV << '-h' if ARGV.length != 6
# rubocop:disable Style/GlobalVars
$options = {}
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('-k', '--kniferb PATH', '/path/to/pivotal.rb') { |v| $options[:knife_rb] = v }
  opts.on('-a', '--args PATH', '/path/to/extra-args.txt') { |v| $options[:args_file] = v }
  opts.on_tail('-h', '--help', 'Show this message') do
    puts opts
    exit
  end
end.parse!

EXTRA_ARGS_FILE = $options[:args_file]
KNIFE_RB = $options[:knife_rb]

def extra_args
  File.readlines(EXTRA_ARGS_FILE).map(&:chomp).join(' ')
end

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 orgs_by_size(orgs)
  Hash[orgs.sort_by { |_k, v| v['size'] }.reverse].keys
end

# rubocop:disable Metrics/MethodLength
def write_script(list)
  last_org = list.pop
  penultimate_org = list.pop
  org_with_users = "$KNIFE_CMD ec restore #{$options[:backup_dir]} --yes --purge -c #{KNIFE_RB} \
  --concurrency #{CONCURRENCY} --with-user-sql --skip-useracl --only-org #{last_org} #{extra_args} 2>&1 \
  | while IFS= read -r line; do echo \"ORG:#{last_org} $(date) $line\"; done >> #{BACKUP_LOG}"
  org_with_useracl = "$KNIFE_CMD ec restore #{$options[:backup_dir]} --yes --purge -c #{KNIFE_RB} \
  --concurrency #{CONCURRENCY} --skip-users --only-org #{penultimate_org} #{extra_args} 2>&1 \
  | while IFS= read -r line; do echo \"ORG:#{penultimate_org} $(date) $line\"; done >> #{BACKUP_LOG}"
  script = <<-SCRIPT
#!/bin/bash
ORGS="#{list.join(' ')}"
export PATH=/opt/chefdk/bin:$PATH

echo "Saving logs to #{BACKUP_LOG}"

echo "Step 1. restoring smallest org and all users"
#{org_with_users} || true

echo "Step 2. backgrounding restore job for orgs"

for org in $ORGS; do
  $KNIFE_CMD ec restore #{$options[:backup_dir]} --yes --purge --concurrency #{CONCURRENCY} -c #{KNIFE_RB} --only-org $org --skip-users --skip-useracl #{extra_args} 2>&1 | while IFS= read -r line; do echo \"ORG:$org $(date) $line\"; done >> #{BACKUP_LOG} &
done

while [[ -n $(jobs -r) ]]; do echo -n "."; sleep 60; done

echo
echo "Step 3. restoring user acls along with last org"
#{org_with_useracl} || true

echo "Step 4. gathering results"

ORGS="${ORGS} #{last_org} #{penultimate_org}"
for org in $ORGS; do
  if grep -q "^ORG:$org .*Finished" #{BACKUP_LOG}; then
    echo "$org succeeded"
  else
    echo "$org failed"
  fi
done
SCRIPT
  puts "Writing: #{SCRIPT_OUTPUT}"
  File.open(SCRIPT_OUTPUT, 'w') do |f|
    f.write(script)
  end
  FileUtils.chmod 'u=rwx', SCRIPT_OUTPUT.to_s
end
# rubocop:enable Metrics/MethodLength
# rubocop:enable Style/GlobalVars

write_script orgs_by_size populate_org_hash