#!/usr/bin/env ruby

require 'dbi'
require 'benchmark'

module Bench
  class Unit

    def setup; end
    def teardown; end

    def go
      benches = self.public_methods.select { |m| m.to_s[0,6] == "bench_" }
      return unless benches.size > 0
      Benchmark.bm(20) do |x|
        for b in benches
          name = b[6,b.size - 1]
          name.gsub!(/_([a-z])/) { |m| $1.upcase }
          name[0,1] = name[0,1].upcase
          callable = method(b.to_sym)

          setup
          x.report(name) { for i in 1..2_000; callable.call(); end }
          teardown
        end
      end
    end
  end
end

class BenchDbdDo < Bench::Unit
  DSN = 'dbi:Pg:postgres'
  DBUSER = 'user'
  DBPASS = ''

  def setup
    @dbh = DBI.connect(DSN, DBUSER, DBPASS)
    @dbh.do 'CREATE TABLE bench (date TIMESTAMP, value FLOAT)'
    @today = Date.today
  end

  def teardown
    @dbh.do 'DROP TABLE bench'
    @dbh.disconnect rescue nil
  end

  def bench_insert_do_no_bind
    @dbh.do("INSERT INTO bench VALUES ('#{(@today + rand(100)).to_s}', #{10 + rand * 300})")
  end

  def bench_insert_do_bindvars
    @dbh.do('INSERT INTO bench VALUES (?, ?)', @today + rand(100), 10 + rand * 300)
  end

  def bench_insert_prepare_block
    @dbh.prepare('INSERT INTO bench VALUES (?, ?)') do |sth|
      sth.execute(@today + rand(100), 10 + rand * 300)
    end
  end

  def bench_insert_prepare_finish
    sth = @dbh.prepare('INSERT INTO bench VALUES (?, ?)')
    sth.execute(@today + rand(100), 10 + rand * 300)
    sth.finish
  end

  def bench_update_do_no_bind
    @dbh.do('UPDATE bench SET value = value + 1')
  end

  def bench_update_do_bindvars
    @dbh.do('UPDATE bench SET value = value + ?', 1)
  end
end

class BenchDbdDoAsync < BenchDbdDo
  def setup
    super
    @dbh['pg_async'] = true
  end
end

if $0 == __FILE__
  ObjectSpace.each_object(Class) { |klass|
    next unless klass.ancestors.include? ::Bench::Unit
    next if klass == ::Bench::Unit
    title_bar = '-' * 72
    title_bar[2, klass.name.size + 4] = '[ ' + klass.name + ' ]'
    puts title_bar
    klass.new.go
    puts
  }
end