jmongo 1.0.3 → 1.1.0

data/test/support/keys.rb

@@ -0,0 +1,45 @@
+class Hash
+  # Return a new hash with all keys converted to strings.
+  def stringify_keys
+    dup.stringify_keys!
+  end
+
+  # Destructively convert all keys to strings.
+  def stringify_keys!
+    keys.each do |key|
+      self[key.to_s] = delete(key)
+    end
+    self
+  end
+
+  # Return a new hash with all keys converted to symbols, as long as
+  # they respond to +to_sym+.
+  def symbolize_keys
+    dup.symbolize_keys!
+  end
+
+  # Destructively convert all keys to symbols, as long as they respond
+  # to +to_sym+.
+  def symbolize_keys!
+    keys.each do |key|
+      self[(key.to_sym rescue key) || key] = delete(key)
+    end
+    self
+  end
+
+  alias_method :to_options,  :symbolize_keys
+  #alias_method :to_options!, :symbolize_keys!
+
+  # Validate all keys in a hash match *valid keys, raising ArgumentError on a mismatch.
+  # Note that keys are NOT treated indifferently, meaning if you use strings for keys but assert symbols
+  # as keys, this will fail.
+  #
+  # ==== Examples
+  #   { :name => "Rob", :years => "28" }.assert_valid_keys(:name, :age) # => raises "ArgumentError: Unknown key(s): years"
+  #   { :name => "Rob", :age => "28" }.assert_valid_keys("name", "age") # => raises "ArgumentError: Unknown key(s): name, age"
+  #   { :name => "Rob", :age => "28" }.assert_valid_keys(:name, :age) # => passes, raises nothing
+  def assert_valid_keys(*valid_keys)
+    unknown_keys = keys - [valid_keys].flatten
+    raise(ArgumentError, "Unknown key(s): #{unknown_keys.join(", ")}") unless unknown_keys.empty?
+  end
+end

data/test/support_test.rb

@@ -0,0 +1,19 @@
+__END__
+require './test/test_helper'
+
+class SupportTest < Test::Unit::TestCase
+
+  def test_command_response_succeeds
+    assert Support.ok?('ok' => 1)
+    assert Support.ok?('ok' => 1.0)
+    assert Support.ok?('ok' => true)
+  end
+
+  def test_command_response_fails
+    assert !Support.ok?('ok' => 0)
+    assert !Support.ok?('ok' => 0.0)
+    assert !Support.ok?('ok' => 0.0)
+    assert !Support.ok?('ok' => 'str')
+    assert !Support.ok?('ok' => false)
+  end
+end

data/test/test_helper.rb

@@ -0,0 +1,111 @@
+$:.unshift(File.join(File.dirname(__FILE__), '..', 'lib'))
+require 'rubygems' if RUBY_VERSION < '1.9.0' && ENV['C_EXT']
+require 'jmongo'
+#require 'test/unit'
+require 'awesome_print'
+
+require 'minitest/autorun'
+
+def silently
+  warn_level = $VERBOSE
+  $VERBOSE = nil
+  result = yield
+  $VERBOSE = warn_level
+  result
+end
+
+def apr(obj, prefix = '=====')
+  puts prefix
+  ap obj
+  puts '====='
+end
+
+begin
+  require 'rubygems' if RUBY_VERSION < "1.9.0" && !ENV['C_EXT']
+  silently { require 'shoulda' }
+  silently { require 'mocha' }
+rescue LoadError
+  puts <<MSG
+
+This test suite requires shoulda and mocha.
+You can install them as follows:
+  gem install shoulda
+  gem install mocha
+
+MSG
+
+  exit
+end
+
+require 'bson_ext/cbson' if !(RUBY_PLATFORM =~ /java/) && ENV['C_EXT']
+
+unless defined? MONGO_TEST_DB
+  MONGO_TEST_DB = 'ruby-test-db'
+end
+
+unless defined? TEST_PORT
+  TEST_PORT = ENV['MONGO_RUBY_DRIVER_PORT'] ? ENV['MONGO_RUBY_DRIVER_PORT'].to_i : Mongo::Connection::DEFAULT_PORT
+end
+
+unless defined? TEST_HOST
+  TEST_HOST = ENV['MONGO_RUBY_DRIVER_HOST'] || 'localhost'
+end
+
+class MiniTest::Unit::TestCase
+  include Mongo
+  include BSON
+
+  def self.standard_connection(options={})
+    Connection.new(TEST_HOST, TEST_PORT, options)
+  end
+
+  def standard_connection(options={})
+    self.class.standard_connection(options)
+  end
+
+  def self.host_port
+    "#{mongo_host}:#{mongo_port}"
+  end
+
+  def self.mongo_host
+    TEST_HOST
+  end
+
+  def self.mongo_port
+    TEST_PORT
+  end
+
+  def host_port
+    self.class.host_port
+  end
+
+  def mongo_host
+    self.class.mongo_host
+  end
+
+  def mongo_port
+    self.class.mongo_port
+  end
+
+  def new_mock_socket(host='localhost', port=27017)
+    socket = Object.new
+    socket.stubs(:setsockopt).with(Socket::IPPROTO_TCP, Socket::TCP_NODELAY, 1)
+    socket.stubs(:close)
+    socket
+  end
+
+  def new_mock_db
+    db = Object.new
+  end
+
+  def assert_raise_error(klass, message)
+    begin
+      yield
+    rescue => e
+      assert_equal klass, e.class
+      assert e.message.include?(message), "#{e.message} does not include #{message}."
+    else
+      flunk "Expected assertion #{klass} but none was raised."
+    end
+  end
+end

