Rake tasks for Shoryuken / AWS SQS queue creation and management

queues.rake

# See https://github.com/phstc/shoryuken/wiki/Creating-a-queue

namespace :queues do

  desc "List all SQS queues"
  task :list => :environment do
    queues.each { |queue_url| puts queue_url }
  end

  desc "Get details for a specific SQS Queue"
  task :queue, [:identifier] => :environment do |t, args|
    unless identifier = args.identifier
      fail 'No SQS queue name or url was specified'
    end
    puts queue(identifier)
  end

  desc "Ensure existence of a SQS Dead Letter Queue"
  task :dead_letter_queue => :environment do
    unless name = ENV['SQS_FAILURE_QUEUE']
      fail 'SQS_FAILURE_QUEUE env variable is not set'
    end
    if queue_names.include?(name)
      puts "SQS Dead Letter Queue '#{name}' exists"
    else
      sqs.create_queue(queue_name: name)
      puts "SQS Dead Letter Queue '#{name}' created"
    end
  end

  desc "Create a new SQS queue with dead job support"
  task :create, [:name, :retries] => :dead_letter_queue do |t, args|
    unless name = args.name
      fail "An SQS queue name must be specified"
    end
    if queue_names.include?(name)
      puts "SQS Queue '#{name}' exists"
    else
      Rake::Task['queues:dead_letter_queue'].invoke
      retries = args.retries || 7
      arn = queue(ENV['SQS_FAILURE_QUEUE'])['QueueArn']
      attrs = { 'RedrivePolicy' =>
        %Q{{"maxReceiveCount":"#{retries}", "deadLetterTargetArn":"#{arn}"}"}
      }
      sqs.create_queue(queue_name: name, attributes: attrs)
      puts "SQS Queue '#{name}' created"
    end
  end

  desc "Setup the default SQS queues for a new project"
  task :setup, [:retries] => :environment do |t, args|
    unless name = ENV['SQS_DEFAULT_QUEUE']
      fail 'SQS_DEFAULT_QUEUE env variable is not set'
    end
    retries = args.retries || 7
    Rake::Task['queues:create'].invoke(name, retries)
  end

  desc "Delete an SQS queue"
  task :delete, [:identifier] => :environment do
    unless identifier = args.identifier
      fail 'No SQS queue name or url was specified'
    end
    sqs.delete_queue(queue_url: queue_url(identifier), attributes: attrs)
    puts "SQS Queue '#{name}' deleted"
  end

  desc "Delete all SQS queues"
  task :delete_all => :environment do
    STDOUT.puts "Are sure you want to delete all SQS Queues? Type 'CONFIRM' to confirm:"
    input = STDIN.gets.chomp
    unless input == 'CONFIRM'
      fail "Aborting deletion of SQS Queues. You entered: #{input}"
    end
    queues.each do |queue|
      Rake::Task['queues:delete'].invoke(queue)
    end
  end

  def sqs
    @sqs ||= Aws::SQS::Client.new
  end

  def queues
    sqs.list_queues.inject([]) do |list, page|
      list.concat(page.queue_urls)
    end
  end

  def queue_names
    queues.map { |queue| queue.rpartition('/').last }
  end

  def queue_url(identifier)
    if /^https?:\/\//.match(identifier)
      identifier
    else
      queues.find { |queue| queue.rpartition('/').last == identifier }
    end
  end

  def queue(identifier)
    sqs.get_queue_attributes(queue_url: queue_url(identifier), attribute_names: ['All']).attributes
  end

end