#!/usr/bin/env ruby

require 'active_record'
require 'lhm/sql_helper'
require 'optparse'

module Lhm
  class KillQueue
    def initialize
      @port = 3306
      @grace = 10
      @tiny = 0.1
      @marker = "%#{ SqlHelper.annotation }%"

      OptionParser.new do |opts|
        opts.on('-h', '--hostname HOSTNAME') { |v| @hostname = v }
        opts.on('-u', '--username USERNAME') { |v| @username = v }
        opts.on('-p', '--password PASSWORD') { |v| @password = v }
        opts.on('-d', '--database DATABASE') { |v| @database = v }
        opts.on('-m', '--mode MODE') { |v| @mode = v.to_sym }
        opts.on('-y', '--confirm') { |v| @confirm = true }
      end.parse!

      unless (@hostname && @username && @password && @database)
        abort usage
      end

      unless ([:kill, :master, :slave].include?(@mode))
        abort 'specify -m kill OR -m master OR -m slave'
      end

      connect
    end

    def usage
      <<-desc.gsub(/^      /, '')
        kills queries on the given server after detecting 'lock table#{ @marker }'.
        usage:
          lhm-kill-queue -h hostname -u username -p password -d database \\
          (-m kill | -m master | -m slave) [--confirm]

      desc
    end

    def run
      case @mode
      when :kill then kill
      when :master then master
      when :slave then slave
      end
    end

    def kill
      lock = trip
      kill_process(lock)
    end

    def master
      lock = trip
      puts "starting to kill non lhm processes in #{ @grace } seconds"
      sleep(@grace + @tiny)

      [list_non_lhm].flatten.each do |process|
        kill_process(process)
        sleep(@tiny)
      end
    end

    def slave
      lock = trip
      puts "starting to kill non lhm SELECT processes in #{ @grace } seconds"
      sleep(@grace + @tiny)

      [list_non_lhm].flatten.each do |process|
        if (select?(process))
          kill_process(process)
          sleep(@tiny)
        end
      end
    end

  private

    def connect
      ActiveRecord::Base.establish_connection({
        :adapter => 'mysql',
        :host => @hostname,
        :port => @port,
        :username => @username,
        :password => @password,
        :database => @database
      })
    end

    def connection
      ActiveRecord::Base.connection
    end

    def list_non_lhm
      select_processes %Q(
        info not like '#{ @marker }' and time > #{ @grace } and command = 'Query'
      )
    end

    def trip
      until res = select_processes("info like 'lock table#{ @marker }'").first
        sleep @tiny
        print '.'
      end

      res
    end

    def kill_process(process_id)
      puts "killing #{ select_statement(process_id) }"

      if (@confirm)
        print "confirm ('y' to confirm): "

        if (gets.strip != 'y')
          puts 'skipped.'
          return
        end
      end

      connection.execute("kill #{ process_id }")
      puts "killed #{ process_id }"
    end

    def select?(process)
      if statement = select_statement(process)
        case statement
        when /delete/i then false
        when /update/i then false
        when /insert/i then false
        else
          !!statement.match(/select/i)
        end
      end
    end

    def select_statement(process)
      if process
        value %Q(
          select info from information_schema.processlist where id = #{ process }
        )
      end
    end

    def select_processes(predicate)
      values %Q(
        select id from information_schema.processlist
         where db = '#{ @database }'
           and user = '#{ @username }'
           and #{ predicate }
      )
    end

    def value(statement)
      connection.select_value(statement)
    end

    def values(statement)
      connection.select_values(statement)
    end
  end
end

killer = Lhm::KillQueue.new
killer.run