data/test/threading/threading_with_large_pool_test.rb

@@ -0,0 +1,90 @@
+require './test/test_helper'
+
+# Essentialy the same as test_threading.rb but with an expanded pool for
+# testing multiple connections.
+class TestThreadingLargePool < Test::Unit::TestCase
+
+  include Mongo
+
+  @@db = standard_connection(:pool_size => 50, :timeout => 60).db(MONGO_TEST_DB)
+  @@coll = @@db.collection('thread-test-collection')
+
+  def set_up_safe_data
+    @@db.drop_collection('duplicate')
+    @@db.drop_collection('unique')
+    @duplicate = @@db.collection('duplicate')
+    @unique    = @@db.collection('unique')
+
+    @duplicate.insert("test" => "insert")
+    @duplicate.insert("test" => "update")
+    @unique.insert("test" => "insert")
+    @unique.insert("test" => "update")
+    @unique.create_index("test", :unique => true)
+  end
+
+  def test_safe_update
+    set_up_safe_data
+    threads = []
+    300.times do |i|
+      threads[i] = Thread.new do
+        if i % 2 == 0
+          assert_raise Mongo::OperationFailure do
+            @unique.update({"test" => "insert"}, {"$set" => {"test" => "update"}}, :safe => true)
+          end
+        else
+          @duplicate.update({"test" => "insert"}, {"$set" => {"test" => "update"}}, :safe => true)
+        end
+      end
+    end
+
+    300.times do |i|
+      threads[i].join
+    end
+  end
+
+  def test_safe_insert
+    set_up_safe_data
+    threads = []
+    300.times do |i|
+      threads[i] = Thread.new do
+        if i % 2 == 0
+          assert_raise Mongo::OperationFailure do
+            @unique.insert({"test" => "insert"}, :safe => true)
+          end
+        else
+          @duplicate.insert({"test" => "insert"}, :safe => true)
+        end
+      end
+    end
+
+    300.times do |i|
+      threads[i].join
+    end
+  end
+
+  def test_threading
+    @@coll.drop
+    @@coll = @@db.collection('thread-test-collection')
+
+    1000.times do |i|
+      @@coll.insert("x" => i)
+    end
+
+    threads = []
+
+    10.times do |i|
+      threads[i] = Thread.new do
+        sum = 0
+        @@coll.find().each do |document|
+          sum += document["x"]
+        end
+        assert_equal 499500, sum
+      end
+    end
+
+    10.times do |i|
+      threads[i].join
+    end
+  end
+
+end

data/test/threading_test.rb

