query_subscriber.rb

# To prevent memory leak possibilities with user errors there are constraints.
#
# * Only allows one subscriber to be active at a time. When
#   instantiated it always removes all QuerySubscribers subscribed to the :client.
# * It requires a block of assertions where the instantiated monitor is
#   available for access to the #query_count.
# * It is meant for a single instantiated use, and removes the subscriber from
#   the client's @monitoring#subscribers Hash after the block executes. It also
#   clears the recorded queries from the subscriber.
#   This means you can use the QueryMonitor as many times as you like in one off
#   situations, but you will not be able to use the same monitor across tests,
#   without subscribing again.
# * It is recommended a user only uses the .listen interface.
#
# Example:
#
# describe 'the number of queries when calling first on supplier properties' do
#   it 'queries supplier properties with find once' do
#     QueryMonitor.listen('supplier_properties') do |monitor|
#       Property.first
#       monitor.query_count.should == 1
#     end
#   end
# end

class QueryMonitor
  attr_reader :client, :subscriber

  def initialize(table_name, client_name, command)
    @client = Mongoid.client client_name
    @subscriber = QuerySubscriber.new table_name, command

    # Prevents a user from manually instantiating and overloading monitor
    # with subscribers. This is a redundancy step since #subscribe silently fails if
    # the @subscriber is already subscribed.
    remove_old_subscribers
  end

  delegate :query_count, :queries=, to: :subscriber

  def self.listen(table_name, client_name: :availability, command: 'find')
    new(table_name, client_name, command).tap do |monitor|
      monitor.subscribe
      yield monitor
      monitor.queries = []
      monitor.unsubscribe
    end
  end

  def subscribe
    client.subscribe Mongo::Monitoring::COMMAND, subscriber unless subscribed?
  end

  def unsubscribe
    command_subscribers.delete subscriber
  end

  def subscribed?
    command_subscribers.include? subscriber
  end

  def subscribers
    command_subscribers.select { |s| s.is_a? subscriber.class }
  end

  def remove_old_subscribers
    subscribers.each { |s| command_subscribers.delete s }
  end

  def command_subscribers
    client.instance_variable_get("@monitoring").subscribers[Mongo::Monitoring::COMMAND]
  end

  class QuerySubscriber
    attr_reader :command, :table_name
    attr_accessor :queries

    def initialize(table_name, command)
      @table_name = table_name
      @command = command
      @queries = []
    end

    def started(event)
      queries.push event if matched_query? event
    end

    def failed(event)
    end

    def succeeded(event)
    end

    def query_count
      queries.size
    end

    def matched_query?(event)
      event.command.dig(command) == table_name
    end
  end
end