require 'method_struct'

class CheckForN1Queries < MethodStruct.new(:generator, :operation, :debug)
  IGNORED_SQL = [
    /^SHOW/, /^EXPLAIN/, /^PRAGMA (?!(table_info))/, /^SELECT currval/, /^SELECT CAST/,
    /^SELECT @@IDENTITY/, /^SELECT @@ROWCOUNT/, /^SAVEPOINT/, /^ROLLBACK TO SAVEPOINT/,
    /^RELEASE SAVEPOINT/, /^SHOW max_identifier_length/
  ]

  def call
    2.times { generator.call }

    log("BEFORE:")
    old_count = count_total_queries { operation.call }

    generator.call

    log("AFTER:")
    new_count = count_total_queries { operation.call }

    [old_count, new_count]
  end

  private
  def log(msg)
    puts(["\n", msg, "\n"].join) if debug
  end

  def count_total_queries(&block)
    query_count = 0

    callback = proc do |name, start, finish, message_id, values|
      next if values[:name] == 'CACHE' || IGNORED_SQL.any? { |r| values[:sql] =~ r }

      log(values[:sql])
      query_count += 1
    end

    subscribed(callback, "sql.active_record", &block)

    query_count
  end

  # Backported from Rails 3.2.1
  def subscribed(callback, *args, &block)
    subscriber = ActiveSupport::Notifications.subscribe(*args, &callback)
    yield
  ensure
    ActiveSupport::Notifications.unsubscribe(subscriber)
  end
end