@@ -0,0 +1,88 @@
+__END__
+require './test/test_helper'
+
+class TestThreading < Test::Unit::TestCase
+
+  include Mongo
+
+  @@db = standard_connection(:pool_size => 1, :timeout => 30).db(MONGO_TEST_DB)
+  @@coll = @@db.collection('thread-test-collection')
+
+  def set_up_safe_data
+    @@db.drop_collection('duplicate')
+    @@db.drop_collection('unique')
+    @duplicate = @@db.collection('duplicate')
+    @unique    = @@db.collection('unique')
+
+    @duplicate.insert("test" => "insert")
+    @duplicate.insert("test" => "update")
+    @unique.insert("test" => "insert")
+    @unique.insert("test" => "update")
+    @unique.create_index("test", :unique => true)
+  end
+
+  def test_safe_update
+    set_up_safe_data
+    threads = []
+    100.times do |i|
+      threads[i] = Thread.new do
+        if i % 2 == 0
+          assert_raise Mongo::OperationFailure do
+            @unique.update({"test" => "insert"}, {"$set" => {"test" => "update"}}, :safe => true)
+          end
+        else
+          @duplicate.update({"test" => "insert"}, {"$set" => {"test" => "update"}}, :safe => true)
+        end
+      end
+    end
+
+    100.times do |i|
+      threads[i].join
+    end
+  end
+
+  def test_safe_insert
+    set_up_safe_data
+    threads = []
+    100.times do |i|
+      threads[i] = Thread.new do
+        if i % 2 == 0
+          assert_raise Mongo::OperationFailure do
+            @unique.insert({"test" => "insert"}, :safe => true)
+          end
+        else
+          @duplicate.insert({"test" => "insert"}, :safe => true)
+        end
+      end
+    end
+
+    100.times do |i|
+      threads[i].join
+    end
+  end
+
+  def test_threading
+    @@coll.drop
+    @@coll = @@db.collection('thread-test-collection')
+
+    1000.times do |i|
+      @@coll.insert("x" => i)
+    end
+
+    threads = []
+
+    10.times do |i|
+      threads[i] = Thread.new do
+        sum = 0
+        @@coll.find().each do |document|
+          sum += document["x"]
+        end
+        assert_equal 499500, sum
+      end
+    end
+
+    10.times do |i|
+      threads[i].join
+    end
+  end
+end

data/test/tools/auth_repl_set_manager.rb

@@ -0,0 +1,14 @@
+require File.join((File.expand_path(File.dirname(__FILE__))), 'repl_set_manager')
+
+class AuthReplSetManager < ReplSetManager
+  def initialize(opts={})
+    super(opts)
+
+    @key_path = opts[:key_path] || File.join(File.expand_path(File.dirname(__FILE__)), "keyfile.txt")
+    system("chmod 600 #{@key_path}")
+  end
+
+  def start_cmd(n)
+    super + " --keyFile #{@key_path}"
+  end
+end

data/test/tools/keyfile.txt

@@ -0,0 +1 @@
+THIS IS A SECRET KEYFILE FOR REPLICA SETS BWAHAHAHAH

data/test/tools/repl_set_manager.rb

@@ -0,0 +1,377 @@
+#!/usr/bin/ruby
+
+require 'thread'
+
+STDOUT.sync = true
+
+unless defined? Mongo
+  require File.join(File.dirname(__FILE__), '..', '..', 'lib', 'mongo')
+end
+
+class ReplSetManager
+
+  attr_accessor :host, :start_port, :ports, :name, :mongods, :tags, :version
+
+  def initialize(opts={})
+    @start_port = opts[:start_port] || 30000
+    @ports      = []
+    @name       = opts[:name] || 'replica-set-foo'
+    @host       = opts[:host]  || 'localhost'
+    @retries    = opts[:retries] || 30
+    @config     = {"_id" => @name, "members" => []}
+    @durable    = opts.fetch(:durable, false)
+    @path       = File.join(File.expand_path(File.dirname(__FILE__)), "data")
+    @oplog_size = opts.fetch(:oplog_size, 32)
+    @tags = [{"dc" => "ny", "rack" => "a", "db" => "main"},
+             {"dc" => "ny", "rack" => "b", "db" => "main"},
+             {"dc" => "sf", "rack" => "a", "db" => "main"}]
+
+    @arbiter_count   = opts[:arbiter_count]   || 2
+    @secondary_count = opts[:secondary_count] || 2
+    @passive_count   = opts[:passive_count] || 0
+    @primary_count   = 1
+
+    @count = @primary_count + @passive_count + @arbiter_count + @secondary_count
+    if @count > 7
+      raise StandardError, "Cannot create a replica set with #{node_count} nodes. 7 is the max."
+    end
+
+    @mongods   = {}
+    version_string = `mongod --version`
+    version_string =~ /(\d\.\d\.\d)/
+    @version = $1.split(".").map {|d| d.to_i }
+  end
+
+  def start_set
+    begin
+    con = Mongo::Connection.new(@host, @start_port)
+      rescue Mongo::ConnectionFailure
+    end
+
+    if con && ensure_up(1, con)
+      should_start = false
+      puts "** Replica set already started."
+    else
+      should_start = true
+      system("killall mongod")
+      puts "** Starting a replica set with #{@count} nodes"
+    end
+
+    n = 0
+    (@primary_count + @secondary_count).times do
+      init_node(n, should_start) do |attrs|
+        if @version[0] >= 2
+          attrs['tags'] = @tags[n % @tags.size]
+        end
+      end
+      n += 1
+    end
+
+    @passive_count.times do
+      init_node(n, should_start) do |attrs|
+        attrs['priority'] = 0
+      end
+      n += 1
+    end
+
+    @arbiter_count.times do
+      init_node(n, should_start) do |attrs|
+        attrs['arbiterOnly'] = true
+      end
+      n += 1
+    end
+
+    if con && ensure_up(1, con)
+      @mongods.each do |k, v|
+        v['up'] = true
+        v['pid'] = File.open(File.join(v['db_path'], 'mongod.lock')).read.strip
+      end
+    else
+      initiate
+      ensure_up
+    end
+  end
+
+  def cleanup_set
+    system("killall mongod")
+    @count.times do |n|
+      system("rm -rf #{@mongods[n]['db_path']}")
+    end
+  end
+
+  def init_node(n, should_start=true)
+    @mongods[n] ||= {}
+    port = @start_port + n
+    @ports << port
+    @mongods[n]['port'] = port
+    @mongods[n]['db_path'] = get_path("rs-#{port}")
+    @mongods[n]['log_path'] = get_path("log-#{port}")
+    @mongods[n]['start'] = start_cmd(n)
+
+    if should_start
+      system("rm -rf #{@mongods[n]['db_path']}")
+      system("mkdir -p #{@mongods[n]['db_path']}")
+      start(n)
+    end
+
+    member = {'_id' => n, 'host' => "#{@host}:#{@mongods[n]['port']}"}
+
+    if block_given?
+      custom_attrs = {}
+      yield custom_attrs
+      member.merge!(custom_attrs)
+      @mongods[n].merge!(custom_attrs)
+    end
+
+    @config['members'] << member
+  end
+
+  def journal_switch
+    if @version[0] >= 2
+      if @durable
+        "--journal"
+      else
+        "--nojournal"
+      end
+    elsif @durable
+      "--journal"
+    end
+  end
+
+  def start_cmd(n)
+    @mongods[n]['start'] = "mongod --replSet #{@name} --logpath '#{@mongods[n]['log_path']}' " +
+     "--oplogSize #{@oplog_size} #{journal_switch} --dbpath #{@mongods[n]['db_path']} --port #{@mongods[n]['port']} --fork"
+    @mongods[n]['start'] += " --dur" if @durable
+    @mongods[n]['start']
+  end
+
+  def remove_secondary_node
+    primary = get_node_with_state(1)
+    con = get_connection(primary)
+    config = con['local']['system.replset'].find_one
+    secondary = get_node_with_state(2)
+    host_port = "#{@host}:#{@mongods[secondary]['port']}"
+    kill(secondary)
+    @mongods.delete(secondary)
+    @config['members'].reject! {|m| m['host'] == host_port}
+    @config['version'] = config['version'] + 1
+
+    begin
+    con['admin'].command({'replSetReconfig' => @config})
+    rescue Mongo::ConnectionFailure
+    end
+
+    con.close
+
+    return secondary
+  end
+
+  def add_node(n=nil)
+    primary = get_node_with_state(1)
+    con = get_connection(primary)
+    init_node(n || @mongods.length)
+
+    config = con['local']['system.replset'].find_one
+    @config['version'] = config['version'] + 1
+
+    # We expect a connection failure on reconfigure here.
+    begin
+      con['admin'].command({'replSetReconfig' => @config})
+    rescue Mongo::ConnectionFailure
+    end
+
+    con.close
+    ensure_up
+  end
+
+  def kill(node, signal=2)
+    pid = @mongods[node]['pid']
+    puts "** Killing node with pid #{pid} at port #{@mongods[node]['port']}"
+    system("kill -#{signal} #{@mongods[node]['pid']}")
+    @mongods[node]['up'] = false
+    sleep(1)
+  end
+
+  def kill_primary(signal=2)
+    node = get_node_with_state(1)
+    kill(node, signal)
+    return node
+  end
+
+  # Note that we have to rescue a connection failure
+  # when we run the StepDown command because that
+  # command will close the connection.
+  def step_down_primary
+    primary = get_node_with_state(1)
+    con = get_connection(primary)
+    begin
+      con['admin'].command({'replSetStepDown' => 90})
+    rescue Mongo::ConnectionFailure
+    end
+    con.close
+  end
+
+  def kill_secondary
+    node = get_node_with_state(2)
+    kill(node)
+    return node
+  end
+
+  def kill_all_secondaries
+    nodes = get_all_nodes_with_state(2)
+    if nodes
+      nodes.each do |n|
+        kill(n)
+      end
+    end
+  end
+
+  def restart_killed_nodes
+    nodes = @mongods.keys.select do |key|
+      @mongods[key]['up'] == false
+    end
+
+    nodes.each do |node|
+      start(node)
+    end
+
+    ensure_up
+  end
+
+  def get_node_from_port(port)
+    @mongods.keys.detect { |key| @mongods[key]['port'] == port }
+  end
+
+  def start(node)
+    system(@mongods[node]['start'])
+    @mongods[node]['up'] = true
+    sleep(0.5)
+    @mongods[node]['pid'] = File.open(File.join(@mongods[node]['db_path'], 'mongod.lock')).read.strip
+  end
+  alias :restart :start
+
+  def ensure_up(n=nil, connection=nil)
+    print "** Ensuring members are up..."
+
+    attempt(n) do
+      con = connection || get_connection
+      status = con['admin'].command({'replSetGetStatus' => 1})
+      print "."
+      if status['members'].all? { |m| m['health'] == 1 &&
+         [1, 2, 7].include?(m['state']) } &&
+         status['members'].any? { |m| m['state'] == 1 }
+        print "all members up!\n\n"
+        con.close
+        return status
+      else
+        con.close
+        raise Mongo::OperationFailure
+      end
+    end
+
+    return false
+  end
+
+  def primary
+    nodes = get_all_host_pairs_with_state(1)
+    nodes.empty? ? nil : nodes[0]
+  end
+
+  def secondaries
+    get_all_host_pairs_with_state(2)
+  end
+
+  def arbiters
+    get_all_host_pairs_with_state(7)
+  end
+
+  # String used for adding a shard via mongos
+  # using the addshard command.
+  def shard_string
+    str = "#{@name}/"
+    str << @mongods.map do |k, mongod|
+      "#{@host}:#{mongod['port']}"
+    end.join(',')
+    str
+  end
+
+  private
+
+  def initiate
+    con = get_connection
+
+    attempt do
+      con['admin'].command({'replSetInitiate' => @config})
+    end
+
+    con.close
+  end
+
+  def get_all_nodes_with_state(state)
+    status = ensure_up
+    nodes = status['members'].select {|m| m['state'] == state}
+    nodes = nodes.map do |node|
+      host_port = node['name'].split(':')
+      port = host_port[1] ? host_port[1].to_i : 27017
+      @mongods.keys.detect {|key| @mongods[key]['port'] == port}
+    end
+
+    nodes == [] ? false : nodes
+  end
+
+  def get_node_with_state(state)
+    status = ensure_up
+    node = status['members'].detect {|m| m['state'] == state}
+    if node
+      host_port = node['name'].split(':')
+      port = host_port[1] ? host_port[1].to_i : 27017
+      key = @mongods.keys.detect {|n| @mongods[n]['port'] == port}
+      return key
+    else
+      return false
+    end
+  end
+
+  def get_all_host_pairs_with_state(state)
+    status = ensure_up
+    nodes = status['members'].select {|m| m['state'] == state}
+    nodes.map do |node|
+      host_port = node['name'].split(':')
+      port = host_port[1] ? host_port[1].to_i : 27017
+      [host, port]
+    end
+  end
+
+  def get_connection(node=nil)
+    con = attempt do
+      if !node
+        node = @mongods.keys.detect {|key| !@mongods[key]['arbiterOnly'] && @mongods[key]['up'] }
+      end
+      con = Mongo::Connection.new(@host, @mongods[node]['port'], :slave_ok => true)
+    end
+
+    return con
+  end
+
+  def get_path(name)
+    File.join(@path, name)
+  end
+
+  def attempt(retries=nil)
+    raise "No block given!" unless block_given?
+    count = 0
+
+    while count < (retries || @retries) do
+      begin
+        return yield
+        rescue Mongo::OperationFailure, Mongo::ConnectionFailure => ex
+          sleep(2)
+          count += 1
+      end
+    end
+
+    puts "NO MORE ATTEMPTS"
+    raise ex
+  end
+
+end